chore: refactor based on pr comments

Author: keshavdandeva

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

@@ -27,8 +27,6 @@
 import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
 import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
 import com.google.cloud.bigquery.storage.v1.ArrowSchema;
-import io.opentelemetry.api.trace.Span;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import java.io.IOException;
 import java.math.BigDecimal;
@@ -239,32 +237,31 @@ public boolean next() throws SQLException {
             || this.currentBatchRowIndex == (this.vectorSchemaRoot.getRowCount() - 1)) {
           /* Start of iteration or we have exhausted the current batch */
           // Advance the cursor. Potentially blocking operation.
-          BigQueryArrowBatchWrapper batchWrapper;
-          try (Scope scope = Context.current().with(Span.wrap(originalSpanContext)).makeCurrent()) {
-            batchWrapper = this.buffer.take();
-          }
-          if (batchWrapper.getException() != null) {
-            throw new BigQueryJdbcRuntimeException(batchWrapper.getException());
-          }
-          if (batchWrapper.isLast()) {
-            /* Marks the end of the records */
-            if (this.vectorSchemaRoot != null) {
-              // IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds
-              // the last batch
-              this.vectorSchemaRoot.clear();
+          try (Scope scope = makeOriginalContextCurrent()) {
+            BigQueryArrowBatchWrapper batchWrapper = this.buffer.take();
+            if (batchWrapper.getException() != null) {
+              throw new BigQueryJdbcRuntimeException(batchWrapper.getException());
+            }
+            if (batchWrapper.isLast()) {
+              /* Marks the end of the records */
+              if (this.vectorSchemaRoot != null) {
+                // IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds
+                // the last batch
+                this.vectorSchemaRoot.clear();
+              }
+              this.hasReachedEnd = true;
+              this.rowCount++;
+              return false;
             }
-            this.hasReachedEnd = true;
+            // Valid batch, process it
+            ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch();
+            // Populates vectorSchemaRoot
+            this.arrowDeserializer.deserializeArrowBatch(arrowBatch);
+            // Pointing to the first row in this fresh batch
+            this.currentBatchRowIndex = 0;
             this.rowCount++;
-            return false;
+            return true;
           }
-          // Valid batch, process it
-          ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch();
-          // Populates vectorSchemaRoot
-          this.arrowDeserializer.deserializeArrowBatch(arrowBatch);
-          // Pointing to the first row in this fresh batch
-          this.currentBatchRowIndex = 0;
-          this.rowCount++;
-          return true;
         }
         // There are rows left in the current batch.
         else if (this.currentBatchRowIndex < this.vectorSchemaRoot.getRowCount()) {

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

@@ -29,6 +29,8 @@
 import com.google.cloud.bigquery.exception.BigQueryJdbcCoercionNotFoundException;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import java.io.InputStream;
 import java.io.Reader;
 import java.io.StringReader;
@@ -72,6 +74,10 @@ protected BigQueryBaseResultSet(
     this.originalSpanContext = Span.current().getSpanContext();
   }
 
+  protected Scope makeOriginalContextCurrent() {
+    return Context.current().with(Span.wrap(this.originalSpanContext)).makeCurrent();
+  }
+
   public QueryStatistics getQueryStatistics() {
     if (queryStatistics != null) {
       return queryStatistics;

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

@@ -1745,7 +1745,7 @@ private ResultSet getTablesImpl(
     final String catalogParam = effectiveCatalog;
     final String schemaParam = effectiveSchemaPattern;
 
-    Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection);
+    Tracer tracer = this.connection.getTracer();
     SpanContext parentSpanContext = Span.current().getSpanContext();
     Runnable tableFetcher =
         () -> {
@@ -2139,7 +2139,7 @@ private ResultSet getColumnsImpl(
     final String catalogParam = effectiveCatalog;
     final String schemaParam = effectiveSchemaPattern;
 
-    Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection);
+    Tracer tracer = this.connection.getTracer();
     SpanContext parentSpanContext = Span.current().getSpanContext();
     Runnable columnFetcher =
         () -> {
@@ -3699,7 +3699,7 @@ private ResultSet getSchemasImpl(String catalog, String schemaPattern) throws SQ
     final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
     final String catalogParam = catalog;
 
-    Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection);
+    Tracer tracer = this.connection.getTracer();
     SpanContext parentSpanContext = Span.current().getSpanContext();
     Runnable schemaFetcher =
         () -> {
@@ -5377,7 +5377,7 @@ private interface TracedMetadataOperation<T> {
 
   private <T> T withTracing(String spanName, TracedMetadataOperation<T> operation)
       throws SQLException {
-    Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection);
+    Tracer tracer = this.connection.getTracer();
     Span span = tracer.spanBuilder(spanName).startSpan();
     try (Scope scope = span.makeCurrent()) {
       return operation.run();

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

@@ -49,12 +49,4 @@ public static OpenTelemetry getOpenTelemetry(
   public static Tracer getTracer(OpenTelemetry openTelemetry) {
     return openTelemetry.getTracer(INSTRUMENTATION_SCOPE_NAME);
   }
-
-  /** Gets a Tracer for the JDBC driver, fallback to noop if connection has no tracer. */
-  public static Tracer getSafeTracer(BigQueryConnection connection) {
-    if (connection != null && connection.getTracer() != null) {
-      return connection.getTracer();
-    }
-    return getOpenTelemetry(false, false, null).getTracer(INSTRUMENTATION_SCOPE_NAME);
-  }
 }

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

@@ -25,8 +25,6 @@
 import com.google.cloud.bigquery.FieldValue.Attribute;
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
-import io.opentelemetry.api.trace.Span;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -161,7 +159,7 @@ public boolean next() throws SQLException {
       }
       try {
         // Advance the cursor,Potentially blocking operation
-        try (Scope scope = Context.current().with(Span.wrap(originalSpanContext)).makeCurrent()) {
+        try (Scope scope = makeOriginalContextCurrent()) {
           this.cursor = this.buffer.take();
         }
         if (this.cursor.getException() != null) {

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

@@ -1542,7 +1542,7 @@ private interface TracedOperation<T> {
   }
 
   private <T> T withTracing(String spanName, TracedOperation<T> operation) throws SQLException {
-    Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection);
+    Tracer tracer = this.connection.getTracer();
     Span span = tracer.spanBuilder(spanName).startSpan();
     try (Scope scope = span.makeCurrent()) {
       return operation.run(span);
@@ -1561,7 +1561,7 @@ private void fetchNextPages(
       BlockingQueue<Tuple<TableResult, Boolean>> rpcResponseQueue,
       BlockingQueue<BigQueryFieldValueListWrapper> bigQueryFieldValueListWrapperBlockingQueue,
       TableResult result) {
-    Tracer tracer = BigQueryJdbcOpenTelemetry.getSafeTracer(this.connection);
+    Tracer tracer = this.connection.getTracer();
     SpanContext parentSpanContext = Span.current().getSpanContext();
     String currentPageToken = firstPageToken;
     TableResult currentResults = result;

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

@@ -52,6 +52,7 @@
 import com.google.cloud.bigquery.storage.v1.ReadSession;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Maps;
+import io.opentelemetry.api.OpenTelemetry;
 import java.io.IOException;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -127,6 +128,9 @@ private Job getJobMock(
   @BeforeEach
   public void setUp() throws IOException, SQLException {
     bigQueryConnection = mock(BigQueryConnection.class);
+    doReturn(OpenTelemetry.noop().getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME))
+        .when(bigQueryConnection)
+        .getTracer();
     rpcFactoryMock = mock(BigQueryRpcFactory.class);
     bigquery = mock(BigQuery.class);
     bigQueryConnection.bigQuery = bigquery;