diff --git a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 188b3f9a8..b89c568d1 100644 --- a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -131,6 +131,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { String sslTrustStorePassword; long maxBytesBilled; Map labels; + String requestReason; BigQueryConnection(String url) throws IOException { this.connectionUrl = url; @@ -347,6 +348,12 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { BigQueryJdbcUrlUtility.METADATA_FETCH_THREAD_COUNT_PROPERTY_NAME, BigQueryJdbcUrlUtility.DEFAULT_METADATA_FETCH_THREAD_COUNT_VALUE, this.connectionClassName); + this.requestReason = + BigQueryJdbcUrlUtility.parseStringProperty( + url, + BigQueryJdbcUrlUtility.REQUEST_REASON_PROPERTY_NAME, + null, + this.connectionClassName); HEADER_PROVIDER = createHeaderProvider(); this.bigQuery = getBigQueryConnection(); @@ -383,7 +390,12 @@ HeaderProvider createHeaderProvider() { String partnerToken = buildPartnerToken(this.connectionUrl); String headerToken = DEFAULT_JDBC_TOKEN_VALUE + "/" + getLibraryVersion(this.getClass()) + partnerToken; - return FixedHeaderProvider.create("user-agent", headerToken); + Map headers = new java.util.HashMap<>(); + headers.put("user-agent", headerToken); + if (this.requestReason != null) { + headers.put("x-goog-request-reason", this.requestReason); + } + return FixedHeaderProvider.create(headers); } protected void addOpenStatements(Statement statement) { diff --git a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java index 3b26f7be5..4d11810a2 100644 --- a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java +++ b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java @@ -160,6 +160,7 @@ final class BigQueryJdbcUrlUtility { OAUTH2_TOKEN_URI_PROPERTY_NAME, HTAPI_ENDPOINT_OVERRIDE_PROPERTY_NAME, STS_ENDPOINT_OVERRIDE_PROPERTY_NAME); + static final String REQUEST_REASON_PROPERTY_NAME = "RequestReason"; static final List BYOID_PROPERTIES = Arrays.asList( BYOID_AUDIENCE_URI_PROPERTY_NAME, @@ -249,7 +250,8 @@ final class BigQueryJdbcUrlUtility { BigQueryConnectionProperty.newBuilder() .setName(OAUTH_SA_IMPERSONATION_CHAIN_PROPERTY_NAME) .setDescription( - "Comma separated list of service account emails in the impersonation chain.") + "Comma separated list of service account emails in the impersonation" + + " chain.") .build(), BigQueryConnectionProperty.newBuilder() .setName(OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME) @@ -569,6 +571,12 @@ final class BigQueryJdbcUrlUtility { .setDescription( "The password for accessing the Java TrustStore that is specified using" + " the property SSLTrustStore.") + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(REQUEST_REASON_PROPERTY_NAME) + .setDescription( + "Reason for the request, which is passed as the x-goog-request-reason" + + " header.") .build()))); private BigQueryJdbcUrlUtility() {} diff --git a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java index 6df4724b2..1cc0ad7bb 100644 --- a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java +++ b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java @@ -165,6 +165,23 @@ public void testHeaderProviderWithInvalidPartner() throws IOException, SQLExcept } } + @Test + public void testHeaderProviderWithRequestReason() throws IOException, SQLException { + String requestReason = "Ticket123"; + String url = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "OAuthType=2;ProjectId=MyBigQueryProject;" + + "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;" + + "OAuthClientSecret=redactedToken;RequestReason=" + + requestReason; + try (BigQueryConnection connection = new BigQueryConnection(url)) { + HeaderProvider headerProvider = connection.createHeaderProvider(); + java.util.Map headers = headerProvider.getHeaders(); + assertTrue(headers.containsKey("x-goog-request-reason")); + assertEquals(requestReason, headers.get("x-goog-request-reason")); + } + } + @Test public void testWriteAPIConnectionProperties() throws SQLException { // Test without connection properties. Defaults to default values. diff --git a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java index 86f087bf3..3cd0733a7 100644 --- a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java +++ b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java @@ -800,4 +800,18 @@ public void testParseLabelsEmpty() { Map labels = BigQueryJdbcUrlUtility.parseLabels(connection_uri, null); assertNull(labels); } + + @Test + public void testParseRequestReason() { + String url = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "OAuthType=3;ProjectId=testProject;RequestReason=testingRequestReason;"; + String requestReason = + BigQueryJdbcUrlUtility.parseStringProperty( + url, + BigQueryJdbcUrlUtility.REQUEST_REASON_PROPERTY_NAME, + null, + "testParseRequestReason"); + assertEquals("testingRequestReason", requestReason); + } }