feat(bigquery-jdbc): record Statement execution and ConnectionAttempts

java-bigquery-jdbc/README.MD

@@ -92,27 +92,15 @@ 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
 
 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
 

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;
@@ -258,6 +260,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);
   }
 

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -40,6 +40,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;
@@ -483,6 +485,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;
   }
 
@@ -541,6 +546,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;
   }
 
@@ -681,6 +689,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);
     }
@@ -915,6 +925,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
@@ -972,6 +988,9 @@ public DatabaseMetaData getMetaData() throws SQLException {
     if (databaseMetaData == null) {
       databaseMetaData = new BigQueryDatabaseMetaData(this);
     }
+
+    TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_METADATA_RETRIEVAL);
+
     return databaseMetaData;
   }
 
@@ -1474,6 +1493,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;
   }
 

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,25 @@ 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());
+      TelemetryManager.recordConnectionAttempt(
+          Status.STATUS_ERROR, TelemetryManager.extractErrorCode(e), authType);
+      LOG.warning("Getting a warning: %s", e.getMessage());
+      return null;
+    } catch (Throwable t) {
+      TelemetryManager.recordConnectionAttempt(
+          Status.STATUS_ERROR, TelemetryManager.extractErrorCode(t), authType);
+      throw t;
     }
-    return null;
   }
 
   /**

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.sql.Time;
@@ -157,6 +159,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());
   }
 
@@ -239,6 +243,9 @@ void setParameter(
     if (parameter.getIndex() == -1) {
       parametersList.add(parameter);
     }
+
+    TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_PARAMETER_BINDING);
+
     LOG.finest("Parameter set { %s }", parameter.toString());
   }
 
@@ -272,6 +279,8 @@ void setParameter(
     parameter.setParamType(paramType);
     parameter.setScale(scale);
 
+    TelemetryManager.recordFeatureUsage(DriverFeature.DRIVER_FEATURE_PARAMETER_BINDING);
+
     LOG.finest("Parameter set { %s }", parameter.toString());
   }
 

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();
     }
 

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;
@@ -323,7 +326,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<BigQueryJdbcParameter> currentParameterList = this.batchParameters.peek();
@@ -336,10 +348,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 {
@@ -370,6 +396,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);
       }
     }
   }
@@ -542,6 +570,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;
   }