Added ability to specify a custom gRPC stub for the stackdriver exporter (#1410)

3 files changed, 37 insertions, 4 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4115d38f..d2b0e0b0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,8 @@
 - Add support to handle the Tracestate in the SpanContext.
 - Remove global synchronization from the get current stats state.
 - Add get/from{Byte} methods on TraceOptions and deprecate get/from{Bytes}.
+- Add an API to `StackdriverTraceConfiguration` to allow setting a
+  `TraceServiceStub` instance to be used for export RPC calls.
 
 ## 0.15.1 - 2018-08-28
 - Improve propagation performance by avoiding doing string formatting when calling checkArgument.
diff --git a/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceConfiguration.java b/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceConfiguration.java
index 3733b763..f78832d0 100644
--- a/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceConfiguration.java
+++ b/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceConfiguration.java
@@ -18,6 +18,7 @@ package io.opencensus.exporter.trace.stackdriver;
 
 import com.google.auth.Credentials;
 import com.google.auto.value.AutoValue;
+import com.google.cloud.trace.v2.stub.TraceServiceStub;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.Immutable;
 
@@ -51,6 +52,15 @@ public abstract class StackdriverTraceConfiguration {
   public abstract String getProjectId();
 
   /**
+   * Returns a TraceServiceStub instance used to make RPC calls.
+   *
+   * @return the trace service stub.
+   * @since 0.16
+   */
+  @Nullable
+  public abstract TraceServiceStub getTraceServiceStub();
+
+  /**
    * Returns a new {@link Builder}.
    *
    * @return a {@code Builder}.
@@ -89,6 +99,15 @@ public abstract class StackdriverTraceConfiguration {
     public abstract Builder setProjectId(String projectId);
 
     /**
+     * Sets the trace service stub used to send gRPC calls.
+     *
+     * @param traceServiceStub the {@code TraceServiceStub}.
+     * @return this.
+     * @since 0.16
+     */
+    public abstract Builder setTraceServiceStub(TraceServiceStub traceServiceStub);
+
+    /**
      * Builds a {@link StackdriverTraceConfiguration}.
      *
      * @return a {@code StackdriverTraceConfiguration}.
diff --git a/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceExporter.java b/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceExporter.java
index 566bebc8..0182ae94 100644
--- a/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceExporter.java
+++ b/exporters/trace/stackdriver/src/main/java/io/opencensus/exporter/trace/stackdriver/StackdriverTraceExporter.java
@@ -21,6 +21,8 @@ import static com.google.common.base.Preconditions.checkState;
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.ServiceOptions;
+import com.google.cloud.trace.v2.TraceServiceClient;
+import com.google.cloud.trace.v2.stub.TraceServiceStub;
 import com.google.common.annotations.VisibleForTesting;
 import io.opencensus.trace.Tracing;
 import io.opencensus.trace.export.SpanExporter;
@@ -76,10 +78,20 @@ public final class StackdriverTraceExporter {
       checkState(handler == null, "Stackdriver exporter is already registered.");
       Credentials credentials = configuration.getCredentials();
       String projectId = configuration.getProjectId();
-      registerInternal(
-          StackdriverV2ExporterHandler.createWithCredentials(
-              credentials != null ? credentials : GoogleCredentials.getApplicationDefault(),
-              projectId != null ? projectId : ServiceOptions.getDefaultProjectId()));
+      projectId = projectId != null ? projectId : ServiceOptions.getDefaultProjectId();
+
+      StackdriverV2ExporterHandler handler;
+      TraceServiceStub stub = configuration.getTraceServiceStub();
+      if (stub == null) {
+        handler =
+            StackdriverV2ExporterHandler.createWithCredentials(
+                credentials != null ? credentials : GoogleCredentials.getApplicationDefault(),
+                projectId);
+      } else {
+        handler = new StackdriverV2ExporterHandler(projectId, TraceServiceClient.create(stub));
+      }
+
+      registerInternal(handler);
     }
   }