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

Commit

Permalink
feat: enable setting quota_project_id (#1128)
Browse files Browse the repository at this point in the history
  • Loading branch information
@summer-ji-eng
summer-ji-eng authored Jul 24, 2020
1 parent 89a3dcf commit 20bb200
Show file tree
Hide file tree
Showing 7 changed files with 803 additions and 6 deletions.
53 changes: 47 additions & 6 deletions gax/src/main/java/com/google/api/gax/rpc/ClientContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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();
Expand All @@ -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);
}
Expand Down Expand Up @@ -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 {

Expand All @@ -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);

Expand Down
16 changes: 16 additions & 0 deletions gax/src/main/java/com/google/api/gax/rpc/ClientSettings.java
View file
Original file line number Diff line number Diff line change
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
59 changes: 59 additions & 0 deletions gax/src/main/java/com/google/api/gax/rpc/StubSettings.java
View file
Original file line number Diff line number Diff line change
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-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;
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,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();
Expand All @@ -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();
Expand All @@ -208,6 +239,7 @@ protected Builder(ClientContext clientContext) {
FixedWatchdogProvider.create(clientContext.getStreamWatchdog());
this.streamWatchdogCheckInterval = clientContext.getStreamWatchdogCheckInterval();
this.tracerFactory = clientContext.getTracerFactory();
this.quotaProjectId = getQuotaProjectIdFromClientContext(clientContext);
}
}

Expand All @@ -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();
}

Expand All @@ -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();
}

Expand All @@ -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();
}

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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)
Expand Down
Loading

0 comments on commit 20bb200

Please sign in to comment.