Add header to indicate user tracking suppression (#89)

Author: DibyojyotiS

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

@@ -3,6 +3,7 @@
 import static io.grpc.Metadata.ASCII_STRING_MARSHALLER;
 import static java.util.Objects.requireNonNull;
 import static org.hypertrace.core.grpcutils.context.RequestContextConstants.CACHE_MEANINGFUL_HEADERS;
+import static org.hypertrace.core.grpcutils.context.RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY;
 import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
 
 import com.google.common.collect.ListMultimap;
@@ -120,6 +121,21 @@ public Optional<String> getRequestId() {
     return this.getHeaderValue(RequestContextConstants.REQUEST_ID_HEADER_KEY);
   }
 
+  public boolean isUserTrackingSuppressed() {
+    return this.getHeaderValue(SUPPRESS_USER_TRACKING_HEADER_KEY)
+        .map(Boolean::parseBoolean)
+        .orElse(false);
+  }
+
+  public RequestContext withUserTrackingSuppressed(boolean suppressed) {
+    this.removeHeader(SUPPRESS_USER_TRACKING_HEADER_KEY);
+    return this.put(SUPPRESS_USER_TRACKING_HEADER_KEY, String.valueOf(suppressed));
+  }
+
+  public RequestContext withUserTrackingSuppressed() {
+    return this.withUserTrackingSuppressed(true);
+  }
+
   private Optional<Jwt> getJwt() {
     return this.getHeaderValue(RequestContextConstants.AUTHORIZATION_HEADER)
         .flatMap(JWT_PARSER::fromAuthHeader);

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

@@ -13,6 +13,7 @@ public class RequestContextConstants {
   public static final String TENANT_ID_HEADER_KEY = "x-tenant-id";
   public static final String REQUEST_ID_HEADER_KEY = "request-id";
   public static final String CONTEXT_ID_HEADER_KEY = "context-id";
+  public static final String SUPPRESS_USER_TRACKING_HEADER_KEY = "x-suppress-user-tracking";
 
   public static final Metadata.Key<String> TENANT_ID_METADATA_KEY =
       Metadata.Key.of(TENANT_ID_HEADER_KEY, ASCII_STRING_MARSHALLER);
@@ -24,6 +25,7 @@ public class RequestContextConstants {
       Set.of(
           TENANT_ID_HEADER_KEY,
           CONTEXT_ID_HEADER_KEY,
+          SUPPRESS_USER_TRACKING_HEADER_KEY,
           "X-B3-",
           "grpc-trace-bin",
           "traceparent",

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

@@ -2,7 +2,9 @@
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.google.common.collect.ImmutableList;
 import io.grpc.Metadata;
@@ -251,4 +253,86 @@ void backwardsCompatibilityForAdd() {
     assertEquals(Optional.of("f-v2"), requestContext.getHeaderValue("first"));
     assertEquals(List.of("f-v2"), requestContext.getAllHeaderValues("first"));
   }
+
+  @Test
+  void testUserTrackingSuppressedDefault() {
+    RequestContext requestContext = new RequestContext();
+    assertFalse(requestContext.isUserTrackingSuppressed());
+  }
+
+  @Test
+  void testUserTrackingSuppressedWithTrueValue() {
+    RequestContext requestContext =
+        new RequestContext().put(RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY, "true");
+    assertTrue(requestContext.isUserTrackingSuppressed());
+  }
+
+  @Test
+  void testUserTrackingSuppressedWithFalseValue() {
+    RequestContext requestContext =
+        new RequestContext()
+            .put(RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY, "false");
+    assertFalse(requestContext.isUserTrackingSuppressed());
+  }
+
+  @Test
+  void testWithUserTrackingSuppressedBoolean() {
+    RequestContext originalContext1 = new RequestContext();
+    RequestContext suppressedContext = originalContext1.withUserTrackingSuppressed(true);
+    assertTrue(suppressedContext.isUserTrackingSuppressed());
+
+    RequestContext originalContext2 = new RequestContext();
+    RequestContext unsuppressedContext = originalContext2.withUserTrackingSuppressed(false);
+    assertFalse(unsuppressedContext.isUserTrackingSuppressed());
+  }
+
+  @Test
+  void testWithUserTrackingSuppressedDefault() {
+    RequestContext originalContext = new RequestContext();
+
+    RequestContext suppressedContext = originalContext.withUserTrackingSuppressed();
+    assertTrue(suppressedContext.isUserTrackingSuppressed());
+  }
+
+  @Test
+  void testWithUserTrackingSuppressedPreservesOtherHeaders() {
+    RequestContext originalContext =
+        new RequestContext()
+            .put(RequestContextConstants.TENANT_ID_HEADER_KEY, "tenant-123")
+            .put("other-header", "other-value");
+
+    RequestContext suppressedContext = originalContext.withUserTrackingSuppressed();
+
+    assertTrue(suppressedContext.isUserTrackingSuppressed());
+    assertEquals(Optional.of("tenant-123"), suppressedContext.getTenantId());
+    assertEquals(Optional.of("other-value"), suppressedContext.getHeaderValue("other-header"));
+  }
+
+  @Test
+  void testWithUserTrackingSuppressedOverridesExisting() {
+    RequestContext originalContext =
+        new RequestContext()
+            .put(RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY, "false");
+
+    assertFalse(originalContext.isUserTrackingSuppressed());
+
+    RequestContext suppressedContext = originalContext.withUserTrackingSuppressed();
+    assertTrue(suppressedContext.isUserTrackingSuppressed());
+  }
+
+  @Test
+  void testUserTrackingSuppressedHeaderPropagation() {
+    Metadata metadata = new Metadata();
+    metadata.put(
+        Metadata.Key.of(
+            RequestContextConstants.SUPPRESS_USER_TRACKING_HEADER_KEY,
+            Metadata.ASCII_STRING_MARSHALLER),
+        "true");
+    metadata.put(RequestContextConstants.TENANT_ID_METADATA_KEY, "test-tenant");
+
+    RequestContext requestContext = RequestContext.fromMetadata(metadata);
+
+    assertTrue(requestContext.isUserTrackingSuppressed());
+    assertEquals(Optional.of("test-tenant"), requestContext.getTenantId());
+  }
 }