fix(bigquery-jdbc): scope Storage Read/Write API credentials for TPC

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -26,6 +26,8 @@
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.auth.Credentials;
 import com.google.auth.http.HttpTransportFactory;
+import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.BigQueryOptions;
@@ -1457,7 +1459,10 @@ private BigQueryReadClient getBigQueryReadClientConnection() throws IOException
       bigQueryReadSettings.setQuotaProjectId(this.catalog);
     }
     if (this.credentials != null) {
-      CredentialsProvider fixedProvider = FixedCredentialsProvider.create(credentials);
+      CredentialsProvider fixedProvider =
+          FixedCredentialsProvider.create(
+              scopeStorageCredentials(
+                  this.credentials, BigQueryReadSettings.getDefaultServiceScopes()));
       bigQueryReadSettings.setCredentialsProvider(fixedProvider);
     }
     if (this.overrideProperties.containsKey(
@@ -1519,7 +1524,10 @@ private BigQueryWriteClient getBigQueryWriteClientConnection() throws IOExceptio
       bigQueryWriteSettings.setQuotaProjectId(this.catalog);
     }
     if (this.credentials != null) {
-      CredentialsProvider fixedProvider = FixedCredentialsProvider.create(credentials);
+      CredentialsProvider fixedProvider =
+          FixedCredentialsProvider.create(
+              scopeStorageCredentials(
+                  this.credentials, BigQueryWriteSettings.getDefaultServiceScopes()));
       bigQueryWriteSettings.setCredentialsProvider(fixedProvider);
     }
     // Same endpoint as READ API
@@ -1539,6 +1547,19 @@ private BigQueryWriteClient getBigQueryWriteClientConnection() throws IOExceptio
     return BigQueryWriteClient.create(bigQueryWriteSettings.build());
   }
 
+  static Credentials scopeStorageCredentials(Credentials credentials, List<String> defaultScopes) {
+    Credentials scopedCredentials = credentials;
+    if (scopedCredentials instanceof GoogleCredentials
+        && ((GoogleCredentials) scopedCredentials).createScopedRequired()) {
+      scopedCredentials = ((GoogleCredentials) scopedCredentials).createScoped(defaultScopes);
+    }
+    if (scopedCredentials instanceof ServiceAccountCredentials) {
+      scopedCredentials =
+          ((ServiceAccountCredentials) scopedCredentials).createWithUseJwtAccessWithScope(true);
+    }
+    return scopedCredentials;
+  }
+
   RetrySettings getRetrySettings() {
     RetrySettings.Builder retrySettingsBuilder = null;
 

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java

@@ -41,7 +41,9 @@
 import com.google.api.gax.paging.Page;
 import com.google.api.gax.rpc.HeaderProvider;
 import com.google.api.gax.rpc.TransportChannelProvider;
+import com.google.auth.Credentials;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.DataFormatOptions;
@@ -54,18 +56,26 @@
 import com.google.cloud.bigquery.TableResult;
 import com.google.cloud.bigquery.exception.BigQueryJdbcException;
 import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
+import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings;
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
+import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
 import com.google.cloud.logging.Logging;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import java.io.IOException;
 import java.io.InputStream;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
 import java.sql.SQLException;
 import java.util.Arrays;
+import java.util.Base64;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.logging.Level;
@@ -86,6 +96,8 @@ public class BigQueryConnectionTest extends BigQueryJdbcLoggingBaseTest {
 
   private static final String DEFAULT_VERSION = "0.0.0";
   private static final String DEFAULT_JDBC_TOKEN_VALUE = "Google-BigQuery-JDBC-Driver";
+  private static final String TPC_UNIVERSE = "apis-tpczero.goog";
+  private static final String DEFAULT_UNIVERSE = "googleapis.com";
   private static final String BASE_URL =
       "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
           + "OAuthType=2;OAuthAccessToken=redacted;ProjectId=project;";
@@ -960,4 +972,150 @@ public void testDuplicateSessionIdKeysAreRejected() {
         assertThrows(BigQueryJdbcException.class, () -> new BigQueryConnection(url));
     assertTrue(ex.getMessage().contains("multiple 'session_id' entries"));
   }
+
+  @Test
+  public void testReadAndWriteClientsScopeServiceAccountCredentialsForTpc() throws Exception {
+    String url =
+        getUriOAuthServiceAccount()
+            .append("ProjectId", "project")
+            .append("universeDomain", TPC_UNIVERSE)
+            .toString();
+
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      BigQueryReadClient readClient = connection.getBigQueryReadClient();
+      assertEquals(TPC_UNIVERSE, readClient.getSettings().getUniverseDomain());
+      assertStorageCredentialsScoped(
+          readClient.getSettings().getCredentialsProvider().getCredentials(),
+          TPC_UNIVERSE,
+          storageUri(TPC_UNIVERSE, "BigQueryRead"),
+          BigQueryReadSettings.getDefaultServiceScopes());
+
+      BigQueryWriteClient writeClient = connection.getBigQueryWriteClient();
+      assertEquals(TPC_UNIVERSE, writeClient.getSettings().getUniverseDomain());
+      assertStorageCredentialsScoped(
+          writeClient.getSettings().getCredentialsProvider().getCredentials(),
+          TPC_UNIVERSE,
+          storageUri(TPC_UNIVERSE, "BigQueryWrite"),
+          BigQueryWriteSettings.getDefaultServiceScopes());
+    }
+  }
+
+  @Test
+  public void testReadAndWriteClientsScopeServiceAccountCredentialsForDefaultUniverse()
+      throws Exception {
+    String url = getUriOAuthServiceAccount().append("ProjectId", "project").toString();
+
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      BigQueryReadClient readClient = connection.getBigQueryReadClient();
+      assertEquals(DEFAULT_UNIVERSE, readClient.getSettings().getUniverseDomain());
+      assertStorageCredentialsScoped(
+          readClient.getSettings().getCredentialsProvider().getCredentials(),
+          DEFAULT_UNIVERSE,
+          storageUri(DEFAULT_UNIVERSE, "BigQueryRead"),
+          BigQueryReadSettings.getDefaultServiceScopes());
+
+      BigQueryWriteClient writeClient = connection.getBigQueryWriteClient();
+      assertEquals(DEFAULT_UNIVERSE, writeClient.getSettings().getUniverseDomain());
+      assertStorageCredentialsScoped(
+          writeClient.getSettings().getCredentialsProvider().getCredentials(),
+          DEFAULT_UNIVERSE,
+          storageUri(DEFAULT_UNIVERSE, "BigQueryWrite"),
+          BigQueryWriteSettings.getDefaultServiceScopes());
+    }
+  }
+
+  @Test
+  public void testReadAndWriteClientsPassThroughNonServiceAccountCredentials() throws Exception {
+    String url = getUriOAuthToken().append("ProjectId", "project").toString();
+
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      Credentials readCredentials =
+          connection
+              .getBigQueryReadClient()
+              .getSettings()
+              .getCredentialsProvider()
+              .getCredentials();
+      Credentials writeCredentials =
+          connection
+              .getBigQueryWriteClient()
+              .getSettings()
+              .getCredentialsProvider()
+              .getCredentials();
+
+      assertFalse(readCredentials instanceof ServiceAccountCredentials);
+      assertSame(connection.credentials, readCredentials);
+      assertSame(connection.credentials, writeCredentials);
+      assertEquals(
+          "Bearer RedactedToken",
+          readCredentials
+              .getRequestMetadata(storageUri(DEFAULT_UNIVERSE, "BigQueryRead"))
+              .get("Authorization")
+              .get(0));
+    }
+  }
+
+  @Test
+  public void testScopeStorageCredentialsPreservesExplicitScopes() throws Exception {
+    List<String> explicitScopes =
+        Arrays.asList(
+            "https://www.googleapis.com/auth/bigquery",
+            "https://www.googleapis.com/auth/drive.readonly");
+    Credentials alreadyScoped =
+        ServiceAccountCredentials.newBuilder()
+            .setClientEmail("service@account")
+            .setPrivateKeyString(fake_pkcs8_key)
+            .setUniverseDomain(TPC_UNIVERSE)
+            .build()
+            .createScoped(explicitScopes);
+
+    Credentials result =
+        BigQueryConnection.scopeStorageCredentials(
+            alreadyScoped, BigQueryReadSettings.getDefaultServiceScopes());
+
+    assertStorageCredentialsScoped(
+        result, TPC_UNIVERSE, storageUri(TPC_UNIVERSE, "BigQueryRead"), explicitScopes);
+  }
+
+  private static URI storageUri(String universeDomain, String service) {
+    return URI.create(
+        "https://bigquerystorage."
+            + universeDomain
+            + "/google.cloud.bigquery.storage.v1."
+            + service);
+  }
+
+  /**
+   * Asserts that {@code credentials} are service account credentials that sign a self-signed JWT
+   * carrying {@code expectedScopes} (and no audience) for the given universe.
+   */
+  private static void assertStorageCredentialsScoped(
+      Credentials credentials, String expectedUniverse, URI uri, List<String> expectedScopes)
+      throws IOException {
+    assertTrue(credentials instanceof ServiceAccountCredentials);
+    ServiceAccountCredentials sa = (ServiceAccountCredentials) credentials;
+    assertFalse(sa.createScopedRequired());
+    assertTrue(sa.getUseJwtAccessWithScope());
+    assertEquals(expectedUniverse, sa.getUniverseDomain());
+
+    JsonObject jwtPayload = extractBearerJwtPayload(credentials, uri);
+    assertFalse(jwtPayload.has("aud"));
+    assertTrue(jwtPayload.has("scope"));
+    String scopeClaim = jwtPayload.get("scope").getAsString();
+    for (String expectedScope : expectedScopes) {
+      assertTrue(scopeClaim.contains(expectedScope), "Missing scope: " + expectedScope);
+    }
+  }
+
+  private static JsonObject extractBearerJwtPayload(Credentials credentials, URI uri)
+      throws IOException {
+    Map<String, List<String>> metadata = credentials.getRequestMetadata(uri);
+    String authHeader = metadata.get("Authorization").get(0);
+    assertTrue(authHeader.startsWith("Bearer "));
+    String jwt = authHeader.substring("Bearer ".length());
+    String[] parts = jwt.split("\\.");
+    assertEquals(3, parts.length);
+    String payloadJson =
+        new String(Base64.getUrlDecoder().decode(parts[1]), StandardCharsets.UTF_8);
+    return JsonParser.parseString(payloadJson).getAsJsonObject();
+  }
 }

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBase.java

@@ -284,6 +284,11 @@ public static String getBaseConnectionUrl() {
       getBaseConnectionUrl() + "ProjectId=" + DEFAULT_CATALOG + ";OAuthType=3;Timeout=3600;";
   public static final BigQuery bigQuery = BigQueryJdbcBaseTest.getBigQuery(connectionUrl);
 
+  /**
+   * Query template returning rows 1..N; use with {@code String.format(GENERATE_ARRAY_QUERY, n)}.
+   */
+  public static final String GENERATE_ARRAY_QUERY = "SELECT * FROM UNNEST(GENERATE_ARRAY(1, %s));";
+
   public static final String FORCE_READ_API_PROPERTIES =
       ";EnableHighThroughputAPI=1;HighThroughputActivationRatio=0;HighThroughputMinTableSize=0;";
 
@@ -438,6 +443,11 @@ protected static GoogleCredentials getCredentials() throws IOException {
         new ByteArrayInputStream(authJson.toString().getBytes(StandardCharsets.UTF_8)));
   }
 
+  /** Asserts that the result set was served by the Storage Read (High-Throughput) API. */
+  protected static void assertArrowResultSet(ResultSet resultSet) {
+    assertTrue(resultSet.getClass().getName().contains("BigQueryArrowResultSet"));
+  }
+
   protected int resultSetRowCount(ResultSet resultSet) throws SQLException {
     int rowCount = 0;
     while (resultSet.next()) {

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java

@@ -228,7 +228,6 @@ public void testIterateOrderArrowMultiThread() throws SQLException {
   }
 
   @Test
-  @Tag("disable_tpc")
   public void testReadAPIPathLarge() throws SQLException {
     Properties withReadApi = new Properties();
     withReadApi.setProperty("EnableHighThroughputAPI", "1");
@@ -240,15 +239,15 @@ public void testReadAPIPathLarge() throws SQLException {
         DriverManager.getConnection(String.format(connectionUrl, DEFAULT_CATALOG), withReadApi);
     Statement statement = connection.createStatement();
     int expectedCnt = 5000;
-    String longQuery = String.format(BASE_QUERY, expectedCnt);
+    String longQuery = String.format(GENERATE_ARRAY_QUERY, expectedCnt);
     ResultSet arrowResultSet = statement.executeQuery(longQuery);
+    assertArrowResultSet(arrowResultSet);
     assertEquals(expectedCnt, resultSetRowCount(arrowResultSet));
     arrowResultSet.close();
     connection.close();
   }
 
   @Test
-  @Tag("disable_tpc")
   public void testReadAPIPathLargeWithThresholdParameters() throws SQLException {
     String connectionUri =
         ITBigQueryJDBCTest.connection_uri
@@ -259,15 +258,15 @@ public void testReadAPIPathLargeWithThresholdParameters() throws SQLException {
     Connection connection = DriverManager.getConnection(connectionUri);
     Statement statement = connection.createStatement();
     int expectedCnt = 1000;
-    String longQuery = String.format(BASE_QUERY, expectedCnt);
+    String longQuery = String.format(GENERATE_ARRAY_QUERY, expectedCnt);
     ResultSet arrowResultSet = statement.executeQuery(longQuery);
+    assertArrowResultSet(arrowResultSet);
     assertEquals(expectedCnt, resultSetRowCount(arrowResultSet));
     arrowResultSet.close();
     connection.close();
   }
 
   @Test
-  @Tag("disable_tpc")
   public void testReadAPIPathLargeWithThresholdNotMet() throws SQLException {
     String connectionUri =
         ITBigQueryJDBCTest.connection_uri
@@ -277,10 +276,11 @@ public void testReadAPIPathLargeWithThresholdNotMet() throws SQLException {
     Connection connection = DriverManager.getConnection(connectionUri);
     Statement statement = connection.createStatement();
     int expectedCnt = 5000;
-    String longQuery = String.format(BASE_QUERY, expectedCnt);
-    ResultSet arrowResultSet = statement.executeQuery(longQuery);
-    assertEquals(expectedCnt, resultSetRowCount(arrowResultSet));
-    arrowResultSet.close();
+    String longQuery = String.format(GENERATE_ARRAY_QUERY, expectedCnt);
+    try (ResultSet jsonResultSet = statement.executeQuery(longQuery)) {
+      assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet"));
+      assertEquals(expectedCnt, resultSetRowCount(jsonResultSet));
+    }
     connection.close();
   }
 
@@ -1420,7 +1420,7 @@ public void testIntervalDataTypeWithArrowResultSet() throws SQLException {
     Statement statement = connection.createStatement();
 
     ResultSet resultSet = statement.executeQuery(selectQuery);
-    assertTrue(resultSet.getClass().getName().contains("BigQueryArrowResultSet"));
+    assertArrowResultSet(resultSet);
     resultSet.next();
     assertEquals("0-0 10 -12:30:0.0", resultSet.getString("intervalField"));
 
@@ -2138,7 +2138,7 @@ public void testRangeDataTypeWithArrowResultSet() throws SQLException {
     Connection connection = DriverManager.getConnection(connection_uri);
     Statement statement = connection.createStatement();
     ResultSet resultSet = statement.executeQuery(selectQuery);
-    assertTrue(resultSet.getClass().getName().contains("BigQueryArrowResultSet"));
+    assertArrowResultSet(resultSet);
     resultSet.next();
     assertEquals("[2024-07-14, 2024-09-23)", resultSet.getString("rangeField"));
     connection.close();