refactor: add shortcut to create a tenant id request context (#8)

* refactor: add shortcut to create a tenant id request context

* refactor: add syntax to wrap an existing rx call

* refactor: add syntax to wrap existing maybe call

Author: aaron-steinfeld

grpc-client-rx-utils/src/main/java/org/hypertrace/core/grpcutils/client/rx/DefaultGrpcRxExecutionContext.java

@@ -3,10 +3,12 @@
 import io.grpc.Context;
 import io.grpc.stub.StreamObserver;
 import io.reactivex.rxjava3.core.Completable;
+import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Single;
 import java.util.concurrent.Callable;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import org.hypertrace.core.grpcutils.context.RequestContext;
 
 class DefaultGrpcRxExecutionContext implements GrpcRxExecutionContext {
@@ -22,6 +24,24 @@ public <TResp> Single<TResp> call(Callable<TResp> callable) {
     return Single.fromCallable(buildContext().wrap(callable));
   }
 
+  @Override
+  public <TResp> Single<TResp> wrapSingle(Supplier<Single<TResp>> singleSupplier) {
+    try {
+      return buildContext().call(singleSupplier::get);
+    } catch (Exception e) {
+      return Single.error(e);
+    }
+  }
+
+  @Override
+  public <TResp> Maybe<TResp> wrapMaybe(Supplier<Maybe<TResp>> maybeSupplier) {
+    try {
+      return buildContext().call(maybeSupplier::get);
+    } catch (Exception e) {
+      return Maybe.error(e);
+    }
+  }
+
   @Override
   public Completable run(Runnable runnable) {
     return Completable.fromRunnable(buildContext().wrap(runnable));

grpc-client-rx-utils/src/main/java/org/hypertrace/core/grpcutils/client/rx/GrpcRxExecutionContext.java

@@ -2,10 +2,12 @@
 
 import io.grpc.stub.StreamObserver;
 import io.reactivex.rxjava3.core.Completable;
+import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Single;
 import java.util.concurrent.Callable;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import org.hypertrace.core.grpcutils.context.RequestContext;
 
 /** An execution context that can turn various types of executions into their Rx equivalents. */
@@ -16,6 +18,16 @@ public interface GrpcRxExecutionContext {
    */
   <TResp> Single<TResp> call(Callable<TResp> callable);
 
+  /**
+   * Creates the provided single in this execution context, returning the result.
+   */
+  <TResp> Single<TResp> wrapSingle(Supplier<Single<TResp>> singleSupplier);
+
+  /**
+   * Creates the provided maybe in this execution context, returning the result.
+   */
+  <TResp> Maybe<TResp> wrapMaybe(Supplier<Maybe<TResp>> maybeSupplier);
+
   /**
    * Executes the given runnable in this execution context, triggering completion or error once the
    * call completes.
@@ -34,4 +46,8 @@ static GrpcRxExecutionContext forCurrentContext() {
   static GrpcRxExecutionContext forContext(RequestContext requestContext) {
     return new DefaultGrpcRxExecutionContext(requestContext);
   }
+
+  static GrpcRxExecutionContext forTenantContext(String tenantId) {
+    return forContext(RequestContext.forTenantId(tenantId));
+  }
 }

grpc-client-rx-utils/src/test/java/org/hypertrace/core/grpcutils/client/rx/DefaultGrpcRxExecutionContextTest.java

@@ -8,6 +8,7 @@
 import static org.mockito.Mockito.when;
 
 import io.reactivex.rxjava3.core.Completable;
+import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Observable;
 import io.reactivex.rxjava3.core.Single;
 import io.reactivex.rxjava3.observers.TestObserver;
@@ -76,4 +77,22 @@ void canPropagateErrors() {
     completable.subscribe(testObserver);
     testObserver.assertError(UnsupportedOperationException.class);
   }
+
+  @Test
+  void canWrapSingle() {
+    Single<?> single =
+        new DefaultGrpcRxExecutionContext(this.mockContext)
+            .wrapSingle(() -> Single.just(RequestContext.CURRENT.get().getTenantId()));
+
+    assertEquals(TEST_TENANT_ID_OPTIONAL, single.blockingGet());
+  }
+
+  @Test
+  void canWrapMaybe() {
+    Maybe<?> maybe =
+        new DefaultGrpcRxExecutionContext(this.mockContext)
+            .wrapMaybe(() -> Maybe.just(RequestContext.CURRENT.get().getTenantId()));
+
+    assertEquals(TEST_TENANT_ID_OPTIONAL, maybe.blockingGet());
+  }
 }

grpc-client-rx-utils/src/test/java/org/hypertrace/core/grpcutils/client/rx/GrpcRxExecutionContextTest.java

@@ -1,8 +1,10 @@
 package org.hypertrace.core.grpcutils.client.rx;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertSame;
 
 import io.grpc.Context;
+import java.util.Optional;
 import org.hypertrace.core.grpcutils.context.RequestContext;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -39,4 +41,15 @@ void canCreateExecutionContextWithProvidedContext() throws Exception {
                         .call(RequestContext.CURRENT::get))
             .blockingGet());
   }
+
+  @Test
+  void canCreateExecutionContextForProvidedTenant() throws Exception {
+    final String testTenant = "testTenant";
+    assertEquals(
+        Optional.of(testTenant),
+        GrpcRxExecutionContext.forTenantContext(testTenant)
+            .call(RequestContext.CURRENT::get)
+            .map(RequestContext::getTenantId)
+            .blockingGet());
+  }
 }

grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java

@@ -6,25 +6,26 @@
 import java.util.Optional;
 
 /**
- * Context of the GRPC request that should be carried and can made available to the services
- * so that the service can use them.
- * We use this to propagate headers across services.
+ * Context of the GRPC request that should be carried and can made available to the services so that
+ * the service can use them. We use this to propagate headers across services.
  */
 public class RequestContext {
   public static final Context.Key<RequestContext> CURRENT = Context.key("request_context");
 
+  public static RequestContext forTenantId(String tenantId) {
+    RequestContext requestContext = new RequestContext();
+    requestContext.add(RequestContextConstants.TENANT_ID_HEADER_KEY, tenantId);
+    return requestContext;
+  }
+
   private final Map<String, String> headers = new HashMap<>();
 
-  /**
-   * Reads tenant id from this RequestContext based on the tenant id http header and returns it.
-   */
+  /** Reads tenant id from this RequestContext based on the tenant id http header and returns it. */
   public Optional<String> getTenantId() {
     return get(RequestContextConstants.TENANT_ID_HEADER_KEY);
   }
 
-  /**
-   * Method to read all GRPC request headers from this RequestContext.
-   */
+  /** Method to read all GRPC request headers from this RequestContext. */
   public Map<String, String> getRequestHeaders() {
     return getAll();
   }

grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/RequestContextTest.java

@@ -5,9 +5,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
-/**
- * Unit tests for {@link RequestContext} and utility methods in it.
- */
+/** Unit tests for {@link RequestContext} and utility methods in it. */
 public class RequestContextTest {
   private static final String TENANT_ID = "example-tenant-id";
   private static final String TEST_AUTH_HEADER = "Bearer sample-auth-header";
@@ -34,9 +32,20 @@ public void testGetRequestHeaders() {
 
     Assertions.assertEquals(
         Map.of(
-            RequestContextConstants.AUTHORIZATION_HEADER, TEST_AUTH_HEADER,
-            "x-some-tenant-header", "v1"
-        ),
+            RequestContextConstants.AUTHORIZATION_HEADER,
+            TEST_AUTH_HEADER,
+            "x-some-tenant-header",
+            "v1"),
         requestHeaders);
   }
+
+  @Test
+  public void testCreateForTenantId() {
+    RequestContext requestContext = RequestContext.forTenantId(TENANT_ID);
+    Assertions.assertEquals(Optional.of(TENANT_ID), requestContext.getTenantId());
+    Assertions.assertEquals(
+        Optional.of(TENANT_ID), requestContext.get(RequestContextConstants.TENANT_ID_HEADER_KEY));
+    Assertions.assertEquals(
+        Map.of(RequestContextConstants.TENANT_ID_HEADER_KEY, TENANT_ID), requestContext.getAll());
+  }
 }