Skip to content

Commit

Permalink
Fix: Fix issue where Universe Domain is not correctly set.
Browse files Browse the repository at this point in the history 
This PR fixes an issue where the universe domain is not correctly set. This also adds integration tests to verify the change.
  • Loading branch information
@PhongChuong
PhongChuong committed Feb 28, 2024
1 parent dfa8d53 commit 466cd4c
Show file tree
Hide file tree
Showing 6 changed files with 361 additions and 3 deletions.
1 change: 1 addition & 0 deletions ...age/src/main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ public static EnhancedBigQueryReadStub create(
// Configure the base settings.
BigQueryReadStubSettings.Builder baseSettingsBuilder =
BigQueryReadStubSettings.newBuilder()
.setUniverseDomain(settings.getUniverseDomain())
.setTransportChannelProvider(settings.getTransportChannelProvider())
.setEndpoint(settings.getEndpoint())
.setHeaderProvider(settings.getHeaderProvider())
Expand Down
1 change: 0 additions & 1 deletion ...main/java/com/google/cloud/bigquery/storage/v1/stub/EnhancedBigQueryReadStubSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,6 @@ protected Builder(ClientContext clientContext) {

// Defaults provider
BigQueryReadStubSettings.Builder baseDefaults = BigQueryReadStubSettings.newBuilder();
setEndpoint(baseDefaults.getEndpoint());
setTransportChannelProvider(defaultTransportChannelProvider());
setCredentialsProvider(baseDefaults.getCredentialsProvider());
setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
Expand Down
1 change: 1 addition & 0 deletions ...rc/main/java/com/google/cloud/bigquery/storage/v1beta2/stub/EnhancedBigQueryReadStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ public static EnhancedBigQueryReadStub create(
// Configure the base settings.
BigQueryReadStubSettings.Builder baseSettingsBuilder =
BigQueryReadStubSettings.newBuilder()
.setUniverseDomain(settings.getUniverseDomain())
.setTransportChannelProvider(settings.getTransportChannelProvider())
.setEndpoint(settings.getEndpoint())
.setHeaderProvider(settings.getHeaderProvider())
Expand Down
1 change: 0 additions & 1 deletion ...java/com/google/cloud/bigquery/storage/v1beta2/stub/EnhancedBigQueryReadStubSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,6 @@ protected Builder(ClientContext clientContext) {

// Defaults provider
BigQueryReadStubSettings.Builder baseDefaults = BigQueryReadStubSettings.newBuilder();
setEndpoint(baseDefaults.getEndpoint());
setTransportChannelProvider(defaultTransportChannelProvider());
setCredentialsProvider(baseDefaults.getCredentialsProvider());
setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
Expand Down
180 changes: 179 additions & 1 deletion ...ystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,20 @@

package com.google.cloud.bigquery.storage.v1.it;

import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.core.InstantiatingExecutorProvider;
import com.google.api.gax.rpc.ServerStream;
import com.google.api.gax.rpc.UnauthenticatedException;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.RetryOption;
import com.google.cloud.ServiceOptions;
import com.google.cloud.bigquery.BigQuery;
Expand Down Expand Up @@ -54,7 +59,9 @@
import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
import com.google.common.base.Preconditions;
import com.google.protobuf.Timestamp;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.util.ArrayList;
Expand Down Expand Up @@ -82,7 +89,6 @@

/** Integration tests for BigQuery Storage API. */
public class ITBigQueryStorageTest {

private static final Logger LOG = Logger.getLogger(ITBigQueryStorageTest.class.getName());
private static final String DATASET = RemoteBigQueryHelper.generateDatasetName();
private static final String DESCRIPTION = "BigQuery Storage Java client test dataset";
Expand All @@ -91,6 +97,66 @@ public class ITBigQueryStorageTest {
private static String parentProjectId;
private static BigQuery bigquery;

private static final String FAKE_JSON_CRED_WITH_GOOGLE_DOMAIN =
"{\n"
+ " \"private_key_id\": \"somekeyid\",\n"
+ " \"private_key\": \"-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggS"
+ "kAgEAAoIBAQC+K2hSuFpAdrJI\\nnCgcDz2M7t7bjdlsadsasad+fvRSW6TjNQZ3p5LLQY1kSZRqBqylRkzteMOyHg"
+ "aR\\n0Pmxh3ILCND5men43j3h4eDbrhQBuxfEMalkG92sL+PNQSETY2tnvXryOvmBRwa/\\nQP/9dJfIkIDJ9Fw9N4"
+ "Bhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nknddadwkwewcVxHFhcZJO+XWf6ofLUXpRwiTZakGMn8EE1uVa2"
+ "LgczOjwWHGi99MFjxSer5m9\\n1tCa3/KEGKiS/YL71JvjwX3mb+cewlkcmweBKZHM2JPTk0ZednFSpVZMtycjkbLa"
+ "\\ndYOS8V85AgMBewECggEBAKksaldajfDZDV6nGqbFjMiizAKJolr/M3OQw16K6o3/\\n0S31xIe3sSlgW0+UbYlF"
+ "4U8KifhManD1apVSC3csafaspP4RZUHFhtBywLO9pR5c\\nr6S5aLp+gPWFyIp1pfXbWGvc5VY/v9x7ya1VEa6rXvL"
+ "sKupSeWAW4tMj3eo/64ge\\nsdaceaLYw52KeBYiT6+vpsnYrEkAHO1fF/LavbLLOFJmFTMxmsNaG0tuiJHgjshB\\"
+ "n82DpMCbXG9YcCgI/DbzuIjsdj2JC1cascSP//3PmefWysucBQe7Jryb6NQtASmnv\\nCdDw/0jmZTEjpe4S1lxfHp"
+ "lAhHFtdgYTvyYtaLZiVVkCgYEA8eVpof2rceecw/I6\\n5ng1q3Hl2usdWV/4mZMvR0fOemacLLfocX6IYxT1zA1FF"
+ "JlbXSRsJMf/Qq39mOR2\\nSpW+hr4jCoHeRVYLgsbggtrevGmILAlNoqCMpGZ6vDmJpq6ECV9olliDvpPgWOP+\\nm"
+ "YPDreFBGxWvQrADNbRt2dmGsrsCgYEAyUHqB2wvJHFqdmeBsaacewzV8x9WgmeX\\ngUIi9REwXlGDW0Mz50dxpxcK"
+ "CAYn65+7TCnY5O/jmL0VRxU1J2mSWyWTo1C+17L0\\n3fUqjxL1pkefwecxwecvC+gFFYdJ4CQ/MHHXU81Lwl1iWdF"
+ "Cd2UoGddYaOF+KNeM\\nHC7cmqra+JsCgYEAlUNywzq8nUg7282E+uICfCB0LfwejuymR93CtsFgb7cRd6ak\\nECR"
+ "8FGfCpH8ruWJINllbQfcHVCX47ndLZwqv3oVFKh6pAS/vVI4dpOepP8++7y1u\\ncoOvtreXCX6XqfrWDtKIvv0vjl"
+ "HBhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nkndj5uNl5SiuVxHFhcZJO+XWf6ofLUregtevZakGMn8EE1uVa"
+ "2AY7eafmoU/nZPT\\n00YB0TBATdCbn/nBSuKDESkhSg9s2GEKQZG5hBmL5uCMfo09z3SfxZIhJdlerreP\\nJ7gSi"
+ "dI12N+EZxYd4xIJh/HFDgp7RRO87f+WJkofMQKBgGTnClK1VMaCRbJZPriw\\nEfeFCoOX75MxKwXs6xgrw4W//AYG"
+ "GUjDt83lD6AZP6tws7gJ2IwY/qP7+lyhjEqN\\nHtfPZRGFkGZsdaksdlaksd323423d+15/UvrlRSFPNj1tWQmNKk"
+ "XyRDW4IG1Oa2p\\nrALStNBx5Y9t0/LQnFI4w3aG\\n-----END PRIVATE KEY-----\\n\",\n"
+ " \"project_id\": \"someprojectid\",\n"
+ " \"client_email\": \"someclientid@developer.gserviceaccount.com\",\n"
+ " \"client_id\": \"someclientid.apps.googleusercontent.com\",\n"
+ " \"type\": \"service_account\",\n"
+ " \"universe_domain\": \"googleapis.com\"\n"
+ "}";

private static final String FAKE_JSON_CRED_WITH_INVALID_DOMAIN =
"{\n"
+ " \"private_key_id\": \"somekeyid\",\n"
+ " \"private_key\": \"-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggS"
+ "kAgEAAoIBAQC+K2hSuFpAdrJI\\nnCgcDz2M7t7bjdlsadsasad+fvRSW6TjNQZ3p5LLQY1kSZRqBqylRkzteMOyHg"
+ "aR\\n0Pmxh3ILCND5men43j3h4eDbrhQBuxfEMalkG92sL+PNQSETY2tnvXryOvmBRwa/\\nQP/9dJfIkIDJ9Fw9N4"
+ "Bhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nknddadwkwewcVxHFhcZJO+XWf6ofLUXpRwiTZakGMn8EE1uVa2"
+ "LgczOjwWHGi99MFjxSer5m9\\n1tCa3/KEGKiS/YL71JvjwX3mb+cewlkcmweBKZHM2JPTk0ZednFSpVZMtycjkbLa"
+ "\\ndYOS8V85AgMBewECggEBAKksaldajfDZDV6nGqbFjMiizAKJolr/M3OQw16K6o3/\\n0S31xIe3sSlgW0+UbYlF"
+ "4U8KifhManD1apVSC3csafaspP4RZUHFhtBywLO9pR5c\\nr6S5aLp+gPWFyIp1pfXbWGvc5VY/v9x7ya1VEa6rXvL"
+ "sKupSeWAW4tMj3eo/64ge\\nsdaceaLYw52KeBYiT6+vpsnYrEkAHO1fF/LavbLLOFJmFTMxmsNaG0tuiJHgjshB\\"
+ "n82DpMCbXG9YcCgI/DbzuIjsdj2JC1cascSP//3PmefWysucBQe7Jryb6NQtASmnv\\nCdDw/0jmZTEjpe4S1lxfHp"
+ "lAhHFtdgYTvyYtaLZiVVkCgYEA8eVpof2rceecw/I6\\n5ng1q3Hl2usdWV/4mZMvR0fOemacLLfocX6IYxT1zA1FF"
+ "JlbXSRsJMf/Qq39mOR2\\nSpW+hr4jCoHeRVYLgsbggtrevGmILAlNoqCMpGZ6vDmJpq6ECV9olliDvpPgWOP+\\nm"
+ "YPDreFBGxWvQrADNbRt2dmGsrsCgYEAyUHqB2wvJHFqdmeBsaacewzV8x9WgmeX\\ngUIi9REwXlGDW0Mz50dxpxcK"
+ "CAYn65+7TCnY5O/jmL0VRxU1J2mSWyWTo1C+17L0\\n3fUqjxL1pkefwecxwecvC+gFFYdJ4CQ/MHHXU81Lwl1iWdF"
+ "Cd2UoGddYaOF+KNeM\\nHC7cmqra+JsCgYEAlUNywzq8nUg7282E+uICfCB0LfwejuymR93CtsFgb7cRd6ak\\nECR"
+ "8FGfCpH8ruWJINllbQfcHVCX47ndLZwqv3oVFKh6pAS/vVI4dpOepP8++7y1u\\ncoOvtreXCX6XqfrWDtKIvv0vjl"
+ "HBhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nkndj5uNl5SiuVxHFhcZJO+XWf6ofLUregtevZakGMn8EE1uVa"
+ "2AY7eafmoU/nZPT\\n00YB0TBATdCbn/nBSuKDESkhSg9s2GEKQZG5hBmL5uCMfo09z3SfxZIhJdlerreP\\nJ7gSi"
+ "dI12N+EZxYd4xIJh/HFDgp7RRO87f+WJkofMQKBgGTnClK1VMaCRbJZPriw\\nEfeFCoOX75MxKwXs6xgrw4W//AYG"
+ "GUjDt83lD6AZP6tws7gJ2IwY/qP7+lyhjEqN\\nHtfPZRGFkGZsdaksdlaksd323423d+15/UvrlRSFPNj1tWQmNKk"
+ "XyRDW4IG1Oa2p\\nrALStNBx5Y9t0/LQnFI4w3aG\\n-----END PRIVATE KEY-----\\n\",\n"
+ " \"project_id\": \"someprojectid\",\n"
+ " \"client_email\": \"someclientid@developer.gserviceaccount.com\",\n"
+ " \"client_id\": \"someclientid.apps.googleusercontent.com\",\n"
+ " \"type\": \"service_account\",\n"
+ " \"universe_domain\": \"fake.domain\"\n"
+ "}";

@BeforeClass
public static void beforeClass() throws IOException {
client = BigQueryReadClient.create();
Expand Down Expand Up @@ -859,6 +925,108 @@ public void testSimpleReadWithBackgroundExecutorProvider() throws IOException {
assertEquals(164_656, rowCount);
}

@Test
public void testUniverseDomainWithInvalidUniverseDomain() throws IOException {
BigQueryReadSettings bigQueryReadSettings =
BigQueryReadSettings.newBuilder()
.setCredentialsProvider(
FixedCredentialsProvider.create(loadCredentials(FAKE_JSON_CRED_WITH_GOOGLE_DOMAIN)))
.setUniverseDomain("invalid.domain")
.build();
BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);

String table =
BigQueryResource.FormatTableResource(
/* projectId = */ "bigquery-public-data",
/* datasetId = */ "samples",
/* tableId = */ "shakespeare");

try {
localClient.createReadSession(
/* parent = */ parentProjectId,
/* readSession = */ ReadSession.newBuilder()
.setTable(table)
.setDataFormat(DataFormat.AVRO)
.build(),
/* maxStreamCount = */ 1);
fail("RPCs to invalid universe domain should fail");
} catch (UnauthenticatedException e) {
assertThat(
(e.getMessage()
.contains("does not match the universe domain found in the credentials")))
.isTrue();
}
}

@Test
public void testInvalidUniverseDomainWithMismatchCredentials() throws IOException {
BigQueryReadSettings bigQueryReadSettings =
BigQueryReadSettings.newBuilder()
.setCredentialsProvider(
FixedCredentialsProvider.create(
loadCredentials(FAKE_JSON_CRED_WITH_INVALID_DOMAIN)))
.setUniverseDomain("invalid.domain")
.build();
BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);

String table =
BigQueryResource.FormatTableResource(
/* projectId = */ "bigquery-public-data",
/* datasetId = */ "samples",
/* tableId = */ "shakespeare");

try {
ReadSession session =
localClient.createReadSession(
/* parent = */ parentProjectId,
/* readSession = */ ReadSession.newBuilder()
.setTable(table)
.setDataFormat(DataFormat.AVRO)
.build(),
/* maxStreamCount = */ 1);
fail("RPCs to invalid universe domain should fail");
} catch (UnauthenticatedException e) {
assertThat(
(e.getMessage()
.contains("does not match the universe domain found in the credentials")))
.isTrue();
}
}

@Test
public void testUniverseDomainWithMatchingDomain() throws IOException {
// Test a valid domain using the default credentials and Google default universe domain.
BigQueryReadSettings bigQueryReadSettings =
BigQueryReadSettings.newBuilder().setUniverseDomain("googleapis.com").build();
BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);

String table =
BigQueryResource.FormatTableResource(
/* projectId = */ "bigquery-public-data",
/* datasetId = */ "samples",
/* tableId = */ "shakespeare");

ReadSession session =
localClient.createReadSession(
/* parent = */ parentProjectId,
/* readSession = */ ReadSession.newBuilder()
.setTable(table)
.setDataFormat(DataFormat.AVRO)
.build(),
/* maxStreamCount = */ 1);

ReadRowsRequest readRowsRequest =
ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build();

long rowCount = 0;
ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
for (ReadRowsResponse response : stream) {
rowCount += response.getRowCount();
}

assertEquals(164_656, rowCount);
}

/**
* Reads to the specified row offset within the stream. If the stream does not have the desired
* rows to read, it will read all of them.
Expand Down Expand Up @@ -1015,4 +1183,14 @@ private Job RunQueryJobAndExpectSuccess(QueryJobConfiguration configuration)

return completedJob;
}

static GoogleCredentials loadCredentials(String credentialFile) {
try {
InputStream keyStream = new ByteArrayInputStream(credentialFile.getBytes());
return GoogleCredentials.fromStream(keyStream);
} catch (IOException e) {
fail("Couldn't create fake JSON credentials.");
}
return null;
}
}
Loading

0 comments on commit 466cd4c

Please sign in to comment.