Add spans to query and listPartitions

Author: whuffman36

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -866,6 +866,17 @@ 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")
+              .setAttribute("language", "Java")
+              .setAllAttributes(completeJobId.getOtelAttributes())
+              .startSpan();
+    }
     try {
       return BigQueryRetryHelper.runWithRetries(
           new Callable<Boolean>() {
@@ -881,6 +892,10 @@ public Boolean call() throws IOException {
           EMPTY_RETRY_CONFIG);
     } catch (BigQueryRetryHelperException e) {
       throw BigQueryException.translateAndThrow(e);
+    } finally {
+      if (jobDelete != null) {
+        jobDelete.end();
+      }
     }
   }
 
@@ -1382,26 +1397,43 @@ && getOptions().getOpenTelemetryTracer() != null) {
 
   @Override
   public List<String> listPartitions(TableId tableId) {
-    List<String> partitions = new ArrayList<String>();
-    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")
+              .setAttribute("language", "Java")
+              .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<String> partitions = new ArrayList<String>();
+      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<Table> listTables(
@@ -1886,20 +1918,38 @@ public TableResult query(QueryJobConfiguration configuration, JobOption... optio
               .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(configuration.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(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(
@@ -2008,31 +2058,51 @@ 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(configuration.getOtelAttributes())
+              .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
@@ -2302,8 +2372,6 @@ private static Attributes otelAttributesFromQueryRequest(QueryRequest request) {
         .put("location", getFieldAsString(request.getLocation()))
         .put("maxResults", getFieldAsString(request.getMaxResults()))
         .put("maximumBytesBilled", getFieldAsString(request.getMaximumBytesBilled()))
-        .put("query", getFieldAsString(request.getQuery()))
-        .put("queryParameters", getFieldAsString(request.getQueryParameters()))
         .put("requestId", getFieldAsString(request.getRequestId()))
         .put("reservation", getFieldAsString(request.getReservation()))
         .put("timeoutMs", getFieldAsString(request.getTimeoutMs()))

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java

@@ -199,7 +199,7 @@ public boolean getUseInt64Timestamps() {
    *
    * @return true if tracing is enabled, false if not
    */
-  @BetaApi
+  @BetaApi("Span names and attributes are subject to change without notice")
   public boolean isOpenTelemetryTracingEnabled() {
     return enableOpenTelemetryTracing;
   }
@@ -209,7 +209,7 @@ public boolean isOpenTelemetryTracingEnabled() {
    *
    * @return OpenTelemetry tracer object or {@code null} if not set
    */
-  @BetaApi
+  @BetaApi("Span names and attributes are subject to change without notice")
   public Tracer getOpenTelemetryTracer() {
     return openTelemetryTracer;
   }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetId.java

@@ -90,7 +90,6 @@ protected Attributes getOtelAttributes() {
     return Attributes.builder()
         .put("project", this.getProject())
         .put("dataset", this.getDataset())
-        .put("db.name", this.getDataset())
         .build();
   }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelId.java

@@ -111,7 +111,6 @@ protected Attributes getOtelAttributes() {
     return Attributes.builder()
         .put("project", this.getProject())
         .put("dataset", this.getDataset())
-        .put("db.name", this.getDataset())
         .put("model", this.getModel())
         .build();
   }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryJobConfiguration.java

@@ -32,6 +32,7 @@
 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 java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -1213,4 +1214,25 @@ public QueryParameter apply(Map.Entry<String, QueryParameterValue> entry) {
               return queryParameterPb;
             }
           };
+
+  private static String getFieldAsString(Object field) {
+    return field == null ? "null" : field.toString();
+  }
+
+  public Attributes getOtelAttributes() {
+    return Attributes.builder()
+        .put("destinationTable", getFieldAsString(this.getDestinationTable()))
+        .put("defaultDataset", getFieldAsString(this.getDefaultDataset()))
+        .put("createSession", getFieldAsString(this.createSession()))
+        .put("flattenResults", getFieldAsString(this.flattenResults()))
+        .put("priority", getFieldAsString(this.getPriority()))
+        .put("useQueryCache", getFieldAsString(this.useQueryCache()))
+        .put("dryRun", getFieldAsString(this.dryRun()))
+        .put("jobTimeoutMs", getFieldAsString(this.getJobTimeoutMs()))
+        .put("labels", getFieldAsString(this.getLabels()))
+        .put("connectionProperties", getFieldAsString(this.getConnectionProperties()))
+        .put("jobCreationMode", getFieldAsString(this.getJobCreationMode()))
+        .put("reservation", getFieldAsString(this.getReservation()))
+        .build();
+  }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineId.java

@@ -114,7 +114,6 @@ protected Attributes getOtelAttributes() {
     return Attributes.builder()
         .put("project", this.getProject())
         .put("dataset", this.getDataset())
-        .put("db.name", this.getDataset())
         .put("routine", this.getRoutine())
         .build();
   }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableId.java

@@ -120,7 +120,6 @@ protected Attributes getOtelAttributes() {
     return Attributes.builder()
         .put("project", this.getProject())
         .put("dataset", this.getDataset())
-        .put("db.name", this.getDataset())
         .put("table", this.getTable())
         .build();
   }