feat(bqjdbc): Add BigQueryJdbcMdc context to entry level methods - part 3

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

@@ -63,17 +63,14 @@ static BigQueryArrowBatchWrapper ofError(Exception exception) {
   }
 
   ArrowRecordBatch getCurrentArrowBatch() {
-    LOG.finest("++enter++");
     return this.currentArrowBatch;
   }
 
   JsonStringArrayList getNestedRecords() {
-    LOG.finest("++enter++");
     return this.nestedRecords;
   }
 
   boolean isLast() {
-    LOG.finest("++enter++");
     return this.isLast;
   }
 

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

@@ -181,7 +181,6 @@ private ArrowDeserializer(ArrowSchema arrowSchema) throws IOException {
     }
 
     private void deserializeArrowBatch(ArrowRecordBatch batch) throws SQLException {
-      LOG.finest("++enter++");
       try {
         if (vectorSchemaRoot != null) {
           // Clear vectorSchemaRoot before populating a new batch
@@ -213,6 +212,15 @@ public void close() {
   @Override
   public boolean next() throws SQLException {
     checkClosed();
+    try {
+      return nextImpl();
+    } catch (SQLException | RuntimeException ex) {
+      BigQueryJdbcMdc.setContextPersistent(this.connectionId);
+      throw ex;
+    }
+  }
+
+  private boolean nextImpl() throws SQLException {
     if (this.isNested) {
       if (this.currentNestedBatch == null || this.currentNestedBatch.getNestedRecords() == null) {
         IllegalStateException ex =
@@ -319,10 +327,17 @@ private Object getObjectInternal(int columnIndex) throws SQLException {
 
   @Override
   public Object getObject(int columnIndex) throws SQLException {
-
     // columnIndex is SQL index starting at 1
-    LOG.finest("++enter++");
     checkClosed();
+    try {
+      return getObjectImpl(columnIndex);
+    } catch (SQLException | RuntimeException ex) {
+      BigQueryJdbcMdc.setContextPersistent(this.connectionId);
+      throw ex;
+    }
+  }
+
+  private Object getObjectImpl(int columnIndex) throws SQLException {
     Object value = getObjectInternal(columnIndex);
     if (value == null) {
       return null;
@@ -451,14 +466,22 @@ private String formatRangeElement(Object element, StandardSQLTypeName elementTyp
   }
 
   @Override
-  public void close() {
-    LOG.fine("Closing BigqueryArrowResultSet %s.", this);
-    this.isClosed = true;
-    if (ownedThread != null && !ownedThread.isInterrupted()) {
-      // interrupt the producer thread when result set is closed
-      ownedThread.interrupt();
+  public void close() throws SQLException {
+    if (isClosed()) {
+      return;
+    }
+
+    try (BigQueryJdbcMdc.MdcCloseable mdc =
+        BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) {
+      LOG.finest("++enter++");
+      LOG.fine("Closing BigqueryArrowResultSet %s.", this);
+      this.isClosed = true;
+      if (ownedThread != null && !ownedThread.isInterrupted()) {
+        // interrupt the producer thread when result set is closed
+        ownedThread.interrupt();
+      }
+      super.close();
     }
-    super.close();
   }
 
   @Override
@@ -492,7 +515,6 @@ public boolean isFirst() throws SQLException {
 
   @Override
   public boolean isLast() throws SQLException {
-    LOG.finest("++enter++");
     checkClosed();
     if (this.isNested) {
       return this.nestedRowIndex == this.toIndexExclusive - 1;

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

@@ -46,7 +46,8 @@
 
 public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet
     implements BigQueryResultSet {
-  protected final BigQueryJdbcCustomLogger LOG = new BigQueryJdbcCustomLogger(this.toString());
+  protected static final BigQueryJdbcCustomLogger LOG =
+      new BigQueryJdbcCustomLogger(BigQueryBaseResultSet.class.getName());
   private BigQuery bigQuery;
   private JobId jobId;
   private String queryId;
@@ -58,6 +59,8 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet
   protected boolean isClosed = false;
   protected boolean wasNull = false;
   protected final BigQueryTypeCoercer bigQueryTypeCoercer = BigQueryTypeCoercionUtility.INSTANCE;
+  protected BigQueryConnection connection = null;
+  protected String connectionId = null;
 
   protected BigQueryBaseResultSet(
       BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) {
@@ -66,6 +69,8 @@ protected BigQueryBaseResultSet(
     this.schema = schema;
     this.schemaFieldList = schema != null ? schema.getFields() : null;
     this.isNested = isNested;
+    this.connection = statement != null ? (BigQueryConnection) statement.getConnection() : null;
+    this.connectionId = this.connection != null ? this.connection.getConnectionId() : null;
   }
 
   public QueryStatistics getQueryStatistics() {
@@ -97,7 +102,7 @@ public String getQueryId() {
   }
 
   @Override
-  public void close() {
+  public void close() throws SQLException {
     try {
       if (statement != null && statement.isCloseOnCompletion() && !statement.hasMoreResults()) {
         statement.close();