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

Commit

Permalink
Add quota project Id in clientContext
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 8e7b92f commit c641572
Show file tree
Hide file tree
Showing 3 changed files with 177 additions and 42 deletions.
9 changes: 8 additions & 1 deletion gax/src/main/java/com/google/api/gax/rpc/ClientContext.java
View file
Expand Up @@ -96,6 +96,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 +113,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 Down Expand Up @@ -200,6 +204,7 @@ 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())
Expand Down Expand Up @@ -229,6 +234,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
18 changes: 18 additions & 0 deletions gax/src/main/java/com/google/api/gax/rpc/StubSettings.java
View file
Expand Up @@ -195,6 +195,23 @@ protected Builder(StubSettings settings) {
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 Down Expand Up @@ -222,6 +239,7 @@ protected Builder(ClientContext clientContext) {
FixedWatchdogProvider.create(clientContext.getStreamWatchdog());
this.streamWatchdogCheckInterval = clientContext.getStreamWatchdogCheckInterval();
this.tracerFactory = clientContext.getTracerFactory();
this.quotaProjectID = getQuotaProjectIDFromClientContext(clientContext);
}
}

Expand Down
192 changes: 151 additions & 41 deletions gax/src/test/java/com/google/api/gax/rpc/ClientSettingsTest.java
View file
Expand Up @@ -44,6 +44,7 @@
import com.google.api.gax.rpc.testing.FakeClientSettings;
import com.google.auth.Credentials;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.common.collect.ImmutableMap;
import com.google.common.truth.Truth;
import java.io.ByteArrayInputStream;
import java.io.IOException;
Expand All @@ -61,6 +62,49 @@

@RunWith(JUnit4.class)
public class ClientSettingsTest {
private static final String QUOTA_PROJECT_ID_KEY = "x-google-user-project";
private static final String QUOTA_PROJECT_ID_FROM_HEADER_VALUE = "quota_project_id_from_headers";
private static final String QUOTA_PROJECT_ID_FROM_BUILDERS = "quota_project_id_from_builders";
private static final String QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE =
"quota_project_id_from_internal_headers";
private static final String QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE =
"quota_project_id_from_credentials";
private static final String QUOTA_PROJECT_ID_FROM_CONTEXT =
"quota_project_id_from_client_context";
private static final String JSON_KEY_QUOTA_PROJECT_ID =
"{\n"
+ " \"private_key_id\": \"somekeyid\",\n"
+ " \"private_key\": \"-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggS"
+ "kAgEAAoIBAQC+K2hSuFpAdrJI\\nnCgcDz2M7t7bjdlsadsasad+fvRSW6TjNQZ3p5LLQY1kSZRqBqylRkzteMOyHg"
+ "aR\\n0Pmxh3ILCND5men43j3h4eDbrhQBuxfEMalkG92sL+PNQSETY2tnvXryOvmBRwa/\\nQP/9dJfIkIDJ9Fw9N4"
+ "Bhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nknddadwkwewcVxHFhcZJO+XWf6ofLUXpRwiTZakGMn8EE1uVa2"
+ "LgczOjwWHGi99MFjxSer5m9\\n1tCa3/KEGKiS/YL71JvjwX3mb+cewlkcmweBKZHM2JPTk0ZednFSpVZMtycjkbLa"
+ "\\ndYOS8V85AgMBewECggEBAKksaldajfDZDV6nGqbFjMiizAKJolr/M3OQw16K6o3/\\n0S31xIe3sSlgW0+UbYlF"
+ "4U8KifhManD1apVSC3csafaspP4RZUHFhtBywLO9pR5c\\nr6S5aLp+gPWFyIp1pfXbWGvc5VY/v9x7ya1VEa6rXvL"
+ "sKupSeWAW4tMj3eo/64ge\\nsdaceaLYw52KeBYiT6+vpsnYrEkAHO1fF/LavbLLOFJmFTMxmsNaG0tuiJHgjshB\\"
+ "n82DpMCbXG9YcCgI/DbzuIjsdj2JC1cascSP//3PmefWysucBQe7Jryb6NQtASmnv\\nCdDw/0jmZTEjpe4S1lxfHp"
+ "lAhHFtdgYTvyYtaLZiVVkCgYEA8eVpof2rceecw/I6\\n5ng1q3Hl2usdWV/4mZMvR0fOemacLLfocX6IYxT1zA1FF"
+ "JlbXSRsJMf/Qq39mOR2\\nSpW+hr4jCoHeRVYLgsbggtrevGmILAlNoqCMpGZ6vDmJpq6ECV9olliDvpPgWOP+\\nm"
+ "YPDreFBGxWvQrADNbRt2dmGsrsCgYEAyUHqB2wvJHFqdmeBsaacewzV8x9WgmeX\\ngUIi9REwXlGDW0Mz50dxpxcK"
+ "CAYn65+7TCnY5O/jmL0VRxU1J2mSWyWTo1C+17L0\\n3fUqjxL1pkefwecxwecvC+gFFYdJ4CQ/MHHXU81Lwl1iWdF"
+ "Cd2UoGddYaOF+KNeM\\nHC7cmqra+JsCgYEAlUNywzq8nUg7282E+uICfCB0LfwejuymR93CtsFgb7cRd6ak\\nECR"
+ "8FGfCpH8ruWJINllbQfcHVCX47ndLZwqv3oVFKh6pAS/vVI4dpOepP8++7y1u\\ncoOvtreXCX6XqfrWDtKIvv0vjl"
+ "HBhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nkndj5uNl5SiuVxHFhcZJO+XWf6ofLUregtevZakGMn8EE1uVa"
+ "2AY7eafmoU/nZPT\\n00YB0TBATdCbn/nBSuKDESkhSg9s2GEKQZG5hBmL5uCMfo09z3SfxZIhJdlerreP\\nJ7gSi"
+ "dI12N+EZxYd4xIJh/HFDgp7RRO87f+WJkofMQKBgGTnClK1VMaCRbJZPriw\\nEfeFCoOX75MxKwXs6xgrw4W//AYG"
+ "GUjDt83lD6AZP6tws7gJ2IwY/qP7+lyhjEqN\\nHtfPZRGFkGZsdaksdlaksd323423d+15/UvrlRSFPNj1tWQmNKk"
+ "XyRDW4IG1Oa2p\\nrALStNBx5Y9t0/LQnFI4w3aG\\n-----END PRIVATE KEY-----\\n\",\n"
+ " \"project_id\": \"someprojectid\",\n"
+ " \"client_email\": \"someclientid@developer.gserviceaccount.com\",\n"
+ " \"client_id\": \"someclientid.apps.googleusercontent.com\",\n"
+ " \"type\": \"service_account\",\n"
+ " \"quota_project_id\": \""
+ QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE
+ "\"\n"
+ "}";

private static final GoogleCredentials credentialsWithQuotaProject =
loadCredentials(JSON_KEY_QUOTA_PROJECT_ID);

@Test
public void testEmptyBuilder() throws Exception {
Expand Down Expand Up @@ -152,6 +196,7 @@ public void testBuilder() throws Exception {

@Test
public void testBuilderFromClientContext() throws Exception {
final String QUOTA_PROJECT_ID_FROM_CONTEXT = "some_quota_project_id_from_context";
ApiClock clock = Mockito.mock(ApiClock.class);
ApiCallContext callContext = FakeCallContext.createDefault();
Map<String, String> headers = Collections.singletonMap("spiffykey", "spiffyvalue");
Expand All @@ -172,6 +217,7 @@ public void testBuilderFromClientContext() throws Exception {
.setHeaders(headers)
.setStreamWatchdog(watchdog)
.setStreamWatchdogCheckInterval(watchdogCheckInterval)
.setQuotaProjectID(QUOTA_PROJECT_ID_FROM_CONTEXT)
.build();

FakeClientSettings.Builder builder = new FakeClientSettings.Builder(clientContext);
Expand All @@ -186,6 +232,7 @@ public void testBuilderFromClientContext() throws Exception {
Truth.assertThat(builder.getWatchdogProvider()).isInstanceOf(FixedWatchdogProvider.class);
Truth.assertThat(builder.getWatchdogProvider().getWatchdog()).isSameInstanceAs(watchdog);
Truth.assertThat(builder.getWatchdogCheckInterval()).isEqualTo(watchdogCheckInterval);
Truth.assertThat(builder.getQuotaProjectID()).isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT);
}

@Test
Expand Down Expand Up @@ -266,47 +313,7 @@ static GoogleCredentials loadCredentials(String credentialFile) {
}

@Test
public void testBuilderQuotaProjectID() {
final String QUOTA_PROJECT_ID_KEY = "x-google-user-project";
final String QUOTA_PROJECT_ID_FROM_HEADER_VALUE = "quota_project_id_from_headers";
final String QUOTA_PROJECT_ID_FROM_BUILDERS = "quota_project_id_from_builders";
final String QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE =
"quota_project_id_from_internal_headers";
final String QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE = "quota_project_id_from_credentials";
final String JSON_KEY_QUOTA_PROJECT_ID =
"{\n"
+ " \"private_key_id\": \"somekeyid\",\n"
+ " \"private_key\": \"-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggS"
+ "kAgEAAoIBAQC+K2hSuFpAdrJI\\nnCgcDz2M7t7bjdlsadsasad+fvRSW6TjNQZ3p5LLQY1kSZRqBqylRkzteMOyHg"
+ "aR\\n0Pmxh3ILCND5men43j3h4eDbrhQBuxfEMalkG92sL+PNQSETY2tnvXryOvmBRwa/\\nQP/9dJfIkIDJ9Fw9N4"
+ "Bhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nknddadwkwewcVxHFhcZJO+XWf6ofLUXpRwiTZakGMn8EE1uVa2"
+ "LgczOjwWHGi99MFjxSer5m9\\n1tCa3/KEGKiS/YL71JvjwX3mb+cewlkcmweBKZHM2JPTk0ZednFSpVZMtycjkbLa"
+ "\\ndYOS8V85AgMBewECggEBAKksaldajfDZDV6nGqbFjMiizAKJolr/M3OQw16K6o3/\\n0S31xIe3sSlgW0+UbYlF"
+ "4U8KifhManD1apVSC3csafaspP4RZUHFhtBywLO9pR5c\\nr6S5aLp+gPWFyIp1pfXbWGvc5VY/v9x7ya1VEa6rXvL"
+ "sKupSeWAW4tMj3eo/64ge\\nsdaceaLYw52KeBYiT6+vpsnYrEkAHO1fF/LavbLLOFJmFTMxmsNaG0tuiJHgjshB\\"
+ "n82DpMCbXG9YcCgI/DbzuIjsdj2JC1cascSP//3PmefWysucBQe7Jryb6NQtASmnv\\nCdDw/0jmZTEjpe4S1lxfHp"
+ "lAhHFtdgYTvyYtaLZiVVkCgYEA8eVpof2rceecw/I6\\n5ng1q3Hl2usdWV/4mZMvR0fOemacLLfocX6IYxT1zA1FF"
+ "JlbXSRsJMf/Qq39mOR2\\nSpW+hr4jCoHeRVYLgsbggtrevGmILAlNoqCMpGZ6vDmJpq6ECV9olliDvpPgWOP+\\nm"
+ "YPDreFBGxWvQrADNbRt2dmGsrsCgYEAyUHqB2wvJHFqdmeBsaacewzV8x9WgmeX\\ngUIi9REwXlGDW0Mz50dxpxcK"
+ "CAYn65+7TCnY5O/jmL0VRxU1J2mSWyWTo1C+17L0\\n3fUqjxL1pkefwecxwecvC+gFFYdJ4CQ/MHHXU81Lwl1iWdF"
+ "Cd2UoGddYaOF+KNeM\\nHC7cmqra+JsCgYEAlUNywzq8nUg7282E+uICfCB0LfwejuymR93CtsFgb7cRd6ak\\nECR"
+ "8FGfCpH8ruWJINllbQfcHVCX47ndLZwqv3oVFKh6pAS/vVI4dpOepP8++7y1u\\ncoOvtreXCX6XqfrWDtKIvv0vjl"
+ "HBhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nkndj5uNl5SiuVxHFhcZJO+XWf6ofLUregtevZakGMn8EE1uVa"
+ "2AY7eafmoU/nZPT\\n00YB0TBATdCbn/nBSuKDESkhSg9s2GEKQZG5hBmL5uCMfo09z3SfxZIhJdlerreP\\nJ7gSi"
+ "dI12N+EZxYd4xIJh/HFDgp7RRO87f+WJkofMQKBgGTnClK1VMaCRbJZPriw\\nEfeFCoOX75MxKwXs6xgrw4W//AYG"
+ "GUjDt83lD6AZP6tws7gJ2IwY/qP7+lyhjEqN\\nHtfPZRGFkGZsdaksdlaksd323423d+15/UvrlRSFPNj1tWQmNKk"
+ "XyRDW4IG1Oa2p\\nrALStNBx5Y9t0/LQnFI4w3aG\\n-----END PRIVATE KEY-----\\n\",\n"
+ " \"project_id\": \"someprojectid\",\n"
+ " \"client_email\": \"someclientid@developer.gserviceaccount.com\",\n"
+ " \"client_id\": \"someclientid.apps.googleusercontent.com\",\n"
+ " \"type\": \"service_account\",\n"
+ " \"quota_project_id\": \""
+ QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE
+ "\"\n"
+ "}";

final GoogleCredentials credentialsWithQuotaProject =
loadCredentials(JSON_KEY_QUOTA_PROJECT_ID);
public void testBuilderFromSettings_QuotaProjectID() {

CredentialsProvider credentialsProvider_no_quota = Mockito.mock(CredentialsProvider.class);
HeaderProvider headerProvider_no_quota = Mockito.mock(HeaderProvider.class);
Expand Down Expand Up @@ -422,4 +429,107 @@ public Credentials getCredentials() throws IOException {
Truth.assertThat(builder_setQuotaFromAllSourcesOrder.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_BUILDERS);
}

@Test
public void testBuilderFromClientContext_QuotaProjectID() {
ApiCallContext callContext = FakeCallContext.createDefault();

ClientContext clientContextQuotaOnly =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setQuotaProjectID(QUOTA_PROJECT_ID_FROM_CONTEXT)
.build();
FakeClientSettings.Builder builderQuotaOnly =
new FakeClientSettings.Builder(clientContextQuotaOnly);

ClientContext clientContextCredentialOnly =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setCredentials(credentialsWithQuotaProject)
.build();
FakeClientSettings.Builder builderCredentialOnly =
new FakeClientSettings.Builder(clientContextCredentialOnly);

ClientContext clientContextCredentialAndQuota =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setCredentials(credentialsWithQuotaProject)
.setQuotaProjectID(QUOTA_PROJECT_ID_FROM_CONTEXT)
.build();
FakeClientSettings.Builder builderCredentialAndQuota =
new FakeClientSettings.Builder(clientContextCredentialAndQuota);

ClientContext clientContextHeadersOnly =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setHeaders(ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_HEADER_VALUE))
.build();
FakeClientSettings.Builder builderHeadersOnly =
new FakeClientSettings.Builder(clientContextHeadersOnly);

ClientContext clientContextHeadersAndQuota =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setQuotaProjectID(QUOTA_PROJECT_ID_FROM_CONTEXT)
.setHeaders(ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_HEADER_VALUE))
.build();
FakeClientSettings.Builder builderHeadersAndQuota =
new FakeClientSettings.Builder(clientContextHeadersAndQuota);

ClientContext clientContextInternalHeadersOnly =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setInternalHeaders(
ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE))
.build();
FakeClientSettings.Builder builderInternalHeadersOnly =
new FakeClientSettings.Builder(clientContextInternalHeadersOnly);

ClientContext clientContextInternalHeadersAndQuota =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setInternalHeaders(
ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE))
.setQuotaProjectID(QUOTA_PROJECT_ID_FROM_CONTEXT)
.build();
FakeClientSettings.Builder builderInternalHeadersAndQuota =
new FakeClientSettings.Builder(clientContextInternalHeadersAndQuota);

ClientContext clientContextQuotaFromAllSources =
ClientContext.newBuilder()
.setTransportChannel(Mockito.mock(TransportChannel.class))
.setDefaultCallContext(callContext)
.setHeaders(
ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE))
.setCredentials(credentialsWithQuotaProject)
.setQuotaProjectID(QUOTA_PROJECT_ID_FROM_CONTEXT)
.setInternalHeaders(
ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE))
.build();
FakeClientSettings.Builder builderQuotaFromAllSources =
new FakeClientSettings.Builder(clientContextQuotaFromAllSources);

Truth.assertThat(builderQuotaOnly.getQuotaProjectID()).isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT);
Truth.assertThat(builderCredentialOnly.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE);
Truth.assertThat(builderCredentialAndQuota.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT);
Truth.assertThat(builderHeadersOnly.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_HEADER_VALUE);
Truth.assertThat(builderHeadersAndQuota.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT);
Truth.assertThat(builderInternalHeadersOnly.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE);
Truth.assertThat(builderInternalHeadersAndQuota.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT);
Truth.assertThat(builderQuotaFromAllSources.getQuotaProjectID())
.isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT);
}
}

0 comments on commit c641572

Please sign in to comment.