Merge pull request #176 from looker-open-source/revert-171-shivane/re…

…moveOAuthTokenFromQueryParams Revert "Remove OAuth Token From Query Params"
looker-open-source · Nov 7, 2023 · 7df2172 · 7df2172
2 parents 5b04532 + 1cb4bff
commit 7df2172
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 67 deletions.
diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java b/src/main/java/net/starschema/clouddb/jdbc/BQSupportFuncts.java
@@ -646,36 +646,6 @@ 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)
@@ -691,7 +661,7 @@ static Bigquery.Jobs.Query getSyncQuery(
       qr.setMaxResults(maxResults);
     }
 
-    return bigquery.jobs().query(projectId, qr);
+    return bigquery.jobs().query(projectId, qr).execute();
   }
 
   /**

diff --git a/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java b/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java
@@ -119,6 +119,9 @@ private static Bigquery.Builder createBqBuilderForCredential(
       if (userAgent != null) {
         requestInitializer.setUserAgent(userAgent);
       }
+      if (oauthToken != null) {
+        requestInitializer.setOauthToken(oauthToken);
+      }
 
       bqBuilder.setBigqueryRequestInitializer(requestInitializer);
     }
@@ -503,11 +506,20 @@ 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) {
@@ -517,6 +529,9 @@ 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,7 +3,6 @@
 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;
@@ -173,39 +172,6 @@ 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,8 +75,9 @@ public void Connect() throws Exception {
       PreparedStatementTests.con =
           DriverManager.getConnection(
               BQSupportFuncts.constructUrlFromPropertiesFile(
-                  BQSupportFuncts.readFromPropFile("installedaccount1.properties")),
-              BQSupportFuncts.readFromPropFile("installedaccount1.properties"));
+                  BQSupportFuncts.readFromPropFile(
+                      "src/test/resources/installedaccount1.properties")),
+              BQSupportFuncts.readFromPropFile("src/test/resources/installedaccount1.properties"));
     } catch (Exception e) {
       e.printStackTrace();
     }