diff --git a/google-cloud-bigquery/pom.xml b/google-cloud-bigquery/pom.xml index a9d3bb46b..d5f1de4ef 100644 --- a/google-cloud-bigquery/pom.xml +++ b/google-cloud-bigquery/pom.xml @@ -156,6 +156,16 @@ error_prone_annotations + + + io.opentelemetry + opentelemetry-api + + + io.opentelemetry + opentelemetry-context + + com.google.api @@ -208,6 +218,23 @@ proto-google-cloud-datacatalog-v1 test + + + + io.opentelemetry + opentelemetry-sdk + test + + + io.opentelemetry + opentelemetry-sdk-common + test + + + io.opentelemetry + opentelemetry-sdk-trace + test + diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 9ede02838..2efe8f92a 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -51,6 +51,9 @@ import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.context.Scope; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -268,7 +271,18 @@ public Dataset create(DatasetInfo datasetInfo, DatasetOption... options) { : datasetInfo.getDatasetId().getProject()) .toPb(); final Map optionsMap = optionMap(options); - try { + Span datasetCreate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + datasetCreate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.createDataset") + .setAllAttributes(datasetInfo.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope datasetCreateScope = datasetCreate != null ? datasetCreate.makeCurrent() : null) { return Dataset.fromPb( this, BigQueryRetryHelper.runWithRetries( @@ -284,6 +298,10 @@ public com.google.api.services.bigquery.model.Dataset call() throws IOException EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (datasetCreate != null) { + datasetCreate.end(); + } } } @@ -298,7 +316,18 @@ public Table create(TableInfo tableInfo, TableOption... options) { .toPb(); handleExternalTableSchema(tablePb); final Map optionsMap = optionMap(options); - try { + Span tableCreate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + tableCreate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.createTable") + .setAllAttributes(tableInfo.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope tableCreateScope = tableCreate != null ? tableCreate.makeCurrent() : null) { return Table.fromPb( this, BigQueryRetryHelper.runWithRetries( @@ -314,6 +343,10 @@ public com.google.api.services.bigquery.model.Table call() throws IOException { EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (tableCreate != null) { + tableCreate.end(); + } } } @@ -337,7 +370,18 @@ public Routine create(RoutineInfo routineInfo, RoutineOption... options) { : routineInfo.getRoutineId().getProject()) .toPb(); final Map optionsMap = optionMap(options); - try { + Span routineCreate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + routineCreate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.createRoutine") + .setAllAttributes(routineInfo.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope createRoutineScope = routineCreate != null ? routineCreate.makeCurrent() : null) { return Routine.fromPb( this, BigQueryRetryHelper.runWithRetries( @@ -353,6 +397,10 @@ public com.google.api.services.bigquery.model.Routine call() throws IOException EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (routineCreate != null) { + routineCreate.end(); + } } } @@ -388,6 +436,17 @@ Job create(JobInfo jobInfo, Supplier idProvider, JobOption... options) { final boolean idRandom = (jobInfo.getJobId() == null); final Map optionsMap = optionMap(options); + Span jobCreate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + jobCreate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.createJob") + .setAllAttributes(jobInfo.getJobId().getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } BigQueryException createException; // NOTE(pongad): This double-try structure is admittedly odd. // translateAndThrow itself throws, and pretends to return an exception only @@ -396,90 +455,96 @@ Job create(JobInfo jobInfo, Supplier idProvider, JobOption... options) { // Fixing this entails some work on BaseServiceException.translate. // Since that affects a bunch of APIs, we should fix this as a separate change. final JobId[] finalJobId = new JobId[1]; - try { + try (Scope jobCreateScope = jobCreate != null ? jobCreate.makeCurrent() : null) { try { - return Job.fromPb( - this, - BigQueryRetryHelper.runWithRetries( - new Callable() { - @Override - public com.google.api.services.bigquery.model.Job call() throws IOException { - if (idRandom) { - // re-generate a new random job with the same jobInfo when jobId is not - // provided by the user - JobInfo recreatedJobInfo = - jobInfo.toBuilder().setJobId(idProvider.get()).build(); - com.google.api.services.bigquery.model.Job newJobPb = - recreatedJobInfo.setProjectId(getOptions().getProjectId()).toPb(); - finalJobId[0] = recreatedJobInfo.getJobId(); - return bigQueryRpc.createSkipExceptionTranslation(newJobPb, optionsMap); - } else { - com.google.api.services.bigquery.model.Job jobPb = - jobInfo.setProjectId(getOptions().getProjectId()).toPb(); - return bigQueryRpc.createSkipExceptionTranslation(jobPb, optionsMap); + try { + return Job.fromPb( + this, + BigQueryRetryHelper.runWithRetries( + new Callable() { + @Override + public com.google.api.services.bigquery.model.Job call() throws IOException { + if (idRandom) { + // re-generate a new random job with the same jobInfo when jobId is not + // provided by the user + JobInfo recreatedJobInfo = + jobInfo.toBuilder().setJobId(idProvider.get()).build(); + com.google.api.services.bigquery.model.Job newJobPb = + recreatedJobInfo.setProjectId(getOptions().getProjectId()).toPb(); + finalJobId[0] = recreatedJobInfo.getJobId(); + return bigQueryRpc.createSkipExceptionTranslation(newJobPb, optionsMap); + } else { + com.google.api.services.bigquery.model.Job jobPb = + jobInfo.setProjectId(getOptions().getProjectId()).toPb(); + return bigQueryRpc.createSkipExceptionTranslation(jobPb, optionsMap); + } } - } - }, - getRetryOptions(optionsMap) != null - ? RetryOption.mergeToSettings( - getOptions().getRetrySettings(), getRetryOptions(optionsMap)) - : getOptions().getRetrySettings(), - BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER, - getOptions().getClock(), - getBigQueryRetryConfig(optionsMap) != null - ? getBigQueryRetryConfig(optionsMap) - : DEFAULT_RETRY_CONFIG)); - } catch (BigQueryRetryHelperException e) { - throw BigQueryException.translateAndThrow(e); + }, + getRetryOptions(optionsMap) != null + ? RetryOption.mergeToSettings( + getOptions().getRetrySettings(), getRetryOptions(optionsMap)) + : getOptions().getRetrySettings(), + BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER, + getOptions().getClock(), + getBigQueryRetryConfig(optionsMap) != null + ? getBigQueryRetryConfig(optionsMap) + : DEFAULT_RETRY_CONFIG)); + } catch (BigQueryRetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } catch (BigQueryException e) { + createException = e; } - } catch (BigQueryException e) { - createException = e; - } - - if (!idRandom) { - if (createException instanceof BigQueryException - && createException.getCause() != null - && createException.getCause().getMessage() != null) { - - Pattern pattern = Pattern.compile(".*Already.*Exists:.*Job.*", Pattern.CASE_INSENSITIVE); - Matcher matcher = pattern.matcher(createException.getCause().getMessage()); - - if (matcher.find()) { - // If the Job ALREADY EXISTS, retrieve it. - Job job = this.getJob(jobInfo.getJobId(), JobOption.fields(JobField.STATISTICS)); - - long jobCreationTime = job.getStatistics().getCreationTime(); - long jobMinStaleTime = System.currentTimeMillis(); - long jobMaxStaleTime = - java.time.Instant.ofEpochMilli(jobMinStaleTime) - .minus(1, java.time.temporal.ChronoUnit.DAYS) - .toEpochMilli(); - - // Only return the job if it has been created in the past 24 hours. - // This is assuming any job older than 24 hours is a valid duplicate JobID - // and not a false positive like b/290419183 - if (jobCreationTime >= jobMaxStaleTime && jobCreationTime <= jobMinStaleTime) { - return job; + + if (!idRandom) { + if (createException instanceof BigQueryException + && createException.getCause() != null + && createException.getCause().getMessage() != null) { + + Pattern pattern = Pattern.compile(".*Already.*Exists:.*Job.*", Pattern.CASE_INSENSITIVE); + Matcher matcher = pattern.matcher(createException.getCause().getMessage()); + + if (matcher.find()) { + // If the Job ALREADY EXISTS, retrieve it. + Job job = this.getJob(jobInfo.getJobId(), JobOption.fields(JobField.STATISTICS)); + + long jobCreationTime = job.getStatistics().getCreationTime(); + long jobMinStaleTime = System.currentTimeMillis(); + long jobMaxStaleTime = + java.time.Instant.ofEpochMilli(jobMinStaleTime) + .minus(1, java.time.temporal.ChronoUnit.DAYS) + .toEpochMilli(); + + // Only return the job if it has been created in the past 24 hours. + // This is assuming any job older than 24 hours is a valid duplicate JobID + // and not a false positive like b/290419183 + if (jobCreationTime >= jobMaxStaleTime && jobCreationTime <= jobMinStaleTime) { + return job; + } } } + throw createException; } - throw createException; - } - // If create RPC fails, it's still possible that the job has been successfully - // created, and get might work. - // We can only do this if we randomly generated the ID. Otherwise we might - // mistakenly fetch a job created by someone else. - Job job; - try { - job = getJob(finalJobId[0]); - } catch (BigQueryException e) { - throw createException; - } - if (job == null) { - throw createException; + // If create RPC fails, it's still possible that the job has been successfully + // created, and get might work. + // We can only do this if we randomly generated the ID. Otherwise we might + // mistakenly fetch a job created by someone else. + Job job; + try { + job = getJob(finalJobId[0]); + } catch (BigQueryException e) { + throw createException; + } + if (job == null) { + throw createException; + } + return job; + } finally { + if (jobCreate != null) { + jobCreate.end(); + } } - return job; } @Override @@ -491,7 +556,18 @@ public Dataset getDataset(String datasetId, DatasetOption... options) { public Dataset getDataset(final DatasetId datasetId, DatasetOption... options) { final DatasetId completeDatasetId = datasetId.setProjectId(getOptions().getProjectId()); final Map optionsMap = optionMap(options); - try { + Span datasetGet = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + datasetGet = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.getDataset") + .setAllAttributes(completeDatasetId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope datasetGetScope = datasetGet != null ? datasetGet.makeCurrent() : null) { com.google.api.services.bigquery.model.Dataset answer = BigQueryRetryHelper.runWithRetries( new Callable() { @@ -514,6 +590,10 @@ public com.google.api.services.bigquery.model.Dataset call() throws IOException return null; } throw BigQueryException.translateAndThrow(e); + } finally { + if (datasetGet != null) { + datasetGet.end(); + } } } @@ -524,7 +604,24 @@ public Page listDatasets(DatasetListOption... options) { @Override public Page listDatasets(String projectId, DatasetListOption... options) { - return listDatasets(projectId, getOptions(), optionMap(options)); + Span datasetsList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + datasetsList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listDatasets") + .setAttribute("bq.dataset.project_id", projectId) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope datasetsListScope = datasetsList != null ? datasetsList.makeCurrent() : null) { + return listDatasets(projectId, getOptions(), optionMap(options)); + } finally { + if (datasetsList != null) { + datasetsList.end(); + } + } } private static Page listDatasets( @@ -574,7 +671,18 @@ public boolean delete(String datasetId, DatasetDeleteOption... options) { public boolean delete(DatasetId datasetId, DatasetDeleteOption... options) { final DatasetId completeDatasetId = datasetId.setProjectId(getOptions().getProjectId()); final Map optionsMap = optionMap(options); - try { + Span datasetDelete = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + datasetDelete = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.deleteDataset") + .setAllAttributes(datasetId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope datasetDeleteScope = datasetDelete != null ? datasetDelete.makeCurrent() : null) { return BigQueryRetryHelper.runWithRetries( new Callable() { @Override @@ -592,6 +700,10 @@ public Boolean call() throws IOException { return false; } throw BigQueryException.translateAndThrow(e); + } finally { + if (datasetDelete != null) { + datasetDelete.end(); + } } } @@ -607,7 +719,17 @@ public boolean delete(TableId tableId) { Strings.isNullOrEmpty(tableId.getProject()) ? getOptions().getProjectId() : tableId.getProject()); - try { + Span tableDelete = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + tableDelete = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.deleteTable") + .setAllAttributes(tableId.getOtelAttributes()) + .startSpan(); + } + try (Scope tableDeleteScope = tableDelete != null ? tableDelete.makeCurrent() : null) { return BigQueryRetryHelper.runWithRetries( new Callable() { @Override @@ -627,6 +749,10 @@ public Boolean call() throws IOException { return false; } throw BigQueryException.translateAndThrow(e); + } finally { + if (tableDelete != null) { + tableDelete.end(); + } } } @@ -637,7 +763,17 @@ public boolean delete(ModelId modelId) { Strings.isNullOrEmpty(modelId.getProject()) ? getOptions().getProjectId() : modelId.getProject()); - try { + Span modelDelete = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + modelDelete = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.deleteModel") + .setAllAttributes(modelId.getOtelAttributes()) + .startSpan(); + } + try (Scope modelDeleteScope = modelDelete != null ? modelDelete.makeCurrent() : null) { return BigQueryRetryHelper.runWithRetries( new Callable() { @Override @@ -657,6 +793,10 @@ public Boolean call() throws IOException { return false; } throw BigQueryException.translateAndThrow(e); + } finally { + if (modelDelete != null) { + modelDelete.end(); + } } } @@ -667,7 +807,17 @@ public boolean delete(RoutineId routineId) { Strings.isNullOrEmpty(routineId.getProject()) ? getOptions().getProjectId() : routineId.getProject()); - try { + Span routineDelete = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + routineDelete = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.deleteRoutine") + .setAllAttributes(routineId.getOtelAttributes()) + .startSpan(); + } + try (Scope routineDeleteScope = routineDelete != null ? routineDelete.makeCurrent() : null) { return BigQueryRetryHelper.runWithRetries( new Callable() { @Override @@ -687,6 +837,10 @@ public Boolean call() throws IOException { return false; } throw BigQueryException.translateAndThrow(e); + } finally { + if (routineDelete != null) { + routineDelete.end(); + } } } @@ -697,6 +851,16 @@ public boolean delete(JobId jobId) { Strings.isNullOrEmpty(jobId.getProject()) ? getOptions().getProjectId() : jobId.getProject()); + Span jobDelete = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + jobDelete = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.deleteJob") + .setAllAttributes(completeJobId.getOtelAttributes()) + .startSpan(); + } try { return BigQueryRetryHelper.runWithRetries( new Callable() { @@ -712,6 +876,10 @@ public Boolean call() throws IOException { EMPTY_RETRY_CONFIG); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (jobDelete != null) { + jobDelete.end(); + } } } @@ -720,7 +888,18 @@ public Dataset update(DatasetInfo datasetInfo, DatasetOption... options) { final com.google.api.services.bigquery.model.Dataset datasetPb = datasetInfo.setProjectId(getOptions().getProjectId()).toPb(); final Map optionsMap = optionMap(options); - try { + Span datasetUpdate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + datasetUpdate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.updateDataset") + .setAllAttributes(datasetInfo.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope datasetUpdateScope = datasetUpdate != null ? datasetUpdate.makeCurrent() : null) { return Dataset.fromPb( this, BigQueryRetryHelper.runWithRetries( @@ -736,6 +915,10 @@ public com.google.api.services.bigquery.model.Dataset call() throws IOException EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (datasetUpdate != null) { + datasetUpdate.end(); + } } } @@ -750,7 +933,18 @@ public Table update(TableInfo tableInfo, TableOption... options) { .toPb(); handleExternalTableSchema(tablePb); final Map optionsMap = optionMap(options); - try { + Span tableUpdate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + tableUpdate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.updateTable") + .setAllAttributes(tableInfo.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope tableUpdateScope = tableUpdate != null ? tableUpdate.makeCurrent() : null) { return Table.fromPb( this, BigQueryRetryHelper.runWithRetries( @@ -766,6 +960,10 @@ public com.google.api.services.bigquery.model.Table call() throws IOException { EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (tableUpdate != null) { + tableUpdate.end(); + } } } @@ -779,7 +977,18 @@ public Model update(ModelInfo modelInfo, ModelOption... options) { : modelInfo.getModelId().getProject()) .toPb(); final Map optionsMap = optionMap(options); - try { + Span modelUpdate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + modelUpdate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.updateModel") + .setAllAttributes(modelInfo.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope modelUpdateScope = modelUpdate != null ? modelUpdate.makeCurrent() : null) { return Model.fromPb( this, BigQueryRetryHelper.runWithRetries( @@ -795,6 +1004,10 @@ public com.google.api.services.bigquery.model.Model call() throws IOException { EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (modelUpdate != null) { + modelUpdate.end(); + } } } @@ -808,7 +1021,18 @@ public Routine update(RoutineInfo routineInfo, RoutineOption... options) { : routineInfo.getRoutineId().getProject()) .toPb(); final Map optionsMap = optionMap(options); - try { + Span routineUpdate = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + routineUpdate = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.updateRoutine") + .setAllAttributes(routineInfo.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope routineUpdateScope = routineUpdate != null ? routineUpdate.makeCurrent() : null) { return Routine.fromPb( this, BigQueryRetryHelper.runWithRetries( @@ -824,6 +1048,10 @@ public com.google.api.services.bigquery.model.Routine call() throws IOException EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (routineUpdate != null) { + routineUpdate.end(); + } } } @@ -842,7 +1070,18 @@ public Table getTable(TableId tableId, TableOption... options) { ? getOptions().getProjectId() : tableId.getProject()); final Map optionsMap = optionMap(options); - try { + Span tableGet = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + tableGet = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.getTable") + .setAllAttributes(tableId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope tableGetScope = tableGet != null ? tableGet.makeCurrent() : null) { com.google.api.services.bigquery.model.Table answer = BigQueryRetryHelper.runWithRetries( new Callable() { @@ -868,6 +1107,10 @@ public com.google.api.services.bigquery.model.Table call() throws IOException { return null; } throw BigQueryException.translateAndThrow(e); + } finally { + if (tableGet != null) { + tableGet.end(); + } } } @@ -884,7 +1127,18 @@ public Model getModel(ModelId modelId, ModelOption... options) { ? getOptions().getProjectId() : modelId.getProject()); final Map optionsMap = optionMap(options); - try { + Span modelGet = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + modelGet = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.getModel") + .setAllAttributes(modelId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope modelGetScope = modelGet != null ? modelGet.makeCurrent() : null) { com.google.api.services.bigquery.model.Model answer = BigQueryRetryHelper.runWithRetries( new Callable() { @@ -910,6 +1164,10 @@ public com.google.api.services.bigquery.model.Model call() throws IOException { return null; } throw BigQueryException.translateAndThrow(e); + } finally { + if (modelGet != null) { + modelGet.end(); + } } } @@ -926,7 +1184,18 @@ public Routine getRoutine(RoutineId routineId, RoutineOption... options) { ? getOptions().getProjectId() : routineId.getProject()); final Map optionsMap = optionMap(options); - try { + Span routineGet = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + routineGet = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.getRoutine") + .setAllAttributes(routineId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope routineGetScope = routineGet != null ? routineGet.makeCurrent() : null) { com.google.api.services.bigquery.model.Routine answer = BigQueryRetryHelper.runWithRetries( new Callable() { @@ -952,67 +1221,189 @@ public com.google.api.services.bigquery.model.Routine call() throws IOException return null; } throw BigQueryException.translateAndThrow(e); + } finally { + if (routineGet != null) { + routineGet.end(); + } } } @Override public Page listTables(String datasetId, TableListOption... options) { - return listTables( - DatasetId.of(getOptions().getProjectId(), datasetId), getOptions(), optionMap(options)); + Span tablesList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + tablesList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listTables") + .setAllAttributes(DatasetId.of(datasetId).getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope tablesListScope = tablesList != null ? tablesList.makeCurrent() : null) { + return listTables( + DatasetId.of(getOptions().getProjectId(), datasetId), getOptions(), optionMap(options)); + } finally { + if (tablesList != null) { + tablesList.end(); + } + } } @Override public Page
listTables(DatasetId datasetId, TableListOption... options) { DatasetId completeDatasetId = datasetId.setProjectId(getOptions().getProjectId()); - return listTables(completeDatasetId, getOptions(), optionMap(options)); + Span tablesList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + tablesList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listTables") + .setAllAttributes(completeDatasetId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope tablesListScope = tablesList != null ? tablesList.makeCurrent() : null) { + return listTables(completeDatasetId, getOptions(), optionMap(options)); + } finally { + if (tablesList != null) { + tablesList.end(); + } + } } @Override public Page listModels(String datasetId, ModelListOption... options) { - return listModels( - DatasetId.of(getOptions().getProjectId(), datasetId), getOptions(), optionMap(options)); + Span modelsList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + modelsList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listModels") + .setAllAttributes(DatasetId.of(datasetId).getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope modelsListScope = modelsList != null ? modelsList.makeCurrent() : null) { + return listModels( + DatasetId.of(getOptions().getProjectId(), datasetId), getOptions(), optionMap(options)); + } finally { + if (modelsList != null) { + modelsList.end(); + } + } } @Override public Page listModels(DatasetId datasetId, ModelListOption... options) { DatasetId completeDatasetId = datasetId.setProjectId(getOptions().getProjectId()); - return listModels(completeDatasetId, getOptions(), optionMap(options)); + Span modelsList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + modelsList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listModels") + .setAllAttributes(datasetId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope modelsListScope = modelsList != null ? modelsList.makeCurrent() : null) { + return listModels(completeDatasetId, getOptions(), optionMap(options)); + } finally { + if (modelsList != null) { + modelsList.end(); + } + } } @Override public Page listRoutines(String datasetId, RoutineListOption... options) { - return listRoutines( - DatasetId.of(getOptions().getProjectId(), datasetId), getOptions(), optionMap(options)); + Span routinesList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + routinesList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listRoutines") + .setAllAttributes(DatasetId.of(datasetId).getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope routinesListScope = routinesList != null ? routinesList.makeCurrent() : null) { + return listRoutines( + DatasetId.of(getOptions().getProjectId(), datasetId), getOptions(), optionMap(options)); + } finally { + if (routinesList != null) { + routinesList.end(); + } + } } @Override public Page listRoutines(DatasetId datasetId, RoutineListOption... options) { DatasetId completeDatasetId = datasetId.setProjectId(getOptions().getProjectId()); - return listRoutines(completeDatasetId, getOptions(), optionMap(options)); + Span routinesList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + routinesList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listRoutines") + .setAllAttributes(datasetId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope routinesListScope = routinesList != null ? routinesList.makeCurrent() : null) { + return listRoutines(completeDatasetId, getOptions(), optionMap(options)); + } finally { + if (routinesList != null) { + routinesList.end(); + } + } } @Override public List listPartitions(TableId tableId) { - List partitions = new ArrayList(); - String partitionsTable = tableId.getTable() + "$__PARTITIONS_SUMMARY__"; - TableId metaTableId = - tableId.getProject() == null - ? TableId.of(tableId.getDataset(), partitionsTable) - : TableId.of(tableId.getProject(), tableId.getDataset(), partitionsTable); - Table metaTable = getTable(metaTableId); - Schema metaSchema = metaTable.getDefinition().getSchema(); - String partition_id = null; - for (Field field : metaSchema.getFields()) { - if (field.getName().equals("partition_id")) { - partition_id = field.getName(); - break; - } + Span listPartitions = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + listPartitions = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listPartitions") + .setAllAttributes(tableId.getOtelAttributes()) + .startSpan(); } - TableResult result = metaTable.list(metaSchema); - for (FieldValueList list : result.iterateAll()) { - partitions.add(list.get(partition_id).getStringValue()); + try (Scope listPartitionsScope = listPartitions != null ? listPartitions.makeCurrent() : null) { + List partitions = new ArrayList(); + String partitionsTable = tableId.getTable() + "$__PARTITIONS_SUMMARY__"; + TableId metaTableId = + tableId.getProject() == null + ? TableId.of(tableId.getDataset(), partitionsTable) + : TableId.of(tableId.getProject(), tableId.getDataset(), partitionsTable); + Table metaTable = getTable(metaTableId); + Schema metaSchema = metaTable.getDefinition().getSchema(); + String partition_id = null; + for (Field field : metaSchema.getFields()) { + if (field.getName().equals("partition_id")) { + partition_id = field.getName(); + break; + } + } + TableResult result = metaTable.list(metaSchema); + for (FieldValueList list : result.iterateAll()) { + partitions.add(list.get(partition_id).getStringValue()); + } + return partitions; + } finally { + if (listPartitions != null) { + listPartitions.end(); + } } - return partitions; } private static Page
listTables( @@ -1165,33 +1556,52 @@ public Rows apply(RowToInsert rowToInsert) { requestPb.setRows(rowsPb); TableDataInsertAllResponse responsePb; - if (allInsertIdsSet[0]) { - // allowing retries only if all row insertIds are set (used for deduplication) - try { + Span insertAll = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + insertAll = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.insertAll") + .setAllAttributes(request.getOtelAttributes()) + .startSpan(); + } + try (Scope insertAllScope = insertAll != null ? insertAll.makeCurrent() : null) { + if (allInsertIdsSet[0]) { + // allowing retries only if all row insertIds are set (used for deduplication) + try { + responsePb = + BigQueryRetryHelper.runWithRetries( + new Callable() { + @Override + public TableDataInsertAllResponse call() throws Exception { + return bigQueryRpc.insertAllSkipExceptionTranslation( + tableId.getProject(), + tableId.getDataset(), + tableId.getTable(), + requestPb); + } + }, + getOptions().getRetrySettings(), + BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER, + getOptions().getClock(), + EMPTY_RETRY_CONFIG); + } catch (BigQueryRetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } else { + // Use insertAll that translate the exception as we are not retrying. responsePb = - BigQueryRetryHelper.runWithRetries( - new Callable() { - @Override - public TableDataInsertAllResponse call() throws Exception { - return bigQueryRpc.insertAllSkipExceptionTranslation( - tableId.getProject(), tableId.getDataset(), tableId.getTable(), requestPb); - } - }, - getOptions().getRetrySettings(), - BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER, - getOptions().getClock(), - EMPTY_RETRY_CONFIG); - } catch (BigQueryRetryHelperException e) { - throw BigQueryException.translateAndThrow(e); + bigQueryRpc.insertAll( + tableId.getProject(), tableId.getDataset(), tableId.getTable(), requestPb); } - } else { - // Use insertAll that translate the exception as we are not retrying. - responsePb = - bigQueryRpc.insertAll( - tableId.getProject(), tableId.getDataset(), tableId.getTable(), requestPb); - } - return InsertAllResponse.fromPb(responsePb); + return InsertAllResponse.fromPb(responsePb); + } finally { + if (insertAll != null) { + insertAll.end(); + } + } } @Override @@ -1215,11 +1625,28 @@ public TableResult listTableData( public TableResult listTableData(TableId tableId, Schema schema, TableDataListOption... options) { Tuple, Long> data = listTableData(tableId, schema, getOptions(), optionMap(options)); - return TableResult.newBuilder() - .setSchema(schema) - .setTotalRows(data.y()) - .setPageNoSchema(data.x()) - .build(); + Span tableDataList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + tableDataList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listTableData") + .setAllAttributes(tableId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope tableDataListScope = tableDataList != null ? tableDataList.makeCurrent() : null) { + return TableResult.newBuilder() + .setSchema(schema) + .setTotalRows(data.y()) + .setPageNoSchema(data.x()) + .build(); + } finally { + if (tableDataList != null) { + tableDataList.end(); + } + } } private static Tuple, Long> listTableData( @@ -1295,7 +1722,18 @@ public Job getJob(JobId jobId, JobOption... options) { jobId.getLocation() == null && getOptions().getLocation() != null ? getOptions().getLocation() : jobId.getLocation()); - try { + Span jobGet = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + jobGet = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.getJob") + .setAllAttributes(completeJobId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope jobGetScope = jobGet != null ? jobGet.makeCurrent() : null) { com.google.api.services.bigquery.model.Job answer = BigQueryRetryHelper.runWithRetries( new Callable() { @@ -1321,12 +1759,32 @@ public com.google.api.services.bigquery.model.Job call() throws IOException { return null; } throw BigQueryException.translateAndThrow(e); + } finally { + if (jobGet != null) { + jobGet.end(); + } } } @Override public Page listJobs(JobListOption... options) { - return listJobs(getOptions(), optionMap(options)); + Span jobsList = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + jobsList = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.listJobs") + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope jobsListScope = jobsList != null ? jobsList.makeCurrent() : null) { + return listJobs(getOptions(), optionMap(options)); + } finally { + if (jobsList != null) { + jobsList.end(); + } + } } private static Page listJobs( @@ -1377,7 +1835,17 @@ public boolean cancel(JobId jobId) { jobId.getLocation() == null && getOptions().getLocation() != null ? getOptions().getLocation() : jobId.getLocation()); - try { + Span jobCancel = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + jobCancel = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.cancelJob") + .setAllAttributes(completeJobId.getOtelAttributes()) + .startSpan(); + } + try (Scope jobCancelScope = jobCancel != null ? jobCancel.makeCurrent() : null) { return BigQueryRetryHelper.runWithRetries( new Callable() { @Override @@ -1395,6 +1863,10 @@ public Boolean call() throws IOException { return false; } throw BigQueryException.translateAndThrow(e); + } finally { + if (jobCancel != null) { + jobCancel.end(); + } } } @@ -1408,27 +1880,56 @@ public TableResult query(QueryJobConfiguration configuration, JobOption... optio .setJobCreationMode(getOptions().getDefaultJobCreationMode()) .build(); - // If all parameters passed in configuration are supported by the query() method on the backend, - // put on fast path - QueryRequestInfo requestInfo = - new QueryRequestInfo(configuration, getOptions().getUseInt64Timestamps()); - if (requestInfo.isFastQuerySupported(null)) { - String projectId = getOptions().getProjectId(); - QueryRequest content = requestInfo.toPb(); - if (getOptions().getLocation() != null) { - content.setLocation(getOptions().getLocation()); + Span querySpan = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + querySpan = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.query") + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope queryScope = querySpan != null ? querySpan.makeCurrent() : null) { + // If all parameters passed in configuration are supported by the query() method on the + // backend, + // put on fast path + QueryRequestInfo requestInfo = + new QueryRequestInfo(configuration, getOptions().getUseInt64Timestamps()); + if (requestInfo.isFastQuerySupported(null)) { + String projectId = getOptions().getProjectId(); + QueryRequest content = requestInfo.toPb(); + if (getOptions().getLocation() != null) { + content.setLocation(getOptions().getLocation()); + } + return queryRpc(projectId, content, options); + } + // Otherwise, fall back to the existing create query job logic + return create(JobInfo.of(configuration), options).getQueryResults(); + } finally { + if (querySpan != null) { + querySpan.end(); } - return queryRpc(projectId, content, options); } - // Otherwise, fall back to the existing create query job logic - return create(JobInfo.of(configuration), options).getQueryResults(); } private TableResult queryRpc( final String projectId, final QueryRequest content, JobOption... options) throws InterruptedException { com.google.api.services.bigquery.model.QueryResponse results; - try { + Span queryRpc = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + queryRpc = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.queryRpc") + .setAttribute("bq.query.project_id", projectId) + .setAllAttributes(otelAttributesFromQueryRequest(content)) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope queryRpcScope = queryRpc != null ? queryRpc.makeCurrent() : null) { results = BigQueryRetryHelper.runWithRetries( new Callable() { @@ -1444,6 +1945,10 @@ public com.google.api.services.bigquery.model.QueryResponse call() DEFAULT_RETRY_CONFIG); } catch (BigQueryRetryHelper.BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (queryRpc != null) { + queryRpc.end(); + } } if (results.getErrors() != null) { @@ -1513,37 +2018,74 @@ public TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOp throws InterruptedException, JobException { Job.checkNotDryRun(configuration, "query"); - // If all parameters passed in configuration are supported by the query() method on the backend, - // put on fast path - QueryRequestInfo requestInfo = - new QueryRequestInfo(configuration, getOptions().getUseInt64Timestamps()); - if (requestInfo.isFastQuerySupported(jobId)) { - // Be careful when setting the projectID in JobId, if a projectID is specified in the JobId, - // the job created by the query method will use that project. This may cause the query to - // fail with "Access denied" if the project do not have enough permissions to run the job. - - String projectId = - jobId.getProject() != null ? jobId.getProject() : getOptions().getProjectId(); - QueryRequest content = requestInfo.toPb(); - // Be careful when setting the location, if a location is specified in the BigQueryOption or - // JobId the job created by the query method will be in that location, even if the table to be - // queried is in a different location. This may cause the query to fail with - // "BigQueryException: Not found" - if (jobId.getLocation() != null) { - content.setLocation(jobId.getLocation()); - } else if (getOptions().getLocation() != null) { - content.setLocation(getOptions().getLocation()); - } + Span querySpan = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + querySpan = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.query") + .setAllAttributes(jobId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope queryScope = querySpan != null ? querySpan.makeCurrent() : null) { + // If all parameters passed in configuration are supported by the query() method on the + // backend, + // put on fast path + QueryRequestInfo requestInfo = + new QueryRequestInfo(configuration, getOptions().getUseInt64Timestamps()); + if (requestInfo.isFastQuerySupported(jobId)) { + // Be careful when setting the projectID in JobId, if a projectID is specified in the JobId, + // the job created by the query method will use that project. This may cause the query to + // fail with "Access denied" if the project do not have enough permissions to run the job. + + String projectId = + jobId.getProject() != null ? jobId.getProject() : getOptions().getProjectId(); + QueryRequest content = requestInfo.toPb(); + // Be careful when setting the location, if a location is specified in the BigQueryOption or + // JobId the job created by the query method will be in that location, even if the table to + // be + // queried is in a different location. This may cause the query to fail with + // "BigQueryException: Not found" + if (jobId.getLocation() != null) { + content.setLocation(jobId.getLocation()); + } else if (getOptions().getLocation() != null) { + content.setLocation(getOptions().getLocation()); + } - return queryRpc(projectId, content, options); + return queryRpc(projectId, content, options); + } + return create(JobInfo.of(jobId, configuration), options).getQueryResults(); + } finally { + if (querySpan != null) { + querySpan.end(); + } } - return create(JobInfo.of(jobId, configuration), options).getQueryResults(); } @Override public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) { Map optionsMap = optionMap(options); - return getQueryResults(jobId, getOptions(), optionsMap); + Span getQueryResults = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + getQueryResults = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.getQueryResults") + .setAllAttributes(jobId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope getQueryResultsScope = + getQueryResults != null ? getQueryResults.makeCurrent() : null) { + return getQueryResults(jobId, getOptions(), optionsMap); + } finally { + if (getQueryResults != null) { + getQueryResults.end(); + } + } } private static QueryResponse getQueryResults( @@ -1619,8 +2161,19 @@ public Policy getIamPolicy(TableId tableId, IAMOption... options) { ? getOptions().getProjectId() : tableId.getProject()); - try { - final Map optionsMap = optionMap(options); + final Map optionsMap = optionMap(options); + Span iamPolicyGet = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + iamPolicyGet = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.getIamPolicy") + .setAllAttributes(tableId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope iamPolicyGetScope = iamPolicyGet != null ? iamPolicyGet.makeCurrent() : null) { return convertFromApiPolicy( BigQueryRetryHelper.runWithRetries( new Callable() { @@ -1636,6 +2189,10 @@ public com.google.api.services.bigquery.model.Policy call() throws IOException { EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (iamPolicyGet != null) { + iamPolicyGet.end(); + } } } @@ -1646,8 +2203,21 @@ public Policy setIamPolicy(TableId tableId, final Policy policy, IAMOption... op Strings.isNullOrEmpty(tableId.getProject()) ? getOptions().getProjectId() : tableId.getProject()); - try { - final Map optionsMap = optionMap(options); + + final Map optionsMap = optionMap(options); + Span iamPolicySet = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + iamPolicySet = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.setIamPolicy") + .setAllAttributes(tableId.getOtelAttributes()) + .setAllAttributes(otelAttributesFromPolicy(policy)) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope iamPolicySetScope = iamPolicySet != null ? iamPolicySet.makeCurrent() : null) { return convertFromApiPolicy( BigQueryRetryHelper.runWithRetries( new Callable() { @@ -1663,6 +2233,10 @@ public com.google.api.services.bigquery.model.Policy call() throws IOException { EMPTY_RETRY_CONFIG)); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (iamPolicySet != null) { + iamPolicySet.end(); + } } } @@ -1674,8 +2248,21 @@ public List testIamPermissions( Strings.isNullOrEmpty(tableId.getProject()) ? getOptions().getProjectId() : tableId.getProject()); - try { - final Map optionsMap = optionMap(options); + final Map optionsMap = optionMap(options); + Span testIamPermissions = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + testIamPermissions = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.testIamPermissions") + .setAllAttributes(tableId.getOtelAttributes()) + .setAttribute("bq.iam.permissions", permissions.toString()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope testIamPermissionsScope = + testIamPermissions != null ? testIamPermissions.makeCurrent() : null) { com.google.api.services.bigquery.model.TestIamPermissionsResponse response = BigQueryRetryHelper.runWithRetries( new Callable() { @@ -1695,6 +2282,10 @@ public com.google.api.services.bigquery.model.TestIamPermissionsResponse call() : ImmutableList.copyOf(response.getPermissions()); } catch (BigQueryRetryHelperException e) { throw BigQueryException.translateAndThrow(e); + } finally { + if (testIamPermissions != null) { + testIamPermissions.end(); + } } } @@ -1708,6 +2299,39 @@ public com.google.api.services.bigquery.model.TestIamPermissionsResponse call() return optionMap; } + private static String getFieldAsString(Object field) { + return field == null ? "null" : field.toString(); + } + + private static Attributes otelAttributesFromOptions(Option... options) { + Attributes attributes = Attributes.builder().build(); + for (Option option : options) { + attributes = + attributes.toBuilder() + .put("bq.option." + option.getRpcOption().toString(), option.getValue().toString()) + .build(); + } + return attributes; + } + + private static Attributes otelAttributesFromQueryRequest(QueryRequest request) { + return Attributes.builder() + .put("bq.query.dry_run", getFieldAsString(request.getDryRun())) + .put("bq.query.job_creation_mode", getFieldAsString(request.getJobCreationMode())) + .put("bq.query.kind", getFieldAsString(request.getKind())) + .put("bq.query.location", getFieldAsString(request.getLocation())) + .put("bq.query.request_id", getFieldAsString(request.getRequestId())) + .put("bq.query.use_query_cache", getFieldAsString(request.getUseQueryCache())) + .build(); + } + + private static Attributes otelAttributesFromPolicy(Policy policy) { + return Attributes.builder() + .put("bq.policy.version", getFieldAsString(policy.getVersion())) + .put("bq.policy.bindings", getFieldAsString(policy.getBindings())) + .build(); + } + static BigQueryRetryConfig getBigQueryRetryConfig(Map options) { return (BigQueryRetryConfig) options.getOrDefault(BigQueryRpc.Option.BIGQUERY_RETRY_CONFIG, null); diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java index a1dacc4cf..36eeb2086 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java @@ -16,6 +16,7 @@ package com.google.cloud.bigquery; +import com.google.api.core.BetaApi; import com.google.cloud.ServiceDefaults; import com.google.cloud.ServiceOptions; import com.google.cloud.ServiceRpc; @@ -25,6 +26,7 @@ import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc; import com.google.cloud.http.HttpTransportOptions; import com.google.common.collect.ImmutableSet; +import io.opentelemetry.api.trace.Tracer; import java.util.Set; public class BigQueryOptions extends ServiceOptions { @@ -39,6 +41,8 @@ public class BigQueryOptions extends ServiceOptions { private boolean setThrowNotFound; private boolean useInt64Timestamps; private JobCreationMode defaultJobCreationMode = JobCreationMode.JOB_CREATION_MODE_UNSPECIFIED; + private boolean enableOpenTelemetryTracing; + private Tracer openTelemetryTracer; public static class DefaultBigQueryFactory implements BigQueryFactory { @@ -64,6 +68,8 @@ public static class Builder extends ServiceOptions.Builder { @@ -171,6 +201,26 @@ public JobCreationMode getDefaultJobCreationMode() { return defaultJobCreationMode; } + /** + * Returns whether this BigQuery instance has OpenTelemetry tracing enabled + * + * @return true if tracing is enabled, false if not + */ + @BetaApi("Span names and attributes are subject to change without notice") + public boolean isOpenTelemetryTracingEnabled() { + return enableOpenTelemetryTracing; + } + + /** + * Returns the OpenTelemetry tracer used by this BigQuery instance + * + * @return OpenTelemetry tracer object or {@code null} if not set + */ + @BetaApi("Span names and attributes are subject to change without notice") + public Tracer getOpenTelemetryTracer() { + return openTelemetryTracer; + } + @SuppressWarnings("unchecked") @Override public Builder toBuilder() { diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetId.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetId.java index 0e2ad02b2..c06d257d7 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetId.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetId.java @@ -20,6 +20,7 @@ import static com.google.common.base.Strings.isNullOrEmpty; import com.google.api.services.bigquery.model.DatasetReference; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.Objects; @@ -84,4 +85,11 @@ DatasetReference toPb() { static DatasetId fromPb(DatasetReference datasetRef) { return new DatasetId(datasetRef.getProjectId(), datasetRef.getDatasetId()); } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .put("bq.dataset.project", this.getProject()) + .put("bq.dataset.id", this.getDataset()) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java index 90d6c2cd5..918761995 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.java @@ -27,6 +27,7 @@ import com.google.common.base.MoreObjects; import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.List; import java.util.Map; @@ -746,6 +747,18 @@ public static DatasetInfo of(String datasetId) { return newBuilder(datasetId).build(); } + private static String getFieldAsString(Object field) { + return field == null ? "null" : field.toString(); + } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .putAll(this.getDatasetId().getOtelAttributes()) + .put("bq.dataset.last_modified", getFieldAsString(this.getLastModified())) + .put("bq.dataset.location", getFieldAsString(this.getLocation())) + .build(); + } + static DatasetInfo fromPb(Dataset datasetPb) { return new BuilderImpl(datasetPb).build(); } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/InsertAllRequest.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/InsertAllRequest.java index 53952bc6c..ce9380098 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/InsertAllRequest.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/InsertAllRequest.java @@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.Collections; import java.util.HashMap; @@ -479,4 +480,15 @@ public boolean equals(Object obj) { && Objects.equals(skipInvalidRows, other.skipInvalidRows) && Objects.equals(templateSuffix, other.templateSuffix); } + + private static String getFieldAsString(Object field) { + return field == null ? "null" : field.toString(); + } + + public Attributes getOtelAttributes() { + return Attributes.builder() + .put("bq.insert_all.table", getFieldAsString(this.getTable().getTable())) + .put("bq.insert_all.template_suffix", getFieldAsString(this.getTemplateSuffix())) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java index b966cebe1..4bfc2aa00 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.java @@ -21,6 +21,7 @@ import com.google.api.services.bigquery.model.JobReference; import com.google.auto.value.AutoValue; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.UUID; import javax.annotation.Nullable; @@ -123,4 +124,16 @@ static JobId fromPb(JobReference jobRef) { .setLocation(jobRef.getLocation()) .build(); } + + private static String getFieldAsString(Object field) { + return field == null ? "null" : field.toString(); + } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .put("bq.job.id", getFieldAsString(this.getJob())) + .put("bq.job.location", getFieldAsString(this.getLocation())) + .put("bq.job.project", getFieldAsString(this.getProject())) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelId.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelId.java index 9356c6935..adb497ffa 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelId.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelId.java @@ -21,6 +21,7 @@ import com.google.api.services.bigquery.model.ModelReference; import com.google.common.base.Function; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.Objects; @@ -105,4 +106,12 @@ ModelReference toPb() { static ModelId fromPb(ModelReference modelRef) { return new ModelId(modelRef.getProjectId(), modelRef.getDatasetId(), modelRef.getModelId()); } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .put("bq.model.project", this.getProject()) + .put("bq.model.dataset", this.getDataset()) + .put("bq.model.id", this.getModel()) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java index 922ed441c..3039483b5 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.java @@ -26,6 +26,7 @@ import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.Collections; import java.util.List; @@ -453,4 +454,19 @@ Model toPb() { static ModelInfo fromPb(Model modelPb) { return new BuilderImpl(modelPb).build(); } + + private static String getFieldAsString(Object field) { + return field == null ? "null" : field.toString(); + } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .putAll(this.getModelId().getOtelAttributes()) + .put("bq.model.type", getFieldAsString(this.getModelType())) + .put("bq.model.creation_time", getFieldAsString(this.getCreationTime())) + .put("bq.model.last_modified_time", getFieldAsString(this.getLastModifiedTime())) + .put("bq.model.expiration_time", getFieldAsString(this.getExpirationTime())) + .put("bq.model.location", getFieldAsString(this.getLocation())) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineId.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineId.java index c5c3dadb3..a632fc5a1 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineId.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineId.java @@ -21,6 +21,7 @@ import com.google.api.services.bigquery.model.RoutineReference; import com.google.common.base.Function; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.Objects; @@ -108,4 +109,12 @@ static RoutineId fromPb(RoutineReference routineRef) { return new RoutineId( routineRef.getProjectId(), routineRef.getDatasetId(), routineRef.getRoutineId()); } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .put("bq.routine.project", this.getProject()) + .put("bq.routine.dataset", this.getDataset()) + .put("bq.routine.id", this.getRoutine()) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java index 646a2fdea..c13b90b41 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java @@ -24,6 +24,7 @@ import com.google.common.base.MoreObjects; import com.google.common.base.Strings; import com.google.common.collect.Lists; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.Collections; import java.util.List; @@ -544,4 +545,17 @@ Routine toPb() { static RoutineInfo fromPb(Routine routinePb) { return new BuilderImpl(routinePb).build(); } + + private static String getFieldAsString(Object field) { + return field == null ? "null" : field.toString(); + } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .putAll(this.getRoutineId().getOtelAttributes()) + .put("bq.routine.type", getFieldAsString(this.getRoutineType())) + .put("bq.routine.creation_time", getFieldAsString(this.getCreationTime())) + .put("bq.routine.last_modified_time", getFieldAsString(this.getLastModifiedTime())) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableId.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableId.java index b74055d4f..fcc208453 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableId.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableId.java @@ -22,6 +22,7 @@ import com.google.api.services.bigquery.model.TableReference; import com.google.common.base.Function; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.util.Objects; @@ -114,4 +115,12 @@ TableReference toPb() { static TableId fromPb(TableReference tableRef) { return new TableId(tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()); } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .put("bq.table.project", this.getProject()) + .put("bq.table.dataset", this.getDataset()) + .put("bq.table.id", this.getTable()) + .build(); + } } diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java index 6e050d064..b3236f4c3 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java @@ -25,6 +25,7 @@ import com.google.api.services.bigquery.model.Table; import com.google.common.base.Function; import com.google.common.base.MoreObjects; +import io.opentelemetry.api.common.Attributes; import java.io.Serializable; import java.math.BigInteger; import java.util.Map; @@ -763,4 +764,19 @@ Table toPb() { static TableInfo fromPb(Table tablePb) { return new BuilderImpl(tablePb).build(); } + + private static String getFieldAsString(Object field) { + return field == null ? "null" : field.toString(); + } + + protected Attributes getOtelAttributes() { + return Attributes.builder() + .putAll(this.getTableId().getOtelAttributes()) + .put("bq.table.creation_time", getFieldAsString(this.getCreationTime())) + .put("bq.table.expiration_time", getFieldAsString(this.getExpirationTime())) + .put("bq.table.last_modified_time", getFieldAsString(this.getLastModifiedTime())) + .put("bq.table.num_bytes", getFieldAsString(this.getNumBytes())) + .put("bq.table.num_rows", getFieldAsString(this.getNumRows())) + .build(); + } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index bd17b14fc..eceadcbca 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -161,6 +161,17 @@ import com.google.common.io.BaseEncoding; import com.google.common.util.concurrent.ListenableFuture; import com.google.gson.JsonObject; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.ByteArrayInputStream; import java.io.IOException; import java.io.InputStream; @@ -220,6 +231,11 @@ public class ITBigQueryTest { private static final String STORAGE_BILLING_MODEL = "LOGICAL"; private static final Long MAX_TIME_TRAVEL_HOURS = 120L; private static final Long MAX_TIME_TRAVEL_HOURS_DEFAULT = 168L; + private static final Map, Object>> OTEL_ATTRIBUTES = + new HashMap, Object>>(); + private static final Map OTEL_PARENT_SPAN_IDS = new HashMap<>(); + private static final Map OTEL_SPAN_IDS_TO_NAMES = new HashMap<>(); + private static final String OTEL_PARENT_SPAN_ID = "0000000000000000"; private static final String CLOUD_SAMPLES_DATA = Optional.fromNullable(System.getenv("CLOUD_SAMPLES_DATA_BUCKET")).or("cloud-samples-data"); private static final Map LABELS = @@ -1017,6 +1033,32 @@ public class ITBigQueryTest { private static BigQuery bigquery; private static Storage storage; + private static OpenTelemetry otel; + + private static class TestSpanExporter implements io.opentelemetry.sdk.trace.export.SpanExporter { + @Override + public CompletableResultCode export(Collection collection) { + if (collection.isEmpty()) { + return CompletableResultCode.ofFailure(); + } + for (SpanData data : collection) { + OTEL_ATTRIBUTES.put(data.getName(), data.getAttributes().asMap()); + OTEL_PARENT_SPAN_IDS.put(data.getName(), data.getParentSpanId()); + OTEL_SPAN_IDS_TO_NAMES.put(data.getSpanId(), data.getName()); + } + return CompletableResultCode.ofSuccess(); + } + + @Override + public CompletableResultCode flush() { + return CompletableResultCode.ofSuccess(); + } + + @Override + public CompletableResultCode shutdown() { + return CompletableResultCode.ofSuccess(); + } + } @Rule public Timeout globalTimeout = Timeout.seconds(300); @@ -1025,6 +1067,13 @@ public static void beforeClass() throws InterruptedException, IOException { RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); RemoteStorageHelper storageHelper = RemoteStorageHelper.create(); Map labels = ImmutableMap.of("test-job-name", "test-load-job"); + SdkTracerProvider tracerProvider = + SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(new TestSpanExporter())) + .setSampler(Sampler.alwaysOn()) + .build(); + otel = OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).buildAndRegisterGlobal(); + bigquery = bigqueryHelper.getOptions().getService(); storage = storageHelper.getOptions().getService(); storage.create(BucketInfo.of(BUCKET)); @@ -7493,4 +7542,157 @@ public void testStatementType() throws InterruptedException { assertTrue(remoteTable.getDefinition() instanceof MaterializedViewDefinition); assertTrue(remoteTable.delete()); } + + @Test + public void testOpenTelemetryTracingDatasets() { + Tracer tracer = otel.getTracer("Test Tracer"); + BigQueryOptions otelOptions = + BigQueryOptions.newBuilder() + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .build(); + BigQuery bigquery = otelOptions.getService(); + + Span parentSpan = + tracer + .spanBuilder("Test Parent Span") + .setNoParent() + .setAttribute("test-attribute", "test-value") + .startSpan(); + String billingModelDataset = RemoteBigQueryHelper.generateDatasetName(); + + try (Scope parentScope = parentSpan.makeCurrent()) { + DatasetInfo info = + DatasetInfo.newBuilder(billingModelDataset) + .setDescription(DESCRIPTION) + .setMaxTimeTravelHours(72L) + .setLabels(LABELS) + .build(); + + Dataset dataset = bigquery.create(info); + assertNotNull(dataset); + dataset = bigquery.getDataset(dataset.getDatasetId().getDataset()); + assertNotNull(dataset); + + DatasetInfo updatedInfo = + DatasetInfo.newBuilder(billingModelDataset) + .setDescription("Updated Description") + .setMaxTimeTravelHours(96L) + .setLabels(LABELS) + .build(); + + dataset = bigquery.update(updatedInfo, DatasetOption.accessPolicyVersion(2)); + assertEquals(dataset.getDescription(), "Updated Description"); + assertTrue(bigquery.delete(dataset.getDatasetId())); + } finally { + parentSpan.end(); + Map, Object> createMap = + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.createDataset"); + assertEquals(createMap.get(AttributeKey.stringKey("bq.dataset.location")), "null"); + + Map, Object> getMap = + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.getDataset"); + assertEquals(getMap.get(AttributeKey.stringKey("bq.dataset.id")), billingModelDataset); + + Map, Object> updateMap = + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.updateDataset"); + assertEquals(updateMap.get(AttributeKey.stringKey("bq.option.ACCESS_POLICY_VERSION")), "2"); + + Map, Object> deleteMap = + OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.deleteDataset"); + assertEquals(deleteMap.get(AttributeKey.stringKey("bq.dataset.id")), billingModelDataset); + + // All should be children spans of parentSpan + assertEquals( + OTEL_SPAN_IDS_TO_NAMES.get( + OTEL_PARENT_SPAN_IDS.get("com.google.cloud.bigquery.BigQuery.getDataset")), + "Test Parent Span"); + assertEquals( + OTEL_SPAN_IDS_TO_NAMES.get( + OTEL_PARENT_SPAN_IDS.get("com.google.cloud.bigquery.BigQuery.createDataset")), + "Test Parent Span"); + assertEquals( + OTEL_SPAN_IDS_TO_NAMES.get( + OTEL_PARENT_SPAN_IDS.get("com.google.cloud.bigquery.BigQuery.deleteDataset")), + "Test Parent Span"); + assertEquals(OTEL_PARENT_SPAN_IDS.get("Test Parent Span"), OTEL_PARENT_SPAN_ID); + RemoteBigQueryHelper.forceDelete(bigquery, billingModelDataset); + } + } + + @Test + public void testOpenTelemetryTracingTables() { + Tracer tracer = otel.getTracer("Test Tracer"); + BigQueryOptions otelOptions = + BigQueryOptions.newBuilder() + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .build(); + BigQuery bigquery = otelOptions.getService(); + + String tableName = "test_otel_table"; + StandardTableDefinition tableDefinition = StandardTableDefinition.of(TABLE_SCHEMA); + TableInfo tableInfo = + TableInfo.newBuilder(TableId.of(DATASET, tableName), tableDefinition) + .setDescription("Some Description") + .build(); + Table createdTable = bigquery.create(tableInfo); + assertThat(createdTable.getDescription()).isEqualTo("Some Description"); + + assertEquals( + OTEL_PARENT_SPAN_IDS.get("com.google.cloud.bigquery.BigQuery.createTable"), + OTEL_PARENT_SPAN_ID); + assertEquals( + OTEL_ATTRIBUTES + .get("com.google.cloud.bigquery.BigQuery.createTable") + .get(AttributeKey.stringKey("bq.table.id")), + tableName); + assertEquals( + OTEL_ATTRIBUTES + .get("com.google.cloud.bigquery.BigQuery.createTable") + .get(AttributeKey.stringKey("bq.table.creation_time")), + "null"); + + Table updatedTable = + bigquery.update(createdTable.toBuilder().setDescription("Updated Description").build()); + assertThat(updatedTable.getDescription()).isEqualTo("Updated Description"); + + assertNotNull(OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.updateTable")); + assertEquals( + OTEL_PARENT_SPAN_IDS.get("com.google.cloud.bigquery.BigQuery.updateTable"), + OTEL_PARENT_SPAN_ID); + } + + @Test + public void testOpenTelemetryTracingQuery() throws InterruptedException { + Tracer tracer = otel.getTracer("Test Tracer"); + BigQueryOptions otelOptions = + BigQueryOptions.newBuilder() + .setEnableOpenTelemetryTracing(true) + .setOpenTelemetryTracer(tracer) + .build(); + BigQuery bigquery = otelOptions.getService(); + + // Stateless query + bigquery.getOptions().setDefaultJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL); + TableResult tableResult = executeSimpleQuery(bigquery); + assertNotNull(tableResult.getQueryId()); + assertNull(tableResult.getJobId()); + + assertNotNull(OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.queryRpc")); + + // Query job + String query = "SELECT TimestampField, StringField, BooleanField FROM " + TABLE_ID.getTable(); + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder(query).setDefaultDataset(DatasetId.of(DATASET)).build(); + Job job = bigquery.create(JobInfo.of(JobId.of(), config)); + + TableResult result = job.getQueryResults(); + assertNotNull(result.getJobId()); + assertEquals(QUERY_RESULT_SCHEMA, result.getSchema()); + + assertNotNull(OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.getQueryResults")); + assertNotNull(OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.listTableData")); + assertNotNull(OTEL_ATTRIBUTES.get("com.google.cloud.bigquery.BigQuery.createJob")); + } } diff --git a/pom.xml b/pom.xml index 6d8f73a80..801aea035 100644 --- a/pom.xml +++ b/pom.xml @@ -59,6 +59,15 @@ + + + io.opentelemetry + opentelemetry-bom + 1.48.0 + pom + import + + com.google.cloud google-cloud-shared-dependencies