feat: add support for quotaProjectId header

googleapis · Dec 5, 2019 · 9c16cc1 · 9c16cc1
1 parent e0abb5c
commit 9c16cc1
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/gax/src/main/java/com/google/api/gax/rpc/ApiClientHeaderProvider.java b/gax/src/main/java/com/google/api/gax/rpc/ApiClientHeaderProvider.java
@@ -42,6 +42,7 @@
 @BetaApi("The surface for customizing headers is not stable yet and may change in the future.")
 public class ApiClientHeaderProvider implements HeaderProvider, Serializable {
   private static final long serialVersionUID = -8876627296793342119L;
+  static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project";
 
   private final Map<String, String> headers;
 
@@ -65,6 +66,10 @@ protected ApiClientHeaderProvider(Builder builder) {
       headersBuilder.put(builder.getResourceHeaderKey(), builder.getResourceToken());
     }
 
+    if (builder.getQuotaProjectIdToken() != null) {
+      headersBuilder.put(QUOTA_PROJECT_ID_HEADER_KEY, builder.getQuotaProjectIdToken());
+    }
+
     this.headers = headersBuilder.build();
   }
 
@@ -101,6 +106,7 @@ public static class Builder {
     private String generatedLibToken;
     private String generatedRuntimeToken;
     private String transportToken;
+    private String quotaProjectIdToken;
 
     private String resourceHeaderKey;
     private String resourceToken;
@@ -113,6 +119,7 @@ protected Builder() {
       generatedLibToken = null;
       setClientRuntimeToken(GaxProperties.getGaxVersion());
       transportToken = null;
+      quotaProjectIdToken = null;
 
       resourceHeaderKey = getDefaultResourceHeaderKey();
       resourceToken = null;
@@ -172,6 +179,15 @@ public Builder setTransportToken(String name, String version) {
       return this;
     }
 
+    public String getQuotaProjectIdToken() {
+      return quotaProjectIdToken;
+    }
+
+    public Builder setQuotaProjectIdToken(String quotaProjectIdToken) {
+      this.quotaProjectIdToken = quotaProjectIdToken;
+      return this;
+    }
+
     public String getResourceHeaderKey() {
       return resourceHeaderKey;
     }

diff --git a/gax/src/test/java/com/google/api/gax/rpc/ApiClientHeaderProviderTest.java b/gax/src/test/java/com/google/api/gax/rpc/ApiClientHeaderProviderTest.java
@@ -123,4 +123,18 @@ public void testCustomHeaderKeys() {
         .matches("^gl-java/.* gapic/4\\.5\\.6 gax/.* grpc/1\\.2\\.3$");
     assertThat(provider.getHeaders().get("custom-header2")).isEqualTo("test-prefix");
   }
+
+  @Test
+  public void testQuotaProjectHeader() {
+    String quotaProjectHeaderValue = "quota-project-id-value";
+    ApiClientHeaderProvider provider =
+            ApiClientHeaderProvider.newBuilder().setClientLibToken("gccl", "1.2.3")
+                    .setQuotaProjectIdToken(quotaProjectHeaderValue)
+                    .build();
+    assertThat(provider.getHeaders().size()).isEqualTo(2);
+    assertThat(provider.getHeaders().get(X_GOOG_API_CLIENT))
+            .matches("^gl-java/.* gccl/1\\.2\\.3 gax/.*$");
+    assertThat(provider.getHeaders().get(ApiClientHeaderProvider.QUOTA_PROJECT_ID_HEADER_KEY))
+            .matches(quotaProjectHeaderValue);
+  }
 }