Skip to content
This repository has been archived by the owner on Sep 26, 2023. It is now read-only.

Commit

Permalink
Enable Quota Project ID in Client setting
Browse files Browse the repository at this point in the history
  • Loading branch information
@summer-ji-eng
summer-ji-eng committed Jun 15, 2020
1 parent 1da065a commit 8e7b92f
Show file tree
Hide file tree
Showing 3 changed files with 241 additions and 0 deletions.
16 changes: 16 additions & 0 deletions gax/src/main/java/com/google/api/gax/rpc/ClientSettings.java
View file
Expand Up @@ -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() {
Expand All @@ -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();
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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();
Expand Down
41 changes: 41 additions & 0 deletions gax/src/main/java/com/google/api/gax/rpc/StubSettings.java
View file
Expand Up @@ -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;
Expand All @@ -60,13 +62,16 @@
*/
public abstract class StubSettings<SettingsT extends StubSettings<SettingsT>> {

static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-google-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;
Expand All @@ -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;
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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)
Expand All @@ -164,6 +175,7 @@ public abstract static class Builder<
private TransportChannelProvider transportChannelProvider;
private ApiClock clock;
private String endpoint;
private String quotaProjectID;
@Nullable private WatchdogProvider streamWatchdogProvider;
@Nonnull private Duration streamWatchdogCheckInterval;
@Nonnull private ApiTracerFactory tracerFactory;
Expand All @@ -177,6 +189,7 @@ 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;
Expand All @@ -191,6 +204,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();
Expand Down Expand Up @@ -234,6 +248,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();
}

Expand All @@ -247,6 +269,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();
}

Expand All @@ -260,6 +286,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();
}

Expand Down Expand Up @@ -298,6 +328,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.
Expand Down Expand Up @@ -364,6 +399,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() {
Expand Down Expand Up @@ -396,6 +436,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)
Expand Down

0 comments on commit 8e7b92f

Please sign in to comment.