Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Neenu1995 wants to merge 2 commits into
base: jdbc-telemetry-feature
Choose a base branch
from
Draft

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dbb8f49
register shutdownhook and connection attempt
Neenu1995 Sep 14, 2026
fce8bfe
record Statement Execution
Neenu1995 Sep 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 44 additions & 4 deletions java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDriver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand All @@ -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 {
Expand Down Expand Up @@ -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 = 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;
}

/**
Expand Down Expand Up @@ -284,4 +314,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;
}
}
20 changes: 20 additions & 0 deletions ...bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryPreparedStatement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<BigQueryJdbcParameter> currentParameterList = this.batchParameters.peek();
Expand All @@ -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 {
Expand Down
53 changes: 52 additions & 1 deletion java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -171,6 +176,8 @@ private void resetStatementFields() {
this.parentJobId = null;
this.currentJobIdIndex = -1;
this.currentUpdateCount = -1;

this.currentExecutionBuilder = StatementExecution.newBuilder();
}

private BigQuerySettings generateBigQuerySettings() {
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down
2 changes: 1 addition & 1 deletion ...uery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/telemetry/v1/TelemetryBatcher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<TelemetryKey, TelemetryAccumulator> metricsMap =
new ConcurrentHashMap<>();

Expand Down
Loading
Loading