Skip to content

Commit f274f62

Browse files
pthirunclaude
pthirun
and
claude
committed
[Controller] Add gRPC support for getKeySchema API
Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
1 parent 58b96b9 commit f274f62

8 files changed

Lines changed: 286 additions & 8 deletions

File tree

‎internal/venice-common/src/main/proto/controller/SchemaGrpcService.proto‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ option java_multiple_files = true;
88

99
service SchemaGrpcService {
1010
rpc getValueSchema(GetValueSchemaGrpcRequest) returns (GetValueSchemaGrpcResponse);
11+
rpc getKeySchema(GetKeySchemaGrpcRequest) returns (GetKeySchemaGrpcResponse);
1112
}
1213

1314
message GetValueSchemaGrpcRequest {
@@ -20,3 +21,13 @@ message GetValueSchemaGrpcResponse {
2021
int32 schemaId = 2;
2122
string schemaStr = 3;
2223
}
24+
25+
message GetKeySchemaGrpcRequest {
26+
ClusterStoreGrpcInfo storeInfo = 1;
27+
}
28+
29+
message GetKeySchemaGrpcResponse {
30+
ClusterStoreGrpcInfo storeInfo = 1;
31+
int32 schemaId = 2;
32+
string schemaStr = 3;
33+
}

‎internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestControllerGrpcEndpoints.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse;
2121
import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest;
2222
import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse;
23+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcRequest;
24+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcResponse;
2325
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcRequest;
2426
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcResponse;
2527
import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest;
@@ -386,6 +388,42 @@ public void testGetValueSchemaGrpcEndpoint() {
386388
assertEquals(exception.getStatus().getCode(), io.grpc.Status.Code.INVALID_ARGUMENT);
387389
}
388390

391+
@Test(timeOut = TIMEOUT_MS)
392+
public void testGetKeySchemaGrpcEndpoint() {
393+
String storeName = Utils.getUniqueString("test_get_key_schema_store");
394+
String controllerGrpcUrl = veniceCluster.getLeaderVeniceController().getControllerGrpcUrl();
395+
ManagedChannel channel = Grpc.newChannelBuilder(controllerGrpcUrl, InsecureChannelCredentials.create()).build();
396+
StoreGrpcServiceGrpc.StoreGrpcServiceBlockingStub storeBlockingStub = StoreGrpcServiceGrpc.newBlockingStub(channel);
397+
SchemaGrpcServiceGrpc.SchemaGrpcServiceBlockingStub schemaBlockingStub =
398+
SchemaGrpcServiceGrpc.newBlockingStub(channel);
399+
400+
ClusterStoreGrpcInfo storeGrpcInfo = ClusterStoreGrpcInfo.newBuilder()
401+
.setClusterName(veniceCluster.getClusterName())
402+
.setStoreName(storeName)
403+
.build();
404+
405+
// Step 1: Create the store
406+
CreateStoreGrpcRequest createStoreGrpcRequest = CreateStoreGrpcRequest.newBuilder()
407+
.setStoreInfo(storeGrpcInfo)
408+
.setOwner("owner")
409+
.setKeySchema(DEFAULT_KEY_SCHEMA)
410+
.setValueSchema("\"string\"")
411+
.build();
412+
CreateStoreGrpcResponse createResponse = storeBlockingStub.createStore(createStoreGrpcRequest);
413+
assertNotNull(createResponse, "Response should not be null");
414+
assertEquals(createResponse.getStoreInfo().getStoreName(), storeName);
415+
416+
// Step 2: Get key schema using gRPC endpoint
417+
GetKeySchemaGrpcRequest getKeySchemaRequest =
418+
GetKeySchemaGrpcRequest.newBuilder().setStoreInfo(storeGrpcInfo).build();
419+
GetKeySchemaGrpcResponse getKeySchemaResponse = schemaBlockingStub.getKeySchema(getKeySchemaRequest);
420+
assertNotNull(getKeySchemaResponse, "Response should not be null");
421+
assertEquals(getKeySchemaResponse.getStoreInfo().getStoreName(), storeName);
422+
assertEquals(getKeySchemaResponse.getStoreInfo().getClusterName(), veniceCluster.getClusterName());
423+
assertEquals(getKeySchemaResponse.getSchemaStr(), DEFAULT_KEY_SCHEMA);
424+
assertEquals(getKeySchemaResponse.getSchemaId(), 1);
425+
}
426+
389427
private static class MockDynamicAccessController extends NoOpDynamicAccessController {
390428
private final Set<String> resourcesInAllowList = ConcurrentHashMap.newKeySet();
391429

‎services/venice-controller/src/main/java/com/linkedin/venice/controller/grpc/server/SchemaGrpcServiceImpl.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package com.linkedin.venice.controller.grpc.server;
22

33
import com.linkedin.venice.controller.server.SchemaRequestHandler;
4+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcRequest;
5+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcResponse;
46
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcRequest;
57
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcResponse;
68
import com.linkedin.venice.protocols.controller.SchemaGrpcServiceGrpc;
@@ -30,4 +32,18 @@ public void getValueSchema(
3032
responseObserver,
3133
request.getStoreInfo());
3234
}
35+
36+
/**
37+
* Retrieves the key schema for a store.
38+
* No ACL check is required for this operation as it only reads store metadata.
39+
*/
40+
@Override
41+
public void getKeySchema(GetKeySchemaGrpcRequest request, StreamObserver<GetKeySchemaGrpcResponse> responseObserver) {
42+
LOGGER.debug("Received getKeySchema with args: {}", request);
43+
ControllerGrpcServerUtils.handleRequest(
44+
SchemaGrpcServiceGrpc.getGetKeySchemaMethod(),
45+
() -> schemaRequestHandler.getKeySchema(request),
46+
responseObserver,
47+
request.getStoreInfo());
48+
}
3349
}

‎services/venice-controller/src/main/java/com/linkedin/venice/controller/server/SchemaRequestHandler.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,10 @@
22

33
import com.linkedin.venice.controller.Admin;
44
import com.linkedin.venice.controller.ControllerRequestHandlerDependencies;
5+
import com.linkedin.venice.exceptions.VeniceException;
56
import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo;
7+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcRequest;
8+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcResponse;
69
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcRequest;
710
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcResponse;
811
import com.linkedin.venice.schema.SchemaEntry;
@@ -47,4 +50,26 @@ public GetValueSchemaGrpcResponse getValueSchema(GetValueSchemaGrpcRequest reque
4750
.setSchemaStr(valueSchemaEntry.getSchema().toString())
4851
.build();
4952
}
53+
54+
/**
55+
* Retrieves the key schema for a store.
56+
* @param request the request containing cluster and store name
57+
* @return response containing the key schema id and schema string
58+
*/
59+
public GetKeySchemaGrpcResponse getKeySchema(GetKeySchemaGrpcRequest request) {
60+
ClusterStoreGrpcInfo storeInfo = request.getStoreInfo();
61+
ControllerRequestParamValidator.validateClusterStoreInfo(storeInfo);
62+
String clusterName = storeInfo.getClusterName();
63+
String storeName = storeInfo.getStoreName();
64+
LOGGER.info("Getting key schema for store: {} in cluster: {}", storeName, clusterName);
65+
SchemaEntry keySchemaEntry = admin.getKeySchema(clusterName, storeName);
66+
if (keySchemaEntry == null) {
67+
throw new VeniceException("Key schema doesn't exist for store: " + storeName);
68+
}
69+
return GetKeySchemaGrpcResponse.newBuilder()
70+
.setStoreInfo(storeInfo)
71+
.setSchemaId(keySchemaEntry.getId())
72+
.setSchemaStr(keySchemaEntry.getSchema().toString())
73+
.build();
74+
}
5075
}

‎services/venice-controller/src/main/java/com/linkedin/venice/controller/server/SchemaRoutes.java‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
import com.linkedin.venice.exceptions.VeniceNoStoreException;
3030
import com.linkedin.venice.meta.Store;
3131
import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo;
32+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcRequest;
33+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcResponse;
3234
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcRequest;
3335
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcResponse;
3436
import com.linkedin.venice.schema.GeneratedSchemaID;
@@ -70,14 +72,20 @@ public Route getKeySchema(Admin admin) {
7072
try {
7173
// No ACL check on getting store metadata
7274
AdminSparkServer.validateParams(request, GET_KEY_SCHEMA.getParams(), admin);
73-
responseObject.setCluster(request.queryParams(CLUSTER));
74-
responseObject.setName(request.queryParams(NAME));
75-
SchemaEntry keySchemaEntry = admin.getKeySchema(responseObject.getCluster(), responseObject.getName());
76-
if (keySchemaEntry == null) {
77-
throw new VeniceException("Key schema doesn't exist for store: " + responseObject.getName());
78-
}
79-
responseObject.setId(keySchemaEntry.getId());
80-
responseObject.setSchemaStr(keySchemaEntry.getSchema().toString());
75+
String clusterName = request.queryParams(CLUSTER);
76+
String storeName = request.queryParams(NAME);
77+
78+
// Convert HTTP request to gRPC request
79+
ClusterStoreGrpcInfo storeInfo =
80+
ClusterStoreGrpcInfo.newBuilder().setClusterName(clusterName).setStoreName(storeName).build();
81+
GetKeySchemaGrpcRequest grpcRequest = GetKeySchemaGrpcRequest.newBuilder().setStoreInfo(storeInfo).build();
82+
83+
// Call the handler and convert the response
84+
GetKeySchemaGrpcResponse grpcResponse = schemaRequestHandler.getKeySchema(grpcRequest);
85+
responseObject.setCluster(clusterName);
86+
responseObject.setName(storeName);
87+
responseObject.setId(grpcResponse.getSchemaId());
88+
responseObject.setSchemaStr(grpcResponse.getSchemaStr());
8189
} catch (Throwable e) {
8290
responseObject.setError(e);
8391
AdminSparkServer.handleError(new VeniceException(e), request, response);

‎services/venice-controller/src/test/java/com/linkedin/venice/controller/grpc/server/SchemaGrpcServiceImplTest.java‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
import com.linkedin.venice.exceptions.VeniceException;
1414
import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo;
1515
import com.linkedin.venice.protocols.controller.ControllerGrpcErrorType;
16+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcRequest;
17+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcResponse;
1618
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcRequest;
1719
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcResponse;
1820
import com.linkedin.venice.protocols.controller.SchemaGrpcServiceGrpc;
@@ -148,4 +150,57 @@ public void testGetValueSchemaReturnsGeneralError() {
148150
assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.GENERAL_ERROR);
149151
assertTrue(errorInfo.getErrorMessage().contains("Internal error fetching schema"));
150152
}
153+
154+
@Test
155+
public void testGetKeySchemaReturnsSuccessfulResponse() {
156+
ClusterStoreGrpcInfo storeInfo =
157+
ClusterStoreGrpcInfo.newBuilder().setClusterName(TEST_CLUSTER).setStoreName(TEST_STORE).build();
158+
GetKeySchemaGrpcRequest request = GetKeySchemaGrpcRequest.newBuilder().setStoreInfo(storeInfo).build();
159+
GetKeySchemaGrpcResponse response =
160+
GetKeySchemaGrpcResponse.newBuilder().setStoreInfo(storeInfo).setSchemaId(1).setSchemaStr("\"string\"").build();
161+
when(schemaRequestHandler.getKeySchema(any(GetKeySchemaGrpcRequest.class))).thenReturn(response);
162+
163+
GetKeySchemaGrpcResponse actualResponse = blockingStub.getKeySchema(request);
164+
165+
assertNotNull(actualResponse, "Response should not be null");
166+
assertEquals(actualResponse.getStoreInfo(), storeInfo, "Store info should match");
167+
assertEquals(actualResponse.getSchemaId(), 1, "Schema ID should match");
168+
assertEquals(actualResponse.getSchemaStr(), "\"string\"", "Schema string should match");
169+
}
170+
171+
@Test
172+
public void testGetKeySchemaReturnsErrorResponse() {
173+
ClusterStoreGrpcInfo storeInfo =
174+
ClusterStoreGrpcInfo.newBuilder().setClusterName(TEST_CLUSTER).setStoreName(TEST_STORE).build();
175+
GetKeySchemaGrpcRequest request = GetKeySchemaGrpcRequest.newBuilder().setStoreInfo(storeInfo).build();
176+
when(schemaRequestHandler.getKeySchema(any(GetKeySchemaGrpcRequest.class)))
177+
.thenThrow(new VeniceException("Key schema doesn't exist for store: " + TEST_STORE));
178+
179+
StatusRuntimeException e = expectThrows(StatusRuntimeException.class, () -> blockingStub.getKeySchema(request));
180+
181+
assertNotNull(e.getStatus(), "Status should not be null");
182+
assertEquals(e.getStatus().getCode(), Status.INTERNAL.getCode());
183+
VeniceControllerGrpcErrorInfo errorInfo = GrpcRequestResponseConverter.parseControllerGrpcError(e);
184+
assertNotNull(errorInfo, "Error info should not be null");
185+
assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.GENERAL_ERROR);
186+
assertTrue(errorInfo.getErrorMessage().contains("Key schema doesn't exist for store"));
187+
}
188+
189+
@Test
190+
public void testGetKeySchemaReturnsBadRequestForInvalidArgument() {
191+
ClusterStoreGrpcInfo storeInfo =
192+
ClusterStoreGrpcInfo.newBuilder().setClusterName(TEST_CLUSTER).setStoreName(TEST_STORE).build();
193+
GetKeySchemaGrpcRequest request = GetKeySchemaGrpcRequest.newBuilder().setStoreInfo(storeInfo).build();
194+
when(schemaRequestHandler.getKeySchema(any(GetKeySchemaGrpcRequest.class)))
195+
.thenThrow(new IllegalArgumentException("Cluster name is mandatory parameter"));
196+
197+
StatusRuntimeException e = expectThrows(StatusRuntimeException.class, () -> blockingStub.getKeySchema(request));
198+
199+
assertNotNull(e.getStatus(), "Status should not be null");
200+
assertEquals(e.getStatus().getCode(), Status.INVALID_ARGUMENT.getCode());
201+
VeniceControllerGrpcErrorInfo errorInfo = GrpcRequestResponseConverter.parseControllerGrpcError(e);
202+
assertNotNull(errorInfo, "Error info should not be null");
203+
assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.BAD_REQUEST);
204+
assertTrue(errorInfo.getErrorMessage().contains("Cluster name is mandatory parameter"));
205+
}
151206
}

‎services/venice-controller/src/test/java/com/linkedin/venice/controller/server/SchemaRequestHandlerTest.java‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88

99
import com.linkedin.venice.controller.Admin;
1010
import com.linkedin.venice.controller.ControllerRequestHandlerDependencies;
11+
import com.linkedin.venice.exceptions.VeniceException;
1112
import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo;
13+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcRequest;
14+
import com.linkedin.venice.protocols.controller.GetKeySchemaGrpcResponse;
1215
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcRequest;
1316
import com.linkedin.venice.protocols.controller.GetValueSchemaGrpcResponse;
1417
import com.linkedin.venice.schema.SchemaEntry;
@@ -80,4 +83,52 @@ public void testGetValueSchemaMissingStoreName() {
8083

8184
schemaRequestHandler.getValueSchema(request);
8285
}
86+
87+
@Test
88+
public void testGetKeySchemaSuccess() {
89+
GetKeySchemaGrpcRequest request = GetKeySchemaGrpcRequest.newBuilder()
90+
.setStoreInfo(ClusterStoreGrpcInfo.newBuilder().setClusterName("testCluster").setStoreName("testStore").build())
91+
.build();
92+
93+
Schema schema = Schema.parse("\"string\"");
94+
SchemaEntry schemaEntry = new SchemaEntry(1, schema);
95+
when(admin.getKeySchema("testCluster", "testStore")).thenReturn(schemaEntry);
96+
97+
GetKeySchemaGrpcResponse response = schemaRequestHandler.getKeySchema(request);
98+
99+
verify(admin, times(1)).getKeySchema("testCluster", "testStore");
100+
assertEquals(response.getStoreInfo().getClusterName(), "testCluster");
101+
assertEquals(response.getStoreInfo().getStoreName(), "testStore");
102+
assertEquals(response.getSchemaId(), 1);
103+
assertEquals(response.getSchemaStr(), "\"string\"");
104+
}
105+
106+
@Test(expectedExceptions = VeniceException.class, expectedExceptionsMessageRegExp = "Key schema doesn't exist for store: testStore")
107+
public void testGetKeySchemaWhenSchemaDoesNotExist() {
108+
GetKeySchemaGrpcRequest request = GetKeySchemaGrpcRequest.newBuilder()
109+
.setStoreInfo(ClusterStoreGrpcInfo.newBuilder().setClusterName("testCluster").setStoreName("testStore").build())
110+
.build();
111+
112+
when(admin.getKeySchema("testCluster", "testStore")).thenReturn(null);
113+
114+
schemaRequestHandler.getKeySchema(request);
115+
}
116+
117+
@Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*[Cc]luster.*")
118+
public void testGetKeySchemaWithMissingClusterName() {
119+
GetKeySchemaGrpcRequest request = GetKeySchemaGrpcRequest.newBuilder()
120+
.setStoreInfo(ClusterStoreGrpcInfo.newBuilder().setStoreName("testStore").build())
121+
.build();
122+
123+
schemaRequestHandler.getKeySchema(request);
124+
}
125+
126+
@Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*[Ss]tore.*")
127+
public void testGetKeySchemaWithMissingStoreName() {
128+
GetKeySchemaGrpcRequest request = GetKeySchemaGrpcRequest.newBuilder()
129+
.setStoreInfo(ClusterStoreGrpcInfo.newBuilder().setClusterName("testCluster").build())
130+
.build();
131+
132+
schemaRequestHandler.getKeySchema(request);
133+
}
83134
}

0 commit comments

Comments
 (0)