Inject namespace header (#2400)

Inject namespace header

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/client/NamespaceInjectWorkflowServiceStubs.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ * Copyright (C) 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this material except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.temporal.client;
+
+import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
+import io.grpc.health.v1.HealthCheckResponse;
+import io.temporal.api.workflowservice.v1.GetSystemInfoResponse;
+import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc;
+import io.temporal.serviceclient.GrpcMetadataProviderInterceptor;
+import io.temporal.serviceclient.WorkflowServiceStubs;
+import io.temporal.serviceclient.WorkflowServiceStubsOptions;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+/** Inject the namespace into the gRPC header */
+class NamespaceInjectWorkflowServiceStubs implements WorkflowServiceStubs {
+  private static Metadata.Key<String> TEMPORAL_NAMESPACE_HEADER_KEY =
+      Metadata.Key.of("temporal-namespace", Metadata.ASCII_STRING_MARSHALLER);
+  private final Metadata metadata;
+  private final WorkflowServiceStubs next;
+
+  public NamespaceInjectWorkflowServiceStubs(WorkflowServiceStubs next, String namespace) {
+    this.next = next;
+    this.metadata = new Metadata();
+    metadata.put(TEMPORAL_NAMESPACE_HEADER_KEY, namespace);
+  }
+
+  @Override
+  public WorkflowServiceStubsOptions getOptions() {
+    return next.getOptions();
+  }
+
+  @Override
+  public WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub() {
+    return next.blockingStub()
+        .withInterceptors(
+            new GrpcMetadataProviderInterceptor(Collections.singleton(() -> metadata)));
+  }
+
+  @Override
+  public WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub() {
+    return next.futureStub()
+        .withInterceptors(
+            new GrpcMetadataProviderInterceptor(Collections.singleton(() -> metadata)));
+  }
+
+  @Override
+  public ManagedChannel getRawChannel() {
+    return next.getRawChannel();
+  }
+
+  @Override
+  public void shutdown() {
+    next.shutdown();
+  }
+
+  @Override
+  public void shutdownNow() {
+    next.shutdownNow();
+  }
+
+  @Override
+  public boolean isShutdown() {
+    return next.isShutdown();
+  }
+
+  @Override
+  public boolean isTerminated() {
+    return next.isTerminated();
+  }
+
+  @Override
+  public boolean awaitTermination(long timeout, TimeUnit unit) {
+    return next.awaitTermination(timeout, unit);
+  }
+
+  @Override
+  public void connect(@Nullable Duration timeout) {
+    next.connect(timeout);
+  }
+
+  @Override
+  public HealthCheckResponse healthCheck() {
+    return next.healthCheck();
+  }
+
+  @Override
+  public Supplier<GetSystemInfoResponse.Capabilities> getServerCapabilities() {
+    return next.getServerCapabilities();
+  }
+}

temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java

@@ -89,6 +89,8 @@ public static WorkflowClient newInstance(
   WorkflowClientInternalImpl(
       WorkflowServiceStubs workflowServiceStubs, WorkflowClientOptions options) {
     options = WorkflowClientOptions.newBuilder(options).validateAndBuildWithDefaults();
+    workflowServiceStubs =
+        new NamespaceInjectWorkflowServiceStubs(workflowServiceStubs, options.getNamespace());
     this.options = options;
     this.workflowServiceStubs = workflowServiceStubs;
     this.metricsScope =

temporal-sdk/src/test/java/io/temporal/authorization/AuthorizationTokenTest.java

@@ -33,6 +33,7 @@
 import io.temporal.workflow.shared.TestWorkflows;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import org.junit.After;
@@ -43,6 +44,8 @@
 import org.junit.runner.Description;
 
 public class AuthorizationTokenTest {
+  private static Metadata.Key<String> TEMPORAL_NAMESPACE_HEADER_KEY =
+      Metadata.Key.of("temporal-namespace", Metadata.ASCII_STRING_MARSHALLER);
   private static final String TASK_QUEUE = "test-workflow";
   private static final String AUTH_TOKEN = "Bearer <token>";
 
@@ -89,7 +92,8 @@ public void start(Listener<RespT> responseListener, Metadata headers) {
                           new GrpcRequest(
                               method.getBareMethodName(),
                               headers.get(
-                                  AuthorizationGrpcMetadataProvider.AUTHORIZATION_HEADER_KEY)));
+                                  AuthorizationGrpcMetadataProvider.AUTHORIZATION_HEADER_KEY),
+                              headers.get(TEMPORAL_NAMESPACE_HEADER_KEY)));
                       super.start(responseListener, headers);
                     }
                   }
@@ -121,9 +125,18 @@ public void allRequestsShouldHaveAnAuthToken() {
     assertEquals("TestWorkflow1-input1", result);
 
     assertFalse(loggedRequests.isEmpty());
+    // These methods are not namespace specific
+    List<String> methodsToSkip =
+        Arrays.asList("GetSystemInfo", "Check", "UnlockTimeSkipping", "LockTimeSkipping");
     for (GrpcRequest grpcRequest : loggedRequests) {
       assertEquals(
           "All requests should have an auth token", AUTH_TOKEN, grpcRequest.authTokenValue);
+      if (!methodsToSkip.contains(grpcRequest.methodName)) {
+        assertEquals(
+            "All requests should have a namespace " + grpcRequest.methodName,
+            testEnvironment.getNamespace(),
+            grpcRequest.namespace);
+      }
     }
   }
 
@@ -139,10 +152,12 @@ public String execute(String input) {
   private static class GrpcRequest {
     final String methodName;
     final String authTokenValue;
+    final String namespace;
 
-    GrpcRequest(String methodName, String authTokenValue) {
+    GrpcRequest(String methodName, String authTokenValue, String namespace) {
       this.methodName = methodName;
       this.authTokenValue = authTokenValue;
+      this.namespace = namespace;
     }
   }
 }

temporal-sdk/src/test/java/io/temporal/workflow/updateTest/UpdateWithStartTest.java

@@ -396,6 +396,7 @@ public void retryUntilDurable() {
     WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub =
         mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
     when(blockingStub.withOption(any(), any())).thenReturn(blockingStub);
+    when(blockingStub.withInterceptors(any())).thenReturn(blockingStub);
     when(blockingStub.withDeadline(any())).thenReturn(blockingStub);
 
     Scope scope = mock(Scope.class);
@@ -457,6 +458,7 @@ public void handleSuccessfulStartButUpdateOnlyErr() {
     WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub =
         mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
     when(blockingStub.withOption(any(), any())).thenReturn(blockingStub);
+    when(blockingStub.withInterceptors(any())).thenReturn(blockingStub);
     when(blockingStub.withDeadline(any())).thenReturn(blockingStub);
 
     Scope scope = mock(Scope.class);

temporal-serviceclient/src/main/java/io/temporal/serviceclient/ServiceStubsOptions.java

@@ -634,10 +634,7 @@ public T addGrpcMetadataProvider(GrpcMetadataProvider grpcMetadataProvider) {
      */
     public T addApiKey(AuthorizationTokenSupplier apiKey) {
       addGrpcMetadataProvider(
-          new AuthorizationGrpcMetadataProvider(
-              () -> {
-                return "Bearer " + apiKey.supply();
-              }));
+          new AuthorizationGrpcMetadataProvider(() -> "Bearer " + apiKey.supply()));
       return self();
     }