fix(bqjdbc): add Google Driver scope to all credential types

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

@@ -139,6 +139,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
   Long listenerPoolSize;
   String partnerToken;
   DatabaseMetaData databaseMetaData;
+  Boolean reqGoogleDriveScope;
 
   BigQueryConnection(String url) throws IOException {
     this(url, DataSource.fromUrl(url));
@@ -171,9 +172,15 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
       this.overrideProperties.put(
           BigQueryJdbcUrlUtility.UNIVERSE_DOMAIN_OVERRIDE_PROPERTY_NAME, this.universeDomain);
     }
+
+    this.reqGoogleDriveScope =
+        BigQueryJdbcUrlUtility.convertIntToBoolean(
+            String.valueOf(ds.getRequestGoogleDriveScope()),
+            BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME);
+
     this.credentials =
         BigQueryJdbcOAuthUtility.getCredentials(
-            authProperties, overrideProperties, this.connectionClassName);
+            authProperties, overrideProperties, this.reqGoogleDriveScope, this.connectionClassName);
     String defaultDatasetString = ds.getDefaultDataset();
     if (defaultDatasetString == null || defaultDatasetString.trim().isEmpty()) {
       this.defaultDataset = null;

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

@@ -80,6 +80,13 @@ final class BigQueryJdbcOAuthUtility {
           + "Thank you for using JDBC Driver for Google BigQuery!\n"
           + "You may now close the window.</body></html>";
 
+  static final String BIGQUERY_SCOPE = "https://www.googleapis.com/auth/bigquery";
+  static final String DRIVE_READONLY_SCOPE = "https://www.googleapis.com/auth/drive.readonly";
+
+  static final List<String> DEFAULT_BIGQUERY_SCOPES = Arrays.asList(BIGQUERY_SCOPE);
+  static final List<String> BIGQUERY_WITH_DRIVE_SCOPES =
+      Arrays.asList(BIGQUERY_SCOPE, DRIVE_READONLY_SCOPE);
+
   private static final int USER_AUTH_TIMEOUT_MS = 120000;
   private static final BigQueryJdbcCustomLogger LOG =
       new BigQueryJdbcCustomLogger(BigQueryJdbcOAuthUtility.class.getName());
@@ -117,6 +124,7 @@ static Map<String, String> parseOAuthProperties(DataSource ds, String callerClas
       throw new IllegalArgumentException(OAUTH_TYPE_ERROR_MESSAGE);
     }
     oauthProperties.put(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME, String.valueOf(authType));
+
     switch (authType) {
       case GOOGLE_SERVICE_ACCOUNT:
         // For using a Google Service Account (OAuth Type 0)
@@ -144,11 +152,6 @@ static Map<String, String> parseOAuthProperties(DataSource ds, String callerClas
             BigQueryJdbcUrlUtility.OAUTH_CLIENT_ID_PROPERTY_NAME, ds.getOAuthClientId());
         oauthProperties.put(
             BigQueryJdbcUrlUtility.OAUTH_CLIENT_SECRET_PROPERTY_NAME, ds.getOAuthClientSecret());
-        int reqGoogleDriveScope = ds.getRequestGoogleDriveScope();
-        oauthProperties.put(
-            BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME,
-            String.valueOf(reqGoogleDriveScope));
-        LOG.fine("RequestGoogleDriveScope parsed.");
         break;
       case PRE_GENERATED_TOKEN:
         String refreshToken = ds.getOAuthRefreshToken();
@@ -239,7 +242,7 @@ static Map<String, String> parseOAuthProperties(DataSource ds, String callerClas
           BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME,
           ds.getOAuthSAImpersonationScopes() != null
               ? ds.getOAuthSAImpersonationScopes()
-              : BigQueryJdbcUrlUtility.DEFAULT_OAUTH_SA_IMPERSONATION_SCOPES_VALUE);
+              : BIGQUERY_SCOPE);
       oauthProperties.put(
           BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_PROPERTY_NAME,
           ds.getOAuthSAImpersonationTokenLifetime() != null
@@ -258,6 +261,7 @@ static Map<String, String> parseOAuthProperties(DataSource ds, String callerClas
   static GoogleCredentials getCredentials(
       Map<String, String> authProperties,
       Map<String, String> overrideProperties,
+      Boolean reqGoogleDriveScopeBool,
       String callerClassName) {
     LOG.finest("++enter++\t" + callerClassName);
 
@@ -280,15 +284,24 @@ static GoogleCredentials getCredentials(
         break;
       case APPLICATION_DEFAULT_CREDENTIALS:
         // This auth method doesn't support service account impersonation
-        return getApplicationDefaultCredentials(callerClassName);
+
+        credentials = getApplicationDefaultCredentials(callerClassName);
+        break;
       case EXTERNAL_ACCOUNT_AUTH:
         // This auth method doesn't support service account impersonation
-        return getExternalAccountAuthCredentials(authProperties, callerClassName);
+        credentials = getExternalAccountAuthCredentials(authProperties, callerClassName);
+        break;
       default:
         throw new IllegalStateException(OAUTH_TYPE_ERROR_MESSAGE);
     }
 
-    return getServiceAccountImpersonatedCredentials(credentials, authProperties);
+    if (reqGoogleDriveScopeBool) {
+      credentials = credentials.createScoped(BIGQUERY_WITH_DRIVE_SCOPES);
+      LOG.fine("Added Google Drive read-only scope centrally to GoogleCredentials.");
+    }
+
+    return getServiceAccountImpersonatedCredentials(
+        credentials, reqGoogleDriveScopeBool, authProperties);
   }
 
   private static boolean isFileExists(String filename) {
@@ -388,29 +401,10 @@ static UserAuthorizer getUserAuthorizer(
       String callerClassName)
       throws URISyntaxException {
     LOG.finest("++enter++\t" + callerClassName);
+
     List<String> scopes = new ArrayList<>();
     scopes.add("https://www.googleapis.com/auth/bigquery");
 
-    // Add Google Drive scope conditionally
-    if (authProperties.containsKey(
-        BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME)) {
-      try {
-        int driveScopeValue =
-            Integer.parseInt(
-                authProperties.get(
-                    BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME));
-        if (driveScopeValue == 1) {
-          scopes.add("https://www.googleapis.com/auth/drive.readonly");
-          LOG.fine("Added Google Drive read-only scope. Caller: " + callerClassName);
-        }
-      } catch (NumberFormatException e) {
-        LOG.severe(
-            "Invalid value for RequestGoogleDriveScope, defaulting to not request Drive scope."
-                + " Caller: "
-                + callerClassName);
-      }
-    }
-
     List<String> responseTypes = new ArrayList<>();
     responseTypes.add("code");
 
@@ -500,14 +494,18 @@ private static GoogleCredentials getPreGeneratedAccessTokenCredentials(
       builder.setUniverseDomain(
           overrideProperties.get(BigQueryJdbcUrlUtility.UNIVERSE_DOMAIN_OVERRIDE_PROPERTY_NAME));
     }
+
     LOG.info("Connection established. Auth Method: Pre-generated Access Token.");
-    return builder
-        .setAccessToken(
-            AccessToken.newBuilder()
-                .setTokenValue(
-                    authProperties.get(BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_PROPERTY_NAME))
-                .build())
-        .build();
+    GoogleCredentials credentials =
+        builder
+            .setAccessToken(
+                AccessToken.newBuilder()
+                    .setTokenValue(
+                        authProperties.get(BigQueryJdbcUrlUtility.OAUTH_ACCESS_TOKEN_PROPERTY_NAME))
+                    .build())
+            .build();
+
+    return credentials;
   }
 
   static GoogleCredentials getPreGeneratedTokensCredentials(
@@ -552,6 +550,7 @@ static UserCredentials getPreGeneratedRefreshTokenCredentials(
       userCredentialsBuilder.setUniverseDomain(
           overrideProperties.get(BigQueryJdbcUrlUtility.UNIVERSE_DOMAIN_OVERRIDE_PROPERTY_NAME));
     }
+
     LOG.info("Connection established. Auth Method: Pre-generated Refresh Token.");
     return userCredentialsBuilder.build();
   }
@@ -571,6 +570,7 @@ private static GoogleCredentials getApplicationDefaultCredentials(String callerC
       LOG.info(
           "Connection established. Auth Method: Application Default Credentials, Principal: %s.",
           principal);
+
       return credentials;
     } catch (IOException exception) {
       // TODO throw exception
@@ -616,6 +616,7 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
         }
       }
 
+      GoogleCredentials credentials;
       if (credentialsPath != null) {
         return ExternalAccountCredentials.fromStream(
             Files.newInputStream(Paths.get(credentialsPath)));
@@ -634,9 +635,11 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
   // This function checks if connection string contains configuration for
   // credentials impersonation. If not, it returns regular credentials object.
   // If impersonated service account is provided, returns Credentials object
-  // accomodating this information.
+  // accommodating this information.
   private static GoogleCredentials getServiceAccountImpersonatedCredentials(
-      GoogleCredentials credentials, Map<String, String> authProperties) {
+      GoogleCredentials credentials,
+      Boolean reqGoogleDriveScopeBool,
+      Map<String, String> authProperties) {
 
     String impersonationEmail =
         authProperties.get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME);
@@ -653,10 +656,18 @@ private static GoogleCredentials getServiceAccountImpersonatedCredentials(
 
     // Scopes has a default value, so it should never be null
     List<String> impersonationScopes =
-        Arrays.asList(
-            authProperties
-                .get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME)
-                .split(","));
+        new java.util.ArrayList<>(
+            Arrays.asList(
+                authProperties
+                    .get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME)
+                    .split(",")));
+
+    if (reqGoogleDriveScopeBool) {
+      if (!impersonationScopes.contains(DRIVE_READONLY_SCOPE)) {
+        impersonationScopes.add(DRIVE_READONLY_SCOPE);
+        LOG.fine("Added Google Drive read-only scope to impersonation scopes.");
+      }
+    }
 
     // Token lifetime has a default value, so it should never be null
     String impersonationLifetime =

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

@@ -70,8 +70,7 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
   static final String HTAPI_ACTIVATION_RATIO_PROPERTY_NAME = "HighThroughputActivationRatio";
   static final String KMS_KEY_NAME_PROPERTY_NAME = "KMSKeyName";
   static final String QUERY_PROPERTIES_NAME = "QueryProperties";
-  static final int DEFAULT_HTAPI_ACTIVATION_RATIO_VALUE =
-      2; // TODO: to adjust this value before private preview based on performance testing.
+  static final int DEFAULT_HTAPI_ACTIVATION_RATIO_VALUE = 2;
   static final String HTAPI_MIN_TABLE_SIZE_PROPERTY_NAME = "HighThroughputMinTableSize";
   static final int DEFAULT_HTAPI_MIN_TABLE_SIZE_VALUE = 100;
   static final int DEFAULT_OAUTH_TYPE_VALUE = -1;
@@ -86,8 +85,6 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
   static final String DEFAULT_OAUTH_SA_IMPERSONATION_CHAIN_VALUE = null;
   static final String OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME =
       "ServiceAccountImpersonationScopes";
-  static final String DEFAULT_OAUTH_SA_IMPERSONATION_SCOPES_VALUE =
-      "https://www.googleapis.com/auth/bigquery";
   static final String OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_PROPERTY_NAME =
       "ServiceAccountImpersonationTokenLifetime";
   static final String DEFAULT_OAUTH_SA_IMPERSONATION_TOKEN_LIFETIME_VALUE = "3600";

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

@@ -413,4 +413,28 @@ public void testBigQueryJobCreationMode_default() throws Exception {
           bq.getOptions().getDefaultJobCreationMode(), JobCreationMode.JOB_CREATION_OPTIONAL);
     }
   }
+
+  @Test
+  public void testWithDriveScopeTrue() throws Exception {
+    String url = BASE_URL + "RequestGoogleDriveScope=1;";
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      assertTrue(connection.reqGoogleDriveScope);
+    }
+  }
+
+  @Test
+  public void testWithDriveScopeFalse() throws Exception {
+    String url = BASE_URL + "RequestGoogleDriveScope=0;";
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      assertFalse(connection.reqGoogleDriveScope);
+    }
+  }
+
+  @Test
+  public void testWithDriveScopeDefault() throws Exception {
+    String url = BASE_URL;
+    try (BigQueryConnection connection = new BigQueryConnection(url)) {
+      assertFalse(connection.reqGoogleDriveScope);
+    }
+  }
 }