From dbb8f495284075b399d4597732476687cc7a0a41 Mon Sep 17 00:00:00 2001 From: Neenu1995 Date: Mon, 14 Sep 2026 10:44:55 -0400 Subject: [PATCH 1/2] 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/2] 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(); - } }