Bypass dry-run job for read-only tokens.

Author: logachev

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

@@ -41,8 +41,11 @@
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
 import com.google.cloud.http.HttpTransportOptions;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.sql.CallableStatement;
 import java.sql.Connection;
 import java.sql.DatabaseMetaData;
@@ -53,6 +56,7 @@
 import java.sql.Statement;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.ConcurrentModificationException;
 import java.util.List;
 import java.util.Map;
@@ -143,6 +147,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
   String partnerToken;
   DatabaseMetaData databaseMetaData;
   Boolean reqGoogleDriveScope;
+  private boolean isReadOnlyTokenUsed = false;
 
   BigQueryConnection(String url) throws IOException {
     this(url, DataSource.fromUrl(url));
@@ -172,6 +177,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
       this.jobTimeoutInSeconds = ds.getJobTimeout();
       this.authProperties =
           BigQueryJdbcOAuthUtility.parseOAuthProperties(ds, this.connectionClassName);
+      this.isReadOnlyTokenUsed = checkisReadOnlyTokenUsed(this.authProperties);
       this.catalog = ds.getProjectId();
       this.universeDomain = ds.getUniverseDomain();
 
@@ -1193,4 +1199,45 @@ public CallableStatement prepareCall(
     }
     return prepareCall(sql);
   }
+
+  public boolean isReadOnlyTokenUsed() {
+    return this.isReadOnlyTokenUsed;
+  }
+
+  private boolean checkisReadOnlyTokenUsed(Map<String, String> authProps) {
+    if (authProps == null) {
+      return false;
+    }
+    BigQueryJdbcOAuthUtility.AuthType oauthType =
+        BigQueryJdbcOAuthUtility.AuthType.fromValue(
+            authProps.get(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME));
+    if (oauthType != BigQueryJdbcOAuthUtility.AuthType.PRE_GENERATED_TOKEN) {
+      return false;
+    }
+    String token = authProps.get(BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_PROPERTY_NAME);
+    if (token == null || token.isEmpty()) {
+      return false;
+    }
+
+    // Try to parse scope from the token if it is a JWT
+    try {
+      String[] parts = token.split("\\.");
+      if (parts.length != 3) {
+        return false;
+      }
+      String payloadJson =
+          new String(Base64.getUrlDecoder().decode(parts[1]), StandardCharsets.UTF_8);
+      JsonObject payload = JsonParser.parseString(payloadJson).getAsJsonObject();
+      if (!payload.has("scope")) {
+        return false;
+      }
+      String scope = payload.get("scope").getAsString();
+      if (scope.contains("https://www.googleapis.com/auth/bigquery.readonly")) {
+        return true;
+      }
+    } catch (Exception e) {
+      // Likely invalid token and auth is going to fail later.
+    }
+    return false;
+  }
 }

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

@@ -749,5 +749,16 @@ static AuthType fromValue(int value) {
       LOG.severe(ex.getMessage(), ex);
       throw ex;
     }
+
+    static AuthType fromValue(String value) {
+      for (AuthType authType : values()) {
+        if (authType.name().equalsIgnoreCase(value)) {
+          return authType;
+        }
+      }
+      IllegalStateException ex = new IllegalStateException(OAUTH_TYPE_ERROR_MESSAGE + ": " + value);
+      LOG.severe(ex.getMessage(), ex);
+      throw ex;
+    }
   }
 }

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

@@ -333,6 +333,13 @@ private boolean executeImpl(String sql) throws SQLException {
 
   StatementType getStatementType(QueryJobConfiguration queryJobConfiguration) throws SQLException {
     LOG.finest("++enter++");
+    // BQ Read-only tokens are not recommended to use, they have a lot of known flaws.
+    // We're supporting them in a limited capacity, for pure SELECT statements.
+    if (this.connection.isReadOnlyTokenUsed()) {
+      LOG.warning(
+          "Read-only token detected, skipping dry run and assuming StatementType is SELECT.");
+      return StatementType.SELECT;
+    }
     QueryJobConfiguration dryRunJobConfiguration =
         queryJobConfiguration.toBuilder().setDryRun(true).build();
     Job job;

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

@@ -30,8 +30,12 @@
 import java.io.InputStream;
 import java.sql.SQLException;
 import java.util.Properties;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 
 public class BigQueryConnectionTest {
 
@@ -437,4 +441,26 @@ public void testWithDriveScopeDefault() throws Exception {
       assertFalse(connection.reqGoogleDriveScope);
     }
   }
+
+  @ParameterizedTest
+  @CsvSource({
+    "https://www.googleapis.com/auth/bigquery.readonly, true",
+    "https://www.googleapis.com/auth/bigquery, false"
+  })
+  public void testIsReadOnlyTokenProvided(String scope, boolean expectedIsReadOnly) throws Exception {
+    String payload = "{\"scope\":\"" + scope + "\"}";
+    String encodedPayload =
+        Base64.getUrlEncoder()
+            .encodeToString(payload.getBytes(StandardCharsets.UTF_8));
+    String url =
+        "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
+            + "OAuthType=2;ProjectId=MyBigQueryProject;"
+            + "OAuthAccessToken=header."
+            + encodedPayload
+            + ".signature;";
+
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      assertEquals(expectedIsReadOnly, connection.isReadOnlyTokenUsed());
+    }
+  }
 }

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

@@ -68,6 +68,8 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 
@@ -480,4 +482,22 @@ public void testCancelWithJoblessQuery() throws SQLException, InterruptedExcepti
     // And no backend cancellation was attempted
     verify(bigquery, Mockito.never()).cancel(any(JobId.class));
   }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  public void testGetStatementType(boolean isReadOnlyTokenUsed) throws Exception {
+    doReturn(isReadOnlyTokenUsed).when(bigQueryConnection).isReadOnlyTokenUsed();
+
+    Job dryRunJobMock = getJobMock(null, null, StatementType.SELECT);
+    doReturn(dryRunJobMock).when(bigquery).create(any(JobInfo.class));
+
+    BigQueryStatement statementSpy = Mockito.spy(bigQueryStatement);
+    QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration.newBuilder(query).build();
+
+    StatementType type = statementSpy.getStatementType(queryJobConfiguration);
+
+    assertThat(type).isEqualTo(StatementType.SELECT);
+    verify(bigquery, isReadOnlyTokenUsed ? Mockito.never() : Mockito.times(1))
+        .create(any(JobInfo.class));
+  }
 }