Skip to content

Commit 63719ae

Browse files
aaron-steinfeldaaron-steinfeld
authored
feat: add support for more advanced contextual keys (#25)
1 parent 967f216 commit 63719ae

File tree

3 files changed

+102
-25
lines changed

3 files changed

+102
-25
lines changed

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

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package org.hypertrace.core.grpcutils.context;
22

3+
import java.util.Collection;
34
import java.util.Map;
45
import java.util.Map.Entry;
56
import java.util.Objects;
7+
import java.util.Set;
68
import java.util.function.Consumer;
79
import java.util.function.Function;
810
import java.util.function.Supplier;
@@ -11,12 +13,13 @@
1113
class DefaultContextualKey<T> implements ContextualKey<T> {
1214
private final RequestContext context;
1315
private final T data;
14-
private final Map<String, String> meaningfulContextHeaders;
16+
private final Map<String, String> cacheableContextHeaders;
1517

16-
DefaultContextualKey(RequestContext context, T data) {
18+
DefaultContextualKey(RequestContext context, T data, Collection<String> cacheableHeaderNames) {
1719
this.context = context;
1820
this.data = data;
19-
this.meaningfulContextHeaders = this.extractMeaningfulHeaders(context.getRequestHeaders());
21+
this.cacheableContextHeaders =
22+
this.extractCacheableHeaders(context.getRequestHeaders(), cacheableHeaderNames);
2023
}
2124

2225
@Override
@@ -55,30 +58,32 @@ public boolean equals(Object o) {
5558
if (o == null || getClass() != o.getClass()) return false;
5659
DefaultContextualKey<?> that = (DefaultContextualKey<?>) o;
5760
return Objects.equals(getData(), that.getData())
58-
&& meaningfulContextHeaders.equals(that.meaningfulContextHeaders);
61+
&& cacheableContextHeaders.equals(that.cacheableContextHeaders);
5962
}
6063

6164
@Override
6265
public int hashCode() {
63-
return Objects.hash(getData(), meaningfulContextHeaders);
66+
return Objects.hash(getData(), cacheableContextHeaders);
6467
}
6568

6669
@Override
6770
public String toString() {
6871
return "DefaultContextualKey{"
6972
+ "data="
7073
+ data
71-
+ ", meaningfulContextHeaders="
72-
+ meaningfulContextHeaders
74+
+ ", cacheableContextHeaders="
75+
+ cacheableContextHeaders
7376
+ '}';
7477
}
7578

76-
private Map<String, String> extractMeaningfulHeaders(Map<String, String> allHeaders) {
79+
private Map<String, String> extractCacheableHeaders(
80+
Map<String, String> allHeaders, Collection<String> cacheableHeaderNames) {
81+
Set<String> cacheableHeaderNameSet =
82+
cacheableHeaderNames.stream()
83+
.map(String::toLowerCase)
84+
.collect(Collectors.toUnmodifiableSet());
7785
return allHeaders.entrySet().stream()
78-
.filter(
79-
entry ->
80-
RequestContextConstants.CACHE_MEANINGFUL_HEADERS.contains(
81-
entry.getKey().toLowerCase()))
86+
.filter(entry -> cacheableHeaderNameSet.contains(entry.getKey().toLowerCase()))
8287
.collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
8388
}
8489
}

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

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package org.hypertrace.core.grpcutils.context;
22

3-
import io.grpc.Context;
3+
import static org.hypertrace.core.grpcutils.context.RequestContextConstants.CACHE_MEANINGFUL_HEADERS;
4+
import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
45

6+
import io.grpc.Context;
57
import java.util.Collections;
68
import java.util.HashMap;
79
import java.util.List;
@@ -87,11 +89,53 @@ public void run(@Nonnull Runnable runnable) {
8789
Context.current().withValue(RequestContext.CURRENT, this).run(runnable);
8890
}
8991

92+
/**
93+
* @deprecated - Use {@link #buildInternalContextualKey(Object)} ()} or {@link
94+
* #buildUserContextualKey(Object)} instead, as appropriate. This delegates to {@link
95+
* #buildUserContextualKey(Object)} to match its original implementation for compatibility.
96+
*/
97+
@Deprecated
9098
public <T> ContextualKey<T> buildContextualKey(T data) {
91-
return new DefaultContextualKey<>(this, data);
99+
return this.buildUserContextualKey(data);
92100
}
93101

102+
/**
103+
* @deprecated - Use {@link #buildInternalContextualKey()} or {@link #buildUserContextualKey()}
104+
* instead, as appropriate. This delegates to {@link #buildUserContextualKey()} to match its
105+
* original implementation for compatibility.
106+
*/
107+
@Deprecated
94108
public ContextualKey<Void> buildContextualKey() {
95-
return new DefaultContextualKey<>(this, null);
109+
return this.buildUserContextualKey();
110+
}
111+
112+
/** This returns a cache key based on this user request context. */
113+
public ContextualKey<Void> buildUserContextualKey() {
114+
return new DefaultContextualKey<>(this, null, CACHE_MEANINGFUL_HEADERS);
115+
}
116+
117+
/**
118+
* An extension of {@link #buildUserContextualKey()} that also includes {@code data} as part of
119+
* the cache key.
120+
*/
121+
public <T> ContextualKey<T> buildUserContextualKey(T data) {
122+
return new DefaultContextualKey<>(this, data, CACHE_MEANINGFUL_HEADERS);
123+
}
124+
125+
/**
126+
* This returns a cache key based on this internal request context. It should only be used if this
127+
* request is not on behalf of a specific user - either one initiated by the platform or by the
128+
* agent.
129+
*/
130+
public ContextualKey<Void> buildInternalContextualKey() {
131+
return new DefaultContextualKey<>(this, null, List.of(TENANT_ID_HEADER_KEY));
132+
}
133+
134+
/**
135+
* An extension of {@link @buildInternalContextualKey()} that also includes {@code data} as part
136+
* of the cache key.
137+
*/
138+
public <T> ContextualKey<T> buildInternalContextualKey(T data) {
139+
return new DefaultContextualKey<>(this, data, List.of(TENANT_ID_HEADER_KEY));
96140
}
97141
}

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

Lines changed: 38 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.hypertrace.core.grpcutils.context;
22

3+
import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
34
import static org.junit.jupiter.api.Assertions.assertEquals;
45
import static org.junit.jupiter.api.Assertions.assertNotEquals;
56
import static org.junit.jupiter.api.Assertions.assertSame;
@@ -10,6 +11,7 @@
1011
import static org.mockito.Mockito.times;
1112
import static org.mockito.Mockito.verify;
1213

14+
import java.util.List;
1315
import java.util.function.Consumer;
1416
import java.util.function.Function;
1517
import java.util.function.Supplier;
@@ -20,7 +22,8 @@ class DefaultContextualKeyTest {
2022
@Test
2123
void callsProvidedMethodsInContext() {
2224
RequestContext testContext = RequestContext.forTenantId("test-tenant");
23-
ContextualKey<String> key = new DefaultContextualKey<>(testContext, "input");
25+
ContextualKey<String> key =
26+
new DefaultContextualKey<>(testContext, "input", List.of(TENANT_ID_HEADER_KEY));
2427

2528
Function<String, String> testFunction =
2629
value ->
@@ -40,7 +43,8 @@ void callsProvidedMethodsInContext() {
4043
@Test
4144
void runsProvidedMethodInContext() {
4245
RequestContext testContext = RequestContext.forTenantId("test-tenant");
43-
ContextualKey<String> key = new DefaultContextualKey<>(testContext, "input");
46+
ContextualKey<String> key =
47+
new DefaultContextualKey<>(testContext, "input", List.of(TENANT_ID_HEADER_KEY));
4448

4549
Consumer<String> testConsumer = mock(Consumer.class);
4650

@@ -67,19 +71,43 @@ void matchesEquivalentKeysOnly() {
6771
RequestContext tenant2Context = RequestContext.forTenantId("second");
6872

6973
assertEquals(
70-
new DefaultContextualKey<>(tenant1Context, "input"),
71-
new DefaultContextualKey<>(tenant1Context, "input"));
74+
new DefaultContextualKey<>(tenant1Context, "input", List.of(TENANT_ID_HEADER_KEY)),
75+
new DefaultContextualKey<>(tenant1Context, "input", List.of(TENANT_ID_HEADER_KEY)));
7276

7377
assertEquals(
74-
new DefaultContextualKey<>(tenant1Context, "input"),
75-
new DefaultContextualKey<>(alternateTenant1Context, "input"));
78+
new DefaultContextualKey<>(tenant1Context, "input", List.of(TENANT_ID_HEADER_KEY)),
79+
new DefaultContextualKey<>(
80+
alternateTenant1Context, "input", List.of(TENANT_ID_HEADER_KEY)));
7681

7782
assertNotEquals(
78-
new DefaultContextualKey<>(tenant1Context, "input"),
79-
new DefaultContextualKey<>(tenant2Context, "input"));
83+
new DefaultContextualKey<>(tenant1Context, "input", List.of(TENANT_ID_HEADER_KEY)),
84+
new DefaultContextualKey<>(tenant2Context, "input", List.of(TENANT_ID_HEADER_KEY)));
8085

8186
assertNotEquals(
82-
new DefaultContextualKey<>(tenant1Context, "input"),
83-
new DefaultContextualKey<>(tenant1Context, "other input"));
87+
new DefaultContextualKey<>(tenant1Context, "input", List.of(TENANT_ID_HEADER_KEY)),
88+
new DefaultContextualKey<>(tenant1Context, "other input", List.of(TENANT_ID_HEADER_KEY)));
89+
}
90+
91+
@Test
92+
void matchesMultipleHeaders() {
93+
RequestContext firstContext = RequestContext.forTenantId("tenant");
94+
firstContext.add("secondHeader", "second");
95+
firstContext.add("thirdHeader", "third");
96+
97+
RequestContext secondContext = RequestContext.forTenantId("tenant");
98+
secondContext.add("secondHeader", "second");
99+
secondContext.add("thirdHeader", "fourth");
100+
101+
assertEquals(
102+
new DefaultContextualKey<>(
103+
firstContext, "input", List.of(TENANT_ID_HEADER_KEY, "secondHeader")),
104+
new DefaultContextualKey<>(
105+
secondContext, "input", List.of(TENANT_ID_HEADER_KEY, "secondHeader")));
106+
107+
assertNotEquals(
108+
new DefaultContextualKey<>(
109+
firstContext, "input", List.of(TENANT_ID_HEADER_KEY, "secondHeader", "thirdHeader")),
110+
new DefaultContextualKey<>(
111+
secondContext, "input", List.of(TENANT_ID_HEADER_KEY, "secondHeader", "thirdHeader")));
84112
}
85113
}

0 commit comments

Comments
 (0)