Skip to content

Commit 577ab05

Browse files
aaron-steinfeldaaron-steinfeld
authored
Add Request ID to RequestContext (#35)
* feat: add request id and expose context hydrate from metadata * refactor: initialize request id for tenant-based creation * test: update request context test
1 parent c7e058b commit 577ab05

File tree

5 files changed

+130
-127
lines changed

5 files changed

+130
-127
lines changed

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

+36
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,14 @@
44
import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
55

66
import io.grpc.Context;
7+
import io.grpc.Metadata;
8+
import java.nio.charset.StandardCharsets;
79
import java.util.Collections;
810
import java.util.HashMap;
911
import java.util.List;
1012
import java.util.Map;
1113
import java.util.Optional;
14+
import java.util.UUID;
1215
import java.util.concurrent.Callable;
1316
import javax.annotation.Nonnull;
1417

@@ -22,6 +25,35 @@ public class RequestContext {
2225
public static RequestContext forTenantId(String tenantId) {
2326
RequestContext requestContext = new RequestContext();
2427
requestContext.add(RequestContextConstants.TENANT_ID_HEADER_KEY, tenantId);
28+
requestContext.add(RequestContextConstants.REQUEST_ID_HEADER_KEY, UUID.randomUUID().toString());
29+
return requestContext;
30+
}
31+
32+
public static RequestContext fromMetadata(Metadata metadata) {
33+
RequestContext requestContext = new RequestContext();
34+
35+
// Go over all the headers and copy the allowed headers to the RequestContext.
36+
metadata.keys().stream()
37+
.filter(
38+
k ->
39+
RequestContextConstants.HEADER_PREFIXES_TO_BE_PROPAGATED.stream()
40+
.anyMatch(prefix -> k.toLowerCase().startsWith(prefix.toLowerCase())))
41+
.forEach(
42+
k -> {
43+
String value;
44+
// check if key ends with binary suffix
45+
if (k.toLowerCase().endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
46+
byte[] bytes = metadata.get(Metadata.Key.of(k, Metadata.BINARY_BYTE_MARSHALLER));
47+
value = new String(bytes, StandardCharsets.UTF_8);
48+
} else {
49+
value = metadata.get(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER));
50+
}
51+
// The value could be null or empty for some keys so validate that.
52+
if (value != null && !value.isEmpty()) {
53+
requestContext.add(k, value);
54+
}
55+
});
56+
2557
return requestContext;
2658
}
2759

@@ -53,6 +85,10 @@ public List<String> getRoles(String rolesClaim) {
5385
return getJwt().map(jwt -> jwt.getRoles(rolesClaim)).orElse(Collections.emptyList());
5486
}
5587

88+
public Optional<String> getRequestId() {
89+
return this.get(RequestContextConstants.REQUEST_ID_HEADER_KEY);
90+
}
91+
5692
private Optional<Jwt> getJwt() {
5793
return get(RequestContextConstants.AUTHORIZATION_HEADER).flatMap(jwtParser::fromAuthHeader);
5894
}

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
*/
1212
public class RequestContextConstants {
1313
public static final String TENANT_ID_HEADER_KEY = "x-tenant-id";
14+
public static final String REQUEST_ID_HEADER_KEY = "request-id";
1415

1516
public static final Metadata.Key<String> TENANT_ID_METADATA_KEY =
1617
Metadata.Key.of(TENANT_ID_HEADER_KEY, ASCII_STRING_MARSHALLER);
@@ -25,7 +26,8 @@ public class RequestContextConstants {
2526
"grpc-trace-bin",
2627
"traceparent",
2728
"tracestate",
28-
AUTHORIZATION_HEADER);
29+
AUTHORIZATION_HEADER,
30+
REQUEST_ID_HEADER_KEY);
2931

3032
/**
3133
* These headers may affect returned results and should be accounted for in any cached remote

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

+90-1
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,16 @@
11
package org.hypertrace.core.grpcutils.context;
22

3+
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
34
import static org.junit.jupiter.api.Assertions.assertEquals;
45

56
import com.google.common.collect.ImmutableList;
7+
import io.grpc.Metadata;
68
import java.util.List;
79
import java.util.Map;
810
import java.util.Optional;
11+
import java.util.Set;
12+
import java.util.UUID;
13+
import org.junit.jupiter.api.Assertions;
914
import org.junit.jupiter.api.Test;
1015

1116
/** Unit tests for {@link RequestContext} and utility methods in it. */
@@ -45,11 +50,18 @@ void testGetRequestHeaders() {
4550
@Test
4651
void testCreateForTenantId() {
4752
RequestContext requestContext = RequestContext.forTenantId(TENANT_ID);
53+
String requestId = requestContext.getRequestId().orElseThrow();
54+
assertDoesNotThrow(() -> UUID.fromString(requestId));
4855
assertEquals(Optional.of(TENANT_ID), requestContext.getTenantId());
4956
assertEquals(
5057
Optional.of(TENANT_ID), requestContext.get(RequestContextConstants.TENANT_ID_HEADER_KEY));
5158
assertEquals(
52-
Map.of(RequestContextConstants.TENANT_ID_HEADER_KEY, TENANT_ID), requestContext.getAll());
59+
Optional.of(requestId), requestContext.get(RequestContextConstants.REQUEST_ID_HEADER_KEY));
60+
assertEquals(
61+
Set.of(
62+
RequestContextConstants.TENANT_ID_HEADER_KEY,
63+
RequestContextConstants.REQUEST_ID_HEADER_KEY),
64+
requestContext.getAll().keySet());
5365
}
5466

5567
@Test
@@ -67,4 +79,81 @@ void testRolesArePropagatedInRequestContext() {
6779
List<String> actualRoles = requestContext.getRoles("roles");
6880
assertEquals(expectedRoles, actualRoles);
6981
}
82+
83+
@Test
84+
public void testMetadataKeys() {
85+
Metadata metadata = new Metadata();
86+
metadata.put(
87+
Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER),
88+
"Bearer Some-bearer-auth");
89+
metadata.put(RequestContextConstants.TENANT_ID_METADATA_KEY, "test-tenant-id");
90+
91+
RequestContext requestContext = RequestContext.fromMetadata(metadata);
92+
93+
Assertions.assertEquals(2, requestContext.getAll().size());
94+
// GRPC Metadata keys are lower cased during creation.
95+
Assertions.assertEquals("Bearer Some-bearer-auth", requestContext.get("authorization").get());
96+
Assertions.assertEquals(
97+
"test-tenant-id",
98+
requestContext.get(RequestContextConstants.TENANT_ID_METADATA_KEY.name()).get());
99+
100+
metadata = new Metadata();
101+
metadata.put(
102+
Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER),
103+
"Bearer Some-bearer-auth");
104+
metadata.put(RequestContextConstants.TENANT_ID_METADATA_KEY, "test-tenant-id");
105+
metadata.put(
106+
Metadata.Key.of("x-some-other-header", Metadata.ASCII_STRING_MARSHALLER),
107+
"Some-other-header-val");
108+
109+
requestContext = RequestContext.fromMetadata(metadata);
110+
111+
Assertions.assertEquals(2, requestContext.getAll().size());
112+
Assertions.assertEquals("Bearer Some-bearer-auth", requestContext.get("authorization").get());
113+
Assertions.assertEquals(
114+
"test-tenant-id",
115+
requestContext.get(RequestContextConstants.TENANT_ID_METADATA_KEY.name()).get());
116+
117+
// Test that header keys are lowercased
118+
metadata = new Metadata();
119+
metadata.put(
120+
Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER),
121+
"Bearer Some-bearer-auth-2");
122+
metadata.put(
123+
Metadata.Key.of("X-tenant-Id", Metadata.ASCII_STRING_MARSHALLER), "test-tenant-id-2");
124+
metadata.put(
125+
Metadata.Key.of("X-Some-Other-Header", Metadata.ASCII_STRING_MARSHALLER),
126+
"Some-other-header-val-2");
127+
128+
requestContext = RequestContext.fromMetadata(metadata);
129+
130+
Assertions.assertEquals(2, requestContext.getAll().size());
131+
Assertions.assertEquals("Bearer Some-bearer-auth-2", requestContext.get("authorization").get());
132+
Assertions.assertEquals(
133+
"test-tenant-id-2",
134+
requestContext.get(RequestContextConstants.TENANT_ID_METADATA_KEY.name()).get());
135+
136+
metadata = new Metadata();
137+
metadata.put(
138+
Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER),
139+
"Bearer Some-bearer-auth-3");
140+
metadata.put(
141+
Metadata.Key.of("X-tenant-Id", Metadata.ASCII_STRING_MARSHALLER), "test-tenant-id-3");
142+
metadata.put(
143+
Metadata.Key.of("X-Some-Other-Header", Metadata.ASCII_STRING_MARSHALLER),
144+
"Some-other-header-val-3");
145+
metadata.put(
146+
Metadata.Key.of("grpc-trace-bin", Metadata.BINARY_BYTE_MARSHALLER),
147+
"AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE".getBytes());
148+
149+
requestContext = RequestContext.fromMetadata(metadata);
150+
151+
Assertions.assertEquals(3, requestContext.getAll().size());
152+
Assertions.assertEquals("Bearer Some-bearer-auth-3", requestContext.get("authorization").get());
153+
Assertions.assertEquals(
154+
"test-tenant-id-3",
155+
requestContext.get(RequestContextConstants.TENANT_ID_METADATA_KEY.name()).get());
156+
Assertions.assertEquals(
157+
"AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE", requestContext.get("grpc-trace-bin").get());
158+
}
70159
}

‎grpc-server-utils/src/main/java/org/hypertrace/core/grpcutils/server/RequestContextServerInterceptor.java

+1-38
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,7 @@
66
import io.grpc.ServerCall;
77
import io.grpc.ServerCallHandler;
88
import io.grpc.ServerInterceptor;
9-
import java.nio.charset.StandardCharsets;
109
import org.hypertrace.core.grpcutils.context.RequestContext;
11-
import org.hypertrace.core.grpcutils.context.RequestContextConstants;
1210

1311
/**
1412
* Interceptor which intercepts the request headers to extract request context and sets it in the
@@ -22,44 +20,9 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
2220
ServerCall<ReqT, RespT> serverCall,
2321
Metadata metadata,
2422
ServerCallHandler<ReqT, RespT> serverCallHandler) {
25-
RequestContext requestContext = createRequestContextFromMetadata(metadata);
23+
RequestContext requestContext = RequestContext.fromMetadata(metadata);
2624
Context ctx = Context.current().withValue(RequestContext.CURRENT, requestContext);
2725

2826
return Contexts.interceptCall(ctx, serverCall, metadata, serverCallHandler);
2927
}
30-
31-
/**
32-
* Add headers that our services would need internally or would need to propagate to downstream
33-
* services, from the metadata object into a RequestContext object. object.
34-
*
35-
* @param metadata
36-
* @return
37-
*/
38-
RequestContext createRequestContextFromMetadata(Metadata metadata) {
39-
RequestContext requestContext = new RequestContext();
40-
41-
// Go over all the headers and copy the allowed headers to the RequestContext.
42-
metadata.keys().stream()
43-
.filter(
44-
k ->
45-
RequestContextConstants.HEADER_PREFIXES_TO_BE_PROPAGATED.stream()
46-
.anyMatch(prefix -> k.toLowerCase().startsWith(prefix.toLowerCase())))
47-
.forEach(
48-
k -> {
49-
String value;
50-
// check if key ends with binary suffix
51-
if (k.toLowerCase().endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
52-
byte[] bytes = metadata.get(Metadata.Key.of(k, Metadata.BINARY_BYTE_MARSHALLER));
53-
value = new String(bytes, StandardCharsets.UTF_8);
54-
} else {
55-
value = metadata.get(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER));
56-
}
57-
// The value could be null or empty for some keys so validate that.
58-
if (value != null && !value.isEmpty()) {
59-
requestContext.add(k, value);
60-
}
61-
});
62-
63-
return requestContext;
64-
}
6528
}

‎grpc-server-utils/src/test/java/org/hypertrace/core/grpcutils/server/RequestContextServerInterceptorTests.java

-87
This file was deleted.

0 commit comments

Comments
 (0)