feat(bigquery): Integrate Otel in client lib (#3747)

* feat(bigquery): Integrate Otel in client lib

* Refactor Otel code into separate file

* Add copyright comment

* fix style

* remove unused dependencies

* remove unused test dependencies

* fix dependency issue

* fix it test

* refactor out OtelHelper class

* change test name to start with lowercase

* change datasetId key to dataset

* remove OtelHelper class

* Clean up attribute values

* Add spans to query and listPartitions

* remove extraneous attributes

* add attribute namespaces, change to snake_case

* Resolve merge conflict

* change previewEnable to JobCreationMode

* fix style

Author: whuffman36

google-cloud-bigquery/pom.xml

@@ -156,6 +156,16 @@
       <artifactId>error_prone_annotations</artifactId>
     </dependency>
 
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+    </dependency>
+
     <!-- Test dependencies -->
     <dependency>
       <groupId>com.google.api</groupId>
@@ -208,6 +218,23 @@
       <artifactId>proto-google-cloud-datacatalog-v1</artifactId>
       <scope>test</scope>
     </dependency>
+
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-common</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-trace</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.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<BigQuery, BigQueryOptions> {
@@ -39,6 +41,8 @@ public class BigQueryOptions extends ServiceOptions<BigQuery, BigQueryOptions> {
   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<BigQuery, BigQueryOpt
 
     private String location;
     private boolean useInt64Timestamps;
+    private boolean enableOpenTelemetryTracing;
+    private Tracer openTelemetryTracer;
 
     private Builder() {}
 
@@ -90,6 +96,28 @@ public Builder setUseInt64Timestamps(boolean useInt64Timestamps) {
       return this;
     }
 
+    /**
+     * Enables OpenTelemetry tracing functionality for this BigQuery instance
+     *
+     * @param enableOpenTelemetryTracing enables OpenTelemetry tracing if true
+     */
+    @BetaApi
+    public Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing) {
+      this.enableOpenTelemetryTracing = enableOpenTelemetryTracing;
+      return this;
+    }
+
+    /**
+     * Sets the OpenTelemetry tracer for this BigQuery instance to be tracer.
+     *
+     * @param tracer OpenTelemetry tracer to be used
+     */
+    @BetaApi
+    public Builder setOpenTelemetryTracer(Tracer tracer) {
+      this.openTelemetryTracer = tracer;
+      return this;
+    }
+
     @Override
     public BigQueryOptions build() {
       return new BigQueryOptions(this);
@@ -100,6 +128,8 @@ private BigQueryOptions(Builder builder) {
     super(BigQueryFactory.class, BigQueryRpcFactory.class, builder, new BigQueryDefaults());
     this.location = builder.location;
     this.useInt64Timestamps = builder.useInt64Timestamps;
+    this.enableOpenTelemetryTracing = builder.enableOpenTelemetryTracing;
+    this.openTelemetryTracer = builder.openTelemetryTracer;
   }
 
   private static class BigQueryDefaults implements ServiceDefaults<BigQuery, BigQueryOptions> {
@@ -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() {

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.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();
+  }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetId.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();
   }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/DatasetInfo.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();
+  }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/InsertAllRequest.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();
+  }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobId.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();
+  }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelId.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();
+  }
 }

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelInfo.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();
+  }
 }