Skip to content

Commit 7c9376b

Browse files
Neenu1995Neenu1995
committed
address comments
1 parent 35a55e1 commit 7c9376b

10 files changed

Lines changed: 186 additions & 40 deletions

File tree

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ public void close() {
109109
}
110110

111111
protected SQLException logAndCreateException(SQLException ex) {
112-
if (BigQueryJdbcRootLogger.isFileLoggingEnabled() && this.statement != null) {
112+
if (this.statement != null) {
113113
try (BigQueryJdbcMdc.MdcCloseable mdc =
114114
BigQueryJdbcMdc.registerInstance(this.statement.connectionId)) {
115115
LOG.severe(ex.getMessage(), ex);
@@ -209,10 +209,7 @@ public ResultSetMetaData getMetaData() throws SQLException {
209209
metaData = BigQueryResultSetMetadata.of(this.schema.getFields(), this.statement);
210210
}
211211
}
212-
if (connectionId != null) {
213-
return BigQueryJdbcContextProxy.wrap(metaData, ResultSetMetaData.class, connectionId);
214-
}
215-
return metaData;
212+
return BigQueryJdbcContextProxy.wrap(metaData, ResultSetMetaData.class);
216213
}
217214

218215
@Override

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -173,8 +173,7 @@ public Connection connect(String url, Properties info) throws SQLException {
173173
logLevel,
174174
logPath,
175175
this.toString());
176-
return BigQueryJdbcContextProxy.wrap(
177-
connection, Connection.class, connection.getConnectionId());
176+
return BigQueryJdbcContextProxy.wrap(connection, Connection.class);
178177
} else {
179178
return null;
180179
}

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

Lines changed: 46 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -43,17 +43,53 @@ private BigQueryJdbcContextProxy(Object target, String connectionId, Class<?> in
4343

4444
/** Wraps a target JDBC object with a dynamic proxy carrying the connection context. */
4545
@SuppressWarnings("unchecked")
46-
static <T> T wrap(Object target, Class<T> interfaceType, String connectionId) {
46+
static <T> T wrap(Object target, Class<T> interfaceType) {
4747
if (target == null) {
4848
return null;
4949
}
50+
String connectionId = extractConnectionId(target);
5051
return (T)
5152
Proxy.newProxyInstance(
5253
interfaceType.getClassLoader(),
5354
new Class<?>[] {interfaceType},
5455
new BigQueryJdbcContextProxy(target, connectionId, interfaceType));
5556
}
5657

58+
private static String extractConnectionId(Object target) {
59+
if (target == null) {
60+
return null;
61+
}
62+
if (target instanceof BigQueryConnection) {
63+
return ((BigQueryConnection) target).getConnectionId();
64+
}
65+
if (target instanceof BigQueryStatement) {
66+
return ((BigQueryStatement) target).connectionId;
67+
}
68+
if (target instanceof BigQueryDatabaseMetaData) {
69+
return ((BigQueryDatabaseMetaData) target).connection.getConnectionId();
70+
}
71+
if (target instanceof BigQueryResultSetMetadata) {
72+
java.sql.Statement stmt = ((BigQueryResultSetMetadata) target).getStatement();
73+
if (stmt != null) {
74+
if (Proxy.isProxyClass(stmt.getClass())) {
75+
InvocationHandler handler = Proxy.getInvocationHandler(stmt);
76+
if (handler instanceof BigQueryJdbcContextProxy) {
77+
return ((BigQueryJdbcContextProxy) handler).connectionId;
78+
}
79+
}
80+
if (stmt instanceof BigQueryStatement) {
81+
return ((BigQueryStatement) stmt).connectionId;
82+
}
83+
}
84+
}
85+
// Fallback to thread active context
86+
String activeId = BigQueryJdbcMdc.getConnectionId();
87+
if (activeId != null) {
88+
return activeId;
89+
}
90+
return null;
91+
}
92+
5793
@Override
5894
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
5995
// Handle standard Object methods explicitly
@@ -110,17 +146,17 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
110146

111147
// Automatically cascade proxy wrappers to child JDBC objects returned from calls
112148
if (result instanceof java.sql.CallableStatement) {
113-
return wrap(result, java.sql.CallableStatement.class, connectionId);
149+
return wrap(result, java.sql.CallableStatement.class);
114150
} else if (result instanceof java.sql.PreparedStatement) {
115-
return wrap(result, java.sql.PreparedStatement.class, connectionId);
151+
return wrap(result, java.sql.PreparedStatement.class);
116152
} else if (result instanceof java.sql.Statement) {
117-
return wrap(result, java.sql.Statement.class, connectionId);
153+
return wrap(result, java.sql.Statement.class);
118154
} else if (result instanceof java.sql.DatabaseMetaData) {
119-
return wrap(result, java.sql.DatabaseMetaData.class, connectionId);
155+
return wrap(result, java.sql.DatabaseMetaData.class);
120156
} else if (result instanceof java.sql.ParameterMetaData) {
121-
return wrap(result, java.sql.ParameterMetaData.class, connectionId);
157+
return wrap(result, java.sql.ParameterMetaData.class);
122158
} else if (result instanceof java.sql.ResultSetMetaData) {
123-
return wrap(result, java.sql.ResultSetMetaData.class, connectionId);
159+
return wrap(result, java.sql.ResultSetMetaData.class);
124160
}
125161

126162
return result;
@@ -129,11 +165,9 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
129165

130166
// Unified Context Logger: Captures and logs every exception exactly once with the Connection
131167
// context
132-
if (BigQueryJdbcRootLogger.isFileLoggingEnabled()) {
133-
try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(connectionId)) {
134-
String errMsg = cause.getMessage() != null ? cause.getMessage() : cause.toString();
135-
LOG.severe("Exception occurred during " + method.getName() + ": " + errMsg, cause);
136-
}
168+
try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(connectionId)) {
169+
String errMsg = cause.getMessage() != null ? cause.getMessage() : cause.toString();
170+
LOG.severe("Exception occurred during " + method.getName() + ": " + errMsg, cause);
137171
}
138172

139173
throw cause;

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcRootLogger.java‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -136,10 +136,6 @@ public static Logger getRootLogger() {
136136
return logger;
137137
}
138138

139-
public static boolean isFileLoggingEnabled() {
140-
return fileHandler != null;
141-
}
142-
143139
public static void setLevel(Level level, String logPath) throws IOException {
144140
if (level != Level.OFF) {
145141
setPath(logPath, level);

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetMetadata.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,10 @@ static BigQueryResultSetMetadata of(FieldList schemaFieldList, Statement stateme
4646
return new BigQueryResultSetMetadata(schemaFieldList, statement);
4747
}
4848

49+
Statement getStatement() {
50+
return this.statement;
51+
}
52+
4953
private Field getField(int sqlColumn) {
5054
return this.schemaFieldList.get(sqlColumn - 1);
5155
}

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

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery.jdbc;
18+
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertFalse;
21+
import static org.junit.jupiter.api.Assertions.assertNotNull;
22+
import static org.junit.jupiter.api.Assertions.assertNull;
23+
import static org.junit.jupiter.api.Assertions.assertThrows;
24+
import static org.junit.jupiter.api.Assertions.assertTrue;
25+
import static org.mockito.Mockito.mock;
26+
import static org.mockito.Mockito.when;
27+
28+
import com.google.cloud.bigquery.Field;
29+
import com.google.cloud.bigquery.FieldList;
30+
import com.google.cloud.bigquery.StandardSQLTypeName;
31+
import java.sql.Connection;
32+
import java.sql.DatabaseMetaData;
33+
import java.sql.ResultSetMetaData;
34+
import java.sql.SQLException;
35+
import java.sql.Statement;
36+
import org.junit.jupiter.api.Test;
37+
38+
public class BigQueryJdbcContextProxyTest {
39+
40+
@Test
41+
public void testExtractConnectionIdFromConnection() throws SQLException {
42+
BigQueryConnection mockConn = mock(BigQueryConnection.class);
43+
when(mockConn.getConnectionId()).thenReturn("conn-uuid-123");
44+
when(mockConn.getAutoCommit())
45+
.thenAnswer(
46+
invocation -> {
47+
assertEquals("conn-uuid-123", BigQueryJdbcMdc.getConnectionId());
48+
return true;
49+
});
50+
51+
Connection proxy = BigQueryJdbcContextProxy.wrap(mockConn, Connection.class);
52+
assertNotNull(proxy);
53+
54+
// Verify active thread-local context is set during method invocation
55+
assertTrue(proxy.getAutoCommit());
56+
// Verify context is cleanly cleared after method exit
57+
assertNull(BigQueryJdbcMdc.getConnectionId());
58+
}
59+
60+
@Test
61+
public void testExtractConnectionIdFromStatement() throws SQLException {
62+
BigQueryConnection mockConn = mock(BigQueryConnection.class);
63+
when(mockConn.getBigQuery()).thenReturn(mock(com.google.cloud.bigquery.BigQuery.class));
64+
65+
BigQueryStatement stmt = new BigQueryStatement(mockConn);
66+
stmt.connectionId = "conn-uuid-456";
67+
68+
Statement proxy = BigQueryJdbcContextProxy.wrap(stmt, Statement.class);
69+
assertNotNull(proxy);
70+
71+
// We can call any statement method (like getUpdateCount) and verify context routing
72+
// by asserting the thread-local value during invocation, but since it's a real object,
73+
// it runs cleanly. Let's assert that no thread-local leaks exist.
74+
assertEquals(-1, proxy.getUpdateCount());
75+
assertNull(BigQueryJdbcMdc.getConnectionId());
76+
}
77+
78+
@Test
79+
public void testExtractConnectionIdFromDatabaseMetaData() throws SQLException {
80+
BigQueryConnection mockConn = mock(BigQueryConnection.class);
81+
when(mockConn.getConnectionId()).thenReturn("conn-uuid-789");
82+
when(mockConn.getBigQuery()).thenReturn(mock(com.google.cloud.bigquery.BigQuery.class));
83+
84+
BigQueryDatabaseMetaData meta = new BigQueryDatabaseMetaData(mockConn);
85+
86+
DatabaseMetaData proxy = BigQueryJdbcContextProxy.wrap(meta, DatabaseMetaData.class);
87+
assertNotNull(proxy);
88+
89+
// Assert read-only capability does not leak thread context
90+
assertFalse(proxy.isReadOnly());
91+
assertNull(BigQueryJdbcMdc.getConnectionId());
92+
}
93+
94+
@Test
95+
public void testExtractConnectionIdFromResultSetMetaData() throws SQLException {
96+
BigQueryConnection mockConn = mock(BigQueryConnection.class);
97+
BigQueryStatement stmt = new BigQueryStatement(mockConn);
98+
stmt.connectionId = "conn-uuid-999";
99+
100+
FieldList fields = FieldList.of(Field.of("col", StandardSQLTypeName.STRING));
101+
BigQueryResultSetMetadata meta = BigQueryResultSetMetadata.of(fields, stmt);
102+
103+
ResultSetMetaData proxy = BigQueryJdbcContextProxy.wrap(meta, ResultSetMetaData.class);
104+
assertNotNull(proxy);
105+
106+
assertEquals(1, proxy.getColumnCount());
107+
assertNull(BigQueryJdbcMdc.getConnectionId());
108+
}
109+
110+
@Test
111+
public void testWrapWithNullContextAndExceptionThrown() throws SQLException {
112+
BigQueryStatement mockStmt = mock(BigQueryStatement.class);
113+
mockStmt.connectionId = null;
114+
when(mockStmt.executeQuery("SELECT *")).thenThrow(new SQLException("Database error"));
115+
116+
Statement proxy = BigQueryJdbcContextProxy.wrap(mockStmt, Statement.class);
117+
assertNotNull(proxy);
118+
119+
SQLException ex = assertThrows(SQLException.class, () -> proxy.executeQuery("SELECT *"));
120+
assertEquals("Database error", ex.getMessage());
121+
assertNull(BigQueryJdbcMdc.getConnectionId());
122+
}
123+
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,7 @@ private void validateConnection(String connection_uri) throws SQLException {
7171
Connection connection = DriverManager.getConnection(connection_uri);
7272
assertNotNull(connection);
7373
assertFalse(connection.isClosed());
74-
String query =
75-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 850";
74+
String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850";
7675
Statement statement = connection.createStatement();
7776
ResultSet jsonResultSet = statement.executeQuery(query);
7877
int totalRows = 0;

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

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -149,17 +149,15 @@ public void testValidAllDataTypesSerializationFromSelectQueryArrowDataset() thro
149149

150150
@Test
151151
public void testFastQueryPathSmall() throws SQLException {
152-
String query =
153-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 850";
152+
String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850";
154153
ResultSet jsonResultSet = bigQueryStatement.executeQuery(query);
155154
assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet"));
156155
assertEquals(850, resultSetRowCount(jsonResultSet));
157156
}
158157

159158
@Test
160159
public void testFastQueryPathEmpty() throws SQLException {
161-
String query =
162-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 0";
160+
String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0";
163161
Connection connection =
164162
DriverManager.getConnection(String.format(connectionUrl, DEFAULT_CATALOG));
165163
Statement bigQueryStatement = connection.createStatement();
@@ -285,13 +283,12 @@ public void testStatelessQueryPathSmall() throws SQLException {
285283

286284
Statement statement = connectionUseStateless.createStatement();
287285

288-
String query =
289-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 850";
286+
String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850";
290287
ResultSet jsonResultSet = statement.executeQuery(query);
291288
Assert.assertEquals(850, resultSetRowCount(jsonResultSet));
292289

293290
String queryEmpty =
294-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 0";
291+
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0";
295292
ResultSet jsonResultSetEmpty = statement.executeQuery(queryEmpty);
296293
Assert.assertEquals(0, resultSetRowCount(jsonResultSetEmpty));
297294
connectionUseStateless.close();
@@ -1556,7 +1553,7 @@ public void testDataSourceOAuthPvtKeyPath() throws SQLException, IOException {
15561553
@Test
15571554
public void testPreparedStatementSmallSelect() throws SQLException {
15581555
String query =
1559-
"SELECT * FROM `bigquery-public-data.samples.shakespeare` where corpus=?" + " LIMIT 1000";
1556+
"SELECT * FROM `bigquery-public-data.samples.shakespeare` where corpus=? LIMIT 1000";
15601557
PreparedStatement preparedStatement = bigQueryConnection.prepareStatement(query);
15611558
preparedStatement.setString(1, "hamlet");
15621559

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -321,8 +321,7 @@ private void assertConnectionPoolingResults(String connectionURL, Long connectio
321321
assertFalse(connection.isClosed());
322322

323323
// Execute query with physical connection
324-
String query =
325-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 850";
324+
String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850";
326325
Statement statement = connection.createStatement();
327326
ResultSet jsonResultSet = statement.executeQuery(query);
328327
assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet"));

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1563,14 +1563,13 @@ public void testStatelessQueryPathSmall() throws SQLException {
15631563

15641564
Statement statement = bigQueryConnectionUseStateless.createStatement();
15651565

1566-
String query =
1567-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 850";
1566+
String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850";
15681567
ResultSet jsonResultSet = statement.executeQuery(query);
15691568
assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet"));
15701569
assertEquals(850, resultSetRowCount(jsonResultSet));
15711570

15721571
String queryEmpty =
1573-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 0";
1572+
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0";
15741573
ResultSet jsonResultSetEmpty = statement.executeQuery(queryEmpty);
15751574
assertTrue(jsonResultSetEmpty.getClass().getName().contains("BigQueryJsonResultSet"));
15761575
assertEquals(0, resultSetRowCount(jsonResultSetEmpty));
@@ -1634,8 +1633,7 @@ public void testNonEnabledUseLegacySQLThrowsSyntaxError() throws SQLException {
16341633

16351634
@Test
16361635
public void testFastQueryPathEmpty() throws SQLException {
1637-
String query =
1638-
"SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT" + " 0";
1636+
String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0";
16391637
ResultSet jsonResultSet = bigQueryStatement.executeQuery(query);
16401638
assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet"));
16411639
assertEquals(0, resultSetRowCount(jsonResultSet));

0 commit comments

Comments
 (0)