fix: partner token overwrite of all subsequent or concurrent connections

Author: keshavdandeva

google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -73,7 +73,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
   String connectionClassName = this.toString();
   private static final String DEFAULT_JDBC_TOKEN_VALUE = "Google-BigQuery-JDBC-Driver";
   private static final String DEFAULT_VERSION = "0.0.0";
-  private static HeaderProvider HEADER_PROVIDER;
+  private HeaderProvider headerProvider;
   BigQueryReadClient bigQueryReadClient = null;
   BigQueryWriteClient bigQueryWriteClient = null;
   BigQuery bigQuery;
@@ -136,6 +136,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
   String requestReason;
   Long connectionPoolSize;
   Long listenerPoolSize;
+  String partnerToken;
 
   BigQueryConnection(String url) throws IOException {
     this.connectionUrl = url;
@@ -236,8 +237,9 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
     this.requestReason = ds.getRequestReason();
     this.connectionPoolSize = ds.getConnectionPoolSize();
     this.listenerPoolSize = ds.getListenerPoolSize();
+    this.partnerToken = ds.getPartnerToken();
 
-    HEADER_PROVIDER = createHeaderProvider();
+    this.headerProvider = createHeaderProvider();
     this.bigQuery = getBigQueryConnection();
   }
 
@@ -259,19 +261,11 @@ String getLibraryVersion(Class<?> libraryClass) {
     return version != null ? version : DEFAULT_VERSION;
   }
 
-  private String buildPartnerToken(String url) {
-    String partnerTokenString =
-        BigQueryJdbcUrlUtility.parsePartnerTokenProperty(url, this.connectionClassName);
-    if (partnerTokenString == null || partnerTokenString.isEmpty()) {
-      return "";
-    }
-    return partnerTokenString;
-  }
-
   HeaderProvider createHeaderProvider() {
-    String partnerToken = buildPartnerToken(this.connectionUrl);
-    String headerToken =
-        DEFAULT_JDBC_TOKEN_VALUE + "/" + getLibraryVersion(this.getClass()) + partnerToken;
+    String headerToken = DEFAULT_JDBC_TOKEN_VALUE + "/" + getLibraryVersion(this.getClass());
+    if (this.partnerToken != null && !this.partnerToken.isEmpty()) {
+      headerToken += this.partnerToken;
+    }
     Map<String, String> headers = new java.util.HashMap<>();
     headers.put("user-agent", headerToken);
     if (this.requestReason != null) {
@@ -937,14 +931,14 @@ private BigQuery getBigQueryConnection() {
       bigQueryOptions.setTransportOptions(this.httpTransportOptions);
     }
 
-    BigQueryOptions options = bigQueryOptions.setHeaderProvider(HEADER_PROVIDER).build();
+    BigQueryOptions options = bigQueryOptions.setHeaderProvider(this.headerProvider).build();
     options.setQueryPreviewEnabled(String.valueOf(this.useStatelessQueryMode));
     return options.getService();
   }
 
   private BigQueryReadClient getBigQueryReadClientConnection() throws IOException {
     BigQueryReadSettings.Builder bigQueryReadSettings =
-        BigQueryReadSettings.newBuilder().setHeaderProvider(HEADER_PROVIDER);
+        BigQueryReadSettings.newBuilder().setHeaderProvider(this.headerProvider);
     if (getRetrySettings() != null) {
       bigQueryReadSettings.createReadSessionSettings().setRetrySettings(getRetrySettings());
     }
@@ -986,7 +980,7 @@ private BigQueryReadClient getBigQueryReadClientConnection() throws IOException
 
   private BigQueryWriteClient getBigQueryWriteClientConnection() throws IOException {
     BigQueryWriteSettings.Builder bigQueryWriteSettings =
-        BigQueryWriteSettings.newBuilder().setHeaderProvider(HEADER_PROVIDER);
+        BigQueryWriteSettings.newBuilder().setHeaderProvider(this.headerProvider);
     if (getRetrySettings() != null) {
       bigQueryWriteSettings.createWriteStreamSettings().setRetrySettings(getRetrySettings());
     }

google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java

@@ -771,10 +771,6 @@ static boolean convertIntToBoolean(String value, String propertyName) {
     }
   }
 
-  static String parsePartnerTokenProperty(String url, String callerClassName) {
-    return parseUriProperty(url, PARTNER_TOKEN_PROPERTY_NAME);
-  }
-
   public static Level parseLogLevel(String logLevelString) {
     int logLevel = logLevelString != null ? Integer.parseInt(logLevelString) : DEFAULT_LOG_LEVEL;
     switch (logLevel) {

google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java

@@ -160,41 +160,40 @@ public void testParsePartnerTokenProperty() {
         "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
             + "PartnerToken=(GPN:partner_company; dev);ProjectId=MyBigQueryProject;";
     String expected = " (GPN:partner_company; dev)";
-    String result =
-        BigQueryJdbcUrlUtility.parsePartnerTokenProperty(url, "testParsePartnerTokenProperty");
+    String result = BigQueryJdbcUrlUtility.parseUriProperty(url, "PartnerToken");
     assertThat(result).isEqualTo(expected);
 
     // Case with only partner name
     url =
         "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
             + "PartnerToken=(GPN:another_partner);ProjectId=MyBigQueryProject;";
     expected = " (GPN:another_partner)";
-    result = BigQueryJdbcUrlUtility.parsePartnerTokenProperty(url, "testParsePartnerTokenProperty");
+    result = BigQueryJdbcUrlUtility.parseUriProperty(url, "PartnerToken");
     assertThat(result).isEqualTo(expected);
 
     // Case when PartnerToken property is not present
     url =
         "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
             + "ProjectId=MyBigQueryProject;";
-    result = BigQueryJdbcUrlUtility.parsePartnerTokenProperty(url, "testParsePartnerTokenProperty");
+    result = BigQueryJdbcUrlUtility.parseUriProperty(url, "PartnerToken");
     assertNull(result);
 
     // Case when PartnerToken property is present but empty
     url = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PartnerToken=();";
-    result = BigQueryJdbcUrlUtility.parsePartnerTokenProperty(url, "testParsePartnerTokenProperty");
+    result = BigQueryJdbcUrlUtility.parseUriProperty(url, "PartnerToken");
     assertNull(result);
 
     // Case when PartnerToken property is present but without partner name
     url = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PartnerToken=(env);";
-    result = BigQueryJdbcUrlUtility.parsePartnerTokenProperty(url, "testParsePartnerTokenProperty");
+    result = BigQueryJdbcUrlUtility.parseUriProperty(url, "PartnerToken");
     assertNull(result);
 
     // Case with extra spaces around the values
     url =
         "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
             + "PartnerToken= ( GPN: partner_name ; test_env ) ;";
     expected = " (GPN: partner_name ; test_env)";
-    result = BigQueryJdbcUrlUtility.parsePartnerTokenProperty(url, "testParsePartnerTokenProperty");
+    result = BigQueryJdbcUrlUtility.parseUriProperty(url, "PartnerToken");
     assertThat(result).isEqualTo(expected);
   }