feat: add X-Goog-Request-Reason header support

Author: pareshgoogle

src/main/java/net/starschema/clouddb/jdbc/BQConnection.java

@@ -246,6 +246,9 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor
     // extract UA String
     String userAgent = caseInsensitiveProps.getProperty("useragent");
 
+    // extract requestReason property
+    String requestReason = caseInsensitiveProps.getProperty("requestreason");
+
     // extract any labels
     this.labels = tryParseLabels(caseInsensitiveProps.getProperty("labels"));
     // extract custom endpoint for connections through restricted VPC
@@ -287,7 +290,8 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor
                 rootUrl,
                 httpTransport,
                 targetServiceAccounts,
-                this.getProjectId());
+                this.getProjectId(),
+                requestReason);
         this.logger.info("Authorized with service account");
       } catch (GeneralSecurityException e) {
         throw new BQSQLException(e);
@@ -305,7 +309,8 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor
                 rootUrl,
                 httpTransport,
                 targetServiceAccounts,
-                this.getProjectId());
+                this.getProjectId(),
+                requestReason);
         this.logger.info("Authorized with OAuth access token");
       } catch (SQLException e) {
         throw new BQSQLException(e);
@@ -320,7 +325,8 @@ public BQConnection(String url, Properties loginProp, HttpTransport httpTranspor
                 rootUrl,
                 httpTransport,
                 targetServiceAccounts,
-                this.getProjectId());
+                this.getProjectId(),
+                requestReason);
       } catch (IOException e) {
         throw new BQSQLException(e);
       }

src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java

@@ -22,7 +22,6 @@
  */
 package net.starschema.clouddb.jdbc;
 
-import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpResponse;
 import com.google.api.client.http.HttpTransport;
@@ -31,7 +30,6 @@
 import com.google.api.client.json.gson.GsonFactory;
 import com.google.api.services.bigquery.Bigquery;
 import com.google.api.services.bigquery.Bigquery.Builder;
-import com.google.api.services.bigquery.BigqueryRequest;
 import com.google.api.services.bigquery.BigqueryRequestInitializer;
 import com.google.api.services.bigquery.BigqueryScopes;
 import com.google.api.services.bigquery.MinifiedBigquery;
@@ -101,13 +99,15 @@ private static Bigquery.Builder createBqBuilderForCredential(
       String rootUrl,
       List<String> targetServiceAccounts,
       @Nullable String oauthToken,
-      @Nullable String projectId) {
+      @Nullable String projectId,
+      @Nullable String requestReason) {
 
     // If targetServiceAccounts is empty this returns the original credential
     credential = impersonateServiceAccount(credential, targetServiceAccounts, projectId);
 
     HttpRequestTimeoutInitializer httpRequestInitializer =
-        createRequestTimeoutInitalizer(credential, connectTimeout, readTimeout);
+        createRequestTimeoutInitalizer(
+            credential, connectTimeout, readTimeout, requestReason, userAgent);
 
     Bigquery.Builder bqBuilder =
         new Builder(httpTransport, JSON_FACTORY, httpRequestInitializer)
@@ -140,7 +140,11 @@ private static Bigquery.Builder createBqBuilderForCredential(
    * @return HttpRequestTimeoutInitializer suitable for use with Bigquery.Builder
    */
   private static HttpRequestTimeoutInitializer createRequestTimeoutInitalizer(
-      GoogleCredentials credential, Integer connectTimeout, Integer readTimeout) {
+      GoogleCredentials credential,
+      Integer connectTimeout,
+      Integer readTimeout,
+      String requestReason,
+      String userAgent) {
     HttpRequestTimeoutInitializer httpRequestInitializer =
         new HttpRequestTimeoutInitializer(credential);
     if (connectTimeout != null) {
@@ -149,6 +153,12 @@ private static HttpRequestTimeoutInitializer createRequestTimeoutInitalizer(
     if (readTimeout != null) {
       httpRequestInitializer.setReadTimeout(readTimeout);
     }
+    if (requestReason != null) {
+      httpRequestInitializer.setRequestReason(requestReason);
+    }
+    if (userAgent != null) {
+      httpRequestInitializer.setUserAgent(userAgent);
+    }
 
     return httpRequestInitializer;
   }
@@ -168,7 +178,8 @@ public static Bigquery authorizeViaToken(
       String rootUrl,
       HttpTransport httpTransport,
       List<String> targetServiceAccounts,
-      String projectId)
+      String projectId,
+      String requestReason)
       throws SQLException {
     GoogleCredentials credential =
         GoogleCredentials.create(new AccessToken(oauthToken, null))
@@ -186,7 +197,8 @@ public static Bigquery authorizeViaToken(
             rootUrl,
             targetServiceAccounts,
             oauthToken,
-            projectId);
+            projectId,
+            requestReason);
 
     return new MinifiedBigquery(bqBuilder);
   }
@@ -279,7 +291,8 @@ public static Bigquery authorizeViaService(
       String rootUrl,
       HttpTransport httpTransport,
       List<String> targetServiceAccounts,
-      String projectId)
+      String projectId,
+      String requestReason)
       throws GeneralSecurityException, IOException {
     GoogleCredentials credential =
         createServiceAccountCredential(
@@ -297,7 +310,8 @@ public static Bigquery authorizeViaService(
             rootUrl,
             targetServiceAccounts,
             /* oauthToken= */ null,
-            projectId);
+            projectId,
+            requestReason);
 
     return new MinifiedBigquery(bqBuilder);
   }
@@ -333,7 +347,8 @@ public static Bigquery authorizeViaApplicationDefault(
       String rootUrl,
       HttpTransport httpTransport,
       List<String> targetServiceAccounts,
-      String projectId)
+      String projectId,
+      String requestReason)
       throws IOException {
     GoogleCredentials credential =
         GoogleCredentials.getApplicationDefault().createScoped(GenerateScopes(false));
@@ -350,7 +365,8 @@ public static Bigquery authorizeViaApplicationDefault(
             rootUrl,
             targetServiceAccounts,
             /* oauthToken= */ null,
-            projectId);
+            projectId,
+            requestReason);
 
     return new MinifiedBigquery(bqBuilder);
   }
@@ -459,9 +475,11 @@ private static PrivateKey getPrivateKeyFromCredentials(String keyPath, String pa
     return (PrivateKey) keystore.getKey(keystore.aliases().nextElement(), password.toCharArray());
   }
 
-  private static class HttpRequestTimeoutInitializer extends HttpCredentialsAdapter {
+  static class HttpRequestTimeoutInitializer extends HttpCredentialsAdapter {
     private Integer readTimeout = null;
     private Integer connectTimeout = null;
+    private String requestReason = null;
+    private String userAgent = null;
 
     public HttpRequestTimeoutInitializer(GoogleCredentials credential) {
       super(credential);
@@ -475,6 +493,18 @@ public void setConnectTimeout(Integer timeout) {
       connectTimeout = timeout;
     }
 
+    public void setRequestReason(String requestReason) {
+      this.requestReason = requestReason;
+    }
+
+    public String getRequestReason() {
+      return requestReason;
+    }
+
+    public void setUserAgent(String userAgent) {
+      this.userAgent = userAgent;
+    }
+
     @Override
     public void initialize(HttpRequest httpRequest) throws IOException {
       super.initialize(httpRequest);
@@ -485,6 +515,12 @@ public void initialize(HttpRequest httpRequest) throws IOException {
       if (readTimeout != null) {
         httpRequest.setReadTimeout(readTimeout);
       }
+      if (userAgent != null) {
+        httpRequest.getHeaders().setUserAgent(userAgent);
+      }
+      if (requestReason != null) {
+        httpRequest.getHeaders().set("X-Goog-Request-Reason", requestReason);
+      }
     }
 
     @Override
@@ -520,16 +556,5 @@ public String getOauthToken() {
     public void setOauthToken(String oauthToken) {
       this.oauthToken = oauthToken;
     }
-
-    @Override
-    public void initializeBigqueryRequest(BigqueryRequest<?> request) throws IOException {
-      if (userAgent != null) {
-        HttpHeaders currentHeaders = request.getRequestHeaders();
-
-        currentHeaders.setUserAgent(userAgent);
-
-        request.setRequestHeaders(currentHeaders);
-      }
-    }
   }
 }

src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java

@@ -924,6 +924,30 @@ public void rootUrlOverrideWorks() throws IOException, SQLException {
         request.getUrl().startsWith("https://restricted.googleapis.com/bigquery/v2/"));
   }
 
+  @Test
+  public void urlWithRequestReasonSendsHeader() throws Exception {
+    properties = getProperties("/vpcaccount.properties");
+    String url = getUrl("/vpcaccount.properties", null) + "&requestReason=test_reason";
+    String mockResponse =
+        "{ \"jobComplete\": true, "
+            + "\"totalRows\": \"0\", "
+            + "\"rows\": [], "
+            + "\"totalBytesProcessed\": \"0\", "
+            + "\"cacheHit\": false }";
+    MockHttpTransport mockTransport =
+        new MockHttpTransport.Builder()
+            .setLowLevelHttpResponse(new MockLowLevelHttpResponse().setContent(mockResponse))
+            .build();
+    bq = new BQConnection(url, properties, mockTransport);
+    BQStatement stmt = new BQStatement(properties.getProperty("projectid"), bq);
+    String sqlStmt = "SELECT word from publicdata:samples.shakespeare LIMIT 100";
+
+    stmt.executeQuery(sqlStmt);
+
+    MockLowLevelHttpRequest request = mockTransport.getLowLevelHttpRequest();
+    Assert.assertEquals("test_reason", request.getFirstHeaderValue("X-Goog-Request-Reason"));
+  }
+
   @Test
   public void timeoutMsRejectsBadValues() throws Exception {
     try {