From 467149efde1599dfdb5cd52e37b63cb47725a367 Mon Sep 17 00:00:00 2001 From: Paresh Maniyar Date: Fri, 17 Apr 2026 09:14:21 +0000 Subject: [PATCH] feat: add X-Goog-Request-Reason header support --- .../starschema/clouddb/jdbc/BQConnection.java | 12 +++- .../clouddb/jdbc/Oauth2Bigquery.java | 71 +++++++++++++------ .../starschema/clouddb/jdbc/JdbcUrlTest.java | 24 +++++++ 3 files changed, 81 insertions(+), 26 deletions(-) diff --git a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java b/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java index 9088c162..97f36de7 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/BQConnection.java +++ b/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); } diff --git a/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java b/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java index 034b2623..b28ead34 100644 --- a/src/main/java/net/starschema/clouddb/jdbc/Oauth2Bigquery.java +++ b/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 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 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 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 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); - } - } } } diff --git a/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java b/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java index 644ee5ea..2bb594df 100644 --- a/src/test/java/net/starschema/clouddb/jdbc/JdbcUrlTest.java +++ b/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 {