only connection extracts connectionID

Author: Neenu1995

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

@@ -109,11 +109,12 @@ public void close() {
   }
 
   protected SQLException logAndCreateException(SQLException ex) {
-    if (this.statement != null) {
-      try (BigQueryJdbcMdc.MdcCloseable mdc =
-          BigQueryJdbcMdc.registerInstance(this.statement.connectionId)) {
-        LOG.severe(ex.getMessage(), ex);
-      }
+    if (this.statement == null) {
+      return ex;
+    }
+    try (BigQueryJdbcMdc.MdcCloseable mdc =
+        BigQueryJdbcMdc.registerInstance(this.statement.connectionId)) {
+      LOG.severe(ex.getMessage(), ex);
     }
     return ex;
   }
@@ -209,7 +210,7 @@ public ResultSetMetaData getMetaData() throws SQLException {
         metaData = BigQueryResultSetMetadata.of(this.schema.getFields(), this.statement);
       }
     }
-    return BigQueryJdbcContextProxy.wrap(metaData, ResultSetMetaData.class);
+    return BigQueryJdbcContextProxy.wrap(metaData, ResultSetMetaData.class, connectionId);
   }
 
   @Override

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

@@ -41,7 +41,7 @@ private BigQueryJdbcContextProxy(Object target, String connectionId, Class<?> in
     this.interfaceType = interfaceType;
   }
 
-  /** Wraps a target JDBC object with a dynamic proxy carrying the connection context. */
+  /** Wraps a target Connection JDBC object, auto-extracting its connection ID. */
   @SuppressWarnings("unchecked")
   static <T> T wrap(Object target, Class<T> interfaceType) {
     if (target == null) {
@@ -55,38 +55,26 @@ static <T> T wrap(Object target, Class<T> interfaceType) {
             new BigQueryJdbcContextProxy(target, connectionId, interfaceType));
   }
 
+  /** Wraps a target child JDBC object, propagating the connection ID parameter directly. */
+  @SuppressWarnings("unchecked")
+  static <T> T wrap(Object target, Class<T> interfaceType, String connectionId) {
+    if (target == null) {
+      return null;
+    }
+    return (T)
+        Proxy.newProxyInstance(
+            interfaceType.getClassLoader(),
+            new Class<?>[] {interfaceType},
+            new BigQueryJdbcContextProxy(target, connectionId, interfaceType));
+  }
+
   private static String extractConnectionId(Object target) {
     if (target == null) {
       return null;
     }
     if (target instanceof BigQueryConnection) {
       return ((BigQueryConnection) target).getConnectionId();
     }
-    if (target instanceof BigQueryStatement) {
-      return ((BigQueryStatement) target).connectionId;
-    }
-    if (target instanceof BigQueryDatabaseMetaData) {
-      return ((BigQueryDatabaseMetaData) target).connection.getConnectionId();
-    }
-    if (target instanceof BigQueryResultSetMetadata) {
-      java.sql.Statement stmt = ((BigQueryResultSetMetadata) target).getStatement();
-      if (stmt != null) {
-        if (Proxy.isProxyClass(stmt.getClass())) {
-          InvocationHandler handler = Proxy.getInvocationHandler(stmt);
-          if (handler instanceof BigQueryJdbcContextProxy) {
-            return ((BigQueryJdbcContextProxy) handler).connectionId;
-          }
-        }
-        if (stmt instanceof BigQueryStatement) {
-          return ((BigQueryStatement) stmt).connectionId;
-        }
-      }
-    }
-    // Fallback to thread active context
-    String activeId = BigQueryJdbcMdc.getConnectionId();
-    if (activeId != null) {
-      return activeId;
-    }
     return null;
   }
 
@@ -146,17 +134,17 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
 
       // Automatically cascade proxy wrappers to child JDBC objects returned from calls
       if (result instanceof java.sql.CallableStatement) {
-        return wrap(result, java.sql.CallableStatement.class);
+        return wrap(result, java.sql.CallableStatement.class, connectionId);
       } else if (result instanceof java.sql.PreparedStatement) {
-        return wrap(result, java.sql.PreparedStatement.class);
+        return wrap(result, java.sql.PreparedStatement.class, connectionId);
       } else if (result instanceof java.sql.Statement) {
-        return wrap(result, java.sql.Statement.class);
+        return wrap(result, java.sql.Statement.class, connectionId);
       } else if (result instanceof java.sql.DatabaseMetaData) {
-        return wrap(result, java.sql.DatabaseMetaData.class);
+        return wrap(result, java.sql.DatabaseMetaData.class, connectionId);
       } else if (result instanceof java.sql.ParameterMetaData) {
-        return wrap(result, java.sql.ParameterMetaData.class);
+        return wrap(result, java.sql.ParameterMetaData.class, connectionId);
       } else if (result instanceof java.sql.ResultSetMetaData) {
-        return wrap(result, java.sql.ResultSetMetaData.class);
+        return wrap(result, java.sql.ResultSetMetaData.class, connectionId);
       }
 
       return result;

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

@@ -65,7 +65,7 @@ public void testExtractConnectionIdFromStatement() throws SQLException {
     BigQueryStatement stmt = new BigQueryStatement(mockConn);
     stmt.connectionId = "conn-uuid-456";
 
-    Statement proxy = BigQueryJdbcContextProxy.wrap(stmt, Statement.class);
+    Statement proxy = BigQueryJdbcContextProxy.wrap(stmt, Statement.class, "conn-uuid-456");
     assertNotNull(proxy);
 
     // We can call any statement method (like getUpdateCount) and verify context routing
@@ -83,7 +83,8 @@ public void testExtractConnectionIdFromDatabaseMetaData() throws SQLException {
 
     BigQueryDatabaseMetaData meta = new BigQueryDatabaseMetaData(mockConn);
 
-    DatabaseMetaData proxy = BigQueryJdbcContextProxy.wrap(meta, DatabaseMetaData.class);
+    DatabaseMetaData proxy =
+        BigQueryJdbcContextProxy.wrap(meta, DatabaseMetaData.class, "conn-uuid-789");
     assertNotNull(proxy);
 
     // Assert read-only capability does not leak thread context
@@ -100,7 +101,8 @@ public void testExtractConnectionIdFromResultSetMetaData() throws SQLException {
     FieldList fields = FieldList.of(Field.of("col", StandardSQLTypeName.STRING));
     BigQueryResultSetMetadata meta = BigQueryResultSetMetadata.of(fields, stmt);
 
-    ResultSetMetaData proxy = BigQueryJdbcContextProxy.wrap(meta, ResultSetMetaData.class);
+    ResultSetMetaData proxy =
+        BigQueryJdbcContextProxy.wrap(meta, ResultSetMetaData.class, "conn-uuid-999");
     assertNotNull(proxy);
 
     assertEquals(1, proxy.getColumnCount());
@@ -113,7 +115,7 @@ public void testWrapWithNullContextAndExceptionThrown() throws SQLException {
     mockStmt.connectionId = null;
     when(mockStmt.executeQuery("SELECT *")).thenThrow(new SQLException("Database error"));
 
-    Statement proxy = BigQueryJdbcContextProxy.wrap(mockStmt, Statement.class);
+    Statement proxy = BigQueryJdbcContextProxy.wrap(mockStmt, Statement.class, null);
     assertNotNull(proxy);
 
     SQLException ex = assertThrows(SQLException.class, () -> proxy.executeQuery("SELECT *"));

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

@@ -143,7 +143,7 @@ public void testProxyExceptionLogRouting() throws Exception {
           .thenThrow(new SQLException("Database error"));
 
       // Wrap it using our proxy (which dynamically extracts "c456" as its connection ID!)
-      Statement proxy = BigQueryJdbcContextProxy.wrap(mockStmt, Statement.class);
+      Statement proxy = BigQueryJdbcContextProxy.wrap(mockStmt, Statement.class, "c456");
       assertNotNull(proxy);
 
       // Call the proxy method. It will throw SQLException