Skip to content
This repository was archived by the owner on May 14, 2026. It is now read-only.

Commit ca3ec24

Browse files
zhumin8zhumin8
authored
feat(gax): append cred-type header for auth metrics (#3186)
This is POC change in gax-java for auth metrics requirements on token usage. See go/googleapis-auth-metric-design for context. [Credentials](https://github.com/googleapis/google-auth-library-java/blob/main/credentials/java/com/google/auth/Credentials.java) will expose `getMetricsCredentialType()` method, this change appends it to existing `x-goog-api-client` header Note: Currently implement in gax at client level. There are 2 edge cases not covered and will create followups for: if handwritten library overrides credentials at rpc level; If handwritten library does not build on gax. (ref: b/370039645, b/370038458) related change in `google-auth-library`googleapis/google-auth-library-java#1503 included in 1.28.0.
1 parent 3c00e8c commit ca3ec24

2 files changed

Lines changed: 103 additions & 2 deletions

File tree

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

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
import com.google.api.gax.tracing.ApiTracerFactory;
4545
import com.google.api.gax.tracing.BaseApiTracerFactory;
4646
import com.google.auth.ApiKeyCredentials;
47+
import com.google.auth.CredentialTypeForMetrics;
4748
import com.google.auth.Credentials;
4849
import com.google.auth.oauth2.GdchCredentials;
4950
import com.google.auto.value.AutoValue;
@@ -210,7 +211,8 @@ public static ClientContext create(StubSettings settings) throws IOException {
210211
if (transportChannelProvider.needsExecutor() && settings.getExecutorProvider() != null) {
211212
transportChannelProvider = transportChannelProvider.withExecutor(backgroundExecutor);
212213
}
213-
Map<String, String> headers = getHeadersFromSettings(settings);
214+
215+
Map<String, String> headers = getHeaders(settings, credentials);
214216
if (transportChannelProvider.needsHeaders()) {
215217
transportChannelProvider = transportChannelProvider.withHeaders(headers);
216218
}
@@ -318,8 +320,11 @@ static GdchCredentials getGdchCredentials(
318320
/**
319321
* Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota
320322
* Project Id.
323+
*
324+
* <p>Then if credentials is present and its type for metrics is not {@code
325+
* CredentialTypeForMetrics.DO_NOT_SEND}, append this type info to x-goog-api-client header.
321326
*/
322-
private static Map<String, String> getHeadersFromSettings(StubSettings settings) {
327+
private static Map<String, String> getHeaders(StubSettings settings, Credentials credentials) {
323328
// Resolve conflicts when merging headers from multiple sources
324329
Map<String, String> userHeaders = settings.getHeaderProvider().getHeaders();
325330
Map<String, String> internalHeaders = settings.getInternalHeaderProvider().getHeaders();
@@ -346,6 +351,20 @@ private static Map<String, String> getHeadersFromSettings(StubSettings settings)
346351
effectiveHeaders.putAll(userHeaders);
347352
effectiveHeaders.putAll(conflictResolution);
348353

354+
return appendCredentialTypeToHeaderIfPresent(effectiveHeaders, credentials);
355+
}
356+
357+
private static Map<String, String> appendCredentialTypeToHeaderIfPresent(
358+
Map<String, String> effectiveHeaders, Credentials credentials) {
359+
CredentialTypeForMetrics credentialTypeForMetrics =
360+
credentials == null
361+
? CredentialTypeForMetrics.DO_NOT_SEND
362+
: credentials.getMetricsCredentialType();
363+
if (credentialTypeForMetrics != CredentialTypeForMetrics.DO_NOT_SEND) {
364+
effectiveHeaders.computeIfPresent(
365+
ApiClientHeaderProvider.getDefaultApiClientHeaderKey(),
366+
(key, value) -> value + " cred-type/" + credentialTypeForMetrics.getLabel());
367+
}
349368
return ImmutableMap.copyOf(effectiveHeaders);
350369
}
351370

‎gax-java/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import static org.junit.jupiter.api.Assertions.assertTrue;
4040
import static org.mockito.Mockito.times;
4141
import static org.mockito.Mockito.verify;
42+
import static org.mockito.Mockito.when;
4243

4344
import com.google.api.core.ApiClock;
4445
import com.google.api.gax.core.BackgroundResource;
@@ -53,6 +54,7 @@
5354
import com.google.api.gax.rpc.testing.FakeStubSettings;
5455
import com.google.api.gax.rpc.testing.FakeTransportChannel;
5556
import com.google.auth.ApiKeyCredentials;
57+
import com.google.auth.CredentialTypeForMetrics;
5658
import com.google.auth.Credentials;
5759
import com.google.auth.oauth2.ComputeEngineCredentials;
5860
import com.google.auth.oauth2.GdchCredentials;
@@ -677,6 +679,86 @@ void testUserAgentConcat() throws Exception {
677679
.containsEntry("user-agent", "user-supplied-agent internal-agent");
678680
}
679681

682+
@Test
683+
void testApiClientHeaderAppendsCredType() throws Exception {
684+
GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class);
685+
when(googleCredentials.getMetricsCredentialType())
686+
.thenReturn(CredentialTypeForMetrics.USER_CREDENTIALS);
687+
688+
Map<String, String> headers =
689+
setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
690+
FixedCredentialsProvider.create(googleCredentials),
691+
FixedHeaderProvider.create("x-goog-api-client", "internal-agent"));
692+
693+
assertThat(headers).containsEntry("x-goog-api-client", "internal-agent cred-type/u");
694+
}
695+
696+
@Test
697+
void testApiClientHeaderDoNotAppendsCredType_whenNoApiClientHeader() throws Exception {
698+
GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class);
699+
when(googleCredentials.getMetricsCredentialType())
700+
.thenReturn(CredentialTypeForMetrics.USER_CREDENTIALS);
701+
702+
Map<String, String> headers =
703+
setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
704+
FixedCredentialsProvider.create(googleCredentials),
705+
FixedHeaderProvider.create("some-other-header", "internal-agent"));
706+
707+
assertThat(headers).doesNotContainKey("x-goog-api-client");
708+
assertThat(headers).containsEntry("some-other-header", "internal-agent");
709+
}
710+
711+
@Test
712+
void testApiClientHeaderDoNotAppendsCredType_whenNullCredentials() throws IOException {
713+
Map<String, String> headers =
714+
setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
715+
NoCredentialsProvider.create(),
716+
FixedHeaderProvider.create("x-goog-api-client", "internal-agent"));
717+
718+
assertThat(headers).containsKey("x-goog-api-client");
719+
assertThat(headers.get("x-goog-api-client")).doesNotContain("cred-type/");
720+
}
721+
722+
@Test
723+
void testApiClientHeaderDoNotAppendsCredType_whenCredTypeDoNotSend() throws Exception {
724+
GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class);
725+
when(googleCredentials.getMetricsCredentialType())
726+
.thenReturn(CredentialTypeForMetrics.DO_NOT_SEND);
727+
728+
Map<String, String> headers =
729+
setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
730+
FixedCredentialsProvider.create(googleCredentials),
731+
FixedHeaderProvider.create("x-goog-api-client", "internal-agent"));
732+
733+
assertThat(headers).containsKey("x-goog-api-client");
734+
assertThat(headers.get("x-goog-api-client")).doesNotContain("cred-type/");
735+
}
736+
737+
private Map<String, String> setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
738+
CredentialsProvider credentialsProvider, HeaderProvider headerProvider) throws IOException {
739+
TransportChannelProvider transportChannelProvider =
740+
new FakeTransportProvider(
741+
FakeTransportChannel.create(new FakeChannel()),
742+
null,
743+
true,
744+
null,
745+
null,
746+
DEFAULT_ENDPOINT);
747+
748+
ClientSettings.Builder builder =
749+
new FakeClientSettings.Builder()
750+
.setExecutorProvider(
751+
FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class)))
752+
.setTransportChannelProvider(transportChannelProvider)
753+
.setCredentialsProvider(credentialsProvider)
754+
.setInternalHeaderProvider(headerProvider);
755+
756+
ClientContext clientContext = ClientContext.create(builder.build());
757+
FakeTransportChannel transportChannel =
758+
(FakeTransportChannel) clientContext.getTransportChannel();
759+
return transportChannel.getHeaders();
760+
}
761+
680762
private static String endpoint = "https://foo.googleapis.com";
681763
private static String mtlsEndpoint = "https://foo.mtls.googleapis.com";
682764

0 commit comments

Comments
 (0)