Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
String sslTrustStorePassword;
long maxBytesBilled;
Map<String, String> labels;
String requestReason;

BigQueryConnection(String url) throws IOException {
this.connectionUrl = url;
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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<String, String> headers = new java.util.HashMap<>();
headers.put("user-agent", headerToken);
if (this.requestReason != null) {
headers.put("x-goog-request-reason", this.requestReason);
}
Comment on lines +395 to +397
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

I have a couple of suggestions for this block to improve robustness and maintainability:

  1. Avoid blank headers: The current code sends the x-goog-request-reason header even if requestReason is an empty string or just whitespace. It's better to avoid sending headers with blank values by adding a check like !this.requestReason.trim().isEmpty().
  2. Use a constant for the header name: The string "x-goog-request-reason" is also used in the test file. Defining it as a public constant in BigQueryJdbcUrlUtility would avoid magic strings and improve maintainability.

Here's a suggested implementation that applies the first point:

Suggested change
if (this.requestReason != null) {
headers.put("x-goog-request-reason", this.requestReason);
}
if (this.requestReason != null && !this.requestReason.trim().isEmpty()) {
headers.put("x-goog-request-reason", this.requestReason);
}

return FixedHeaderProvider.create(headers);
}

protected void addOpenStatements(Statement statement) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<String> BYOID_PROPERTIES =
Arrays.asList(
BYOID_AUDIENCE_URI_PROPERTY_NAME,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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() {}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String> headers = headerProvider.getHeaders();
assertTrue(headers.containsKey("x-goog-request-reason"));
assertEquals(requestReason, headers.get("x-goog-request-reason"));
}
}
Comment on lines +168 to +183
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This test covers the happy path well. To improve test coverage, please consider adding another test case for when RequestReason is provided but is empty or blank in the connection URL (e.g., RequestReason= or RequestReason= ). This would ensure the driver handles this edge case correctly and doesn't send the header, especially if you implement my other suggestion to avoid sending blank headers.


@Test
public void testWriteAPIConnectionProperties() throws SQLException {
// Test without connection properties. Defaults to default values.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -800,4 +800,18 @@ public void testParseLabelsEmpty() {
Map<String, String> 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);
}
}
Loading