feat: set quota project when using impersonation (#131)

When using SA impersonation we need to explicitly set a quota project that will be used for billing services. Here we add the projectId specified in the JDBC URL as the quota project when constructing the ImpersonatedCredentials to use. This effectively makes the projectId declared in the URL the billing project that will be used for all BQ jobs.
looker-open-source · May 13, 2022 · 2bd0a3c · 2bd0a3c
1 parent 1ee44c6
commit 2bd0a3c
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 20 deletions.
diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java
@@ -229,7 +229,8 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor
                 connectTimeout,
                 rootUrl,
                 httpTransport,
-                targetServiceAccounts);
+                targetServiceAccounts,
+                this.getProjectId());
         this.logger.info("Authorized with service account");
       } catch (GeneralSecurityException e) {
         throw new BQSQLException(e);
@@ -246,7 +247,8 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor
                 readTimeout,
                 rootUrl,
                 httpTransport,
-                targetServiceAccounts);
+                targetServiceAccounts,
+                this.getProjectId());
         this.logger.info("Authorized with OAuth access token");
       } catch (SQLException e) {
         throw new BQSQLException(e);
@@ -260,7 +262,8 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor
                 readTimeout,
                 rootUrl,
                 httpTransport,
-                targetServiceAccounts);
+                targetServiceAccounts,
+                this.getProjectId());
       } catch (IOException e) {
         throw new BQSQLException(e);
       }

diff --git a/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java b/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java
@@ -104,10 +104,11 @@ private static Bigquery.Builder createBqBuilderForCredential(
       String userAgent,
       String rootUrl,
       List<String> targetServiceAccounts,
-      @Nullable String oauthToken) {
+      @Nullable String oauthToken,
+      @Nullable String projectId) {
 
     // If targetServiceAccounts is empty this returns the original credential
-    credential = impersonateServiceAccount(credential, targetServiceAccounts);
+    credential = impersonateServiceAccount(credential, targetServiceAccounts, projectId);
 
     HttpRequestTimeoutInitializer httpRequestInitializer =
         createRequestTimeoutInitalizer(credential, connectTimeout, readTimeout);
@@ -170,7 +171,8 @@ public static Bigquery authorizeViaToken(
       Integer readTimeout,
       String rootUrl,
       HttpTransport httpTransport,
-      List<String> targetServiceAccounts)
+      List<String> targetServiceAccounts,
+      String projectId)
       throws SQLException {
     GoogleCredentials credential = GoogleCredentials.create(new AccessToken(oauthToken, null));
 
@@ -185,7 +187,8 @@ public static Bigquery authorizeViaToken(
             userAgent,
             rootUrl,
             targetServiceAccounts,
-            oauthToken);
+            oauthToken,
+            projectId);
 
     return new MinifiedBigquery(bqBuilder);
   }
@@ -277,7 +280,8 @@ public static Bigquery authorizeViaService(
       Integer connectTimeout,
       String rootUrl,
       HttpTransport httpTransport,
-      List<String> targetServiceAccounts)
+      List<String> targetServiceAccounts,
+      String projectId)
       throws GeneralSecurityException, IOException {
     GoogleCredentials credential =
         createServiceAccountCredential(
@@ -294,7 +298,8 @@ public static Bigquery authorizeViaService(
             userAgent,
             rootUrl,
             targetServiceAccounts,
-            null);
+            /* oauthToken= */ null,
+            projectId);
 
     return new MinifiedBigquery(bqBuilder);
   }
@@ -329,7 +334,8 @@ public static Bigquery authorizeViaApplicationDefault(
       Integer readTimeout,
       String rootUrl,
       HttpTransport httpTransport,
-      List<String> targetServiceAccounts)
+      List<String> targetServiceAccounts,
+      String projectId)
       throws IOException {
     GoogleCredentials credential = GoogleCredentials.getApplicationDefault();
 
@@ -344,7 +350,8 @@ public static Bigquery authorizeViaApplicationDefault(
             userAgent,
             rootUrl,
             targetServiceAccounts,
-            null);
+            /* oauthToken= */ null,
+            projectId);
 
     return new MinifiedBigquery(bqBuilder);
   }
@@ -398,21 +405,30 @@ public static String generateAccessToken(
    * @return GoogleCredentials
    */
   private static GoogleCredentials impersonateServiceAccount(
-      GoogleCredentials sourceCredentials, List<String> targetServiceAccounts) {
+      GoogleCredentials sourceCredentials,
+      List<String> targetServiceAccounts,
+      @Nullable String quotaProjectId) {
     if (targetServiceAccounts.isEmpty()) {
       return sourceCredentials;
     }
 
+    // Get target principle at end of delegate chain
     int lastIdx = targetServiceAccounts.size() - 1;
     String targetServiceAccount = targetServiceAccounts.get(lastIdx);
     List<String> delegates = targetServiceAccounts.subList(0, lastIdx);
 
-    return ImpersonatedCredentials.create(
-        sourceCredentials,
-        targetServiceAccount,
-        delegates,
-        GenerateScopes(false),
-        DEFAULT_IMPERSONATION_LIFETIME);
+    ImpersonatedCredentials.Builder builder = ImpersonatedCredentials.newBuilder();
+    builder.setSourceCredentials(sourceCredentials);
+    builder.setTargetPrincipal(targetServiceAccount);
+    builder.setDelegates(delegates);
+    builder.setLifetime(DEFAULT_IMPERSONATION_LIFETIME);
+    builder.setScopes(GenerateScopes(false));
+
+    if (quotaProjectId != null) {
+      builder.setQuotaProjectId(quotaProjectId);
+    }
+
+    return builder.build();
   }
 
   private static GoogleCredentials createServiceAccountCredential(

diff --git a/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java b/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java
@@ -269,7 +269,7 @@ public void impersonatedServiceAccountWithDelegations() throws IOException, SQLE
   }
 
   @Test
-  public void delegationsThowWithBadDelegationChain() throws IOException, SQLException {
+  public void delegationsThrowWithBadDelegationChain() throws IOException, SQLException {
     Properties testProps = getProperties("/applicationdefault.properties");
     String url =
         BQDriver.getURLPrefix()

diff --git a/src/test/resources/vpcaccount.properties b/src/test/resources/vpcaccount.properties
@@ -2,5 +2,4 @@ projectid=super-party-888
 type=service
 user=697117590302-76cr6q3217nck6gks0kf4r151j4d9f8e@developer.gserviceaccount.com
 password=src/test/resources/bigquery_credentials.p12
-dataset=looker_test
 transformquery=true