feat: enable setting quota_project_id (#1128)

gax/src/main/java/com/google/api/gax/rpc/ClientContext.java

@@ -35,6 +35,7 @@
 import com.google.api.gax.core.BackgroundResource;
 import com.google.api.gax.core.ExecutorAsBackgroundResource;
 import com.google.api.gax.core.ExecutorProvider;
+import com.google.api.gax.rpc.internal.QuotaProjectIdHidingCredentials;
 import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.api.gax.tracing.NoopApiTracerFactory;
 import com.google.auth.Credentials;
@@ -60,6 +61,7 @@
  */
 @AutoValue
 public abstract class ClientContext {
+  private static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project";
 
   /**
    * The objects that need to be closed in order to clean up the resources created in the process of
@@ -96,6 +98,9 @@ public abstract class ClientContext {
   @Nullable
   public abstract String getEndpoint();
 
+  @Nullable
+  public abstract String getQuotaProjectId();
+
   /** Gets the {@link ApiTracerFactory} that will be used to generate traces for operations. */
   @BetaApi("The surface for tracing is not stable yet and may change in the future.")
   @Nonnull
@@ -110,7 +115,8 @@ public static Builder newBuilder() {
         .setClock(NanoClock.getDefaultClock())
         .setStreamWatchdog(null)
         .setStreamWatchdogCheckInterval(Duration.ZERO)
-        .setTracerFactory(NoopApiTracerFactory.getInstance());
+        .setTracerFactory(NoopApiTracerFactory.getInstance())
+        .setQuotaProjectId(null);
   }
 
   public abstract Builder toBuilder();
@@ -135,15 +141,19 @@ public static ClientContext create(StubSettings settings) throws IOException {
 
     Credentials credentials = settings.getCredentialsProvider().getCredentials();
 
+    if (settings.getQuotaProjectId() != null) {
+      // If the quotaProjectId is set, wrap original credentials with correct quotaProjectId as
+      // QuotaProjectIdHidingCredentials.
+      // Ensure that a custom set quota project id takes priority over one detected by credentials.
+      // Avoid the backend receiving possibly conflict values of quotaProjectId
+      credentials = new QuotaProjectIdHidingCredentials(credentials);
+    }
+
     TransportChannelProvider transportChannelProvider = settings.getTransportChannelProvider();
     if (transportChannelProvider.needsExecutor()) {
       transportChannelProvider = transportChannelProvider.withExecutor((Executor) executor);
     }
-    Map<String, String> headers =
-        ImmutableMap.<String, String>builder()
-            .putAll(settings.getHeaderProvider().getHeaders())
-            .putAll(settings.getInternalHeaderProvider().getHeaders())
-            .build();
+    Map<String, String> headers = getHeadersFromSettings(settings);
     if (transportChannelProvider.needsHeaders()) {
       transportChannelProvider = transportChannelProvider.withHeaders(headers);
     }
@@ -200,12 +210,41 @@ public static ClientContext create(StubSettings settings) throws IOException {
         .setClock(clock)
         .setDefaultCallContext(defaultCallContext)
         .setEndpoint(settings.getEndpoint())
+        .setQuotaProjectId(settings.getQuotaProjectId())
         .setStreamWatchdog(watchdog)
         .setStreamWatchdogCheckInterval(settings.getStreamWatchdogCheckInterval())
         .setTracerFactory(settings.getTracerFactory())
         .build();
   }
 
+  /**
+   * Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota
+   * Project Id.
+   */
+  private static Map<String, String> getHeadersFromSettings(StubSettings settings) {
+    ImmutableMap.Builder<String, String> headersBuilder = ImmutableMap.builder();
+    if (settings.getQuotaProjectId() != null) {
+      headersBuilder.put(QUOTA_PROJECT_ID_HEADER_KEY, settings.getQuotaProjectId());
+      for (Map.Entry<String, String> entry : settings.getHeaderProvider().getHeaders().entrySet()) {
+        if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) {
+          continue;
+        }
+        headersBuilder.put(entry);
+      }
+      for (Map.Entry<String, String> entry :
+          settings.getInternalHeaderProvider().getHeaders().entrySet()) {
+        if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) {
+          continue;
+        }
+        headersBuilder.put(entry);
+      }
+    } else {
+      headersBuilder.putAll(settings.getHeaderProvider().getHeaders());
+      headersBuilder.putAll(settings.getInternalHeaderProvider().getHeaders());
+    }
+    return headersBuilder.build();
+  }
+
   @AutoValue.Builder
   public abstract static class Builder {
 
@@ -229,6 +268,8 @@ public abstract static class Builder {
 
     public abstract Builder setEndpoint(String endpoint);
 
+    public abstract Builder setQuotaProjectId(String QuotaProjectId);
+
     @BetaApi("The surface for streaming is not stable yet and may change in the future.")
     public abstract Builder setStreamWatchdog(Watchdog watchdog);
 

gax/src/main/java/com/google/api/gax/rpc/ClientSettings.java

@@ -93,6 +93,10 @@ public final String getEndpoint() {
     return stubSettings.getEndpoint();
   }
 
+  public final String getQuotaProjectId() {
+    return stubSettings.getQuotaProjectId();
+  }
+
   @BetaApi("The surface for streaming is not stable yet and may change in the future.")
   @Nullable
   public final WatchdogProvider getWatchdogProvider() {
@@ -114,6 +118,7 @@ public String toString() {
         .add("internalHeaderProvider", getInternalHeaderProvider())
         .add("clock", getClock())
         .add("endpoint", getEndpoint())
+        .add("quotaProjectId", getQuotaProjectId())
         .add("watchdogProvider", getWatchdogProvider())
         .add("watchdogCheckInterval", getWatchdogCheckInterval())
         .toString();
@@ -216,6 +221,11 @@ public B setEndpoint(String endpoint) {
       return self();
     }
 
+    public B setQuotaProjectId(String quotaProjectId) {
+      stubSettings.setQuotaProjectId(quotaProjectId);
+      return self();
+    }
+
     @BetaApi("The surface for streaming is not stable yet and may change in the future.")
     public B setWatchdogProvider(@Nullable WatchdogProvider watchdogProvider) {
       stubSettings.setStreamWatchdogProvider(watchdogProvider);
@@ -264,6 +274,11 @@ public String getEndpoint() {
       return stubSettings.getEndpoint();
     }
 
+    /** Gets the QuotaProjectId that was previously set on this Builder. */
+    public String getQuotaProjectId() {
+      return stubSettings.getQuotaProjectId();
+    }
+
     @BetaApi("The surface for streaming is not stable yet and may change in the future.")
     @Nullable
     public WatchdogProvider getWatchdogProvider() {
@@ -294,6 +309,7 @@ public String toString() {
           .add("internalHeaderProvider", getInternalHeaderProvider())
           .add("clock", getClock())
           .add("endpoint", getEndpoint())
+          .add("quotaProjectId", getQuotaProjectId())
           .add("watchdogProvider", getWatchdogProvider())
           .add("watchdogCheckInterval", getWatchdogCheckInterval())
           .toString();

gax/src/main/java/com/google/api/gax/rpc/StubSettings.java

@@ -41,6 +41,8 @@
 import com.google.api.gax.core.NoCredentialsProvider;
 import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.api.gax.tracing.NoopApiTracerFactory;
+import com.google.auth.Credentials;
+import com.google.auth.oauth2.QuotaProjectIdProvider;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import java.io.IOException;
@@ -60,13 +62,16 @@
  */
 public abstract class StubSettings<SettingsT extends StubSettings<SettingsT>> {
 
+  static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project";
+
   private final ExecutorProvider executorProvider;
   private final CredentialsProvider credentialsProvider;
   private final HeaderProvider headerProvider;
   private final HeaderProvider internalHeaderProvider;
   private final TransportChannelProvider transportChannelProvider;
   private final ApiClock clock;
   private final String endpoint;
+  private final String quotaProjectId;
   @Nullable private final WatchdogProvider streamWatchdogProvider;
   @Nonnull private final Duration streamWatchdogCheckInterval;
   @Nonnull private final ApiTracerFactory tracerFactory;
@@ -80,6 +85,7 @@ protected StubSettings(Builder builder) {
     this.internalHeaderProvider = builder.internalHeaderProvider;
     this.clock = builder.clock;
     this.endpoint = builder.endpoint;
+    this.quotaProjectId = builder.quotaProjectId;
     this.streamWatchdogProvider = builder.streamWatchdogProvider;
     this.streamWatchdogCheckInterval = builder.streamWatchdogCheckInterval;
     this.tracerFactory = builder.tracerFactory;
@@ -115,6 +121,10 @@ public final String getEndpoint() {
     return endpoint;
   }
 
+  public final String getQuotaProjectId() {
+    return quotaProjectId;
+  }
+
   @BetaApi("The surface for streaming is not stable yet and may change in the future.")
   @Nullable
   public final WatchdogProvider getStreamWatchdogProvider() {
@@ -146,6 +156,7 @@ public String toString() {
         .add("internalHeaderProvider", internalHeaderProvider)
         .add("clock", clock)
         .add("endpoint", endpoint)
+        .add("quotaProjectId", quotaProjectId)
         .add("streamWatchdogProvider", streamWatchdogProvider)
         .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)
         .add("tracerFactory", tracerFactory)
@@ -164,6 +175,7 @@ public String toString() {
     private TransportChannelProvider transportChannelProvider;
     private ApiClock clock;
     private String endpoint;
+    private String quotaProjectId;
     @Nullable private WatchdogProvider streamWatchdogProvider;
     @Nonnull private Duration streamWatchdogCheckInterval;
     @Nonnull private ApiTracerFactory tracerFactory;
@@ -177,11 +189,29 @@ protected Builder(StubSettings settings) {
       this.internalHeaderProvider = settings.internalHeaderProvider;
       this.clock = settings.clock;
       this.endpoint = settings.endpoint;
+      this.quotaProjectId = settings.quotaProjectId;
       this.streamWatchdogProvider = settings.streamWatchdogProvider;
       this.streamWatchdogCheckInterval = settings.streamWatchdogCheckInterval;
       this.tracerFactory = settings.tracerFactory;
     }
 
+    /** Get Quota Project ID from Client Context * */
+    private static String getQuotaProjectIdFromClientContext(ClientContext clientContext) {
+      if (clientContext.getQuotaProjectId() != null) {
+        return clientContext.getQuotaProjectId();
+      }
+      if (clientContext.getCredentials() instanceof QuotaProjectIdProvider) {
+        return ((QuotaProjectIdProvider) clientContext.getCredentials()).getQuotaProjectId();
+      }
+      if (clientContext.getHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) {
+        return clientContext.getHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY);
+      }
+      if (clientContext.getInternalHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) {
+        return clientContext.getInternalHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY);
+      }
+      return null;
+    }
+
     protected Builder(ClientContext clientContext) {
       if (clientContext == null) {
         this.executorProvider = InstantiatingExecutorProvider.newBuilder().build();
@@ -191,6 +221,7 @@ protected Builder(ClientContext clientContext) {
         this.internalHeaderProvider = new NoHeaderProvider();
         this.clock = NanoClock.getDefaultClock();
         this.endpoint = null;
+        this.quotaProjectId = null;
         this.streamWatchdogProvider = InstantiatingWatchdogProvider.create();
         this.streamWatchdogCheckInterval = Duration.ofSeconds(10);
         this.tracerFactory = NoopApiTracerFactory.getInstance();
@@ -208,6 +239,7 @@ protected Builder(ClientContext clientContext) {
             FixedWatchdogProvider.create(clientContext.getStreamWatchdog());
         this.streamWatchdogCheckInterval = clientContext.getStreamWatchdogCheckInterval();
         this.tracerFactory = clientContext.getTracerFactory();
+        this.quotaProjectId = getQuotaProjectIdFromClientContext(clientContext);
       }
     }
 
@@ -234,6 +266,14 @@ public B setExecutorProvider(ExecutorProvider executorProvider) {
     /** Sets the CredentialsProvider to use for getting the credentials to make calls with. */
     public B setCredentialsProvider(CredentialsProvider credentialsProvider) {
       this.credentialsProvider = Preconditions.checkNotNull(credentialsProvider);
+      try {
+        Credentials credentials = credentialsProvider.getCredentials();
+        if (this.quotaProjectId == null && credentials instanceof QuotaProjectIdProvider) {
+          this.quotaProjectId = ((QuotaProjectIdProvider) credentials).getQuotaProjectId();
+        }
+      } catch (IOException e) {
+        System.out.println("fail to fetch credentials");
+      }
       return self();
     }
 
@@ -247,6 +287,10 @@ public B setCredentialsProvider(CredentialsProvider credentialsProvider) {
     @BetaApi("The surface for customizing headers is not stable yet and may change in the future.")
     public B setHeaderProvider(HeaderProvider headerProvider) {
       this.headerProvider = headerProvider;
+      if (this.quotaProjectId == null
+          && headerProvider.getHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) {
+        this.quotaProjectId = headerProvider.getHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY);
+      }
       return self();
     }
 
@@ -260,6 +304,10 @@ public B setHeaderProvider(HeaderProvider headerProvider) {
     @BetaApi("The surface for customizing headers is not stable yet and may change in the future.")
     protected B setInternalHeaderProvider(HeaderProvider internalHeaderProvider) {
       this.internalHeaderProvider = internalHeaderProvider;
+      if (this.quotaProjectId == null
+          && internalHeaderProvider.getHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) {
+        this.quotaProjectId = internalHeaderProvider.getHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY);
+      }
       return self();
     }
 
@@ -298,6 +346,11 @@ public B setEndpoint(String endpoint) {
       return self();
     }
 
+    public B setQuotaProjectId(String quotaProjectId) {
+      this.quotaProjectId = quotaProjectId;
+      return self();
+    }
+
     /**
      * Sets how often the {@link Watchdog} will check ongoing streaming RPCs. Defaults to 10 secs.
      * Use {@link Duration#ZERO} to disable.
@@ -364,6 +417,11 @@ public String getEndpoint() {
       return endpoint;
     }
 
+    /** Gets the QuotaProjectId that was previously set on this Builder. */
+    public String getQuotaProjectId() {
+      return quotaProjectId;
+    }
+
     @BetaApi("The surface for streaming is not stable yet and may change in the future.")
     @Nonnull
     public Duration getStreamWatchdogCheckInterval() {
@@ -396,6 +454,7 @@ public String toString() {
           .add("internalHeaderProvider", internalHeaderProvider)
           .add("clock", clock)
           .add("endpoint", endpoint)
+          .add("quotaProjectId", quotaProjectId)
           .add("streamWatchdogProvider", streamWatchdogProvider)
           .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)
           .add("tracerFactory", tracerFactory)