chore: address pr feedback and refactor

Author: keshavdandeva

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

@@ -42,6 +42,7 @@
 import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
 import com.google.cloud.http.HttpTransportOptions;
 import com.google.cloud.logging.Logging;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSortedSet;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.baggage.Baggage;
@@ -438,6 +439,7 @@ String getConnectionUrl() {
     return connectionUrl;
   }
 
+  @VisibleForTesting
   public String getConnectionId() {
     return this.connectionId;
   }

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

@@ -912,8 +912,7 @@ private void processArrowStream(
       enqueueError(arrowBatchWrapperBlockingQueue, e);
       Thread.currentThread().interrupt();
     } catch (Exception e) {
-      if (e.getCause() instanceof InterruptedException
-          || Thread.currentThread().isInterrupted()) {
+      if (e.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) {
         LOG.log(
             Level.WARNING,
             "\n" + Thread.currentThread().getName() + " Interrupted @ arrowStreamProcessor",
@@ -1684,8 +1683,7 @@ private void parseAndPopulateRpcData(
       }
 
     } catch (Exception ex) {
-      if (ex.getCause() instanceof InterruptedException
-          || Thread.currentThread().isInterrupted()) {
+      if (ex.getCause() instanceof InterruptedException || Thread.currentThread().isInterrupted()) {
         LOG.log(
             Level.WARNING,
             "\n" + Thread.currentThread().getName() + " Interrupted @ populateBufferAsync",

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java

@@ -19,26 +19,26 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
 import com.google.api.gax.paging.Page;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.bigquery.jdbc.BigQueryConnection;
+import com.google.cloud.bigquery.jdbc.DataSource;
 import com.google.cloud.logging.LogEntry;
 import com.google.cloud.logging.Logging;
 import com.google.cloud.logging.LoggingOptions;
 import com.google.cloud.trace.v1.TraceServiceClient;
 import com.google.devtools.cloudtrace.v1.Trace;
 import com.google.devtools.cloudtrace.v1.TraceSpan;
 import java.sql.Connection;
-import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Properties;
 import org.junit.jupiter.api.Test;
 
 public class ITOpenTelemetryTest {
@@ -49,112 +49,89 @@ public class ITOpenTelemetryTest {
           "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;ProjectId=%s;OAuthType=3;Timeout=3600;",
           PROJECT_ID);
 
+  private static class TelemetryContext {
+    final String traceId;
+    final String spanId;
+
+    TelemetryContext(String traceId, String spanId) {
+      this.traceId = traceId;
+      this.spanId = spanId;
+    }
+  }
+
   @Test
   public void testExecute_withOpenTelemetryGcpExporter() throws Exception {
     assumeTrue(
         PROJECT_ID != null && !PROJECT_ID.trim().isEmpty(),
         "Skipping OTel E2E tests because no default Project ID is configured.");
 
-    // Step 1: Connect with GCP Exporters enabled
-    Properties props = new Properties();
-    props.setProperty("enableGcpTraceExporter", "true");
-    props.setProperty("enableGcpLogExporter", "true");
-    props.setProperty("LogLevel", "3"); // Triggers FINE log generation
-    props.setProperty("gcpTelemetryProjectId", PROJECT_ID);
-    props.setProperty("EnableHighThroughputAPI", "0");
-    props.setProperty("MaxResults", "50"); // Forces small page size (50) to trigger pagination
+    // Step 1: Connect with GCP Exporters enabled via DataSource
+    DataSource ds = DataSource.fromUrl(CONNECTION_URL);
+    ds.setEnableGcpTraceExporter(true);
+    ds.setEnableGcpLogExporter(true);
+    ds.setLogLevel("5"); // Triggers FINE log generation
+    ds.setGcpTelemetryProjectId(PROJECT_ID);
+    ds.setEnableHighThroughputAPI(false);
+    ds.setMaxResults(50L); // Forces small page size (50) to trigger pagination
 
     String connectionUuid = null;
 
-    try (Connection connection = DriverManager.getConnection(CONNECTION_URL, props);
+    try (Connection connection = ds.getConnection();
         Statement statement = connection.createStatement()) {
 
       // Retrieve the Connection UUID programmatically
       BigQueryConnection bqConnection = connection.unwrap(BigQueryConnection.class);
       connectionUuid = bqConnection.getConnectionId();
       assertNotNull(connectionUuid, "Connection UUID should be generated");
 
-      // Execute an in-memory array query (scans 0 bytes, extremely fast) and force pagination
-      String paginationQuery = "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS id;";
+      // Execute an in-memory array query (scans 0 bytes, extremely fast) and force pagination (3
+      // pages)
+      String paginationQuery = "SELECT * FROM UNNEST(GENERATE_ARRAY(1, 150)) AS id;";
       try (ResultSet paginatedRs = statement.executeQuery(paginationQuery)) {
         int rowCount = 0;
-        while (paginatedRs.next() && rowCount < 1000) {
+        while (paginatedRs.next() && rowCount < 150) {
           rowCount++;
         }
       }
     }
 
-    // Step 2: Retrieve logs from Cloud Logging and extract TraceId
-    String traceId = null;
-    String hexSpanId = null;
+    // Step 2: Retrieve and assert logs
+    TelemetryContext telCtx = verifyAndFetchLogs(connectionUuid);
 
-    try (Logging logging =
-        LoggingOptions.newBuilder().setProjectId(PROJECT_ID).build().getService()) {
-      String filter =
-          "logName:\"projects/"
-              + PROJECT_ID
-              + "/logs/com.google.cloud.bigquery\" AND labels.\"jdbc.connection_id\"=\""
-              + connectionUuid
-              + "\"";
+    // Step 3: Query Cloud Trace and assert parent-child hierarchy
+    Trace trace = verifyAndFetchTrace(telCtx.traceId);
 
-      List<LogEntry> entries = fetchLogsWithRetry(logging, filter);
-      assertFalse(entries.isEmpty(), "Telemetry logs should be exported to GCP");
+    boolean foundParentExecuteQuery = false;
+    boolean foundChildSdkSpans = false;
+    boolean foundPaginationSpans = false;
+    long parentSpanId = 0;
 
-      LogEntry sampleEntry = entries.get(0);
-      traceId = sampleEntry.getTrace();
-      hexSpanId = sampleEntry.getSpanId();
-
-      assertNotNull(traceId, "Log entry must contain TraceId");
-      assertNotNull(hexSpanId, "Log entry must contain SpanId");
-
-      // Verify Connection UUID label correlation on all entries
-      for (LogEntry entry : entries) {
-        assertEquals(connectionUuid, entry.getLabels().get("jdbc.connection_id"));
+    for (TraceSpan span : trace.getSpansList()) {
+      String spanName = span.getName();
+      if (spanName.equals("BigQueryStatement.executeQuery")) {
+        foundParentExecuteQuery = true;
+        parentSpanId = span.getSpanId();
       }
     }
 
-    // Step 3: Query Cloud Trace using TraceId and assert parent-child hierarchy
-    String hexTraceId = traceId;
-    if (traceId.contains("/traces/")) {
-      hexTraceId = traceId.substring(traceId.lastIndexOf("/traces/") + 8);
-    }
-
-    try (TraceServiceClient traceClient = TraceServiceClient.create()) {
-      Trace trace = fetchTraceWithRetry(traceClient, PROJECT_ID, hexTraceId);
-      assertNotNull(trace, "Trace must be found in Cloud Trace API: " + hexTraceId);
-
-      boolean foundParentExecuteQuery = false;
-      boolean foundChildSdkSpans = false;
-      boolean foundPaginationSpans = false;
-      long parentSpanId = 0;
+    assertTrue(
+        foundParentExecuteQuery,
+        "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'");
 
-      for (TraceSpan span : trace.getSpansList()) {
-        String spanName = span.getName();
-        if (spanName.equals("BigQueryStatement.executeQuery")) {
-          foundParentExecuteQuery = true;
-          parentSpanId = span.getSpanId();
-        }
+    // Verify that we captured child spans or linked pagination spans
+    for (TraceSpan span : trace.getSpansList()) {
+      if (span.getParentSpanId() == parentSpanId && parentSpanId != 0) {
+        foundChildSdkSpans = true;
       }
-
-      assertTrue(
-          foundParentExecuteQuery,
-          "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'");
-
-      // Verify that we captured child spans or linked pagination spans
-      for (TraceSpan span : trace.getSpansList()) {
-        if (span.getParentSpanId() == parentSpanId && parentSpanId != 0) {
-          foundChildSdkSpans = true;
-        }
-        if (span.getName().equals("BigQueryStatement.pagination")) {
-          foundPaginationSpans = true;
-        }
+      if (span.getName().equals("BigQueryStatement.pagination")) {
+        foundPaginationSpans = true;
       }
-
-      assertTrue(foundPaginationSpans, "OTel pagination must generate pagination spans");
-      assertTrue(
-          foundChildSdkSpans,
-          "OTel context must propagate parent to downstream pagination child spans");
     }
+
+    assertTrue(foundPaginationSpans, "OTel pagination must generate pagination spans");
+    assertTrue(
+        foundChildSdkSpans,
+        "OTel context must propagate parent to downstream pagination child spans");
   }
 
   @Test
@@ -163,37 +140,48 @@ public void testExecute_withErrorCorrelation() throws Exception {
         PROJECT_ID != null && !PROJECT_ID.trim().isEmpty(),
         "Skipping OTel E2E tests because no default Project ID is configured.");
 
-    // Step 1: Connect with GCP Exporters enabled
-    Properties props = new Properties();
-    props.setProperty("enableGcpTraceExporter", "true");
-    props.setProperty("enableGcpLogExporter", "true");
-    props.setProperty("LogLevel", "3"); // Triggers FINE log generation
-    props.setProperty("gcpTelemetryProjectId", PROJECT_ID);
+    // Step 1: Connect with GCP Exporters enabled via DataSource
+    DataSource ds = DataSource.fromUrl(CONNECTION_URL);
+    ds.setEnableGcpTraceExporter(true);
+    ds.setEnableGcpLogExporter(true);
+    ds.setLogLevel("5"); // Triggers FINE log generation
+    ds.setGcpTelemetryProjectId(PROJECT_ID);
 
     String connectionUuid = null;
 
-    try (Connection connection = DriverManager.getConnection(CONNECTION_URL, props);
+    try (Connection connection = ds.getConnection();
         Statement statement = connection.createStatement()) {
 
       // Retrieve the Connection UUID programmatically
       BigQueryConnection bqConnection = connection.unwrap(BigQueryConnection.class);
       connectionUuid = bqConnection.getConnectionId();
       assertNotNull(connectionUuid, "Connection UUID should be generated");
 
-      // Execute a query designed to fail due to non-existent table
-      boolean caughtException = false;
-      try {
-        statement.executeQuery("SELECT * FROM invalid_dataset.invalid_table;");
-      } catch (SQLException e) {
-        caughtException = true;
+      // Execute a query designed to fail instantly due to syntax error (compiler-level failure)
+      assertThrows(SQLException.class, () -> statement.executeQuery("SELECT * FROM;"));
+    }
+
+    // Step 2: Retrieve and assert logs
+    TelemetryContext telCtx = verifyAndFetchLogs(connectionUuid);
+
+    // Step 3: Query Cloud Trace and assert span status is ERROR
+    Trace trace = verifyAndFetchTrace(telCtx.traceId);
+
+    boolean foundParentExecuteQuery = false;
+
+    for (TraceSpan span : trace.getSpansList()) {
+      String spanName = span.getName();
+      if (spanName.equals("BigQueryStatement.executeQuery")) {
+        foundParentExecuteQuery = true;
       }
-      assertTrue(caughtException, "Expected SQLException to be thrown");
     }
 
-    // Step 2: Retrieve logs from Cloud Logging and assert error logs
-    String traceId = null;
-    String hexSpanId = null;
+    assertTrue(
+        foundParentExecuteQuery,
+        "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'");
+  }
 
+  private TelemetryContext verifyAndFetchLogs(String connectionUuid) throws Exception {
     try (Logging logging =
         LoggingOptions.newBuilder().setProjectId(PROJECT_ID).build().getService()) {
       String filter =
@@ -207,14 +195,22 @@ public void testExecute_withErrorCorrelation() throws Exception {
       assertFalse(entries.isEmpty(), "Telemetry logs should be exported to GCP");
 
       LogEntry sampleEntry = entries.get(0);
-      traceId = sampleEntry.getTrace();
-      hexSpanId = sampleEntry.getSpanId();
+      String traceId = sampleEntry.getTrace();
+      String hexSpanId = sampleEntry.getSpanId();
 
       assertNotNull(traceId, "Log entry must contain TraceId");
       assertNotNull(hexSpanId, "Log entry must contain SpanId");
+
+      // Verify Connection UUID label correlation on all entries
+      for (LogEntry entry : entries) {
+        assertEquals(connectionUuid, entry.getLabels().get("jdbc.connection_id"));
+      }
+
+      return new TelemetryContext(traceId, hexSpanId);
     }
+  }
 
-    // Step 3: Query Cloud Trace using TraceId and assert span status is ERROR
+  private Trace verifyAndFetchTrace(String traceId) throws Exception {
     String hexTraceId = traceId;
     if (traceId.contains("/traces/")) {
       hexTraceId = traceId.substring(traceId.lastIndexOf("/traces/") + 8);
@@ -223,19 +219,7 @@ public void testExecute_withErrorCorrelation() throws Exception {
     try (TraceServiceClient traceClient = TraceServiceClient.create()) {
       Trace trace = fetchTraceWithRetry(traceClient, PROJECT_ID, hexTraceId);
       assertNotNull(trace, "Trace must be found in Cloud Trace API: " + hexTraceId);
-
-      boolean foundParentExecuteQuery = false;
-
-      for (TraceSpan span : trace.getSpansList()) {
-        String spanName = span.getName();
-        if (spanName.equals("BigQueryStatement.executeQuery")) {
-          foundParentExecuteQuery = true;
-        }
-      }
-
-      assertTrue(
-          foundParentExecuteQuery,
-          "Traces must contain JDBC parent span 'BigQueryStatement.executeQuery'");
+      return trace;
     }
   }
 
@@ -246,7 +230,6 @@ private <T> T pollWithRetry(java.util.concurrent.Callable<T> task) throws Interr
 
     while (attempts < maxAttempts) {
       attempts++;
-      Thread.sleep(delayMs);
       try {
         T result = task.call();
         if (result != null) {
@@ -258,6 +241,9 @@ private <T> T pollWithRetry(java.util.concurrent.Callable<T> task) throws Interr
       } catch (Exception e) {
         // Ignore exceptions during remote lookup and retry
       }
+      if (attempts < maxAttempts) {
+        Thread.sleep(delayMs);
+      }
     }
     return null;
   }