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 8c748b1f52bc..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 @@ -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 = AuthenticationType.AUTHENTICATION_TYPE_UNSPECIFIED; try { if (acceptsURL(url)) { Properties connectInfo = info == null ? new Properties() : (Properties) info.clone(); @@ -132,6 +136,17 @@ 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); + } + TelemetryManager.getInstance(connectInfo); + Level logLevel; String logPath; try { @@ -200,14 +215,29 @@ 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 = TelemetryManager.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; } /** 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 014717fd0646..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,9 @@ 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; import com.google.cloud.bigquery.storage.v1.BatchCommitWriteStreamsResponse; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; @@ -318,7 +321,16 @@ public int[] executeBatch() throws SQLException { if (this.batchParameters.isEmpty()) { 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 +343,24 @@ 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)); + 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 { @@ -365,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); } } } @@ -536,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 d1348f6d8091..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,10 @@ 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; 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 +153,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 +177,8 @@ private void resetStatementFields() { this.parentJobId = null; this.currentJobIdIndex = -1; this.currentUpdateCount = -1; + + this.currentExecutionBuilder = StatementExecution.newBuilder(); } private BigQuerySettings generateBigQuerySettings() { @@ -487,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 @@ -574,6 +584,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 +615,7 @@ ExecuteResult executeJob(QueryJobConfiguration jobConfiguration) job = refreshedJob; } } + this.currentExecutionBuilder.setQueryApiType(QueryApiType.QUERY_API_TYPE_STANDARD_REST_API); return new ExecuteResult(tableResult, job); } @@ -655,19 +667,55 @@ 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)); + Thread.currentThread().interrupt(); 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 +879,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 +1124,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 +1136,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; @@ -1745,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/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..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 @@ -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; @@ -34,6 +36,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 +73,7 @@ public static TelemetryManager getInstance(Properties properties) { } TelemetryManager localRef = instance; + if (localRef == null) { synchronized (TelemetryManager.class) { if (globallyDisabled) { @@ -82,6 +87,7 @@ public static TelemetryManager getInstance(Properties properties) { TelemetryBatcher batcher = new TelemetryBatcher(config, transport); localRef = new TelemetryManager(batcher); instance = localRef; + registerShutdownHook(); } } } @@ -150,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; } @@ -161,24 +167,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 +201,7 @@ static void recordConnectionAttempt(Status status, int errorCode, Authentication }); } - static void recordStatementExecution( + public static void recordStatementExecution( StatementType statementType, QueryApiType apiType, Status status, @@ -217,7 +224,21 @@ static void recordStatementExecution( }); } - static void recordFeatureUsage(DriverFeature feature, String customFeatureName) { + 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( () -> { TelemetryManager mgr = instance; @@ -231,4 +252,76 @@ static void recordFeatureUsage(DriverFeature feature, String customFeatureName) } }); } + + public static void recordFeatureUsage(DriverFeature feature) { + recordFeatureUsage(feature, null); + } + + 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()); + } + }); + } + + /** + * 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) { + int depth = 0; + while (t != null && depth++ < 20) { + 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) { + 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..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; @@ -187,4 +188,48 @@ public void testInvalidLogLevelExceptionIsLogged() { && r.getMessage().contains("Failed to parse connection URL properties")); assertThat(foundSevere).isTrue(); } + + @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(); + } } 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..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));