diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryConversionException.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryConversionException.java index 2be7b40ddc6a..ea78a4f11330 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryConversionException.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryConversionException.java @@ -16,18 +16,14 @@ package com.google.cloud.bigquery.exception; -import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; import java.sql.SQLException; /** * Exception for errors that occur when the driver cannot convert a value from one type to another. */ public class BigQueryConversionException extends SQLException { - private static final BigQueryJdbcCustomLogger LOG = - new BigQueryJdbcCustomLogger(BigQueryConversionException.class.getName()); public BigQueryConversionException(String message, Throwable cause) { super(BigQueryJdbcExceptionUtils.formatMessage(message, cause), cause); - LOG.severe(this.getMessage(), this); } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionException.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionException.java index 3241b7a1993e..defbbfcca784 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionException.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionException.java @@ -17,7 +17,6 @@ package com.google.cloud.bigquery.exception; import com.google.api.core.InternalApi; -import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; /** * Thrown to indicate that the coercion was attempted but couldn't be performed successfully because @@ -25,8 +24,6 @@ */ @InternalApi public class BigQueryJdbcCoercionException extends RuntimeException { - private static final BigQueryJdbcCustomLogger LOG = - new BigQueryJdbcCustomLogger(BigQueryJdbcCoercionException.class.getName()); /** * Construct a new exception with the specified cause. @@ -35,6 +32,5 @@ public class BigQueryJdbcCoercionException extends RuntimeException { */ public BigQueryJdbcCoercionException(Exception cause) { super(BigQueryJdbcExceptionUtils.formatMessage("Coercion error", cause), cause); - LOG.severe(this.getMessage(), this); } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionNotFoundException.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionNotFoundException.java index 8a12da311c21..b4eafb2ee583 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionNotFoundException.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcCoercionNotFoundException.java @@ -17,7 +17,6 @@ package com.google.cloud.bigquery.exception; import com.google.api.core.InternalApi; -import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; /** * Thrown to indicate that the current TypeCoercer can not perform the coercion as the Coercion @@ -25,8 +24,6 @@ */ @InternalApi public class BigQueryJdbcCoercionNotFoundException extends RuntimeException { - private static final BigQueryJdbcCustomLogger LOG = - new BigQueryJdbcCustomLogger(BigQueryJdbcCoercionNotFoundException.class.getName()); /** * Construct a new exception. @@ -39,6 +36,5 @@ public BigQueryJdbcCoercionNotFoundException(Class source, Class target) { String.format( "Coercion not found for [%s -> %s] conversion", source.getCanonicalName(), target.getCanonicalName())); - LOG.severe(this.getMessage(), this); } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcException.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcException.java index 4cdc37773e59..a6262b167db4 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcException.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcException.java @@ -17,12 +17,9 @@ package com.google.cloud.bigquery.exception; import com.google.cloud.bigquery.BigQueryException; -import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; import java.sql.SQLException; public class BigQueryJdbcException extends SQLException { - private static final BigQueryJdbcCustomLogger LOG = - new BigQueryJdbcCustomLogger(BigQueryJdbcException.class.getName()); private BigQueryException bigQueryException = null; /** @@ -32,7 +29,6 @@ public class BigQueryJdbcException extends SQLException { */ public BigQueryJdbcException(String message) { super(message); - LOG.severe(message, this); } /** @@ -42,7 +38,6 @@ public BigQueryJdbcException(String message) { */ public BigQueryJdbcException(InterruptedException ex) { super(ex); - LOG.severe(ex.getMessage(), this); } /** @@ -54,7 +49,6 @@ public BigQueryJdbcException(InterruptedException ex) { public BigQueryJdbcException(String message, BigQueryException ex) { super(BigQueryJdbcExceptionUtils.formatMessage(message, ex), ex); this.bigQueryException = ex; - LOG.severe(this.getMessage(), this); } /** @@ -65,7 +59,6 @@ public BigQueryJdbcException(String message, BigQueryException ex) { */ public BigQueryJdbcException(String message, Throwable cause) { super(BigQueryJdbcExceptionUtils.formatMessage(message, cause), cause); - LOG.severe(this.getMessage(), this); } /** @@ -76,7 +69,6 @@ public BigQueryJdbcException(String message, Throwable cause) { */ public BigQueryJdbcException(Throwable cause) { super(cause); - LOG.severe(cause.getMessage(), this); } public BigQueryException getBigQueryException() { diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcRuntimeException.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcRuntimeException.java index b5caad826986..2fe21ab2a7c4 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcRuntimeException.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcRuntimeException.java @@ -16,13 +16,8 @@ package com.google.cloud.bigquery.exception; -import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; - public class BigQueryJdbcRuntimeException extends RuntimeException { - private static final BigQueryJdbcCustomLogger LOG = - new BigQueryJdbcCustomLogger(BigQueryJdbcRuntimeException.class.getName()); - /** * Constructs a new BigQueryJdbcRuntimeException with the given message. * @@ -30,7 +25,6 @@ public class BigQueryJdbcRuntimeException extends RuntimeException { */ public BigQueryJdbcRuntimeException(String message) { super(message); - LOG.severe(message, this); } /** @@ -40,7 +34,6 @@ public BigQueryJdbcRuntimeException(String message) { */ public BigQueryJdbcRuntimeException(Throwable ex) { super(ex); - LOG.severe(ex.getMessage(), this); } /** @@ -50,12 +43,10 @@ public BigQueryJdbcRuntimeException(Throwable ex) { * @param ex Throwable to be thrown. */ public BigQueryJdbcRuntimeException(String message, InterruptedException ex) { - super(BigQueryJdbcExceptionUtils.formatMessage(message, ex), ex); - LOG.severe(this.getMessage(), this); + super(message, ex); } public BigQueryJdbcRuntimeException(String message, Throwable ex) { super(BigQueryJdbcExceptionUtils.formatMessage(message, ex), ex); - LOG.severe(this.getMessage(), this); } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlFeatureNotSupportedException.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlFeatureNotSupportedException.java index 0c4493b29dc3..8039dcd53495 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlFeatureNotSupportedException.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlFeatureNotSupportedException.java @@ -17,12 +17,9 @@ package com.google.cloud.bigquery.exception; import com.google.cloud.bigquery.BigQueryException; -import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; import java.sql.SQLFeatureNotSupportedException; public class BigQueryJdbcSqlFeatureNotSupportedException extends SQLFeatureNotSupportedException { - private static final BigQueryJdbcCustomLogger LOG = - new BigQueryJdbcCustomLogger(BigQueryJdbcSqlFeatureNotSupportedException.class.getName()); /** * Constructs a new BigQueryJdbcSqlFeatureNotSupportedException with the given message. @@ -31,7 +28,6 @@ public class BigQueryJdbcSqlFeatureNotSupportedException extends SQLFeatureNotSu */ public BigQueryJdbcSqlFeatureNotSupportedException(String message) { super(message); - LOG.severe(message, this); } /** @@ -41,6 +37,5 @@ public BigQueryJdbcSqlFeatureNotSupportedException(String message) { */ public BigQueryJdbcSqlFeatureNotSupportedException(BigQueryException ex) { super(ex); - LOG.severe(ex.getMessage(), this); } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlSyntaxErrorException.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlSyntaxErrorException.java index f3b6f525add9..6d51e656542a 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlSyntaxErrorException.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/exception/BigQueryJdbcSqlSyntaxErrorException.java @@ -17,7 +17,6 @@ package com.google.cloud.bigquery.exception; import com.google.cloud.bigquery.BigQueryException; -import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; import java.sql.SQLSyntaxErrorException; /** @@ -26,8 +25,6 @@ * rules. */ public class BigQueryJdbcSqlSyntaxErrorException extends SQLSyntaxErrorException { - private static final BigQueryJdbcCustomLogger LOG = - new BigQueryJdbcCustomLogger(BigQueryJdbcSqlSyntaxErrorException.class.getName()); /** * Constructs a new BigQueryJdbcSqlSyntaxErrorException from BigQueryException @@ -36,11 +33,9 @@ public class BigQueryJdbcSqlSyntaxErrorException extends SQLSyntaxErrorException */ public BigQueryJdbcSqlSyntaxErrorException(BigQueryException ex) { super(ex.getMessage(), "Incorrect SQL syntax."); - LOG.severe(ex.getMessage(), this); } public BigQueryJdbcSqlSyntaxErrorException(String message, BigQueryException ex) { super(BigQueryJdbcExceptionUtils.formatMessage(message, ex), ex); - LOG.severe(this.getMessage(), this); } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index af041dc2a649..0736cd309f57 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java @@ -215,11 +215,8 @@ public boolean next() throws SQLException { checkClosed(); if (this.isNested) { if (this.currentNestedBatch == null || this.currentNestedBatch.getNestedRecords() == null) { - IllegalStateException ex = - new IllegalStateException( - "currentNestedBatch/JsonStringArrayList can not be null working with the nested record"); - LOG.severe(ex.getMessage(), ex); - throw ex; + throw new IllegalStateException( + "currentNestedBatch/JsonStringArrayList can not be null working with the nested record"); } if (this.nestedRowIndex < (this.toIndexExclusive - 1)) { /* Check if there's a next record in the array which can be read */ diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java index d63b72bb6c93..0e72a5fccf71 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java +++ b/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 final BigQueryJdbcCustomLogger LOG = + new BigQueryJdbcCustomLogger(this.getClass().getName()); private BigQuery bigQuery; private JobId jobId; private String queryId; @@ -107,6 +108,25 @@ public void close() { } } + protected SQLException logAndCreateException(SQLException ex) { + if (this.statement == null) { + return ex; + } + try (BigQueryJdbcMdc.MdcCloseable mdc = + BigQueryJdbcMdc.registerInstance(this.statement.connectionId)) { + LOG.severe(ex.getMessage(), ex); + } + return ex; + } + + protected SQLException logAndCreateException(String message) { + return logAndCreateException(new SQLException(message)); + } + + protected SQLException logAndCreateException(String message, Throwable cause) { + return logAndCreateException(new SQLException(message, cause)); + } + protected SQLException createCoercionException( int columnIndex, Class targetClass, Exception cause) throws SQLException { checkClosed(); @@ -115,29 +135,31 @@ protected SQLException createCoercionException( if (isNested) { if (columnIndex == 1) { - return new BigQueryConversionException( - String.format("Cannot convert index column to type %s.", targetClass.getSimpleName()), - cause); + return logAndCreateException( + new BigQueryConversionException( + String.format( + "Cannot convert index column to type %s.", targetClass.getSimpleName()), + cause)); } else if (columnIndex == 2) { Field arrayField = this.schema.getFields().get(0); type = arrayField.getType().getStandardType(); typeName = type.name(); } else { - SQLException ex = + throw logAndCreateException( new SQLException( - "For a nested ResultSet from an Array, columnIndex must be 1 or 2.", cause); - LOG.severe(ex.getMessage(), ex); - throw ex; + "For a nested ResultSet from an Array, columnIndex must be 1 or 2.", cause)); } } else { Field field = this.schemaFieldList.get(columnIndex - 1); type = field.getType().getStandardType(); typeName = type.name(); } - return new BigQueryConversionException( - String.format( - "Cannot convert value of type %s to type %s.", typeName, targetClass.getSimpleName()), - cause); + return logAndCreateException( + new BigQueryConversionException( + String.format( + "Cannot convert value of type %s to type %s.", + typeName, targetClass.getSimpleName()), + cause)); } private StandardSQLTypeName getStandardSQLTypeName(int columnIndex) throws SQLException { @@ -147,25 +169,19 @@ private StandardSQLTypeName getStandardSQLTypeName(int columnIndex) throws SQLEx return StandardSQLTypeName.INT64; } else if (columnIndex == 2) { if (this.schema == null || this.schema.getFields().isEmpty()) { - SQLException ex = new SQLException("Schema not available for nested result set."); - LOG.severe("Schema not available for nested result set.", ex); - throw ex; + throw logAndCreateException("Schema not available for nested result set."); } Field arrayField = this.schema.getFields().get(0); return arrayField.getType().getStandardType(); } else { - SQLException ex = - new SQLException("For a nested ResultSet from an Array, columnIndex must be 1 or 2."); - LOG.severe("For a nested ResultSet from an Array, columnIndex must be 1 or 2.", ex); - throw ex; + throw logAndCreateException( + "For a nested ResultSet from an Array, columnIndex must be 1 or 2."); } } else { if (this.schemaFieldList == null || columnIndex > this.schemaFieldList.size() || columnIndex < 1) { - SQLException ex = new SQLException("Invalid column index: " + columnIndex); - LOG.severe("Invalid column index: " + columnIndex, ex); - throw ex; + throw logAndCreateException("Invalid column index: " + columnIndex); } Field field = this.schemaFieldList.get(columnIndex - 1); return field.getType().getStandardType(); @@ -185,11 +201,16 @@ public boolean wasNull() throws SQLException { @Override public ResultSetMetaData getMetaData() throws SQLException { checkClosed(); - if (this.isNested) { - return BigQueryResultSetMetadata.of(this.schemaFieldList, this.statement); - } else { - return BigQueryResultSetMetadata.of(this.schema.getFields(), this.statement); + String connectionId = this.statement != null ? this.statement.connectionId : null; + ResultSetMetaData metaData; + try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(connectionId)) { + if (this.isNested) { + metaData = BigQueryResultSetMetadata.of(this.schemaFieldList, this.statement); + } else { + metaData = BigQueryResultSetMetadata.of(this.schema.getFields(), this.statement); + } } + return BigQueryJdbcContextProxy.wrap(metaData, ResultSetMetaData.class, connectionId); } @Override @@ -227,9 +248,7 @@ protected int getColumnIndex(String columnLabel) throws SQLException { LOG.finest("++enter++"); checkClosed(); if (columnLabel == null) { - SQLException ex = new SQLException("Column label cannot be null"); - LOG.severe(ex.getMessage(), ex); - throw ex; + throw logAndCreateException("Column label cannot be null"); } // use schema to get the column index, add 1 for SQL index return this.schemaFieldList.getIndex(columnLabel) + 1; diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 4f70de993177..3c64e3a073f8 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -150,8 +150,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { BigQueryConnection(String url, DataSource ds) throws IOException { this.connectionId = UUID.randomUUID().toString(); - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { + try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(this.connectionId)) { LOG.finest("++enter++"); this.connectionUrl = url; @@ -351,15 +350,12 @@ String getConnectionId() { */ @Override public Statement createStatement() throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - BigQueryStatement currentStatement = new BigQueryStatement(this); - LOG.fine("Statement %s created.", currentStatement); - addOpenStatements(currentStatement); - return currentStatement; - } + LOG.finest("++enter++"); + checkClosed(); + BigQueryStatement currentStatement = new BigQueryStatement(this); + LOG.fine("Statement %s created.", currentStatement); + addOpenStatements(currentStatement); + return currentStatement; } /** @@ -378,17 +374,13 @@ public Statement createStatement() throws SQLException { @Override public Statement createStatement(int resultSetType, int resultSetConcurrency) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - if (resultSetType != ResultSet.TYPE_FORWARD_ONLY - || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY) { - throw new BigQueryJdbcSqlFeatureNotSupportedException( - "Unsupported createStatement feature."); - } - return createStatement(); + LOG.finest("++enter++"); + checkClosed(); + if (resultSetType != ResultSet.TYPE_FORWARD_ONLY + || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY) { + throw new BigQueryJdbcSqlFeatureNotSupportedException("Unsupported createStatement feature."); } + return createStatement(); } /** @@ -407,43 +399,33 @@ public Statement createStatement(int resultSetType, int resultSetConcurrency) @Override public Statement createStatement( int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - if (resultSetType != ResultSet.TYPE_FORWARD_ONLY - || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY - || resultSetHoldability != ResultSet.CLOSE_CURSORS_AT_COMMIT) { - throw new BigQueryJdbcSqlFeatureNotSupportedException( - "Unsupported createStatement feature"); - } - return createStatement(); + LOG.finest("++enter++"); + checkClosed(); + if (resultSetType != ResultSet.TYPE_FORWARD_ONLY + || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY + || resultSetHoldability != ResultSet.CLOSE_CURSORS_AT_COMMIT) { + throw new BigQueryJdbcSqlFeatureNotSupportedException("Unsupported createStatement feature"); } + return createStatement(); } @Override public PreparedStatement prepareStatement(String sql) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - PreparedStatement currentStatement = new BigQueryPreparedStatement(this, sql); - LOG.fine("Prepared Statement %s created.", currentStatement); - addOpenStatements(currentStatement); - return currentStatement; - } + LOG.finest("++enter++"); + checkClosed(); + PreparedStatement currentStatement = new BigQueryPreparedStatement(this, sql); + LOG.fine("Prepared Statement %s created.", currentStatement); + addOpenStatements(currentStatement); + return currentStatement; } @Override public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) { - throw new BigQueryJdbcSqlFeatureNotSupportedException("autoGeneratedKeys is not supported"); - } - return prepareStatement(sql); + LOG.finest("++enter++"); + if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) { + throw new BigQueryJdbcSqlFeatureNotSupportedException("autoGeneratedKeys is not supported"); } + return prepareStatement(sql); } @Override @@ -455,32 +437,24 @@ public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throw public PreparedStatement prepareStatement( String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - if (resultSetType != ResultSet.TYPE_FORWARD_ONLY - || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY - || resultSetHoldability != ResultSet.CLOSE_CURSORS_AT_COMMIT) { - throw new BigQueryJdbcSqlFeatureNotSupportedException( - "Unsupported prepareStatement feature"); - } - return prepareStatement(sql); + LOG.finest("++enter++"); + if (resultSetType != ResultSet.TYPE_FORWARD_ONLY + || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY + || resultSetHoldability != ResultSet.CLOSE_CURSORS_AT_COMMIT) { + throw new BigQueryJdbcSqlFeatureNotSupportedException("Unsupported prepareStatement feature"); } + return prepareStatement(sql); } @Override public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - if (resultSetType != ResultSet.TYPE_FORWARD_ONLY - || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY) { - throw new BigQueryJdbcSqlFeatureNotSupportedException( - "Unsupported prepareStatement feature"); - } - return prepareStatement(sql); + LOG.finest("++enter++"); + if (resultSetType != ResultSet.TYPE_FORWARD_ONLY + || resultSetConcurrency != ResultSet.CONCUR_READ_ONLY) { + throw new BigQueryJdbcSqlFeatureNotSupportedException("Unsupported prepareStatement feature"); } + return prepareStatement(sql); } public DatasetId getDefaultDataset() { @@ -694,36 +668,30 @@ Long getListenerPoolSize() { @Override public boolean isValid(int timeout) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - if (timeout < 0) { - throw new BigQueryJdbcException("timeout must be >= 0"); - } - if (!isClosed()) { - try (Statement statement = createStatement(); - ResultSet rs = statement.executeQuery("SELECT 1")) { - LOG.finest("Running validation query"); - if (rs.next()) { - if (rs.getInt(1) == 1) { - return true; - } + LOG.finest("++enter++"); + if (timeout < 0) { + throw new BigQueryJdbcException("timeout must be >= 0"); + } + if (!isClosed()) { + try (Statement statement = createStatement(); + ResultSet rs = statement.executeQuery("SELECT 1")) { + LOG.finest("Running validation query"); + if (rs.next()) { + if (rs.getInt(1) == 1) { + return true; } - } catch (SQLException ex) { - // Ignore } + } catch (SQLException ex) { + // Ignore } - return false; } + return false; } @Override public void abort(Executor executor) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - close(); - } + LOG.finest("++enter++"); + close(); } @Override @@ -763,75 +731,63 @@ public void clearWarnings() { @Override public boolean getAutoCommit() { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - return this.autoCommit; - } + LOG.finest("++enter++"); + checkClosed(); + return this.autoCommit; } @Override public void setAutoCommit(boolean autoCommit) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - checkIfEnabledSession("setAutoCommit"); - if (this.autoCommit == autoCommit) { - return; - } + LOG.finest("++enter++"); + checkClosed(); + checkIfEnabledSession("setAutoCommit"); + if (this.autoCommit == autoCommit) { + return; + } - if (isTransactionStarted()) { - commitTransaction(); - } + if (isTransactionStarted()) { + commitTransaction(); + } - this.autoCommit = autoCommit; - if (!this.autoCommit) { - beginTransaction(); - } + this.autoCommit = autoCommit; + if (!this.autoCommit) { + beginTransaction(); } } @Override public void commit() { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - checkIfEnabledSession("commit"); - if (!isTransactionStarted()) { - IllegalStateException ex = - new IllegalStateException( - "Cannot commit without an active transaction. Please set setAutoCommit to false to start" - + " a transaction."); - LOG.severe(ex.getMessage(), ex); - throw ex; - } - commitTransaction(); - if (!getAutoCommit()) { - beginTransaction(); - } + LOG.finest("++enter++"); + checkClosed(); + checkIfEnabledSession("commit"); + if (!isTransactionStarted()) { + IllegalStateException ex = + new IllegalStateException( + "Cannot commit without an active transaction. Please set setAutoCommit to false to start" + + " a transaction."); + LOG.severe(ex.getMessage(), ex); + throw ex; + } + commitTransaction(); + if (!getAutoCommit()) { + beginTransaction(); } } @Override public void rollback() throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - checkIfEnabledSession("rollback"); - if (!isTransactionStarted()) { - IllegalStateException ex = - new IllegalStateException( - "Cannot rollback without an active transaction. Please set setAutoCommit to false to" - + " start a transaction."); - LOG.severe(ex.getMessage(), ex); - throw ex; - } - rollbackImpl(); + LOG.finest("++enter++"); + checkClosed(); + checkIfEnabledSession("rollback"); + if (!isTransactionStarted()) { + IllegalStateException ex = + new IllegalStateException( + "Cannot rollback without an active transaction. Please set setAutoCommit to false to" + + " start a transaction."); + LOG.severe(ex.getMessage(), ex); + throw ex; } + rollbackImpl(); } private void rollbackImpl() throws SQLException { @@ -853,14 +809,11 @@ private void rollbackImpl() throws SQLException { @Override public DatabaseMetaData getMetaData() throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - if (databaseMetaData == null) { - databaseMetaData = new BigQueryDatabaseMetaData(this); - } - return databaseMetaData; + LOG.finest("++enter++"); + if (databaseMetaData == null) { + databaseMetaData = new BigQueryDatabaseMetaData(this); } + return databaseMetaData; } @Override @@ -870,15 +823,12 @@ public int getTransactionIsolation() { @Override public void setTransactionIsolation(int level) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - if (level != Connection.TRANSACTION_SERIALIZABLE) { - throw new BigQueryJdbcSqlFeatureNotSupportedException( - "Unsupported transaction isolation level"); - } - this.transactionIsolation = level; + LOG.finest("++enter++"); + if (level != Connection.TRANSACTION_SERIALIZABLE) { + throw new BigQueryJdbcSqlFeatureNotSupportedException( + "Unsupported transaction isolation level"); } + this.transactionIsolation = level; } @Override @@ -888,14 +838,11 @@ public int getHoldability() { @Override public void setHoldability(int holdability) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - if (holdability != ResultSet.CLOSE_CURSORS_AT_COMMIT) { - throw new BigQueryJdbcSqlFeatureNotSupportedException( - "CLOSE_CURSORS_AT_COMMIT not supported"); - } - this.holdability = holdability; + if (holdability != ResultSet.CLOSE_CURSORS_AT_COMMIT) { + throw new BigQueryJdbcSqlFeatureNotSupportedException( + "CLOSE_CURSORS_AT_COMMIT not supported"); } + this.holdability = holdability; } /** @@ -911,12 +858,9 @@ public void close() throws SQLException { return; } - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - LOG.finest("++enter++"); - LOG.fine("Closing Connection " + this); - closeImpl(); - } + LOG.finest("++enter++"); + LOG.fine("Closing Connection " + this); + closeImpl(); } private void closeImpl() throws SQLException { @@ -942,7 +886,7 @@ private void closeImpl() throws SQLException { } catch (InterruptedException e) { throw new BigQueryJdbcRuntimeException("Interrupted during close", e); } finally { - BigQueryJdbcMdc.removeInstance(this); + BigQueryJdbcMdc.clear(); BigQueryJdbcRootLogger.closeConnectionHandler(this.connectionId); } this.isClosed = true; @@ -1157,14 +1101,11 @@ private void commitTransaction() { @Override public CallableStatement prepareCall(String sql) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this, this.connectionId)) { - checkClosed(); - CallableStatement currentStatement = new BigQueryCallableStatement(this, sql); - LOG.fine("Callable Statement %s created.", currentStatement); - addOpenStatements(currentStatement); - return currentStatement; - } + checkClosed(); + CallableStatement currentStatement = new BigQueryCallableStatement(this, sql); + LOG.fine("Callable Statement %s created.", currentStatement); + addOpenStatements(currentStatement); + return currentStatement; } @Override diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java index e62d93fca0ed..074c196eecaf 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java @@ -173,7 +173,7 @@ public Connection connect(String url, Properties info) throws SQLException { logLevel, logPath, this.toString()); - return connection; + return BigQueryJdbcContextProxy.wrap(connection, Connection.class); } else { return null; } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxy.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxy.java new file mode 100644 index 000000000000..c686e8070a76 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxy.java @@ -0,0 +1,166 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; + +/** + * Dynamic InvocationHandler that transparently wraps JDBC operations. Sets the connection context + * on the executing thread for connection routing, and cleans it up upon method exit to prevent + * memory leaks and context bleeding. Acts as the unified exception logger for Statement and + * Connection methods. + */ +class BigQueryJdbcContextProxy implements InvocationHandler { + private static final BigQueryJdbcCustomLogger LOG = + new BigQueryJdbcCustomLogger(BigQueryJdbcContextProxy.class.getName()); + + private final Object target; + private final String connectionId; + private final Class interfaceType; + + private BigQueryJdbcContextProxy(Object target, String connectionId, Class interfaceType) { + this.target = target; + this.connectionId = connectionId; + this.interfaceType = interfaceType; + } + + /** Wraps a target Connection JDBC object, auto-extracting its connection ID. */ + @SuppressWarnings("unchecked") + static T wrap(Object target, Class interfaceType) { + if (target == null) { + return null; + } + String connectionId = extractConnectionId(target); + return (T) + Proxy.newProxyInstance( + interfaceType.getClassLoader(), + new Class[] {interfaceType}, + new BigQueryJdbcContextProxy(target, connectionId, interfaceType)); + } + + /** Wraps a target child JDBC object, propagating the connection ID parameter directly. */ + @SuppressWarnings("unchecked") + static T wrap(Object target, Class 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(); + } + return null; + } + + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + // Handle standard Object methods explicitly + if (method.getDeclaringClass() == Object.class) { + String methodName = method.getName(); + if (methodName.equals("equals")) { + Object other = args[0]; + if (other == null) { + return false; + } + if (Proxy.isProxyClass(other.getClass())) { + InvocationHandler handler = Proxy.getInvocationHandler(other); + if (handler instanceof BigQueryJdbcContextProxy) { + return target.equals(((BigQueryJdbcContextProxy) handler).target); + } + } + return target.equals(other); + } + if (methodName.equals("hashCode")) { + return target.hashCode(); + } + if (methodName.equals("toString")) { + return target.toString(); + } + } + + // Support standard JDBC Wrapper unwrap operations + if (method.getName().equals("unwrap") && args != null && args.length == 1) { + Class iface = (Class) args[0]; + if (iface.isInstance(target)) { + return target; + } + try { + return method.invoke(target, args); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + } + if (method.getName().equals("isWrapperFor") && args != null && args.length == 1) { + Class iface = (Class) args[0]; + if (iface.isInstance(target)) { + return true; + } + try { + return method.invoke(target, args); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + } + + // Wrap execution in the context of the active connection for all non-bypassed methods + try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(connectionId)) { + Object result = method.invoke(target, args); + + // Symmetrical Cascade: Dynamic ResultSet concrete classes are deliberately unproxied here. + // Bypassing proxies on high-frequency ResultSet iterations avoids dynamic invocation + // and argument array allocations, allowing the JIT compiler to natively inline getters. + if (result instanceof java.sql.CallableStatement) { + return wrap(result, java.sql.CallableStatement.class, connectionId); + } else if (result instanceof java.sql.PreparedStatement) { + return wrap(result, java.sql.PreparedStatement.class, connectionId); + } else if (result instanceof java.sql.Statement) { + return wrap(result, java.sql.Statement.class, connectionId); + } else if (result instanceof java.sql.DatabaseMetaData) { + return wrap(result, java.sql.DatabaseMetaData.class, connectionId); + } else if (result instanceof java.sql.ParameterMetaData) { + return wrap(result, java.sql.ParameterMetaData.class, connectionId); + } else if (result instanceof java.sql.ResultSetMetaData) { + return wrap(result, java.sql.ResultSetMetaData.class, connectionId); + } + + return result; + } catch (InvocationTargetException e) { + Throwable cause = e.getCause(); + + // Unified Context Logger: Captures and logs every exception exactly once with the Connection + // context + try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(connectionId)) { + String errMsg = cause.getMessage() != null ? cause.getMessage() : cause.toString(); + LOG.severe("Exception occurred during " + method.getName() + ": " + errMsg, cause); + } + + throw cause; + } + } +} diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLogger.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLogger.java index 6932f9b1a2a8..5195539439f2 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLogger.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLogger.java @@ -67,6 +67,21 @@ private void logWithCaller(Level level, Throwable thrown, Supplier msgSu } } + @Override + public void finest(String msg) { + logWithCaller(Level.FINEST, () -> msg); + } + + @Override + public void finer(String msg) { + logWithCaller(Level.FINER, () -> msg); + } + + @Override + public void fine(String msg) { + logWithCaller(Level.FINE, () -> msg); + } + void finest(String format, Object... args) { logWithCaller(Level.FINEST, () -> String.format(format, args)); } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java index 12ea04c30628..36a315b31e93 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java @@ -16,38 +16,15 @@ package com.google.cloud.bigquery.jdbc; -import java.util.UUID; -import java.util.concurrent.ConcurrentHashMap; - -/** - * Lightweight MDC implementation for the BigQuery JDBC driver using InheritableThreadLocal. - * Allocates a dedicated, independent InheritableThreadLocal object per concrete BigQueryConnection - * instance. - */ +/** Lightweight MDC implementation for the BigQuery JDBC driver using InheritableThreadLocal. */ class BigQueryJdbcMdc { - private static final ConcurrentHashMap> - instanceLocals = new ConcurrentHashMap<>(); - private static final ConcurrentHashMap instanceIds = - new ConcurrentHashMap<>(); - - /** Allocates an exclusive InheritableThreadLocal and registers the connection mapping. */ private static final InheritableThreadLocal currentConnectionId = new InheritableThreadLocal<>(); - static MdcCloseable registerInstance(BigQueryConnection connection, String id) { - if (connection != null) { - String cleanId = - instanceIds.computeIfAbsent( - connection, - k -> { - String baseId = (id != null && !id.isEmpty()) ? id : UUID.randomUUID().toString(); - return baseId; - }); - - currentConnectionId.set(cleanId); - InheritableThreadLocal threadLocal = - instanceLocals.computeIfAbsent(connection, k -> new InheritableThreadLocal<>()); - threadLocal.set(cleanId); + /** Sets the current connection context on the executing thread. */ + static MdcCloseable registerInstance(String connectionId) { + if (connectionId != null) { + currentConnectionId.set(connectionId); } return () -> clear(); } @@ -59,22 +36,8 @@ static String getConnectionId() { return currentConnectionId.get(); } - /** Clears the connection ID context from all active connection contexts on the current thread. */ - static void removeInstance(BigQueryConnection connection) { - if (connection != null) { - InheritableThreadLocal local = instanceLocals.remove(connection); - if (local != null) { - local.remove(); - } - instanceIds.remove(connection); - } - } - static void clear() { currentConnectionId.remove(); - for (InheritableThreadLocal local : instanceLocals.values()) { - local.remove(); - } } /** diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java index c59061b25467..17198a753393 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonResultSet.java @@ -137,11 +137,8 @@ public boolean next() throws SQLException { // We are working with the nested record, the cursor would have been // populated. if (this.cursor == null || this.cursor.getArrayFieldValueList() == null) { - IllegalStateException ex = - new IllegalStateException( - "Cursor/ArrayFieldValueList can not be null working with the nested record"); - LOG.severe(ex.getMessage(), ex); - throw ex; + throw new IllegalStateException( + "Cursor/ArrayFieldValueList can not be null working with the nested record"); } // Check if there's a next record in the array which can be read if (this.nestedRowIndex < (this.toIndexExclusive - 1)) { diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetMetadata.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetMetadata.java index d18c689333a4..eb088567a937 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetMetadata.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryResultSetMetadata.java @@ -46,6 +46,10 @@ static BigQueryResultSetMetadata of(FieldList schemaFieldList, Statement stateme return new BigQueryResultSetMetadata(schemaFieldList, statement); } + Statement getStatement() { + return this.statement; + } + private Field getField(int sqlColumn) { return this.schemaFieldList.get(sqlColumn - 1); } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 2c04747f8863..e2dab7b31678 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -235,12 +235,9 @@ private BigQuerySettings generateBigQuerySettings() { */ @Override public ResultSet executeQuery(String sql) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - return executeQueryImpl(sql); - } + LOG.finest("++enter++"); + checkClosed(); + return executeQueryImpl(sql); } private ResultSet executeQueryImpl(String sql) throws SQLException { @@ -262,12 +259,9 @@ private ResultSet executeQueryImpl(String sql) throws SQLException { @Override public long executeLargeUpdate(String sql) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - return executeLargeUpdateImpl(sql); - } + LOG.finest("++enter++"); + checkClosed(); + return executeLargeUpdateImpl(sql); } private long executeLargeUpdateImpl(String sql) throws SQLException { @@ -287,13 +281,8 @@ private long executeLargeUpdateImpl(String sql) throws SQLException { @Override public int executeUpdate(String sql) throws SQLException { - try { - BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId); - LOG.finest("++enter++"); - return checkUpdateCount(executeLargeUpdate(sql)); - } finally { - BigQueryJdbcMdc.clear(); - } + LOG.finest("++enter++"); + return checkUpdateCount(executeLargeUpdate(sql)); } int checkUpdateCount(long updateCount) { @@ -308,12 +297,9 @@ int checkUpdateCount(long updateCount) { @Override public boolean execute(String sql) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - return executeImpl(sql); - } + LOG.finest("++enter++"); + checkClosed(); + return executeImpl(sql); } private boolean executeImpl(String sql) throws SQLException { @@ -402,23 +388,20 @@ public void close() throws SQLException { if (isClosed()) { return; } - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) { - LOG.fine("Closing Statement %s.", this); + LOG.fine("Closing Statement %s.", this); - boolean cancelSucceeded = false; - try { - cancel(); // This attempts to cancel jobs and calls closeStatementResources() - cancelSucceeded = true; - } catch (SQLException e) { - LOG.warning("Failed to cancel statement during close().", e); - } finally { - if (!cancelSucceeded) { - closeStatementResources(); - } - this.connection = null; - this.isClosed = true; + boolean cancelSucceeded = false; + try { + cancel(); // This attempts to cancel jobs and calls closeStatementResources() + cancelSucceeded = true; + } catch (SQLException e) { + LOG.warning("Failed to cancel statement during close().", e); + } finally { + if (!cancelSucceeded) { + closeStatementResources(); } + this.connection = null; + this.isClosed = true; } } @@ -470,31 +453,28 @@ public void setQueryTimeout(int seconds) { */ @Override public void cancel() throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) { - LOG.finest("Statement %s cancelled", this); - synchronized (cancelLock) { - this.isCanceled = true; - for (JobId jobId : this.jobIds) { - try { - this.bigQuery.cancel(jobId); - LOG.info("Job " + jobId + "cancelled."); - } catch (BigQueryException e) { - if (e.getMessage() != null - && (e.getMessage().contains("Job is already in state DONE") - || e.getMessage().contains("Error: 3848323"))) { - LOG.warning("Attempted to cancel a job that was already done: " + jobId); - } else { - throw new BigQueryJdbcException(e); - } + LOG.finest("Statement %s cancelled", this); + synchronized (cancelLock) { + this.isCanceled = true; + for (JobId jobId : this.jobIds) { + try { + this.bigQuery.cancel(jobId); + LOG.info("Job " + jobId + "cancelled."); + } catch (BigQueryException e) { + if (e.getMessage() != null + && (e.getMessage().contains("Job is already in state DONE") + || e.getMessage().contains("Error: 3848323"))) { + LOG.warning("Attempted to cancel a job that was already done: " + jobId); + } else { + throw new BigQueryJdbcException(e); } } - jobIds.clear(); } - // If a ResultSet exists, then it will be closed as well, closing the - // ownedThreads - closeStatementResources(); + jobIds.clear(); } + // If a ResultSet exists, then it will be closed as well, closing the + // ownedThreads + closeStatementResources(); } @Override @@ -1484,12 +1464,9 @@ public boolean hasMoreResults() { @Override public boolean getMoreResults(int current) throws SQLException { - try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(this.connection, this.connectionId)) { - LOG.finest("++enter++"); - checkClosed(); - return getMoreResultsImpl(current); - } + LOG.finest("++enter++"); + checkClosed(); + return getMoreResultsImpl(current); } private boolean getMoreResultsImpl(int current) throws SQLException { diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandler.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandler.java index 1d6c177537e8..46cf647ef1f4 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandler.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandler.java @@ -55,6 +55,9 @@ class PerConnectionFileHandler extends Handler { } private String getLogFilePath(String id) { + if ("Jdbc-default".equals(id)) { + return baseLogPath.resolve("BQ-JDBC-GLOBAL.log").toString(); + } String uuid = id; if (id.startsWith("BQ-JDBC-")) { uuid = id.substring("BQ-JDBC-".length()); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java index 8a86c8805cdb..7de4516427c3 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java @@ -42,12 +42,17 @@ public PooledConnection getPooledConnection() throws SQLException { throw new BigQueryJdbcRuntimeException( "Cannot get pooled connection: unable to get underlying physical connection"); } - Long connectionPoolSize = ((BigQueryConnection) bqConnection).getConnectionPoolSize(); + BigQueryConnection physicalConnection; + if (bqConnection.isWrapperFor(BigQueryConnection.class)) { + physicalConnection = bqConnection.unwrap(BigQueryConnection.class); + } else { + physicalConnection = (BigQueryConnection) bqConnection; + } + Long connectionPoolSize = physicalConnection.getConnectionPoolSize(); if (connectionPoolManager == null) { connectionPoolManager = new PooledConnectionListener(connectionPoolSize); } - BigQueryPooledConnection bqPooledConnection = - new BigQueryPooledConnection((BigQueryConnection) bqConnection); + BigQueryPooledConnection bqPooledConnection = new BigQueryPooledConnection(physicalConnection); bqPooledConnection.addConnectionEventListener(connectionPoolManager); return bqPooledConnection; } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxyTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxyTest.java new file mode 100644 index 000000000000..6cb46cb8de21 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxyTest.java @@ -0,0 +1,125 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery.jdbc; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.bigquery.Field; +import com.google.cloud.bigquery.FieldList; +import com.google.cloud.bigquery.StandardSQLTypeName; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import org.junit.jupiter.api.Test; + +public class BigQueryJdbcContextProxyTest { + + @Test + public void testExtractConnectionIdFromConnection() throws SQLException { + BigQueryConnection mockConn = mock(BigQueryConnection.class); + when(mockConn.getConnectionId()).thenReturn("conn-uuid-123"); + when(mockConn.getAutoCommit()) + .thenAnswer( + invocation -> { + assertEquals("conn-uuid-123", BigQueryJdbcMdc.getConnectionId()); + return true; + }); + + Connection proxy = BigQueryJdbcContextProxy.wrap(mockConn, Connection.class); + assertNotNull(proxy); + + // Verify active thread-local context is set during method invocation + assertTrue(proxy.getAutoCommit()); + // Verify context is cleanly cleared after method exit + assertNull(BigQueryJdbcMdc.getConnectionId()); + } + + @Test + public void testExtractConnectionIdFromStatement() throws SQLException { + BigQueryConnection mockConn = mock(BigQueryConnection.class); + when(mockConn.getBigQuery()).thenReturn(mock(com.google.cloud.bigquery.BigQuery.class)); + + BigQueryStatement stmt = new BigQueryStatement(mockConn); + stmt.connectionId = "conn-uuid-456"; + + Statement proxy = BigQueryJdbcContextProxy.wrap(stmt, Statement.class, "conn-uuid-456"); + assertNotNull(proxy); + + // We can call any statement method (like getUpdateCount) and verify context routing + // by asserting the thread-local value during invocation, but since it's a real object, + // it runs cleanly. Let's assert that no thread-local leaks exist. + assertEquals(-1, proxy.getUpdateCount()); + assertNull(BigQueryJdbcMdc.getConnectionId()); + } + + @Test + public void testExtractConnectionIdFromDatabaseMetaData() throws SQLException { + BigQueryConnection mockConn = mock(BigQueryConnection.class); + when(mockConn.getConnectionId()).thenReturn("conn-uuid-789"); + when(mockConn.getBigQuery()).thenReturn(mock(com.google.cloud.bigquery.BigQuery.class)); + + BigQueryDatabaseMetaData meta = new BigQueryDatabaseMetaData(mockConn); + + DatabaseMetaData proxy = + BigQueryJdbcContextProxy.wrap(meta, DatabaseMetaData.class, "conn-uuid-789"); + assertNotNull(proxy); + + // Assert read-only capability does not leak thread context + assertFalse(proxy.isReadOnly()); + assertNull(BigQueryJdbcMdc.getConnectionId()); + } + + @Test + public void testExtractConnectionIdFromResultSetMetaData() throws SQLException { + BigQueryConnection mockConn = mock(BigQueryConnection.class); + BigQueryStatement stmt = new BigQueryStatement(mockConn); + stmt.connectionId = "conn-uuid-999"; + + FieldList fields = FieldList.of(Field.of("col", StandardSQLTypeName.STRING)); + BigQueryResultSetMetadata meta = BigQueryResultSetMetadata.of(fields, stmt); + + ResultSetMetaData proxy = + BigQueryJdbcContextProxy.wrap(meta, ResultSetMetaData.class, "conn-uuid-999"); + assertNotNull(proxy); + + assertEquals(1, proxy.getColumnCount()); + assertNull(BigQueryJdbcMdc.getConnectionId()); + } + + @Test + public void testWrapWithNullContextAndExceptionThrown() throws SQLException { + BigQueryStatement mockStmt = mock(BigQueryStatement.class); + mockStmt.connectionId = null; + when(mockStmt.executeQuery("SELECT *")).thenThrow(new SQLException("Database error")); + + Statement proxy = BigQueryJdbcContextProxy.wrap(mockStmt, Statement.class, null); + assertNotNull(proxy); + + SQLException ex = assertThrows(SQLException.class, () -> proxy.executeQuery("SELECT *")); + assertEquals("Database error", ex.getMessage()); + assertNull(BigQueryJdbcMdc.getConnectionId()); + } +} diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java index ce192baece10..56d05d6b56c0 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java @@ -79,6 +79,38 @@ public void testLogWithCallerInference() { assertEquals(BigQueryJdbcCustomLoggerTest.class.getName(), record.getSourceClassName()); } + @Test + public void testHotPathLoggerLogToDefaultWhenContextIsNull() { + BigQueryJdbcCustomLogger hotpathLogger = + new BigQueryJdbcCustomLogger("com.google.cloud.bigquery.jdbc.BigQueryArrowResultSet"); + TestHandler hotpathHandler = new TestHandler(); + hotpathLogger.addHandler(hotpathHandler); + hotpathLogger.setLevel(Level.ALL); + + BigQueryJdbcMdc.clear(); // Ensure context is null + hotpathLogger.fine("This should log to default"); + + List records = hotpathHandler.getRecords(); + assertEquals(1, records.size()); // Logged successfully, not dropped! + assertEquals("This should log to default", records.get(0).getMessage()); + } + + @Test + public void testHotPathLoggerNotSilencedWhenContextIsPresent() { + BigQueryJdbcCustomLogger hotpathLogger = + new BigQueryJdbcCustomLogger("com.google.cloud.bigquery.jdbc.BigQueryArrowResultSet"); + TestHandler hotpathHandler = new TestHandler(); + hotpathLogger.addHandler(hotpathHandler); + hotpathLogger.setLevel(Level.ALL); + + BigQueryJdbcMdc.registerInstance("TestConnection"); // Set active context + hotpathLogger.fine("This should not be silenced"); + + List records = hotpathHandler.getRecords(); + assertEquals(1, records.size()); // Allowed! + assertEquals("This should not be silenced", records.get(0).getMessage()); + } + @Test public void testLogWithException() { Exception ex = new Exception("Test exception"); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java index 63c5883fd9c5..55edf6efb89a 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdcTest.java @@ -24,21 +24,10 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.mockito.Mockito; public class BigQueryJdbcMdcTest { - private BigQueryConnection mockConnection1; - private BigQueryConnection mockConnection2; - - @BeforeEach - public void setUp() { - mockConnection1 = Mockito.mock(BigQueryConnection.class); - mockConnection2 = Mockito.mock(BigQueryConnection.class); - } - @AfterEach public void tearDown() { BigQueryJdbcMdc.clear(); @@ -46,37 +35,22 @@ public void tearDown() { @Test public void testRegisterAndRetrieveConnectionId() { - BigQueryJdbcMdc.registerInstance(mockConnection1, "123"); - assertEquals("123", BigQueryJdbcMdc.getConnectionId()); - } - - @Test - public void testRemoveInstance() { - BigQueryJdbcMdc.registerInstance(mockConnection1, "1"); - assertEquals("1", BigQueryJdbcMdc.getConnectionId()); - - BigQueryJdbcMdc.removeInstance(mockConnection1); - // Note: removeInstance does not clear currentConnectionId on the current thread - // based on current implementation. - assertEquals("1", BigQueryJdbcMdc.getConnectionId()); - - BigQueryJdbcMdc.clear(); - assertNull(BigQueryJdbcMdc.getConnectionId()); + BigQueryJdbcMdc.registerInstance("JdbcConnection-123"); + assertEquals("JdbcConnection-123", BigQueryJdbcMdc.getConnectionId()); } @Test public void testClearContext() { - BigQueryJdbcMdc.registerInstance(mockConnection1, "456"); - assertEquals("456", BigQueryJdbcMdc.getConnectionId()); - + BigQueryJdbcMdc.registerInstance("JdbcConnection-456"); + assertEquals("JdbcConnection-456", BigQueryJdbcMdc.getConnectionId()); BigQueryJdbcMdc.clear(); assertNull(BigQueryJdbcMdc.getConnectionId()); } @Test public void testThreadInheritance() throws InterruptedException { - BigQueryJdbcMdc.registerInstance(mockConnection1, "parent"); - assertEquals("parent", BigQueryJdbcMdc.getConnectionId()); + BigQueryJdbcMdc.registerInstance("JdbcConnection-parent"); + assertEquals("JdbcConnection-parent", BigQueryJdbcMdc.getConnectionId()); AtomicReference childConnectionId = new AtomicReference<>(); CountDownLatch latch = new CountDownLatch(1); @@ -90,13 +64,12 @@ public void testThreadInheritance() throws InterruptedException { childThread.start(); assertTrue(latch.await(5, TimeUnit.SECONDS)); - assertEquals("parent", childConnectionId.get()); + assertEquals("JdbcConnection-parent", childConnectionId.get()); } @Test public void testThreadIsolation() throws InterruptedException { CountDownLatch threadARegistered = new CountDownLatch(1); - CountDownLatch threadBChecked = new CountDownLatch(1); CountDownLatch threadBRegistered = new CountDownLatch(1); CountDownLatch testFinished = new CountDownLatch(2); @@ -109,7 +82,7 @@ public void testThreadIsolation() throws InterruptedException { new Thread( () -> { try { - BigQueryJdbcMdc.registerInstance(mockConnection1, "A"); + BigQueryJdbcMdc.registerInstance("JdbcConnection-A"); threadAIdBeforeB.set(BigQueryJdbcMdc.getConnectionId()); threadARegistered.countDown(); @@ -129,7 +102,7 @@ public void testThreadIsolation() throws InterruptedException { threadARegistered.await(); threadBIdBeforeRegister.set(BigQueryJdbcMdc.getConnectionId()); - BigQueryJdbcMdc.registerInstance(mockConnection2, "B"); + BigQueryJdbcMdc.registerInstance("JdbcConnection-B"); threadBIdAfterRegister.set(BigQueryJdbcMdc.getConnectionId()); threadBRegistered.countDown(); } catch (InterruptedException e) { @@ -144,17 +117,17 @@ public void testThreadIsolation() throws InterruptedException { assertTrue(testFinished.await(5, TimeUnit.SECONDS)); - assertEquals("A", threadAIdBeforeB.get()); + assertEquals("JdbcConnection-A", threadAIdBeforeB.get()); assertNull(threadBIdBeforeRegister.get()); - assertEquals("B", threadBIdAfterRegister.get()); - assertEquals("A", threadAIdAfterB.get()); + assertEquals("JdbcConnection-B", threadBIdAfterRegister.get()); + assertEquals("JdbcConnection-A", threadAIdAfterB.get()); } @Test public void testMdcCloseableClearsContext() { try (BigQueryJdbcMdc.MdcCloseable mdc = - BigQueryJdbcMdc.registerInstance(mockConnection1, "789")) { - assertEquals("789", BigQueryJdbcMdc.getConnectionId()); + BigQueryJdbcMdc.registerInstance("JdbcConnection-789")) { + assertEquals("JdbcConnection-789", BigQueryJdbcMdc.getConnectionId()); } assertNull(BigQueryJdbcMdc.getConnectionId()); } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java index f905fa80fe5a..fda5112703fd 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/PerConnectionFileHandlerTest.java @@ -16,11 +16,20 @@ package com.google.cloud.bigquery.jdbc; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import com.google.cloud.bigquery.Field; +import com.google.cloud.bigquery.FieldList; +import com.google.cloud.bigquery.StandardSQLTypeName; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; +import java.sql.SQLException; +import java.sql.Statement; import java.util.Optional; import java.util.logging.Level; import java.util.logging.LogRecord; @@ -63,7 +72,8 @@ private Optional findLogFile(String suffix) throws IOException { @Test public void testInitialization() throws IOException { - assertTrue(findLogFile("-Jdbc.log").isPresent()); + Path defaultLog = tempDir.resolve("BQ-JDBC-GLOBAL.log"); + assertTrue(Files.exists(defaultLog)); } @Test @@ -72,38 +82,137 @@ public void testPublishDefault() throws IOException { handler.publish(record); handler.flush(); - Optional defaultLog = findLogFile("-Jdbc.log"); - assertTrue(defaultLog.isPresent()); - String content = new String(Files.readAllBytes(defaultLog.get())); + Path defaultLog = tempDir.resolve("BQ-JDBC-GLOBAL.log"); + assertTrue(Files.exists(defaultLog)); + String content = new String(Files.readAllBytes(defaultLog)); assertTrue(content.contains("Test message default")); } @Test public void testPublishConnectionSpecific() throws IOException { - BigQueryJdbcMdc.registerInstance(mockConnection, "123"); + BigQueryJdbcMdc.registerInstance("c123"); LogRecord record = new LogRecord(Level.INFO, "Test message connection 123"); handler.publish(record); handler.flush(); - Optional connLog = findLogFile("-123.log"); + Optional connLog = findLogFile("-c123.log"); assertTrue(connLog.isPresent()); - String content = new String(Files.readAllBytes(connLog.get())); - assertTrue(content.contains("Test message connection 123")); } @Test public void testCloseHandler() throws IOException { - BigQueryJdbcMdc.registerInstance(mockConnection, "456"); + BigQueryJdbcMdc.registerInstance("c123"); - LogRecord record = new LogRecord(Level.INFO, "Test message connection 456"); + LogRecord record = new LogRecord(Level.INFO, "Test message connection 123"); handler.publish(record); handler.flush(); - Optional connLog = findLogFile("-456.log"); + Optional connLog = findLogFile("-c123.log"); + assertTrue(connLog.isPresent()); + + handler.closeHandler("c123"); + // File remains on disk, but handler is closed. assertTrue(connLog.isPresent()); + } + + /** + * Verifies that when an exception is thrown from within a dynamic proxy wrapper method while the + * thread-local connection ID context is completely missing (null), the proxy's catch block + * dynamically extracts the connection ID from the target instance, registers it on the executing + * thread, and writes the severe exception log record directly to the connection-specific log file + * (e.g. "BQ-JDBC-timestamp-c456.log") instead of the global log file. + */ + @Test + public void testProxyExceptionLogRouting() throws Exception { + // Register the temp file handler to the root logger so proxy logs are routed to it + java.util.logging.Logger rootLogger = BigQueryJdbcRootLogger.getRootLogger(); + rootLogger.addHandler(handler); + + try { + // Ensure thread context is completely missing (null) before query + BigQueryJdbcMdc.clear(); + assertNull(BigQueryJdbcMdc.getConnectionId()); + + // Create a mock statement with connectionId = "c456" + BigQueryStatement mockStmt = mock(BigQueryStatement.class); + mockStmt.connectionId = "c456"; + + // Mock executeQuery to throw an exception + Mockito.when(mockStmt.executeQuery(Mockito.anyString())) + .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, "c456"); + assertNotNull(proxy); + + // Call the proxy method. It will throw SQLException + assertThrows(SQLException.class, () -> proxy.executeQuery("SELECT *")); + + // Flush the handler to ensure logs are written to disk + handler.flush(); + + // Verify that the exception log got registered and written directly to c456.log! + Optional connLog = findLogFile("-c456.log"); + assertTrue(connLog.isPresent()); + + String content = new String(Files.readAllBytes(connLog.get())); + assertTrue(content.contains("Database error")); + assertTrue(content.contains("Exception occurred during executeQuery")); - handler.closeHandler("BQ-JDBC-456"); - assertTrue(Files.exists(connLog.get())); + } finally { + // Cleanup + rootLogger.removeHandler(handler); + } + } + + /** + * Verifies that when an exception is thrown from within an unproxied, raw ResultSet concrete + * class while the thread-local connection ID context is completely missing (null), the internal + * logAndCreateException() builder dynamically extracts the connection ID from its parent + * statement, registers it on the executing thread, and writes the severe exception log record + * directly to the connection-specific log file (e.g. "BQ-JDBC-timestamp-c789.log") instead of the + * global log file. + */ + @Test + public void testResultSetExceptionLogRouting() throws Exception { + // Register the temp file handler to the root logger so ResultSet logs are captured + java.util.logging.Logger rootLogger = BigQueryJdbcRootLogger.getRootLogger(); + rootLogger.addHandler(handler); + + try { + // Ensure thread context is completely missing (null) before call + BigQueryJdbcMdc.clear(); + assertNull(BigQueryJdbcMdc.getConnectionId()); + + // Create a mock statement with connectionId = "c789" + BigQueryStatement mockStmt = mock(BigQueryStatement.class); + mockStmt.connectionId = "c789"; + + // Create a mock FieldList and schema for the ResultSet + FieldList fields = FieldList.of(Field.of("col", StandardSQLTypeName.STRING)); + com.google.cloud.bigquery.Schema schema = com.google.cloud.bigquery.Schema.of(fields); + + // Instantiate a real BigQueryJsonResultSet (which extends BigQueryBaseResultSet) + // passing the mock statement carrying connectionId "c789" + BigQueryJsonResultSet rs = BigQueryJsonResultSet.of(schema, 0, null, mockStmt, new Thread[0]); + + // Calling findColumn(null) throws SQLException because column label is null + assertThrows(SQLException.class, () -> rs.findColumn(null)); + + // Flush the handler to ensure logs are written to disk + handler.flush(); + + // Verify that the ResultSet exception log got registered and written directly to c789.log! + Optional connLog = findLogFile("-c789.log"); + assertTrue(connLog.isPresent()); + + String content = new String(Files.readAllBytes(connLog.get())); + assertTrue(content.contains("Column label cannot be null")); + + } finally { + // Cleanup + rootLogger.removeHandler(handler); + } } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java index 1f4397e417b3..0877553e42c0 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITAuthTests.java @@ -71,9 +71,7 @@ private void validateConnection(String connection_uri) throws SQLException { Connection connection = DriverManager.getConnection(connection_uri); assertNotNull(connection); assertFalse(connection.isClosed()); - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); int totalRows = 0; @@ -205,7 +203,7 @@ public void testValidGoogleUserAccountAuthentication() throws SQLException { Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery( - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 50"); + "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 50"); assertEquals(50, resultSetRowCount(resultSet)); connection.close(); @@ -229,7 +227,7 @@ public void testValidExternalAccountAuthentication() throws SQLException { Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery( - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 50"); + "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 50"); assertEquals(50, resultSetRowCount(resultSet)); connection.close(); @@ -251,7 +249,7 @@ public void testValidExternalAccountAuthenticationFromFile() throws SQLException Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery( - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 50"); + "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 50"); assertEquals(50, resultSetRowCount(resultSet)); connection.close(); @@ -283,7 +281,7 @@ public void testValidExternalAccountAuthenticationRawJson() throws SQLException Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery( - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 50"); + "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 50"); assertEquals(50, resultSetRowCount(resultSet)); connection.close(); @@ -333,7 +331,7 @@ public void testValidRefreshTokenAuthentication() throws SQLException { Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery( - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 50"); + "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 50"); assertEquals(50, resultSetRowCount(resultSet)); connection.close(); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java index f8c75a2acc9e..26df941651e9 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java @@ -149,9 +149,7 @@ public void testValidAllDataTypesSerializationFromSelectQueryArrowDataset() thro @Test public void testFastQueryPathSmall() throws SQLException { - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; ResultSet jsonResultSet = bigQueryStatement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); assertEquals(850, resultSetRowCount(jsonResultSet)); @@ -159,9 +157,7 @@ public void testFastQueryPathSmall() throws SQLException { @Test public void testFastQueryPathEmpty() throws SQLException { - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 0"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0"; Connection connection = DriverManager.getConnection(String.format(connectionUrl, DEFAULT_CATALOG)); Statement bigQueryStatement = connection.createStatement(); @@ -173,8 +169,8 @@ public void testFastQueryPathEmpty() throws SQLException { @Test public void testSmallSelectAndVerifyResults() throws SQLException { String query = - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` WHERE" - + " repository_name LIKE 'X%' LIMIT 10"; + "SELECT word FROM `bigquery-public-data.samples.shakespeare` WHERE" + + " word LIKE 'X%' LIMIT 10"; ResultSet resultSet = bigQueryStatement.executeQuery(query); int rowCount = 0; @@ -287,15 +283,12 @@ public void testStatelessQueryPathSmall() throws SQLException { Statement statement = connectionUseStateless.createStatement(); - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; ResultSet jsonResultSet = statement.executeQuery(query); Assert.assertEquals(850, resultSetRowCount(jsonResultSet)); String queryEmpty = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 0"; + "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0"; ResultSet jsonResultSetEmpty = statement.executeQuery(queryEmpty); Assert.assertEquals(0, resultSetRowCount(jsonResultSetEmpty)); connectionUseStateless.close(); @@ -326,8 +319,7 @@ public void testDriver() throws SQLException { Connection connection = driver.connect(connection_uri, new Properties()); assertNotNull(connection); Statement st = connection.createStatement(); - boolean rs = - st.execute("Select * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"); + boolean rs = st.execute("Select * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"); assertTrue(rs); connection.close(); } @@ -345,7 +337,8 @@ public void testDefaultDataset() throws SQLException { Connection connection = driver.connect(connection_uri, new Properties()); assertNotNull(connection); assertEquals( - DatasetId.of("testDataset"), ((BigQueryConnection) connection).getDefaultDataset()); + DatasetId.of("testDataset"), + connection.unwrap(BigQueryConnection.class).getDefaultDataset()); String connection_uri_null_default_dataset = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PROJECTID=" @@ -356,7 +349,7 @@ public void testDefaultDataset() throws SQLException { Connection connection2 = driver.connect(connection_uri_null_default_dataset, new Properties()); assertNotNull(connection2); - assertNull(((BigQueryConnection) connection2).getDefaultDataset()); + assertNull(connection2.unwrap(BigQueryConnection.class).getDefaultDataset()); connection.close(); connection2.close(); } @@ -377,7 +370,7 @@ public void testDefaultDatasetWithProject() throws SQLException { assertNotNull(connection); assertEquals( DatasetId.of(PROJECT_ID, "testDataset"), - ((BigQueryConnection) connection).getDefaultDataset()); + connection.unwrap(BigQueryConnection.class).getDefaultDataset()); connection.close(); } @@ -392,7 +385,7 @@ public void testLocation() throws SQLException { assertTrue(driver.acceptsURL(connection_uri)); Connection connection = driver.connect(connection_uri, new Properties()); - assertEquals(((BigQueryConnection) connection).getLocation(), "EU"); + assertEquals(connection.unwrap(BigQueryConnection.class).getLocation(), "EU"); Statement statement = connection.createStatement(); @@ -411,7 +404,7 @@ public void testLocation() throws SQLException { Connection connection2 = driver.connect(connection_uri_null_location, new Properties()); assertNotNull(connection2); - assertNull(((BigQueryConnection) connection2).getLocation()); + assertNull(connection2.unwrap(BigQueryConnection.class).getLocation()); connection.close(); connection2.close(); } @@ -426,11 +419,11 @@ public void testIncorrectLocation() throws SQLException { Driver driver = BigQueryDriver.getRegisteredDriver(); Connection connection = driver.connect(connection_uri, new Properties()); - assertEquals(((BigQueryConnection) connection).getLocation(), "europe-west3"); + assertEquals(connection.unwrap(BigQueryConnection.class).getLocation(), "europe-west3"); // Query a dataset in the US Statement statement = connection.createStatement(); - String query = "SELECT * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"; + String query = "SELECT * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"; BigQueryJdbcException ex = assertThrows(BigQueryJdbcException.class, () -> statement.executeQuery(query)); BigQueryError error = ex.getBigQueryException().getError(); @@ -524,7 +517,8 @@ public void testCreateStatementWithResultSetConcurrencyWhenConnectionClosedThrow @Test public void testSetAutoCommitWithClosedConnectionThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.close(); assertThrows(IllegalStateException.class, () -> connection.setAutoCommit(true)); } @@ -532,7 +526,7 @@ public void testSetAutoCommitWithClosedConnectionThrowsIllegalState() throws SQL @Test public void testSetCommitToFalseWithoutSessionEnabledThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(connection_uri); + DriverManager.getConnection(connection_uri).unwrap(BigQueryConnection.class); assertThrows(IllegalStateException.class, () -> connection.setAutoCommit(false)); connection.close(); } @@ -540,7 +534,8 @@ public void testSetCommitToFalseWithoutSessionEnabledThrowsIllegalState() throws @Test public void testCommitWithConnectionClosedThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.close(); assertThrows(IllegalStateException.class, connection::commit); } @@ -548,7 +543,7 @@ public void testCommitWithConnectionClosedThrowsIllegalState() throws SQLExcepti @Test public void testCommitToFalseWithoutSessionEnabledThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(connection_uri); + DriverManager.getConnection(connection_uri).unwrap(BigQueryConnection.class); assertThrows(IllegalStateException.class, connection::commit); connection.close(); } @@ -556,7 +551,8 @@ public void testCommitToFalseWithoutSessionEnabledThrowsIllegalState() throws SQ @Test public void testCommitWithNoTransactionStartedThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); assertThrows(IllegalStateException.class, connection::commit); connection.close(); } @@ -606,7 +602,8 @@ public void testRollbackOnConnectionClosed() throws SQLException { @Test public void testRollbackWithConnectionClosedThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.close(); assertThrows(IllegalStateException.class, connection::rollback); } @@ -614,7 +611,7 @@ public void testRollbackWithConnectionClosedThrowsIllegalState() throws SQLExcep @Test public void testRollbackToFalseWithoutSessionEnabledThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(connection_uri); + DriverManager.getConnection(connection_uri).unwrap(BigQueryConnection.class); assertThrows(IllegalStateException.class, connection::rollback); connection.close(); } @@ -622,7 +619,8 @@ public void testRollbackToFalseWithoutSessionEnabledThrowsIllegalState() throws @Test public void testRollbackWithoutTransactionStartedThrowsIllegalState() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); assertThrows(IllegalStateException.class, connection::rollback); connection.close(); } @@ -743,7 +741,8 @@ public void testGetLocationWhenConnectionClosedThrows() throws SQLException { connection.close(); assertThrows( - IllegalStateException.class, () -> ((BigQueryConnection) connection).getLocation()); + IllegalStateException.class, + () -> connection.unwrap(BigQueryConnection.class).getLocation()); connection.close(); } @@ -754,7 +753,8 @@ public void testGetDefaultDatasetWhenConnectionClosedThrows() throws SQLExceptio connection.close(); assertThrows( - IllegalStateException.class, () -> ((BigQueryConnection) connection).getDefaultDataset()); + IllegalStateException.class, + () -> connection.unwrap(BigQueryConnection.class).getDefaultDataset()); } @Test @@ -959,8 +959,7 @@ public void testExecuteQueryWithInsert() throws SQLException { @Test public void testExecuteQueryWithMultipleReturns() throws SQLException { - String query = - String.format("SELECT * FROM bigquery-public-data.samples.github_timeline LIMIT 1;"); + String query = String.format("SELECT * FROM bigquery-public-data.samples.shakespeare LIMIT 1;"); assertThrows(BigQueryJdbcException.class, () -> bigQueryStatement.executeQuery(query + query)); } @@ -968,7 +967,7 @@ public void testExecuteQueryWithMultipleReturns() throws SQLException { @Test public void testExecuteUpdateWithSelect() throws SQLException { String selectQuery = - String.format("SELECT * FROM bigquery-public-data.samples.github_timeline LIMIT 1;"); + String.format("SELECT * FROM bigquery-public-data.samples.shakespeare LIMIT 1;"); assertThrows(BigQueryJdbcException.class, () -> bigQueryStatement.executeUpdate(selectQuery)); } @@ -1120,8 +1119,8 @@ public void testExecuteBatchQueryTypeSelectThrowsUnsupported() throws SQLExcepti Driver driver = BigQueryDriver.getRegisteredDriver(); Connection connection = driver.connect(connection_uri, new Properties()); String query = - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` WHERE" - + " repository_name LIKE 'X%' LIMIT 10"; + "SELECT word FROM `bigquery-public-data.samples.shakespeare` WHERE" + + " word LIKE 'X%' LIMIT 10"; Statement statement = connection.createStatement(); assertThrows(IllegalArgumentException.class, () -> statement.addBatch(query)); @@ -1467,7 +1466,7 @@ public void testREPEndpointDataNotFoundThrows() throws SQLException { @Test public void testCloseStatement() throws SQLException { - String query = "SELECT * FROM `bigquery-public-data.samples.github_timeline` LIMIT 10"; + String query = "SELECT * FROM `bigquery-public-data.samples.shakespeare` LIMIT 10"; Statement statement = bigQueryConnection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); assertEquals(10, resultSetRowCount(jsonResultSet)); @@ -1477,7 +1476,7 @@ public void testCloseStatement() throws SQLException { @Test public void testCloseableStatementSingleResult() throws SQLException { - String query = "SELECT * FROM `bigquery-public-data.samples.github_timeline` LIMIT 10"; + String query = "SELECT * FROM `bigquery-public-data.samples.shakespeare` LIMIT 10"; Statement statement = bigQueryConnection.createStatement(); statement.closeOnCompletion(); assertTrue(statement.isCloseOnCompletion()); @@ -1489,7 +1488,7 @@ public void testCloseableStatementSingleResult() throws SQLException { @Test public void testCloseableStatementMultiResult() throws SQLException { - String query = "SELECT * FROM `bigquery-public-data.samples.github_timeline` LIMIT 10;"; + String query = "SELECT * FROM `bigquery-public-data.samples.shakespeare` LIMIT 10;"; Statement statement = bigQueryConnection.createStatement(); statement.closeOnCompletion(); assertTrue(statement.isCloseOnCompletion()); @@ -1507,7 +1506,7 @@ public void testCloseableStatementMultiResult() throws SQLException { @Test public void testCloseableStatementMultiResultExplicitClose() throws SQLException { - String query = "SELECT * FROM `bigquery-public-data.samples.github_timeline` LIMIT 10;"; + String query = "SELECT * FROM `bigquery-public-data.samples.shakespeare` LIMIT 10;"; Statement statement = bigQueryConnection.createStatement(); statement.closeOnCompletion(); assertTrue(statement.isCloseOnCompletion()); @@ -1554,10 +1553,9 @@ public void testDataSourceOAuthPvtKeyPath() throws SQLException, IOException { @Test public void testPreparedStatementSmallSelect() throws SQLException { String query = - "SELECT * FROM `bigquery-public-data.samples.github_timeline` where repository_language=?" - + " LIMIT 1000"; + "SELECT * FROM `bigquery-public-data.samples.shakespeare` where corpus=? LIMIT 1000"; PreparedStatement preparedStatement = bigQueryConnection.prepareStatement(query); - preparedStatement.setString(1, "Java"); + preparedStatement.setString(1, "hamlet"); ResultSet jsonResultSet = preparedStatement.executeQuery(); @@ -2121,7 +2119,7 @@ public void testQueryPropertyTimeZoneQueries() throws SQLException { + "ProjectId=" + PROJECT_ID + ";QueryProperties=time_zone=America/New_York;"; - String query = "SELECT * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"; + String query = "SELECT * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"; Driver driver = BigQueryDriver.getRegisteredDriver(); Connection connection = driver.connect(connection_uri, new Properties()); Statement statement = connection.createStatement(); @@ -2265,7 +2263,8 @@ public void testValidLegacySQLStatement() throws SQLException { @Test public void testMultipleTransactionsThrowsUnsupported() throws SQLException { BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.setAutoCommit(false); Statement statement = connection.createStatement(); assertThrows(BigQueryJdbcException.class, () -> statement.execute("BEGIN TRANSACTION;")); @@ -2757,4 +2756,80 @@ public void validateGetObjectNullValues() throws Exception { } } } + + @Test + public void testPerConnectionLoggingE2E() throws SQLException, IOException { + File tempDir = File.createTempFile("bq-jdbc-e2e-logs", ""); + tempDir.delete(); + tempDir.mkdirs(); + tempDir.deleteOnExit(); + + String logPath = tempDir.getAbsolutePath(); + String targetUri = connection_uri + ";LogLevel=6;LogPath=" + logPath; + + String connectionId = null; + try (Connection conn = DriverManager.getConnection(targetUri)) { + assertNotNull(conn); + + // Extract connection ID using reflection to bypass package-private encapsulation limits in + // test + java.lang.reflect.Method method = + conn.unwrap(BigQueryConnection.class).getClass().getDeclaredMethod("getConnectionId"); + method.setAccessible(true); + connectionId = (String) method.invoke(conn.unwrap(BigQueryConnection.class)); + assertNotNull(connectionId); + + // 1. Execute a local operational method (like close) which does not create cloud jobs + Statement stmt = conn.createStatement(); + stmt.close(); + + // 2. Execute an empty query to trigger a local SQL syntax SQLException + assertThrows(SQLException.class, () -> stmt.executeQuery("")); + } catch (Exception e) { + throw new SQLException("Reflection lookup or E2E execution failed", e); + } finally { + // Cleanly close and remove all file handlers from com.google.cloud.bigquery logger using + // public APIs + java.util.logging.Logger bqLogger = + java.util.logging.Logger.getLogger("com.google.cloud.bigquery"); + for (java.util.logging.Handler h : bqLogger.getHandlers()) { + h.close(); + bqLogger.removeHandler(h); + } + + // Verify physical connection-specific log file creation + // Verify physical connection-specific log file creation (uses first 4 chars of connectionId) + final String shortId = + connectionId != null + ? connectionId.substring(0, Math.min(connectionId.length(), 4)) + : null; + File[] files = + tempDir.listFiles( + (dir, name) -> shortId != null && name.endsWith("-" + shortId + ".log")); + assertNotNull(files); + assertEquals(1, files.length); + + File actualLog = files[0]; + byte[] encoded = java.nio.file.Files.readAllBytes(actualLog.toPath()); + String content = new String(encoded, StandardCharsets.UTF_8); + + // Asserts that the connection ID prefix is formatted inside log entries + assertTrue(content.contains("[" + connectionId + "]")); + // Asserts that operational actions are logged inside the connection log + assertTrue(content.contains("close")); + // Asserts that the exception is connection-routed and logged + assertTrue( + content.contains("Exception occurred during executeQuery"), + "Log content did not contain expected exception! Content: \n" + content); + + // Clean up log files inside temp directory but keep directory alive to avoid afterClass + // NoSuchFileException + File[] remaining = tempDir.listFiles(); + if (remaining != null) { + for (File f : remaining) { + f.delete(); + } + } + } + } } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITConnectionPoolingTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITConnectionPoolingTest.java index ea53aeb84a29..f310ac8c986d 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITConnectionPoolingTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITConnectionPoolingTest.java @@ -321,9 +321,7 @@ private void assertConnectionPoolingResults(String connectionURL, Long connectio assertFalse(connection.isClosed()); // Execute query with physical connection - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITNightlyBigQueryTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITNightlyBigQueryTest.java index 2deefdd345cd..4c067fa41b08 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITNightlyBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITNightlyBigQueryTest.java @@ -781,7 +781,8 @@ public void testMultiStatementTransactionRollbackByUser() throws SQLException { bigQueryStatement.execute(createTransactionTable); BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.setAutoCommit(false); Statement statement = connection.createStatement(); assertTrue(connection.isTransactionStarted()); @@ -825,7 +826,8 @@ public void testMultiStatementTransactionDoesNotCommitWithoutCommit() throws SQL bigQueryStatement.execute(createTransactionTable); BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.setAutoCommit(false); Statement statement = connection.createStatement(); assertTrue(connection.isTransactionStarted()); @@ -866,7 +868,8 @@ public void testValidMultiStatementTransactionCommits() throws SQLException { bigQueryStatement.execute(createTransactionTable); BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.setAutoCommit(false); Statement statement = connection.createStatement(); assertTrue(connection.isTransactionStarted()); @@ -909,7 +912,8 @@ public void testConnectionWithMultipleTransactionCommits() throws SQLException { bigQueryStatement.execute(createTransactionTable); BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.setAutoCommit(false); Statement statement = connection.createStatement(); @@ -968,7 +972,8 @@ public void testTransactionRollbackOnError() throws SQLException { + " ROLLBACK TRANSACTION;\n" + "END;"; BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); Statement statement = connection.createStatement(); statement.execute(transactionOnError); @@ -1421,7 +1426,8 @@ public void testRollbackOnConnectionClosed() throws SQLException { bigQueryStatement.execute(createTransactionTable); BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.setAutoCommit(false); Statement statement = connection.createStatement(); assertTrue(connection.isTransactionStarted()); @@ -1526,7 +1532,8 @@ public void testConnectionClosedRollsBackStartedTransactions() throws SQLExcepti bigQueryStatement.execute(createTransactionTable); BigQueryConnection connection = - (BigQueryConnection) DriverManager.getConnection(session_enabled_connection_uri); + DriverManager.getConnection(session_enabled_connection_uri) + .unwrap(BigQueryConnection.class); connection.setAutoCommit(false); Statement statement = connection.createStatement(); assertTrue(connection.isTransactionStarted()); @@ -1556,16 +1563,13 @@ public void testStatelessQueryPathSmall() throws SQLException { Statement statement = bigQueryConnectionUseStateless.createStatement(); - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; ResultSet jsonResultSet = statement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); assertEquals(850, resultSetRowCount(jsonResultSet)); String queryEmpty = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 0"; + "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0"; ResultSet jsonResultSetEmpty = statement.executeQuery(queryEmpty); assertTrue(jsonResultSetEmpty.getClass().getName().contains("BigQueryJsonResultSet")); assertEquals(0, resultSetRowCount(jsonResultSetEmpty)); @@ -1574,8 +1578,7 @@ public void testStatelessQueryPathSmall() throws SQLException { @Test public void testFastQueryPathMedium() throws SQLException { - String query = - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 9000"; + String query = "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 9000"; ResultSet jsonResultSet = bigQueryStatement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); assertEquals(9000, resultSetRowCount(jsonResultSet)); @@ -1583,8 +1586,7 @@ public void testFastQueryPathMedium() throws SQLException { @Test public void testFastQueryPathLarge() throws SQLException { - String query = - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 18000"; + String query = "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 18000"; ResultSet jsonResultSet = bigQueryStatement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); assertEquals(18000, resultSetRowCount(jsonResultSet)); @@ -1631,9 +1633,7 @@ public void testNonEnabledUseLegacySQLThrowsSyntaxError() throws SQLException { @Test public void testFastQueryPathEmpty() throws SQLException { - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT" - + " 0"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 0"; ResultSet jsonResultSet = bigQueryStatement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); assertEquals(0, resultSetRowCount(jsonResultSet)); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITPSCBigQueryTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITPSCBigQueryTest.java index 2943f25e6818..55ceeb5d7dd5 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITPSCBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITPSCBigQueryTest.java @@ -68,10 +68,9 @@ public void testNoOverrideTimesOut() throws SQLException { assertFalse(connection.isClosed()); assertEquals( "APPLICATION_DEFAULT_CREDENTIALS", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; Statement statement = connection.createStatement(); assertThrows(BigQueryException.class, () -> statement.executeQuery(query)); } @@ -90,10 +89,9 @@ public void testValidADCAuthenticationInPSC() throws SQLException { assertFalse(connection.isClosed()); assertEquals( "APPLICATION_DEFAULT_CREDENTIALS", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); @@ -115,10 +113,9 @@ public void testValidOAuthType2AuthenticationInPSC() throws SQLException { assertFalse(connection.isClosed()); assertEquals( "PRE_GENERATED_TOKEN", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 850"; + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); @@ -143,7 +140,7 @@ public void testValidServiceAccountAuthenticationKeyFileInPSC() throws SQLExcept assertFalse(connection.isClosed()); assertEquals( "GOOGLE_SERVICE_ACCOUNT", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); connection.close(); } @@ -167,9 +164,8 @@ public void testValidServiceAccountAuthenticationViaEmailInPSC() throws SQLExcep assertFalse(connection.isClosed()); assertEquals( "GOOGLE_SERVICE_ACCOUNT", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); - String query = - "SELECT DISTINCT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 850"; + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); + String query = "SELECT DISTINCT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 850"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); assertTrue(jsonResultSet.getClass().getName().contains("BigQueryJsonResultSet")); @@ -287,12 +283,12 @@ public void testValidExternalAccountAuthenticationInPSC() throws SQLException { assertFalse(connection.isClosed()); assertEquals( "EXTERNAL_ACCOUNT_AUTH", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery( - "SELECT repository_name FROM `bigquery-public-data.samples.github_timeline` LIMIT 50"); + "SELECT word FROM `bigquery-public-data.samples.shakespeare` LIMIT 50"); assertNotNull(resultSet); connection.close(); diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITProxyBigQueryTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITProxyBigQueryTest.java index ed019cbe2bf5..e0b02ce66bd8 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITProxyBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITProxyBigQueryTest.java @@ -67,15 +67,14 @@ public void testValidAuthenticatedProxy() throws SQLException { assertFalse(connection.isClosed()); Statement statement = connection.createStatement(); boolean result = - statement.execute( - "Select * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"); + statement.execute("Select * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"); assertTrue(result); connection.close(); } @Test public void testAuthenticatedProxyWithOutAuthDetailsThrows() throws SQLException { - String query = "Select * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"; + String query = "Select * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"; String connection_uri = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + "ProjectId=" @@ -95,7 +94,7 @@ public void testAuthenticatedProxyWithOutAuthDetailsThrows() throws SQLException @Test public void testNonExistingProxyTimesOut() throws SQLException { - String query = "Select * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"; + String query = "Select * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"; String connection_uri = "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + "ProjectId=" @@ -131,8 +130,7 @@ public void testNonAuthenticatedProxy() throws SQLException { assertFalse(connection.isClosed()); Statement statement = connection.createStatement(); boolean result = - statement.execute( - "Select * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"); + statement.execute("Select * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"); assertTrue(result); connection.close(); } @@ -150,8 +148,7 @@ public void testValidNonProxyConnectionQueries() throws SQLException { assertFalse(connection.isClosed()); Statement statement = connection.createStatement(); boolean result = - statement.execute( - "Select * FROM `bigquery-public-data.samples.github_timeline` LIMIT 180"); + statement.execute("Select * FROM `bigquery-public-data.samples.shakespeare` LIMIT 180"); assertTrue(result); connection.close(); } diff --git a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITTPCBigQueryTest.java b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITTPCBigQueryTest.java index 32bddb001903..a3838e690741 100644 --- a/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITTPCBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITTPCBigQueryTest.java @@ -74,7 +74,7 @@ public void testServiceAccountAuthenticationViaEmail() throws SQLException { assertFalse(connection.isClosed()); assertEquals( "GOOGLE_SERVICE_ACCOUNT", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); String query = "SELECT 1"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); @@ -103,7 +103,7 @@ public void testValidApplicationDefaultCredentialsAuthentication() throws SQLExc assertFalse(connection.isClosed()); assertEquals( "APPLICATION_DEFAULT_CREDENTIALS", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); String query = "SELECT * FROM test.test;"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); @@ -163,7 +163,7 @@ public void testSimpleQueryReturns() throws SQLException { assertFalse(connection.isClosed()); assertEquals( "GOOGLE_SERVICE_ACCOUNT", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); String query = "SELECT * FROM test.test;"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query); @@ -192,7 +192,7 @@ public void testServiceAccountKeyFileReturns() throws SQLException { assertFalse(connection.isClosed()); assertEquals( "GOOGLE_SERVICE_ACCOUNT", - ((BigQueryConnection) connection).getAuthProperties().get("OAuthType")); + connection.unwrap(BigQueryConnection.class).getAuthProperties().get("OAuthType")); String query = "SELECT * FROM test.test;"; Statement statement = connection.createStatement(); ResultSet jsonResultSet = statement.executeQuery(query);