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/BigQueryArrowResultSet.java

@@ -212,7 +212,15 @@ public void close() {
 
   @Override
   public boolean next() throws SQLException {
-    checkClosed();
+    try (BigQueryJdbcMdc.MdcCloseable mdc =
+        BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) {
+      LOG.finest("++enter++");
+      checkClosed();
+      return nextImpl();
+    }
+  }
+
+  private boolean nextImpl() throws SQLException {
     if (this.isNested) {
       if (this.currentNestedBatch == null || this.currentNestedBatch.getNestedRecords() == null) {
         IllegalStateException ex =
@@ -319,10 +327,16 @@ private Object getObjectInternal(int columnIndex) throws SQLException {
 
   @Override
   public Object getObject(int columnIndex) throws SQLException {
+    try (BigQueryJdbcMdc.MdcCloseable mdc =
+        BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) {
+      // columnIndex is SQL index starting at 1
+      LOG.finest("++enter++");
+      checkClosed();
+      return getObjectImpl(columnIndex);
+    }
+  }
 
-    // columnIndex is SQL index starting at 1
-    LOG.finest("++enter++");
-    checkClosed();
+  private Object getObjectImpl(int columnIndex) throws SQLException {
     Object value = getObjectInternal(columnIndex);
     if (value == null) {
       return null;
@@ -451,14 +465,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

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

@@ -58,6 +58,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 +68,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 +101,7 @@ public String getQueryId() {
   }
 
   @Override
-  public void close() {
+  public void close() throws SQLException {
     try {
       if (statement != null && statement.isCloseOnCompletion() && !statement.hasMoreResults()) {
         statement.close();