Skip to content

Commit

Permalink
Fix: Fix issue where Universe Domain is not correctly set. (#2423)
Browse files Browse the repository at this point in the history
* Fix: Fix issue where Universe Domain is not correctly set.

This PR fixes an issue where the universe domain is not correctly set. This also adds integration tests to verify the change.

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* Add universe domain integration tests for v1beta1

* Fix issue with universe domain for v1beta and add functional integration test for universe domain.

* Fix integration test failure due to commented out code. Also close newly created clients.

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
  • Loading branch information
PhongChuong and gcf-owl-bot[bot] committed Feb 29, 2024
1 parent 07d98ab commit b7ebd73
Show file tree
Hide file tree
Showing 9 changed files with 666 additions and 5 deletions.
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
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
Expand Up @@ -74,8 +74,8 @@ public static EnhancedBigQueryStorageStub create(
// Configure the base settings.
BigQueryStorageStubSettings.Builder baseSettingsBuilder =
BigQueryStorageStubSettings.newBuilder()
.setUniverseDomain(settings.getUniverseDomain())
.setTransportChannelProvider(settings.getTransportChannelProvider())
.setEndpoint(settings.getEndpoint())
.setHeaderProvider(settings.getHeaderProvider())
.setCredentialsProvider(settings.getCredentialsProvider())
.setStreamWatchdogCheckInterval(settings.getStreamWatchdogCheckInterval())
Expand Down
Expand Up @@ -186,7 +186,6 @@ protected Builder(ClientContext clientContext) {

// Defaults provider
BigQueryStorageStubSettings.Builder baseDefaults = BigQueryStorageStubSettings.newBuilder();
setEndpoint(baseDefaults.getEndpoint());
setTransportChannelProvider(defaultTransportChannelProvider());
setCredentialsProvider(baseDefaults.getCredentialsProvider());
setStreamWatchdogCheckInterval(baseDefaults.getStreamWatchdogCheckInterval());
Expand Down
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
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
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,147 @@ 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();
}
localClient.close();
}

@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();
}
localClient.close();
}

@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);
localClient.close();
}

public void testUniverseDomain() throws IOException {
// This test is not yet part presubmit integration test as it requires the apis-tpclp.goog
// universe domain credentials.
// Test a valid read session in the universe domain gdutst.
BigQueryReadSettings bigQueryReadSettings =
BigQueryReadSettings.newBuilder().setUniverseDomain("apis-tpclp.goog").build();
BigQueryReadClient localClient = BigQueryReadClient.create(bigQueryReadSettings);

String table =
BigQueryResource.FormatTableResource(
/* projectId = */ "google-tpc-testing-environment:cloudsdk-test-project",
/* datasetId = */ "tpc_demo_dataset",
/* tableId = */ "new_table");

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 = localClient.readRowsCallable().call(readRowsRequest);
for (ReadRowsResponse response : stream) {
rowCount += response.getRowCount();
}

assertEquals(1, rowCount);
localClient.close();
}

/**
* 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 +1222,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;
}
}

0 comments on commit b7ebd73

Please sign in to comment.