From dbb8f495284075b399d4597732476687cc7a0a41 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Mon, 14 Sep 2026 10:44:55 -0400 Subject: [PATCH 1/3] register shutdownhook and connection attempt --- .../cloud/bigquery/jdbc/BigQueryDriver.java | 48 +++++++++++++- .../jdbc/telemetry/v1/TelemetryBatcher.java | 2 +- .../jdbc/telemetry/v1/TelemetryManager.java | 62 ++++++++++++++++--- .../bigquery/jdbc/BigQueryDriverTest.java | 45 ++++++++++++++ .../telemetry/v1/TelemetryManagerTest.java | 44 +++++++++++++ 5 files changed, 190 insertions(+), 11 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java index 8c748b1f52bc..d2e66f8b25af 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java @@ -18,6 +18,9 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.AuthenticationType; +import com.google.cloud.bigquery.jdbc.telemetry.v1.Status; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; import io.grpc.LoadBalancerRegistry; import io.grpc.internal.PickFirstLoadBalancerProvider; @@ -124,6 +127,7 @@ public static BigQueryDriver getRegisteredDriver() throws IllegalStateException @Override public Connection connect(String url, Properties info) throws SQLException { LOG.finest("++enter++"); + AuthenticationType authType; try { if (acceptsURL(url)) { Properties connectInfo = info == null ? new Properties() : (Properties) info.clone(); @@ -132,6 +136,18 @@ public Connection connect(String url, Properties info) throws SQLException { String connectionUri = BigQueryJdbcUrlUtility.appendPropertiesToURL( url.substring(5), this.toString(), connectInfo); + + String telemetryOptOut = + BigQueryJdbcUrlUtility.parseUriPropertyWithoutValidation( + connectionUri, BigQueryJdbcUrlUtility.ENABLE_DIAGNOSTIC_TELEMETRY_PROPERTY_NAME); + + if (telemetryOptOut != null) { + connectInfo.setProperty( + BigQueryJdbcUrlUtility.ENABLE_DIAGNOSTIC_TELEMETRY_PROPERTY_NAME, telemetryOptOut); + } + authType = AuthenticationType.AUTHENTICATION_TYPE_UNSPECIFIED; + TelemetryManager.getInstance(connectInfo); + Level logLevel; String logPath; try { @@ -200,12 +216,28 @@ public Connection connect(String url, Properties info) throws SQLException { logLevel, logPath, this.toString()); - return BigQueryJdbcContextProxy.wrap(connection, Connection.class); + + Connection wrapped = BigQueryJdbcContextProxy.wrap(connection, Connection.class); + + authType = TelemetryManager.toAuthenticationType(ds.getOAuthType()); + TelemetryManager.recordConnectionAttempt(Status.STATUS_SUCCESS, 0, authType); + return wrapped; } else { return null; } - } catch (IOException e) { - LOG.warning("Getting a warning: " + e.getMessage()); + } catch (Throwable t) { + int errorCode = extractErrorCode(t); + TelemetryManager.recordConnectionAttempt(Status.STATUS_ERROR, errorCode, authType); + if (t instanceof SQLException) { + throw (SQLException) t; + } else if (t instanceof RuntimeException) { + throw (RuntimeException) t; + } else if (t instanceof IOException) { + LOG.warning("Getting a warning: " + t.getMessage()); + return null; + } else { + throw new BigQueryJdbcException("Failed to establish BigQuery connection", t); + } } return null; } @@ -284,4 +316,14 @@ public Logger getParentLogger() { private static class LazyHolder { static final BigQueryDriver INSTANCE = new BigQueryDriver(); } + + private static int extractErrorCode(Throwable t) { + if (t instanceof SQLException) { + int code = ((SQLException) t).getErrorCode(); + if (code != 0) { + return code; + } + } + return 1000; + } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryBatcher.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryBatcher.java index 7a9e804aa2be..39e28963cd02 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryBatcher.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryBatcher.java @@ -48,7 +48,7 @@ final class TelemetryBatcher implements AutoCloseable { private final boolean ownsExecutor; private final ReentrantLock flushLock = new ReentrantLock(); - // Live telemetry accumulators. Lock-free to eliminate object allocation and GC overhead. + // Live telemetry accumulator. Lock-free to eliminate object allocation and GC overhead. private ConcurrentHashMap metricsMap = new ConcurrentHashMap<>(); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java index 2837257b3534..c2bc05ef3c06 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java @@ -34,6 +34,8 @@ public final class TelemetryManager implements AutoCloseable { private static final Logger logger = new BigQueryJdbcCustomLogger(TelemetryManager.class.getName()); + private static volatile boolean shutdownHookRegistered = false; + private static volatile TelemetryManager instance; private static volatile boolean globallyDisabled = false; @@ -69,6 +71,7 @@ public static TelemetryManager getInstance(Properties properties) { } TelemetryManager localRef = instance; + if (localRef == null) { synchronized (TelemetryManager.class) { if (globallyDisabled) { @@ -82,6 +85,7 @@ public static TelemetryManager getInstance(Properties properties) { TelemetryBatcher batcher = new TelemetryBatcher(config, transport); localRef = new TelemetryManager(batcher); instance = localRef; + registerShutdownHook(); } } } @@ -161,24 +165,25 @@ static StatementType toStatementType(QueryStatistics.StatementType bqStatementTy return desc != null ? StatementType.valueOf(desc) : StatementType.STATEMENT_TYPE_OTHER; } - static AuthenticationType toAuthenticationType(int oauthType) { + public static AuthenticationType toAuthenticationType(int oauthType) { switch (oauthType) { case 0: return AuthenticationType.AUTHENTICATION_TYPE_SERVICE_ACCOUNT; case 1: return AuthenticationType.AUTHENTICATION_TYPE_USER_AUTHENTICATION; case 2: - return AuthenticationType.AUTHENTICATION_TYPE_APPLICATION_DEFAULT_CREDENTIALS; + return AuthenticationType.AUTHENTICATION_TYPE_TOKEN; case 3: - return AuthenticationType.AUTHENTICATION_TYPE_EXTERNAL; + return AuthenticationType.AUTHENTICATION_TYPE_APPLICATION_DEFAULT_CREDENTIALS; case 4: - return AuthenticationType.AUTHENTICATION_TYPE_TOKEN; + return AuthenticationType.AUTHENTICATION_TYPE_EXTERNAL; default: return AuthenticationType.AUTHENTICATION_TYPE_CUSTOM; } } - static void recordConnectionAttempt(Status status, int errorCode, AuthenticationType authType) { + public static void recordConnectionAttempt( + Status status, int errorCode, AuthenticationType authType) { runSafely( () -> { TelemetryManager mgr = instance; @@ -194,7 +199,7 @@ static void recordConnectionAttempt(Status status, int errorCode, Authentication }); } - static void recordStatementExecution( + public static void recordStatementExecution( StatementType statementType, QueryApiType apiType, Status status, @@ -217,7 +222,7 @@ static void recordStatementExecution( }); } - static void recordFeatureUsage(DriverFeature feature, String customFeatureName) { + public static void recordFeatureUsage(DriverFeature feature, String customFeatureName) { runSafely( () -> { TelemetryManager mgr = instance; @@ -231,4 +236,47 @@ static void recordFeatureUsage(DriverFeature feature, String customFeatureName) } }); } + + public static void recordError(int errorCode, int errorXdbcCode, String methodName) { + runSafely( + () -> { + TelemetryManager mgr = instance; + if (mgr != null && mgr.getBatcher() != null) { + mgr.getBatcher() + .offer( + ErrorMetric.newBuilder() + .setErrorCode(errorCode) + .setErrorXdbcCode(errorXdbcCode) + .setMethodName(methodName == null ? "" : methodName) + .build()); + } + }); + } + + private static void registerShutdownHook() { + if (!shutdownHookRegistered) { + synchronized (TelemetryManager.class) { + if (!shutdownHookRegistered) { + try { + Runtime.getRuntime() + .addShutdownHook( + new Thread( + () -> { + try { + closeInstance(); + } catch (Throwable t) { + logger.warning("Error closing TelemetryManager during JVM shutdown"); + } + }, + "bigquery-jdbc-telemetry-shutdown-hook")); + shutdownHookRegistered = true; + } catch (IllegalStateException e) { + // Thrown if the JVM is already in the process of shutting down + } catch (SecurityException e) { + logger.warning("SecurityManager prevented registering telemetry shutdown hook"); + } + } + } + } + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java index 8acbc5abb8dc..38c5a49bc77d 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java @@ -187,4 +187,49 @@ public void testInvalidLogLevelExceptionIsLogged() { && r.getMessage().contains("Failed to parse connection URL properties")); assertThat(foundSevere).isTrue(); } + + @Test + public void testConnect_recordsFailedConnectionWithErrorCode() throws Exception { + ByteArrayOutputStream capturedBytes = new ByteArrayOutputStream(); + MockHttpTransport mockTransport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) { + return new MockLowLevelHttpRequest(url) { + @Override + public LowLevelHttpResponse execute() throws IOException { + if (getStreamingContent() != null) { + getStreamingContent().writeTo(capturedBytes); + } + MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); + response.setStatusCode(200); + return response; + } + }; + } + }; + + TelemetryConfiguration config = TelemetryConfiguration.newBuilder().setEnabled(true).build(); + TelemetryManager.init(config, new ClearcutTransport(mockTransport, config)); + + // Calling connect with invalid OAuthType triggers SQLException + Assertions.assertThrows( + SQLException.class, + () -> + bigQueryDriver.connect( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=invalid;", + new Properties())); + + // Flush and verify error was recorded + TelemetryManager.getInstance().getBatcher().flush(); + + LogRequest logRequest = LogRequest.parseFrom(capturedBytes.toByteArray()); + TelemetryPayload payload = + TelemetryPayload.parseFrom(logRequest.getLogEvent(0).getSourceExtension()); + + assertThat(payload.getConnectionAttemptsCount()).isEqualTo(1); + ConnectionAttempt attempt = payload.getConnectionAttempts(0); + assertThat(attempt.getStatus()).isEqualTo(Status.STATUS_ERROR); + assertThat(attempt.getErrorCode()).isEqualTo(1000); // generic fallback code + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java index 007236f25962..73b4ddd5aa7c 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java @@ -215,4 +215,48 @@ public void testGlobalKillSwitch() { assertNull(TelemetryManager.getInstance()); assertNull(TelemetryManager.getInstance(props1)); } + + @Test + public void testConnect_recordsSuccessfulConnectionTelemetry() throws SQLException { + TelemetryManager.closeInstance(); + Connection connection = + bigQueryDriver.connect( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "OAuthType=2;ProjectId=MyBigQueryProject;" + + "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;" + + "OAuthClientSecret=redactedToken;", + new Properties()); + assertThat(connection).isNotNull(); + assertThat(connection.isClosed()).isFalse(); + // Verify TelemetryManager is initialized and recorded the connection + assertThat(TelemetryManager.isInitialized()).isTrue(); + } + + @Test + public void testConnect_recordsFailedConnectionTelemetry() { + TelemetryManager.closeInstance(); + // Malformed URL causing DataSource parsing failure + Assertions.assertThrows( + SQLException.class, + () -> + bigQueryDriver.connect( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=invalid;", + new Properties())); + assertThat(TelemetryManager.isInitialized()).isTrue(); + } + + @Test + public void testConnect_optOut_noTelemetryRecorded() throws SQLException { + TelemetryManager.closeInstance(); + Connection connection = + bigQueryDriver.connect( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "OAuthType=2;ProjectId=MyBigQueryProject;" + + "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;" + + "OAuthClientSecret=redactedToken;EnableDiagnosticTelemetry=0;", + new Properties()); + assertThat(connection).isNotNull(); + // Since opt-out was requested, TelemetryManager should NOT be initialized + assertThat(TelemetryManager.isInitialized()).isFalse(); + } } From fce8bfe015192291d400901a777223545387184d Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Mon, 14 Sep 2026 16:05:59 -0400 Subject: [PATCH 2/3] record Statement Execution --- .../cloud/bigquery/jdbc/BigQueryDriver.java | 4 +- .../jdbc/BigQueryPreparedStatement.java | 20 ++++++ .../bigquery/jdbc/BigQueryStatement.java | 53 +++++++++++++- .../jdbc/telemetry/v1/TelemetryManager.java | 42 ++++++++++- .../bigquery/jdbc/BigQueryDriverTest.java | 70 +++++++++---------- .../telemetry/v1/TelemetryManagerTest.java | 52 ++------------ 6 files changed, 153 insertions(+), 88 deletions(-) diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java index d2e66f8b25af..a6c945aec66a 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java @@ -127,7 +127,7 @@ public static BigQueryDriver getRegisteredDriver() throws IllegalStateException @Override public Connection connect(String url, Properties info) throws SQLException { LOG.finest("++enter++"); - AuthenticationType authType; + AuthenticationType authType = AuthenticationType.AUTHENTICATION_TYPE_UNSPECIFIED; try { if (acceptsURL(url)) { Properties connectInfo = info == null ? new Properties() : (Properties) info.clone(); @@ -145,7 +145,6 @@ public Connection connect(String url, Properties info) throws SQLException { connectInfo.setProperty( BigQueryJdbcUrlUtility.ENABLE_DIAGNOSTIC_TELEMETRY_PROPERTY_NAME, telemetryOptOut); } - authType = AuthenticationType.AUTHENTICATION_TYPE_UNSPECIFIED; TelemetryManager.getInstance(connectInfo); Level logLevel; @@ -239,7 +238,6 @@ public Connection connect(String url, Properties info) throws SQLException { throw new BigQueryJdbcException("Failed to establish BigQuery connection", t); } } - return null; } /** diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java index 014717fd0646..55cd7545810b 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java @@ -28,6 +28,8 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.StatementExecution; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.cloud.bigquery.storage.v1.BatchCommitWriteStreamsRequest; import com.google.cloud.bigquery.storage.v1.BatchCommitWriteStreamsResponse; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; @@ -319,6 +321,14 @@ public int[] executeBatch() throws SQLException { return result; } if (useWriteAPI()) { + long startTime = System.currentTimeMillis(); + StatementExecution.Builder writeApiExecutionBuilder = + StatementExecution.newBuilder() + .setStatementType( + com.google.cloud.bigquery.jdbc.telemetry.v1.StatementType.STATEMENT_TYPE_INSERT) + .setQueryApiType( + com.google.cloud.bigquery.jdbc.telemetry.v1.QueryApiType + .QUERY_API_TYPE_WRITE_API); try (BigQueryWriteClient writeClient = this.connection.getBigQueryWriteClient()) { LOG.info("Using Write API for bulk INSERT operation."); ArrayList currentParameterList = this.batchParameters.peek(); @@ -331,10 +341,20 @@ public int[] executeBatch() throws SQLException { long rowCount = bulkInsertWithWriteAPI(writeClient); int[] insertArray = new int[Math.toIntExact(rowCount)]; Arrays.fill(insertArray, 1); + + writeApiExecutionBuilder.setStatus( + com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_SUCCESS); + return insertArray; } catch (DescriptorValidationException | IOException | InterruptedException e) { + writeApiExecutionBuilder + .setStatus(com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_ERROR) + .setErrorCode(TelemetryManager.extractErrorCode(e)); throw new BigQueryJdbcRuntimeException("Failed to execute batch with Write API", e); + } finally { + long durationMs = System.currentTimeMillis() - startTime; + TelemetryManager.recordStatementExecution(writeApiExecutionBuilder, durationMs); } } else { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index d1348f6d8091..437bc72d9ec3 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -48,6 +48,9 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlSyntaxErrorException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.QueryApiType; +import com.google.cloud.bigquery.jdbc.telemetry.v1.StatementExecution; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch; import com.google.cloud.bigquery.storage.v1.ArrowSchema; import com.google.cloud.bigquery.storage.v1.ArrowSerializationOptions; @@ -149,6 +152,8 @@ public class BigQueryStatement extends BigQueryNoOpsStatement { private static final ThreadFactory JDBC_THREAD_FACTORY = new BigQueryThreadFactory("BigQuery-Thread-"); + protected StatementExecution.Builder currentExecutionBuilder = StatementExecution.newBuilder(); + static { BigQueryDaemonPollingTask.startGcDaemonTask( referenceQueueArrowRs, @@ -171,6 +176,8 @@ private void resetStatementFields() { this.parentJobId = null; this.currentJobIdIndex = -1; this.currentUpdateCount = -1; + + this.currentExecutionBuilder = StatementExecution.newBuilder(); } private BigQuerySettings generateBigQuerySettings() { @@ -574,6 +581,7 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration) if (result instanceof TableResult) { TableResult tableResult = (TableResult) result; saveSessionIdIfPresent(tableResult); + this.currentExecutionBuilder.setQueryApiType(QueryApiType.QUERY_API_TYPE_JOBLESS_QUERY); return new ExecuteResult(tableResult, null); } @@ -604,6 +612,7 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration) job = refreshedJob; } } + this.currentExecutionBuilder.setQueryApiType(QueryApiType.QUERY_API_TYPE_STANDARD_REST_API); return new ExecuteResult(tableResult, job); } @@ -655,19 +664,54 @@ void runQuery(String query, QueryJobConfiguration jobConfiguration) jobConfiguration.toBuilder().setJobTimeoutMs(Long.valueOf(queryTimeout) * 1000).build(); } + long startTime = System.currentTimeMillis(); + try { resetStatementFields(); ExecuteResult executeResult = executeJob(jobConfiguration); StatementType statementType = getStatementType(executeResult); + + this.currentExecutionBuilder.setStatementType( + TelemetryManager.toStatementType(statementType)); + SqlType queryType = getQueryType(jobConfiguration, statementType); handleQueryResult(query, executeResult.tableResult, queryType, executeResult.job); + + this.currentExecutionBuilder.setStatus( + com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_SUCCESS); + } catch (InterruptedException ex) { + this.currentExecutionBuilder + .setStatus(com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_ERROR) + .setErrorCode(TelemetryManager.extractErrorCode(ex)); throw new BigQueryJdbcRuntimeException("Interrupted during runQuery", ex); } catch (BigQueryException ex) { + this.currentExecutionBuilder + .setStatus( + isCanceled + ? com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_CANCELLED + : com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_ERROR) + .setErrorCode(TelemetryManager.extractErrorCode(ex)); if (ex.getMessage().contains("Syntax error")) { throw new BigQueryJdbcSqlSyntaxErrorException("BigQueryException during runQuery", ex); } throw new BigQueryJdbcException("BigQueryException during runQuery", ex); + } finally { + long durationMs = System.currentTimeMillis() - startTime; + + // Safety net: If an uncaught RuntimeException occurred before setting STATUS_SUCCESS, + // mark it as an ERROR so failures are never reported as UNSPECIFIED. + if (this.currentExecutionBuilder.getStatus() + == com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_UNSPECIFIED) { + this.currentExecutionBuilder + .setStatus( + isCanceled + ? com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_CANCELLED + : com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_ERROR) + .setErrorCode(1000); + } + + TelemetryManager.recordStatementExecution(this.currentExecutionBuilder, durationMs); } } @@ -831,7 +875,6 @@ private QueryStatistics getQueryStatisticsFromJob(TableResult results, Job job) } private void updateAffectedRowCount(Long count) throws SQLException { - // TODO(neenu): check if this need to be closed vs removed) if (this.currentResultSet != null) { try { this.currentResultSet.close(); @@ -1077,6 +1120,7 @@ void processQueryResponse(String query, TableResult results, Job job) throws SQL try { LOG.info("Using ReadAPI to read the data."); resultSet = processArrowResultSet(results, job); + this.currentExecutionBuilder.setQueryApiType(QueryApiType.QUERY_API_TYPE_READ_API); } catch (SQLException e) { if (!isPermissionDeniedException(e)) { throw e; @@ -1088,6 +1132,13 @@ void processQueryResponse(String query, TableResult results, Job job) throws SQL if (resultSet == null) { LOG.info("Using Standard API to read the data."); resultSet = processJsonResultSet(results, job); + + // Jobless vs Standard REST + if (jobId == null) { + this.currentExecutionBuilder.setQueryApiType(QueryApiType.QUERY_API_TYPE_JOBLESS_QUERY); + } else { + this.currentExecutionBuilder.setQueryApiType(QueryApiType.QUERY_API_TYPE_STANDARD_REST_API); + } } this.currentResultSet = resultSet; this.currentUpdateCount = -1; diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java index c2bc05ef3c06..a695d70435a5 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java @@ -16,9 +16,11 @@ package com.google.cloud.bigquery.jdbc.telemetry.v1; +import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.JobStatistics.QueryStatistics; import com.google.cloud.bigquery.jdbc.BigQueryJdbcCustomLogger; import com.google.protobuf.Descriptors.EnumValueDescriptor; +import java.sql.SQLException; import java.util.Properties; import java.util.logging.Level; import java.util.logging.Logger; @@ -154,7 +156,7 @@ static synchronized void resetGlobalDisableForTest() { globallyDisabled = false; } - static StatementType toStatementType(QueryStatistics.StatementType bqStatementType) { + public static StatementType toStatementType(QueryStatistics.StatementType bqStatementType) { if (bqStatementType == null) { return StatementType.STATEMENT_TYPE_UNSPECIFIED; } @@ -222,6 +224,20 @@ public static void recordStatementExecution( }); } + public static void recordStatementExecution( + StatementExecution.Builder statementExecutionBuilder, long durationMs) { + if (statementExecutionBuilder == null) { + return; + } + runSafely( + () -> { + TelemetryManager mgr = instance; + if (mgr != null && mgr.getBatcher() != null) { + mgr.getBatcher().offer(statementExecutionBuilder.build(), durationMs); + } + }); + } + public static void recordFeatureUsage(DriverFeature feature, String customFeatureName) { runSafely( () -> { @@ -253,6 +269,30 @@ public static void recordError(int errorCode, int errorXdbcCode, String methodNa }); } + /** + * Extracts the numeric error code from the throwable chain. Traverses causes to unpack + * BigQueryException (HTTP status codes) or SQLException error codes. Returns 1000 as the fallback + * driver error code. + */ + public static int extractErrorCode(Throwable t) { + while (t != null) { + if (t instanceof BigQueryException) { + int code = ((BigQueryException) t).getCode(); + if (code != 0) { + return code; + } + } + if (t instanceof SQLException) { + int code = ((SQLException) t).getErrorCode(); + if (code != 0) { + return code; + } + } + t = t.getCause(); + } + return 1000; + } + private static void registerShutdownHook() { if (!shutdownHookRegistered) { synchronized (TelemetryManager.class) { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java index 38c5a49bc77d..9f0b07ace74d 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDriverTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.mockito.Mockito.mock; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility; import io.opentelemetry.api.OpenTelemetry; import java.sql.Connection; @@ -189,47 +190,46 @@ public void testInvalidLogLevelExceptionIsLogged() { } @Test - public void testConnect_recordsFailedConnectionWithErrorCode() throws Exception { - ByteArrayOutputStream capturedBytes = new ByteArrayOutputStream(); - MockHttpTransport mockTransport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) { - return new MockLowLevelHttpRequest(url) { - @Override - public LowLevelHttpResponse execute() throws IOException { - if (getStreamingContent() != null) { - getStreamingContent().writeTo(capturedBytes); - } - MockLowLevelHttpResponse response = new MockLowLevelHttpResponse(); - response.setStatusCode(200); - return response; - } - }; - } - }; - - TelemetryConfiguration config = TelemetryConfiguration.newBuilder().setEnabled(true).build(); - TelemetryManager.init(config, new ClearcutTransport(mockTransport, config)); - - // Calling connect with invalid OAuthType triggers SQLException + public void testConnect_recordsSuccessfulConnectionTelemetry() throws SQLException { + TelemetryManager.closeInstance(); + Connection connection = + bigQueryDriver.connect( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "OAuthType=2;ProjectId=MyBigQueryProject;" + + "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;" + + "OAuthClientSecret=redactedToken;", + new Properties()); + assertThat(connection).isNotNull(); + assertThat(connection.isClosed()).isFalse(); + // Verify TelemetryManager is initialized and recorded the connection + assertThat(TelemetryManager.isInitialized()).isTrue(); + } + + @Test + public void testConnect_recordsFailedConnectionTelemetry() { + TelemetryManager.closeInstance(); + // Malformed URL causing DataSource parsing failure Assertions.assertThrows( SQLException.class, () -> bigQueryDriver.connect( "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=invalid;", new Properties())); + assertThat(TelemetryManager.isInitialized()).isTrue(); + } - // Flush and verify error was recorded - TelemetryManager.getInstance().getBatcher().flush(); - - LogRequest logRequest = LogRequest.parseFrom(capturedBytes.toByteArray()); - TelemetryPayload payload = - TelemetryPayload.parseFrom(logRequest.getLogEvent(0).getSourceExtension()); - - assertThat(payload.getConnectionAttemptsCount()).isEqualTo(1); - ConnectionAttempt attempt = payload.getConnectionAttempts(0); - assertThat(attempt.getStatus()).isEqualTo(Status.STATUS_ERROR); - assertThat(attempt.getErrorCode()).isEqualTo(1000); // generic fallback code + @Test + public void testConnect_optOut_noTelemetryRecorded() throws SQLException { + TelemetryManager.closeInstance(); + Connection connection = + bigQueryDriver.connect( + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" + + "OAuthType=2;ProjectId=MyBigQueryProject;" + + "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;" + + "OAuthClientSecret=redactedToken;EnableDiagnosticTelemetry=0;", + new Properties()); + assertThat(connection).isNotNull(); + // Since opt-out was requested, TelemetryManager should NOT be initialized + assertThat(TelemetryManager.isInitialized()).isFalse(); } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java index 73b4ddd5aa7c..cf467afaf319 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManagerTest.java @@ -162,12 +162,12 @@ public void testToAuthenticationType() { AuthenticationType.AUTHENTICATION_TYPE_USER_AUTHENTICATION, TelemetryManager.toAuthenticationType(1)); assertEquals( - AuthenticationType.AUTHENTICATION_TYPE_APPLICATION_DEFAULT_CREDENTIALS, - TelemetryManager.toAuthenticationType(2)); + AuthenticationType.AUTHENTICATION_TYPE_TOKEN, TelemetryManager.toAuthenticationType(2)); assertEquals( - AuthenticationType.AUTHENTICATION_TYPE_EXTERNAL, TelemetryManager.toAuthenticationType(3)); + AuthenticationType.AUTHENTICATION_TYPE_APPLICATION_DEFAULT_CREDENTIALS, + TelemetryManager.toAuthenticationType(3)); assertEquals( - AuthenticationType.AUTHENTICATION_TYPE_TOKEN, TelemetryManager.toAuthenticationType(4)); + AuthenticationType.AUTHENTICATION_TYPE_EXTERNAL, TelemetryManager.toAuthenticationType(4)); assertEquals( AuthenticationType.AUTHENTICATION_TYPE_CUSTOM, TelemetryManager.toAuthenticationType(5)); @@ -215,48 +215,4 @@ public void testGlobalKillSwitch() { assertNull(TelemetryManager.getInstance()); assertNull(TelemetryManager.getInstance(props1)); } - - @Test - public void testConnect_recordsSuccessfulConnectionTelemetry() throws SQLException { - TelemetryManager.closeInstance(); - Connection connection = - bigQueryDriver.connect( - "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" - + "OAuthType=2;ProjectId=MyBigQueryProject;" - + "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;" - + "OAuthClientSecret=redactedToken;", - new Properties()); - assertThat(connection).isNotNull(); - assertThat(connection.isClosed()).isFalse(); - // Verify TelemetryManager is initialized and recorded the connection - assertThat(TelemetryManager.isInitialized()).isTrue(); - } - - @Test - public void testConnect_recordsFailedConnectionTelemetry() { - TelemetryManager.closeInstance(); - // Malformed URL causing DataSource parsing failure - Assertions.assertThrows( - SQLException.class, - () -> - bigQueryDriver.connect( - "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=invalid;", - new Properties())); - assertThat(TelemetryManager.isInitialized()).isTrue(); - } - - @Test - public void testConnect_optOut_noTelemetryRecorded() throws SQLException { - TelemetryManager.closeInstance(); - Connection connection = - bigQueryDriver.connect( - "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;" - + "OAuthType=2;ProjectId=MyBigQueryProject;" - + "OAuthAccessToken=redactedToken;OAuthClientId=redactedToken;" - + "OAuthClientSecret=redactedToken;EnableDiagnosticTelemetry=0;", - new Properties()); - assertThat(connection).isNotNull(); - // Since opt-out was requested, TelemetryManager should NOT be initialized - assertThat(TelemetryManager.isInitialized()).isFalse(); - } } From c1229ea2b10230fb44790845593227f3b1f2a466 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Wed, 16 Sep 2026 14:05:50 -0400 Subject: [PATCH 3/3] add DriverFeature instrumentation --- java-bigquery-jdbc/README.MD | 16 ++------------ .../bigquery/jdbc/BigQueryBaseResultSet.java | 7 ++++++ .../bigquery/jdbc/BigQueryConnection.java | 22 +++++++++++++++++++ .../cloud/bigquery/jdbc/BigQueryDriver.java | 12 +--------- .../jdbc/BigQueryParameterHandler.java | 12 ++++++++-- .../jdbc/BigQueryPooledConnection.java | 6 ++++- .../jdbc/BigQueryPreparedStatement.java | 13 +++++++++++ .../bigquery/jdbc/BigQueryStatement.java | 7 ++++++ .../jdbc/PooledConnectionDataSource.java | 4 ++++ .../jdbc/telemetry/v1/TelemetryManager.java | 7 +++++- 10 files changed, 77 insertions(+), 29 deletions(-) diff --git a/java-bigquery-jdbc/README.MD b/java-bigquery-jdbc/README.MD index cd5cde8e795f..34bbb7b38c54 100644 --- a/java-bigquery-jdbc/README.MD +++ b/java-bigquery-jdbc/README.MD @@ -92,10 +92,6 @@ a connection property (`EnableDiagnosticTelemetry=0`), environment variable, or system property, the service immediately shuts down and remains permanently off for the remainder of the JVM lifecycle. It will only restart when the JVM restarts. -* Global Worker Settings: Because the telemetry batcher is a shared global -resource, properties that configure its behavior (`TelemetryUploadInterval` -and `TelemetryBatchSize`) are exclusively evaluated upon establishing the first -connection. Subsequent overrides for these specific properties are ignored. #### Configuration Keys @@ -103,16 +99,8 @@ To configure these properties, you can use the following keys: * Opt-Out Control: * Connection Property: `EnableDiagnosticTelemetry=0;` - * Environment Variable: `GOOGLE_BIGQUERY_JDBC_TELEMETRY_ENABLED=false` - * System Property: `-Dgoogle.bigquery.jdbc.telemetry.enabled=false` -* Upload Interval: - * Connection Property: `TelemetryUploadInterval=300000;` - * Environment Variable: `GOOGLE_BIGQUERY_JDBC_TELEMETRY_INTERVAL_MS=300000` - * System Property: `-Dgoogle.bigquery.jdbc.telemetry.interval_ms=300000` -* Batch Size: - * Connection Property: `TelemetryBatchSize=5000;` - * Environment Variable: `GOOGLE_BIGQUERY_JDBC_TELEMETRY_BATCH_SIZE=5000` - * System Property: `-Dgoogle.bigquery.jdbc.telemetry.batch_size=5000` + * Environment Variable: `BIGQUERY_JDBC_TELEMETRY_ENABLED=false` + * System Property: `-DBIGQUERY_JDBC_TELEMETRY_ENABLED=false` ## Developer Guide diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java index de137e57933e..d1f05679263c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java @@ -28,6 +28,8 @@ import com.google.cloud.bigquery.StandardSQLTypeName; import com.google.cloud.bigquery.exception.BigQueryConversionException; import com.google.cloud.bigquery.exception.BigQueryJdbcException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.context.Context; @@ -254,6 +256,11 @@ public ResultSetMetaData getMetaData() throws SQLException { metaData = BigQueryResultSetMetadata.of(this.schema.getFields(), this.statement); } } + + TelemetryManager.recordFeatureUsage( + DriverFeature.DRIVER_FEATURE_METADATA_RETRIEVAL, + "DRIVER_FEATURE_RESULTSET_METADATA_RETRIEVAL"); + return BigQueryJdbcContextProxy.wrap(metaData, ResultSetMetaData.class, connectionId); } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 3d2ae4e9e02e..2d25840a6454 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -39,6 +39,8 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; @@ -482,6 +484,9 @@ public Statement createStatement() throws SQLException { BigQueryStatement currentStatement = new BigQueryStatement(this); LOG.fine("Statement %s created.", currentStatement); addOpenStatements(currentStatement); + + TelemetryManager.recordFeatureUsage( + DriverFeature.DRIVER_FEATURE_CUSTOM, "DRIVER_FEATURE_REGULAR_STATEMENT"); return currentStatement; } @@ -540,6 +545,9 @@ public PreparedStatement prepareStatement(String sql) throws SQLException { PreparedStatement currentStatement = new BigQueryPreparedStatement(this, sql); LOG.fine("Prepared Statement %s created.", currentStatement); addOpenStatements(currentStatement); + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_PREPARED_STATEMENT); + return currentStatement; } @@ -680,6 +688,8 @@ private void beginTransaction() { updateSessionInfo(transactionBeginJob.getStatistics().getSessionInfo().getSessionId()); } this.transactionStarted = true; + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_TRANSACTIONS); } catch (InterruptedException ex) { throw new BigQueryJdbcRuntimeException("Failed to begin transaction", ex); } @@ -914,6 +924,12 @@ public void setAutoCommit(boolean autoCommit) throws SQLException { if (!this.autoCommit) { beginTransaction(); } + + if (autoCommit) { + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_AUTOCOMMIT_ENABLED); + } else { + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_AUTOCOMMIT_DISABLED); + } } @Override @@ -971,6 +987,9 @@ public DatabaseMetaData getMetaData() throws SQLException { if (databaseMetaData == null) { databaseMetaData = new BigQueryDatabaseMetaData(this); } + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_METADATA_RETRIEVAL); + return databaseMetaData; } @@ -1463,6 +1482,9 @@ public CallableStatement prepareCall(String sql) throws SQLException { CallableStatement currentStatement = new BigQueryCallableStatement(this, sql); LOG.fine("Callable Statement %s created.", currentStatement); addOpenStatements(currentStatement); + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_CALLABLE_STATEMENT); + return currentStatement; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java index a6c945aec66a..2509ed527c4c 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java @@ -225,7 +225,7 @@ public Connection connect(String url, Properties info) throws SQLException { return null; } } catch (Throwable t) { - int errorCode = extractErrorCode(t); + int errorCode = TelemetryManager.extractErrorCode(t); TelemetryManager.recordConnectionAttempt(Status.STATUS_ERROR, errorCode, authType); if (t instanceof SQLException) { throw (SQLException) t; @@ -314,14 +314,4 @@ public Logger getParentLogger() { private static class LazyHolder { static final BigQueryDriver INSTANCE = new BigQueryDriver(); } - - private static int extractErrorCode(Throwable t) { - if (t instanceof SQLException) { - int code = ((SQLException) t).getErrorCode(); - if (code != 0) { - return code; - } - } - return 1000; - } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryParameterHandler.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryParameterHandler.java index e40e3fcf28ef..c320f7249bda 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryParameterHandler.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryParameterHandler.java @@ -21,6 +21,8 @@ import com.google.cloud.bigquery.StandardSQLTypeName; import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import java.math.BigInteger; import java.sql.SQLException; import java.util.ArrayList; @@ -108,8 +110,7 @@ static Object formatValueForQueryParameter(Object parameterValue, StandardSQLTyp return parameterValue; } - void setParameter(int parameterIndex, Object value, Class type) - throws BigQueryJdbcSqlFeatureNotSupportedException { + void setParameter(int parameterIndex, Object value, Class type) { LOG.finest("++enter++"); LOG.finest("setParameter called by : %s", type.getName()); checkValidIndex(parameterIndex); @@ -133,6 +134,8 @@ void setParameter(int parameterIndex, Object value, Class type) parameter.setParamType(BigQueryStatementParameterType.UNSPECIFIED); parameter.setScale(-1); + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_PARAMETER_BINDING); + LOG.finest("Parameter set { %s }", parameter.toString()); } @@ -215,6 +218,9 @@ void setParameter( if (parameter.getIndex() == -1) { parametersList.add(parameter); } + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_PARAMETER_BINDING); + LOG.finest("Parameter set { %s }", parameter.toString()); } @@ -248,6 +254,8 @@ void setParameter( parameter.setParamType(paramType); parameter.setScale(scale); + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_PARAMETER_BINDING); + LOG.finest("Parameter set { %s }", parameter.toString()); } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPooledConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPooledConnection.java index 99dea20e21ce..ad62ee76bcdb 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPooledConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPooledConnection.java @@ -16,8 +16,11 @@ package com.google.cloud.bigquery.jdbc; +import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.common.annotations.VisibleForTesting; import java.sql.Connection; +import java.sql.DatabaseMetaData; import java.sql.SQLException; import java.util.UUID; import java.util.concurrent.Executor; @@ -231,7 +234,8 @@ public void rollback() throws SQLException { } @Override - public java.sql.DatabaseMetaData getMetaData() throws SQLException { + public DatabaseMetaData getMetaData() throws SQLException { + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_METADATA_RETRIEVAL); return bqConnectionDelegate.getMetaData(); } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java index 55cd7545810b..0c82b0956763 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java @@ -28,6 +28,7 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcException; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; import com.google.cloud.bigquery.jdbc.telemetry.v1.StatementExecution; import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.cloud.bigquery.storage.v1.BatchCommitWriteStreamsRequest; @@ -320,6 +321,7 @@ public int[] executeBatch() throws SQLException { if (this.batchParameters.isEmpty()) { return result; } + if (useWriteAPI()) { long startTime = System.currentTimeMillis(); StatementExecution.Builder writeApiExecutionBuilder = @@ -351,10 +353,14 @@ public int[] executeBatch() throws SQLException { writeApiExecutionBuilder .setStatus(com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_ERROR) .setErrorCode(TelemetryManager.extractErrorCode(e)); + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } throw new BigQueryJdbcRuntimeException("Failed to execute batch with Write API", e); } finally { long durationMs = System.currentTimeMillis() - startTime; TelemetryManager.recordStatementExecution(writeApiExecutionBuilder, durationMs); + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_BATCH_OPERATIONS); } } else { @@ -385,6 +391,8 @@ public int[] executeBatch() throws SQLException { throw new BigQueryJdbcRuntimeException("Interrupted during individual INSERT batch", ex); } catch (SQLException e) { throw new BigQueryJdbcException("SQL error during individual INSERT batch", e); + } finally { + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_BATCH_OPERATIONS); } } } @@ -556,6 +564,11 @@ public ResultSetMetaData getMetaData() throws SQLException { if (this.insertSchema != null) { return BigQueryResultSetMetadata.of(this.insertSchema.getFields(), this); } + + TelemetryManager.recordFeatureUsage( + DriverFeature.DRIVER_FEATURE_METADATA_RETRIEVAL, + "DRIVER_FEATURE_RESULTSET_METADATA_RETRIEVAL"); + return null; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 437bc72d9ec3..2bce2ec57359 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java @@ -48,6 +48,7 @@ import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlSyntaxErrorException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; import com.google.cloud.bigquery.jdbc.telemetry.v1.QueryApiType; import com.google.cloud.bigquery.jdbc.telemetry.v1.StatementExecution; import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; @@ -494,6 +495,8 @@ public void cancel() throws SQLException { // If a ResultSet exists, then it will be closed as well, closing the // ownedThreads closeStatementResources(); + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_STATEMENT_CANCEL); } @Override @@ -684,6 +687,7 @@ void runQuery(String query, QueryJobConfiguration jobConfiguration) this.currentExecutionBuilder .setStatus(com.google.cloud.bigquery.jdbc.telemetry.v1.Status.STATUS_ERROR) .setErrorCode(TelemetryManager.extractErrorCode(ex)); + Thread.currentThread().interrupt(); throw new BigQueryJdbcRuntimeException("Interrupted during runQuery", ex); } catch (BigQueryException ex) { this.currentExecutionBuilder @@ -1796,6 +1800,9 @@ private int[] executeBatchImpl(String combinedQueries) throws SQLException { } clearBatch(); + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_BATCH_OPERATIONS); + return result; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java index 7de4516427c3..0a986bd860b5 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/PooledConnectionDataSource.java @@ -17,6 +17,8 @@ package com.google.cloud.bigquery.jdbc; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; +import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; +import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager; import com.google.common.annotations.VisibleForTesting; import java.sql.Connection; import java.sql.SQLException; @@ -54,6 +56,8 @@ public PooledConnection getPooledConnection() throws SQLException { } BigQueryPooledConnection bqPooledConnection = new BigQueryPooledConnection(physicalConnection); bqPooledConnection.addConnectionEventListener(connectionPoolManager); + + TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_CONNECTION_POOLING); return bqPooledConnection; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java index a695d70435a5..1df96f58dfe2 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryManager.java @@ -253,6 +253,10 @@ public static void recordFeatureUsage(DriverFeature feature, String customFeatur }); } + public static void recordFeatureUsage(DriverFeature feature) { + recordFeatureUsage(feature, null); + } + public static void recordError(int errorCode, int errorXdbcCode, String methodName) { runSafely( () -> { @@ -275,7 +279,8 @@ public static void recordError(int errorCode, int errorXdbcCode, String methodNa * driver error code. */ public static int extractErrorCode(Throwable t) { - while (t != null) { + int depth = 0; + while (t != null && depth++ < 20) { if (t instanceof BigQueryException) { int code = ((BigQueryException) t).getCode(); if (code != 0) {