Revert "Revert "Remove OAuth Token From Query Params""

looker-open-source · Nov 23, 2023 · c2dcf5d · c2dcf5d
1 parent 8f9b83b
commit c2dcf5d
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 19 deletions.
diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java b/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java
@@ -646,6 +646,36 @@ static QueryResponse runSyncQuery(
       Map<String, String> labels,
       boolean useQueryCache)
       throws IOException {
+    return getSyncQuery(
+            bigquery,
+            projectId,
+            querySql,
+            dataSet,
+            useLegacySql,
+            maxBillingBytes,
+            queryTimeoutMs,
+            maxResults,
+            labels,
+            useQueryCache)
+        .execute();
+  }
+
+  /*
+   * Gets a query as specified, but does not execute it.
+   * Introduced for assertions on the property of the query.
+   * */
+  static Bigquery.Jobs.Query getSyncQuery(
+      Bigquery bigquery,
+      String projectId,
+      String querySql,
+      String dataSet,
+      Boolean useLegacySql,
+      Long maxBillingBytes,
+      Long queryTimeoutMs,
+      Long maxResults,
+      Map<String, String> labels,
+      boolean useQueryCache)
+      throws IOException {
     QueryRequest qr =
         new QueryRequest()
             .setLabels(labels)
@@ -661,7 +691,7 @@ static QueryResponse runSyncQuery(
       qr.setMaxResults(maxResults);
     }
 
-    return bigquery.jobs().query(projectId, qr).execute();
+    return bigquery.jobs().query(projectId, qr);
   }
 
   /**

diff --git a/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java b/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java
@@ -119,9 +119,6 @@ private static Bigquery.Builder createBqBuilderForCredential(
       if (userAgent != null) {
         requestInitializer.setUserAgent(userAgent);
       }
-      if (oauthToken != null) {
-        requestInitializer.setOauthToken(oauthToken);
-      }
 
       bqBuilder.setBigqueryRequestInitializer(requestInitializer);
     }
@@ -506,20 +503,11 @@ public boolean handleResponse(
   private static class BigQueryRequestUserAgentInitializer extends BigqueryRequestInitializer {
 
     String userAgent = null;
-    String oauthToken = null;
 
     public void setUserAgent(String userAgent) {
       this.userAgent = userAgent;
     }
 
-    public void setOauthToken(String oauthToken) {
-      this.oauthToken = oauthToken;
-    }
-
-    public String getOauthToken() {
-      return this.oauthToken;
-    }
-
     @Override
     public void initializeBigqueryRequest(BigqueryRequest<?> request) throws IOException {
       if (userAgent != null) {
@@ -529,9 +517,6 @@ public void initializeBigqueryRequest(BigqueryRequest<?> request) throws IOExcep
 
         request.setRequestHeaders(currentHeaders);
       }
-      if (oauthToken != null) {
-        request.setOauthToken(oauthToken);
-      }
     }
   }
 }
diff --git a/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java b/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java
@@ -3,6 +3,7 @@
 import com.google.api.client.testing.http.MockHttpTransport;
 import com.google.api.client.testing.http.MockLowLevelHttpRequest;
 import com.google.api.client.testing.http.MockLowLevelHttpResponse;
+import com.google.api.services.bigquery.Bigquery.Jobs.Query;
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
 import java.nio.file.Files;
@@ -172,6 +173,39 @@ public void canConnectWithOAuthAccessToken()
     stmt.executeQuery("SELECT * FROM orders limit 1");
   }
 
+  @Test
+  public void oAuthAccessTokenOnlyInHeader()
+      throws SQLException, IOException, GeneralSecurityException {
+    // generate access token from service account credentials
+    Properties serviceProps = getProperties("/protectedaccount.properties");
+    String accessToken =
+        Oauth2Bigquery.generateAccessToken(
+            serviceProps.getProperty("user"),
+            serviceProps.getProperty("path"),
+            serviceProps.getProperty("password"),
+            null);
+
+    Properties oauthProps = getProperties("/oauthaccount.properties");
+    oauthProps.setProperty("oauthaccesstoken", accessToken);
+    String url = BQSupportFuncts.constructUrlFromPropertiesFile(oauthProps, true, null);
+    BQConnection bqConn = new BQConnection(url, new Properties());
+    BQStatement stmt = new BQStatement(oauthProps.getProperty("projectid"), bqConn);
+    Query query =
+        BQSupportFuncts.getSyncQuery(
+            bqConn.getBigquery(),
+            oauthProps.getProperty("projectid"),
+            "SELECT * FROM orders limit 1",
+            bqConn.getDataSet(),
+            bqConn.getUseLegacySql(),
+            null,
+            stmt.getSyncTimeoutMillis(),
+            (long) stmt.getMaxRows(),
+            stmt.getAllLabels(),
+            bqConn.getUseQueryCache());
+    String oAuthToken = query.getOauthToken();
+    Assert.assertTrue(oAuthToken == null);
+  }
+
   @Test
   public void unauthorizedResponseForInvalidOAuthAccessToken()
       throws SQLException, IOException, GeneralSecurityException {

diff --git a/src/test/java/net/starschema/clouddb/jdbc/PreparedStatementTests.java b/src/test/java/net/starschema/clouddb/jdbc/PreparedStatementTests.java
@@ -75,9 +75,8 @@ public void Connect() throws Exception {
       PreparedStatementTests.con =
           DriverManager.getConnection(
               BQSupportFuncts.constructUrlFromPropertiesFile(
-                  BQSupportFuncts.readFromPropFile(
-                      "src/test/resources/installedaccount1.properties")),
-              BQSupportFuncts.readFromPropFile("src/test/resources/installedaccount1.properties"));
+                  BQSupportFuncts.readFromPropFile("installedaccount1.properties")),
+              BQSupportFuncts.readFromPropFile("installedaccount1.properties"));
     } catch (Exception e) {
       e.printStackTrace();
     }