From 4375ce9e935e57f2dcc332539931ae08b554a34d Mon Sep 17 00:00:00 2001 From: sjmignot Date: Thu, 10 Sep 2026 10:59:52 -0700 Subject: [PATCH 1/2] Add gRPC multi-query support (onlineQueryMulti) Adds ChalkClient.onlineQueryMulti, which runs several independent online queries in parallel over a single request using the OnlineQueryMulti RPC. The generated stubs for that RPC were already vendored in this repo but had never been wired up, so this is a client-side change only. Callers with N unrelated queries -- different roots, different outputs -- previously paid N round trips. Details: - New OnlineQueryMultiResult (AutoCloseable) exposes the per-query results positionally, plus the request-level errors and the shared trace id. Each sub-result owns its own Arrow child allocator, because OnlineQueryResult.close() closes the allocator it is handed; sharing one would release it while other sub-results' tables still held buffers. Closing the wrapper closes them all. - GRPCClient.onlineQuery is decomposed into serializeInputs, buildBulkRequest, queryStub and decodeBulkResponse, now shared by both paths so single and multi build byte-identical sub-requests. A golden test asserts buildBulkRequest field by field to hold that line. - Every sub-query is sent as a bulk_request: the engine rejects single_request outright ("Please use OnlineQueryBulkRequest w/ a single row"), matching the reference Python client. - A request carries one deadline and one set of headers, so the queries must agree on environmentId, branch, queryName and queryNameVersion -- the API server resolves one branch to one deployment and routes named queries by an exact match on x-chalk-query-name, so a mixed batch could not be honored. The deadline is the longest per-query timeout. - decodeBulkResponse now rethrows on failure. It previously closed its tables and allocator but fell through, returning an OnlineQueryResult wrapping already-closed Arrow objects. The multi path cannot tolerate that: a mid-batch failure would leak the sibling allocators. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 64 ++ .../java/ai/chalk/client/ChalkClient.java | 102 ++++ src/main/java/ai/chalk/client/GRPCClient.java | 425 +++++++++++-- .../models/OnlineQueryBulkResponse.java | 8 + .../chalk/models/OnlineQueryMultiResult.java | 106 ++++ .../java/ai/chalk/client/TestAllClients.java | 17 + .../java/ai/chalk/client/TestGrpcClient.java | 129 ++++ .../ai/chalk/client/TestGrpcMultiQuery.java | 572 ++++++++++++++++++ 8 files changed, 1386 insertions(+), 37 deletions(-) create mode 100644 src/main/java/ai/chalk/models/OnlineQueryMultiResult.java create mode 100644 src/test/java/ai/chalk/client/TestGrpcMultiQuery.java diff --git a/README.md b/README.md index 8d543e7e..f958562e 100644 --- a/README.md +++ b/README.md @@ -378,3 +378,67 @@ public class Main { } } ``` + +### Multi-query + +`onlineQueryMulti` runs several *independent* queries in parallel, using a single request. Reach for +it when the queries are rooted in different feature classes, or otherwise ask for unrelated outputs, +and you would rather not pay a round trip for each. To run one query over many rows of input, use +`onlineQuery` — it is already a bulk operation. + +Multi-query requires the gRPC client. + +```java +import ai.chalk.client.ChalkClient; +import ai.chalk.exceptions.ServerError; +import ai.chalk.models.OnlineQueryMultiResult; +import ai.chalk.models.OnlineQueryParams; +import ai.chalk.models.OnlineQueryResult; + +import java.util.Arrays; +import java.util.List; + +ChalkClient client = ChalkClient.createGrpc(); + +var users = OnlineQueryParams.builder() + .withInput("user.id", Arrays.asList(1, 2, 3)) + .withOutputs("user.email") + .build(); +var merchants = OnlineQueryParams.builder() + .withInput("merchant.id", Arrays.asList("a", "b")) + .withOutputs("merchant.risk_score") + .build(); + +try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(users, merchants))) { + for (ServerError error : multi.getGlobalErrors()) { + // errors that do not belong to any one query + } + for (OnlineQueryResult result : multi.getResults()) { + if (result.getErrors().length > 0) { + // this query failed; the others are unaffected + continue; + } + // do something with the result + } +} +``` + +A few things to know: + +- **Results are positional.** `getResults()` holds one result per query, in the order you passed + them in. +- **Queries fail independently.** A resolver that fails while computing one query is reported on + that query's own result, and leaves the others intact. `getGlobalErrors()` holds the errors the + engine did not attach to a specific result. Note that a query which failed to *run* at all — it + was rate limited, or the engine hit an internal error on it — is reported there, naming the query + by its 1-based position in the message, while its slot in `getResults()` holds an empty result + with no errors of its own. So check the global errors too, rather than reading an empty result as + a successful one. +- **Close the wrapper, not the results.** Closing the `OnlineQueryMultiResult` closes every one of + its results and releases all of their Arrow memory. +- **Some settings apply to the whole request,** which carries one deadline and one set of headers. + Every query must agree on `environmentId`, `branch`, `queryName` and `queryNameVersion`, or the + call throws `ClientException` — a request resolves one branch to one deployment and routes named + queries by an exact match on the query name, so a mixed batch could not be honored. Run those as + separate queries. The deadline is the longest of the per-query timeouts, falling back to the + client-level timeout, so per-query timeouts are not independently enforced. diff --git a/src/main/java/ai/chalk/client/ChalkClient.java b/src/main/java/ai/chalk/client/ChalkClient.java index 3ba53d32..f8af98ed 100644 --- a/src/main/java/ai/chalk/client/ChalkClient.java +++ b/src/main/java/ai/chalk/client/ChalkClient.java @@ -1,11 +1,14 @@ package ai.chalk.client; import ai.chalk.exceptions.ChalkException; +import ai.chalk.exceptions.ClientException; import ai.chalk.models.*; import java.net.http.HttpClient; import java.nio.file.Path; import java.time.Duration; +import java.util.Arrays; +import java.util.List; import java.util.Optional; import java.util.concurrent.ThreadLocalRandom; @@ -129,6 +132,105 @@ default int ping() throws ChalkException { */ OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws ChalkException; + /** + * OnlineQueryMulti computes several independent online queries in parallel, + * using a single request to the Chalk engine. Use it when the queries are + * rooted in different feature classes, or otherwise ask for unrelated + * outputs, and you would rather not pay a round trip for each one. To run + * one query over many rows of input, use {@link #onlineQuery} instead, + * which is already a bulk operation. + * + *

This is supported only by the gRPC client; see {@link ChalkClient#createGrpc()}. + * + *

+ * Example usage: + * 

+     *         {@code
+     *         OnlineQueryParamsComplete users = OnlineQueryParams.builder()
+     *             .withInput("user.id", Arrays.asList(1, 2, 3))
+     *             .withOutputs("user.email")
+     *             .build();
+     *         OnlineQueryParamsComplete merchants = OnlineQueryParams.builder()
+     *             .withInput("merchant.id", Arrays.asList("a", "b"))
+     *             .withOutputs("merchant.risk_score")
+     *             .build();
+     *
+     *         try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(users, merchants))) {
+     *             for (OnlineQueryResult result : multi.getResults()) {
+     *                 if (result.getErrors().length > 0) {
+     *                     // handle this query's failure; the other queries are unaffected
+     *                     continue;
+     *                 }
+     *                 // do something with the result
+     *             }
+     *         }
+     *         }
+     *
+ *

+ * + *

Results are positional. {@link OnlineQueryMultiResult#getResults()} + * holds one result per query, in the order the queries were passed in. + * + *

Queries fail independently. A resolver that fails while computing + * one query is reported on that query's own result, via + * {@link OnlineQueryResult#getErrors()}, and leaves the other queries' results + * intact. {@link OnlineQueryMultiResult#getGlobalErrors()} holds the errors the + * engine did not attach to a specific result -- including the case where a query + * failed to run at all, which is reported there with the query's 1-based position + * in the message rather than on the result. That query still occupies its slot, + * but its result is empty and reports no errors of its own, so check the global + * errors as well before treating an empty result as a successful one. + * + *

Resources. Closing the returned {@link OnlineQueryMultiResult} + * closes every one of its results and releases all of their Arrow memory. + * Do not close the individual results yourself. + * + *

Some settings apply to the whole request. A request carries one + * deadline and one set of headers, so: + *

+ * + * @param params the queries to run, one result per query + * @return {@link OnlineQueryMultiResult} + * @throws ChalkException if the request fails, or if the queries disagree on a + * setting that applies to the whole request + * @see query basics + */ + default OnlineQueryMultiResult onlineQueryMulti(List params) + throws ChalkException { + throw new ClientException( + "onlineQueryMulti is only supported by the gRPC client. Build one with " + + "ChalkClient.builder().withGrpc().build() or ChalkClient.createGrpc()." + ); + } + + /** + * OnlineQueryMulti computes several independent online queries in parallel, + * using a single request to the Chalk engine. + * + *

See {@link #onlineQueryMulti(List)} for details and for the settings that + * every query must agree on. + * + * @return {@link OnlineQueryMultiResult} + * @throws ChalkException + */ + default OnlineQueryMultiResult onlineQueryMulti(OnlineQueryParamsComplete... params) + throws ChalkException { + return onlineQueryMulti(Arrays.asList(params)); + } + /** * UploadFeatures synchronously persists feature values to the online store and * offline store. diff --git a/src/main/java/ai/chalk/client/GRPCClient.java b/src/main/java/ai/chalk/client/GRPCClient.java index f2cdeaae..5ea4519e 100644 --- a/src/main/java/ai/chalk/client/GRPCClient.java +++ b/src/main/java/ai/chalk/client/GRPCClient.java @@ -7,6 +7,7 @@ import ai.chalk.internal.arrow.FeatherProcessor; import ai.chalk.internal.config.Loader; import ai.chalk.internal.config.models.ProjectToken; +import ai.chalk.models.OnlineQueryMultiResult; import ai.chalk.models.OnlineQueryParamsComplete; import ai.chalk.models.OnlineQueryResult; import ai.chalk.models.UploadFeaturesParams; @@ -22,6 +23,7 @@ import io.grpc.*; import io.grpc.stub.MetadataUtils; import lombok.NonNull; +import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.table.Table; @@ -31,6 +33,7 @@ import java.time.Duration; import java.util.*; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import static ai.chalk.internal.arrow.FeatherProcessor.inputsToArrowBytes; @@ -205,19 +208,97 @@ public int ping(int num) { public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws ChalkException { - byte[] bodyBytes; + byte[] bodyBytes = serializeInputs(params, this.allocator, "grpc_online_query_params"); + OnlineQueryBulkRequest request = buildBulkRequest(params, bodyBytes, this.branchId); + + AtomicReference trailersRef = new AtomicReference<>(); + OnlineQueryBulkResponse response = queryStub( + params.getTimeout(), + params.getEnvironmentId(), + params.getQueryName(), + trailersRef + ).onlineQueryBulk(request); + + return decodeBulkResponse( + response, + traceIdFrom(trailersRef), + newResponseAllocator("grpc_online_query_response") + ); + } + + /** + * OnlineQueryMulti computes several independent online queries in parallel, + * using a single request to the Chalk engine. + */ + @Override + public OnlineQueryMultiResult onlineQueryMulti(List paramsList) + throws ChalkException { + validateMultiParams(paramsList); + if (paramsList.isEmpty()) { + return new OnlineQueryMultiResult(List.of(), new ServerError[0], null); + } + + // Resolve the request-level settings before serializing anything, so a batch that + // cannot be satisfied fails before doing the work of encoding its inputs. + Duration deadline = resolveDeadline(paramsList).orElse(null); + String environmentId = resolveEnvironment(paramsList); + String queryName = resolveQueryName(paramsList); + // Validated for agreement but deliberately unused: the gRPC client does not send the + // "x-chalk-branch-id" header that the API server routes branches on, so branch + // selection is inert on this path for single and multi queries alike. See the + // branch/preview-deployment follow-up. + resolveBranch(paramsList, this.branchId); + + OnlineQueryMultiRequest request = buildMultiRequest(paramsList, this.allocator, this.branchId); + + AtomicReference trailersRef = new AtomicReference<>(); + OnlineQueryMultiResponse response = queryStub(deadline, environmentId, queryName, trailersRef) + .onlineQueryMulti(request); + + return decodeMultiResponse( + response, + paramsList.size(), + traceIdFrom(trailersRef), + this.allocator + ); + } + + /** + * Serializes the query's inputs to Arrow IPC bytes, using a short-lived + * child allocator that is released before returning. + */ + static byte[] serializeInputs( + OnlineQueryParamsComplete params, + BufferAllocator parentAllocator, + String allocatorName + ) throws ClientException { try ( - var childAllocator = allocator.newChildAllocator( - "grpc_online_query_params", + var childAllocator = parentAllocator.newChildAllocator( + allocatorName, 0, FeatherProcessor.ALLOCATOR_SIZE_REQUEST ) ) { - bodyBytes = inputsToArrowBytes(params.getInputs(), childAllocator); + return inputsToArrowBytes(params.getInputs(), childAllocator); } catch (Exception e) { throw new ClientException("Failed to serialize OnlineQueryParams", e); } + } + /* + * Builds the bulk request for a single query. Shared by the single-query and + * multi-query paths, so both produce byte-identical sub-requests. + * + * TODO: params.getStaleness() and params.isStorePlanStages() are not sent on the + * gRPC path, even though OnlineQueryBulkRequest has a `staleness` field. Wiring + * staleness up requires deciding how to encode a Duration as the bucket string + * the server expects. + */ + static OnlineQueryBulkRequest buildBulkRequest( + OnlineQueryParamsComplete params, + byte[] inputsFeather, + @Nullable String clientBranchId + ) { List resolvedOutputs = params.getOutputs(); if (resolvedOutputs == null) { resolvedOutputs = new ArrayList<>(); @@ -243,8 +324,8 @@ public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws Ch var context = OnlineQueryContext.newBuilder(); if (params.getBranch() != null && !params.getBranch().isEmpty()) { context.setBranchId(params.getBranch()); - } else if (this.branchId != null && !this.branchId.isEmpty()) { - context.setBranchId(this.branchId); + } else if (clientBranchId != null && !clientBranchId.isEmpty()) { + context.setBranchId(clientBranchId); } if (params.getCorrelationId() != null) { context.setCorrelationId(params.getCorrelationId()); @@ -289,44 +370,92 @@ public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws Ch options.putAllMetadata(params.getMeta()); } - var request = OnlineQueryBulkRequest.newBuilder() - .setInputsFeather(ByteString.copyFrom(bodyBytes)) + return OnlineQueryBulkRequest.newBuilder() + .setInputsFeather(ByteString.copyFrom(inputsFeather)) .addAllOutputs(outputs) .addAllNow(now) .setBodyType(FeatherBodyType.FEATHER_BODY_TYPE_TABLE) .setContext(context) .setResponseOptions(options) .build(); + } - AtomicReference trailersRef = new AtomicReference<>(); - OnlineQueryBulkResponse response = this.stubsProvider.getQueryStub(Optional.ofNullable(params.getTimeout())) + /* + * Wraps each set of params as a bulk sub-query, preserving order. The engine + * rejects `single_request` outright ("OnlineQueryRequest is unsupported; Please + * use OnlineQueryBulkRequest w/ a single row"), so every sub-query is a bulk + * request, exactly as the reference Python client does. + */ + static OnlineQueryMultiRequest buildMultiRequest( + List paramsList, + BufferAllocator parentAllocator, + @Nullable String clientBranchId + ) throws ClientException { + var request = OnlineQueryMultiRequest.newBuilder(); + for (int i = 0; i < paramsList.size(); i++) { + OnlineQueryParamsComplete params = paramsList.get(i); + byte[] inputsFeather = serializeInputs( + params, parentAllocator, "grpc_online_query_multi_params[" + i + "]" + ); + request.addQueries( + GenericSingleQuery.newBuilder() + .setBulkRequest(buildBulkRequest(params, inputsFeather, clientBranchId)) + .build() + ); + } + return request.build(); + } + + private QueryServiceGrpc.QueryServiceBlockingStub queryStub( + @Nullable Duration timeout, + @Nullable String environmentIdOverride, + @Nullable String queryName, + AtomicReference trailersRef + ) { + return this.stubsProvider.getQueryStub(Optional.ofNullable(timeout)) .withInterceptors( - MetadataUtils.newCaptureMetadataInterceptor(new AtomicReference<>(), trailersRef), - this.getRequestHeaderInterceptor(params.getEnvironmentId(), params.getQueryName()) - ) - .onlineQueryBulk(request); + MetadataUtils.newCaptureMetadataInterceptor(new AtomicReference<>(), trailersRef), + this.getRequestHeaderInterceptor(environmentIdOverride, queryName) + ); + } - var meta = GrpcSerializer.toQueryMeta( - response.getResponseMeta(), - trailersRef.get().get(CHALK_TRACE_ID_KEY) - ); + private static @Nullable String traceIdFrom(AtomicReference trailersRef) { + Metadata trailers = trailersRef.get(); + return trailers == null ? null : trailers.get(CHALK_TRACE_ID_KEY); + } - ServerError[] errors = new ServerError[response.getErrorsCount()]; - for (int i = 0; i < response.getErrorsCount(); i++) { - errors[i] = GrpcSerializer.toServerError(response.getErrors(i)); - } + private BufferAllocator newResponseAllocator(String name) { + return this.allocator.newChildAllocator(name, 0, FeatherProcessor.ALLOCATOR_SIZE_RESPONSE); + } + /* + * Decodes one bulk response. + * + * Takes ownership of responseAllocator: on success it is handed to the returned + * OnlineQueryResult, which releases it when closed. On any failure this method + * closes the tables it decoded along with the allocator, and rethrows, so no + * Arrow memory is retained. + */ + static OnlineQueryResult decodeBulkResponse( + OnlineQueryBulkResponse response, + @Nullable String traceId, + BufferAllocator responseAllocator + ) throws ChalkException { Table scalars = null; Map groups = new HashMap<>(); - var responseAlloc = this.allocator.newChildAllocator( - "grpc_online_query_response", 0, FeatherProcessor.ALLOCATOR_SIZE_RESPONSE - ); try { + var meta = GrpcSerializer.toQueryMeta(response.getResponseMeta(), traceId); + + ServerError[] errors = new ServerError[response.getErrorsCount()]; + for (int i = 0; i < response.getErrorsCount(); i++) { + errors[i] = GrpcSerializer.toServerError(response.getErrors(i)); + } + if (!response.getScalarsData().isEmpty()) { try { scalars = FeatherProcessor.convertBytesToTable( response.getScalarsData().toByteArray(), - responseAlloc + responseAllocator ); } catch (Exception e) { throw new ClientException("Failed to convert scalar data bytes to table", e); @@ -338,7 +467,7 @@ public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws Ch try { groups.put( fqn, - FeatherProcessor.convertBytesToTable(entry.getValue().toByteArray(), responseAlloc) + FeatherProcessor.convertBytesToTable(entry.getValue().toByteArray(), responseAllocator) ); } catch (Exception e) { throw new ClientException( @@ -346,23 +475,245 @@ public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws Ch ); } } + + return new OnlineQueryResult(scalars, groups, errors, meta, responseAllocator); + } catch (Exception e) { + releaseQuietly(e, scalars, groups.values(), responseAllocator); + throw asChalkException(e, "Failed to decode online query response"); + } + } + + /* + * Decodes a multi-query response into one OnlineQueryResult per sub-query. + * + * Each sub-result gets its own child allocator of rootAllocator, because + * OnlineQueryResult.close() closes the allocator it is handed: sharing one + * allocator across sub-results would release it while the other sub-results' + * tables still held buffers. On any failure, every result decoded so far is + * closed and the exception is rethrown, so no Arrow memory is retained. + */ + static OnlineQueryMultiResult decodeMultiResponse( + OnlineQueryMultiResponse response, + int expectedCount, + @Nullable String traceId, + BufferAllocator rootAllocator + ) throws ChalkException { + ServerError[] globalErrors = new ServerError[response.getErrorsCount()]; + for (int i = 0; i < response.getErrorsCount(); i++) { + globalErrors[i] = GrpcSerializer.toServerError(response.getErrors(i)); + } + + // The engine returns no sub-responses at all when the request failed before any + // sub-query was dispatched; the global errors explain why. This is the one case + // where the results are not 1:1 with the submitted queries. + if (response.getResponsesCount() == 0 && globalErrors.length > 0) { + return new OnlineQueryMultiResult(List.of(), globalErrors, traceId); + } + + // Sub-responses are correlated with sub-queries only by position, so a count + // mismatch leaves every result unattributable. Fail rather than guess. + if (response.getResponsesCount() != expectedCount) { + throw new ClientException(String.format( + "Expected %d sub-responses for %d queries, but got %d%s", + expectedCount, + expectedCount, + response.getResponsesCount(), + globalErrors.length == 0 ? "" : "; errors: " + describeErrors(globalErrors) + )); + } + + List decoded = new ArrayList<>(expectedCount); + try { + for (int i = 0; i < response.getResponsesCount(); i++) { + GenericSingleResponse single = response.getResponses(i); + // The engine sets the bulk_response case up front for every bulk sub-query, so + // this only fires if it rejected the sub-query's type outright -- which it does + // for single_request, and this client only ever sends bulk_request. + if (single.getQueryCase() != GenericSingleResponse.QueryCase.BULK_RESPONSE) { + throw new ClientException(String.format( + "Sub-response at index %d is not a bulk response (got %s); " + + "the Java client only issues bulk sub-queries", + i, single.getQueryCase() + )); + } + BufferAllocator subAllocator = rootAllocator.newChildAllocator( + "grpc_online_query_multi_response[" + i + "]", + 0, + FeatherProcessor.ALLOCATOR_SIZE_RESPONSE + ); + try { + // decodeBulkResponse releases subAllocator itself if it throws. + decoded.add(decodeBulkResponse(single.getBulkResponse(), traceId, subAllocator)); + } catch (Exception e) { + throw new ClientException( + String.format("Failed to decode sub-response at index %d", i), e + ); + } + } } catch (Exception e) { - if (scalars != null) { + for (OnlineQueryResult result : decoded) { + try { + result.close(); + } catch (Exception closeError) { + e.addSuppressed(closeError); + } + } + throw asChalkException(e, "Failed to decode multi-query response"); + } + + return new OnlineQueryMultiResult(decoded, globalErrors, traceId); + } + + /** + * Validates the per-query parameters of a multi-query. + */ + static void validateMultiParams(List paramsList) throws ClientException { + if (paramsList == null) { + throw new ClientException("onlineQueryMulti requires a non-null list of queries"); + } + for (int i = 0; i < paramsList.size(); i++) { + OnlineQueryParamsComplete params = paramsList.get(i); + if (params == null) { + throw new ClientException( + String.format("onlineQueryMulti received a null query at index %d", i) + ); + } + if (!isBlank(params.getQueryNameVersion()) && isBlank(params.getQueryName())) { + throw new ClientException(String.format( + "The query at index %d specifies a queryNameVersion without a queryName", i + )); + } + } + } + + /** + * Resolves the deadline for the whole request as the longest of the per-query + * timeouts, so that no sub-query is cancelled before the time its caller + * allowed. Returns empty when no query sets one, which lets the client-level + * timeout apply. Per-query timeouts are not independently enforced: a single + * request carries a single deadline. + */ + static Optional resolveDeadline(List paramsList) { + Duration longest = null; + for (OnlineQueryParamsComplete params : paramsList) { + Duration timeout = params.getTimeout(); + if (timeout == null) { + continue; + } + if (longest == null || timeout.compareTo(longest) > 0) { + longest = timeout; + } + } + return Optional.ofNullable(longest); + } + + static @Nullable String resolveEnvironment(List paramsList) + throws ClientException { + return requireAgreement(paramsList, OnlineQueryParamsComplete::getEnvironmentId, "environmentId"); + } + + static @Nullable String resolveBranch( + List paramsList, + @Nullable String clientBranchId + ) throws ClientException { + return requireAgreement( + paramsList, + params -> isBlank(params.getBranch()) ? clientBranchId : params.getBranch(), + "branch" + ); + } + + static @Nullable String resolveQueryName(List paramsList) + throws ClientException { + String queryName = requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryName, "queryName"); + requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryNameVersion, "queryNameVersion"); + return queryName; + } + + /* + * Returns the value that every query agrees on, or throws if two queries disagree. + * + * Branch, environment and query name are all carried by per-request headers, of + * which a single request has exactly one set: the API server resolves one branch + * to one deployment, and named queries are routed by an exact match on the + * "x-chalk-query-name" header. A batch that mixed these values could not be + * satisfied, and silently picking one would run a sub-query somewhere its caller + * never asked for. + */ + private static @Nullable String requireAgreement( + List paramsList, + Function getter, + String settingName + ) throws ClientException { + if (paramsList.isEmpty()) { + return null; + } + String resolved = blankToNull(getter.apply(paramsList.get(0))); + for (int i = 1; i < paramsList.size(); i++) { + String candidate = blankToNull(getter.apply(paramsList.get(i))); + if (!Objects.equals(resolved, candidate)) { + throw new ClientException(String.format( + "onlineQueryMulti requires every query to agree on %s, but found %s at " + + "index 0 and %s at index %d. A single request carries one %s for " + + "all of its queries, so run these as separate queries instead.", + settingName, describe(resolved), describe(candidate), i, settingName + )); + } + } + return resolved; + } + + private static ChalkException asChalkException(Exception e, String fallbackMessage) { + return e instanceof ChalkException ? (ChalkException) e : new ClientException(fallbackMessage, e); + } + + private static void releaseQuietly( + Throwable primary, + @Nullable Table scalars, + Collection groups, + @Nullable BufferAllocator allocator + ) { + if (scalars != null) { + try { scalars.close(); + } catch (Exception e) { + primary.addSuppressed(e); } - for (var table : groups.values()) { + } + for (Table table : groups) { + try { table.close(); + } catch (Exception e) { + primary.addSuppressed(e); + } + } + if (allocator != null) { + try { + allocator.close(); + } catch (Exception e) { + primary.addSuppressed(e); } - responseAlloc.close(); } + } - return new OnlineQueryResult( - scalars, - groups, - errors, - meta, - responseAlloc - ); + private static String describeErrors(ServerError[] errors) { + List messages = new ArrayList<>(); + for (ServerError error : errors) { + messages.add(error.getMessage() == null ? "(no message)" : error.getMessage()); + } + return String.join("; ", messages); + } + + private static String describe(@Nullable String value) { + return value == null ? "none" : "'" + value + "'"; + } + + private static boolean isBlank(@Nullable String value) { + return value == null || value.isEmpty(); + } + + private static @Nullable String blankToNull(@Nullable String value) { + return isBlank(value) ? null : value; } public UploadFeaturesResult uploadFeatures(UploadFeaturesParams params) throws ChalkException { diff --git a/src/main/java/ai/chalk/internal/request/models/OnlineQueryBulkResponse.java b/src/main/java/ai/chalk/internal/request/models/OnlineQueryBulkResponse.java index 776d43cc..8805b2ad 100644 --- a/src/main/java/ai/chalk/internal/request/models/OnlineQueryBulkResponse.java +++ b/src/main/java/ai/chalk/internal/request/models/OnlineQueryBulkResponse.java @@ -43,6 +43,14 @@ public static OnlineQueryBulkResponse fromBytes(byte[] bytes, BufferAllocator al return new OnlineQueryBulkResponse(resultFeatherMap); } + /* + * TODO: the REST wire format is already multi-query shaped -- `query_results_bytes` + * is a map keyed by query index, and this reads only key "0". Generalizing this is + * what a REST implementation of ChalkClient.onlineQueryMulti would need. Note that + * OnlineQueryResultFeather.close() closes a *shared* allocator once per map entry, + * which must be fixed first: with more than one entry it would release the allocator + * while the other entries' tables still hold buffers. + */ public OnlineQueryResult toResult() throws ChalkException { if (!(this.queryResults.containsKey("0"))) { throw new ClientException("malformed online query bulk response"); diff --git a/src/main/java/ai/chalk/models/OnlineQueryMultiResult.java b/src/main/java/ai/chalk/models/OnlineQueryMultiResult.java new file mode 100644 index 00000000..652815ce --- /dev/null +++ b/src/main/java/ai/chalk/models/OnlineQueryMultiResult.java @@ -0,0 +1,106 @@ +package ai.chalk.models; + +import ai.chalk.exceptions.ServerError; +import lombok.Getter; + +import javax.annotation.Nullable; +import java.util.Collections; +import java.util.List; + +/** + * OnlineQueryMultiResult holds the results of a multi-query, i.e. several + * independent online queries executed in parallel by a single request. + * + *

Closing this object closes every sub-result and releases all Arrow + * memory associated with the request. Callers should use try-with-resources + * and should not close the individual {@link OnlineQueryResult}s + * returned by {@link #getResults()}. + */ +public class OnlineQueryMultiResult implements AutoCloseable { + + /** + * results holds one {@link OnlineQueryResult} per sub-query, in the same + * order the queries were passed in. Errors that pertain to a specific + * sub-query are found on that sub-query's result, via + * {@link OnlineQueryResult#getErrors()}. + * + *

Each result owns its own Arrow allocator, which is released when + * this object is closed. + * + *

This list is 1:1 with the queries that were submitted, with one + * exception: if the request failed before any sub-query was dispatched, + * this list is empty and {@link #getGlobalErrors()} explains why. + */ + @Getter + private final List results; + + /** + * globalErrors holds the errors that the engine did not attach to a specific + * sub-query's result. A resolver that fails while computing a sub-query is + * reported on that sub-query's own result instead, so most errors are not + * global. + * + *

An error here is not necessarily unrelated to a particular sub-query. + * When a sub-query fails to execute at all -- it was rate limited, or + * the engine hit a transport or internal error running it -- the engine + * records the failure here and names the sub-query by its 1-based position in + * the message text, e.g. {@code "... while executing query 2/3: ..."}. That + * sub-query still occupies its slot in {@link #getResults()}, but the result + * is empty and reports no errors of its own. So an empty result is not proof + * of success: check this list too before concluding that every sub-query ran. + */ + @Getter + private final ServerError[] globalErrors; + + /** + * traceId is the trace ID of the single underlying request, and is + * therefore shared by every sub-query. + */ + @Getter + private final String traceId; + + private boolean closed = false; + + public OnlineQueryMultiResult( + List results, + ServerError[] globalErrors, + @Nullable String traceId + ) { + this.results = Collections.unmodifiableList(results); + this.globalErrors = globalErrors; + this.traceId = traceId; + } + + /** + * close releases the resources associated with every sub-result. It is + * idempotent. If closing a sub-result fails, the remaining sub-results + * are still closed and the first failure is thrown with the rest + * attached as suppressed exceptions. + */ + @Override + public void close() { + if (this.closed) { + return; + } + this.closed = true; + + RuntimeException failure = null; + for (OnlineQueryResult result : this.results) { + if (result == null) { + continue; + } + try { + result.close(); + } catch (RuntimeException e) { + if (failure == null) { + failure = e; + } else { + failure.addSuppressed(e); + } + } + } + if (failure != null) { + throw failure; + } + } +} diff --git a/src/test/java/ai/chalk/client/TestAllClients.java b/src/test/java/ai/chalk/client/TestAllClients.java index aa4d7ac2..7026e168 100644 --- a/src/test/java/ai/chalk/client/TestAllClients.java +++ b/src/test/java/ai/chalk/client/TestAllClients.java @@ -2,6 +2,7 @@ import ai.chalk.client.e2e.FraudTemplateFeatures; import ai.chalk.client.e2e.User; +import ai.chalk.exceptions.ClientException; import ai.chalk.models.*; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; @@ -13,6 +14,8 @@ import java.time.Duration; import java.util.*; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -284,5 +287,19 @@ public void testNamedQuery(String clientType) throws Exception { } } + /** + * Multi-query is a gRPC-only capability today. The REST client inherits the + * rejecting default on ChalkClient; make that contract explicit. + */ + @Test + public void testOnlineQueryMultiUnsupportedOnRest() { + var params = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, List.of("1")) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .build(); + var thrown = assertThrows( + ClientException.class, () -> restClient.onlineQueryMulti(List.of(params))); + assertTrue(thrown.getMessage().contains("gRPC")); + } } diff --git a/src/test/java/ai/chalk/client/TestGrpcClient.java b/src/test/java/ai/chalk/client/TestGrpcClient.java index f0c60094..7cf4e791 100644 --- a/src/test/java/ai/chalk/client/TestGrpcClient.java +++ b/src/test/java/ai/chalk/client/TestGrpcClient.java @@ -5,7 +5,10 @@ import ai.chalk.client.e2e.User; import ai.chalk.models.ErrorCode; import ai.chalk.models.ErrorCodeCategory; +import ai.chalk.exceptions.ClientException; +import ai.chalk.models.OnlineQueryMultiResult; import ai.chalk.models.OnlineQueryParams; +import ai.chalk.models.OnlineQueryParamsComplete; import ai.chalk.models.OnlineQueryResult; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; @@ -160,5 +163,131 @@ public void testOnlineQueryOptionalParamsSanity() throws Exception { } } + @Test + public void testOnlineQueryMultiScalars() throws Exception { + var firstIds = List.of("1", "2", "3"); + var secondIds = List.of("4", "5"); + var first = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, firstIds) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .build(); + var second = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, secondIds) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .build(); + + try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(first, second))) { + assert multi.getGlobalErrors().length == 0; + assert multi.getResults().size() == 2; + + var firstUsers = multi.getResults().get(0).unmarshal(User.class); + assert firstUsers.length == firstIds.size(); + assert firstUsers[0].socure_score.getValue().equals(123.0); + + var secondUsers = multi.getResults().get(1).unmarshal(User.class); + assert secondUsers.length == secondIds.size(); + } + } + + /** + * The engine executes the sub-queries independently, so one failing query must + * leave the others' results intact and must not produce a global error. + */ + @Test + public void testOnlineQueryMultiPartialFailure() throws Exception { + var userIds = List.of("1", "2", "3"); + var healthy = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, userIds) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .build(); + var crashing = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, userIds) + .withOutputs(FraudTemplateFeatures.user.socure_score, FraudTemplateFeatures.user.crashingFeature) + .build(); + + try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(healthy, crashing))) { + // the batch itself succeeded; the failure belongs to one sub-query + assert multi.getGlobalErrors().length == 0; + assert multi.getResults().size() == 2; + + var healthyResult = multi.getResults().get(0); + assert healthyResult.getErrors().length == 0; + var users = healthyResult.unmarshal(User.class); + assert users.length == userIds.size(); + assert users[0].socure_score.getValue().equals(123.0); + + var crashingResult = multi.getResults().get(1); + assert crashingResult.getErrors().length > 0; + assert crashingResult.getErrors()[0].getCode().equals(ErrorCode.RESOLVER_FAILED); + } + } + + @Test + public void testOnlineQueryMultiTraceId() throws Exception { + var params = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, List.of("1")) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .build(); + + try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(params, params))) { + assert multi.getTraceId() != null && !multi.getTraceId().isEmpty(); + // one request means one trace id, shared by every sub-query + for (var result : multi.getResults()) { + assert multi.getTraceId().equals(result.getMeta().getTraceId()); + } + } + } + + /** + * A one-element batch must agree with the equivalent single query, which + * cross-checks the request builders and response decoders the two paths share. + */ + @Test + public void testOnlineQueryMultiMatchesSingleQuery() throws Exception { + var userIds = List.of("1", "2", "3"); + var params = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, userIds) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .build(); + + try ( + OnlineQueryResult single = client.onlineQuery(params); + OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(params)) + ) { + assert multi.getResults().size() == 1; + var singleUsers = single.unmarshal(User.class); + var multiUsers = multi.getResults().get(0).unmarshal(User.class); + assert singleUsers.length == multiUsers.length; + for (int i = 0; i < singleUsers.length; i++) { + assert singleUsers[i].socure_score.getValue() + .equals(multiUsers[i].socure_score.getValue()); + } + } + } + + @Test + public void testOnlineQueryMultiRejectsConflictingQueryName() throws Exception { + var named = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, List.of("1")) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .withQueryName("chalk-java::multiA") + .build(); + var otherName = OnlineQueryParams.builder() + .withInput(FraudTemplateFeatures.user.id, List.of("1")) + .withOutputs(FraudTemplateFeatures.user.socure_score) + .withQueryName("chalk-java::multiB") + .build(); + + // Query name is routed on per request, so a batch cannot mix names. + org.junit.jupiter.api.Assertions.assertThrows( + ClientException.class, () -> client.onlineQueryMulti(List.of(named, otherName))); + } + @Test + public void testOnlineQueryMultiEmptyList() throws Exception { + try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of())) { + assert multi.getResults().isEmpty(); + assert multi.getGlobalErrors().length == 0; + } + } } diff --git a/src/test/java/ai/chalk/client/TestGrpcMultiQuery.java b/src/test/java/ai/chalk/client/TestGrpcMultiQuery.java new file mode 100644 index 00000000..1fb58e88 --- /dev/null +++ b/src/test/java/ai/chalk/client/TestGrpcMultiQuery.java @@ -0,0 +1,572 @@ +package ai.chalk.client; + +import ai.chalk.exceptions.ClientException; +import ai.chalk.internal.arrow.FeatherProcessor; +import ai.chalk.models.OnlineQueryMultiResult; +import ai.chalk.models.OnlineQueryParams; +import ai.chalk.models.OnlineQueryParamsComplete; +import ai.chalk.models.OnlineQueryResult; +import ai.chalk.protos.chalk.common.v1.ChalkError; +import ai.chalk.protos.chalk.common.v1.ErrorCode; +import ai.chalk.protos.chalk.common.v1.ErrorCodeCategory; +import ai.chalk.protos.chalk.common.v1.FeatherBodyType; +import ai.chalk.protos.chalk.common.v1.GenericSingleResponse; +import ai.chalk.protos.chalk.common.v1.OnlineQueryBulkResponse; +import ai.chalk.protos.chalk.common.v1.OnlineQueryMetadata; +import ai.chalk.protos.chalk.common.v1.OnlineQueryMultiResponse; +import ai.chalk.protos.chalk.common.v1.OnlineQueryResponse; +import com.google.protobuf.ByteString; +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.time.ZonedDateTime; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Offline tests for the gRPC multi-query request builders and response decoders. + * These exercise the pure helpers in {@link GRPCClient} without a live engine. + * + *

Every test inherits {@link AllocatorTest}'s teardown, which closes the root + * allocator and therefore fails any case that leaked Arrow memory. + */ +public class TestGrpcMultiQuery extends AllocatorTest { + + private static final String TRACE_ID = "trace-abc"; + + // ---------------------------------------------------------------- helpers + + private static OnlineQueryParamsComplete params(String queryName) { + var builder = OnlineQueryParams.builder() + .withInput("user.id", List.of(1, 2, 3)) + .withOutputs("user.socure_score"); + if (queryName != null) { + builder = builder.withQueryName(queryName); + } + return builder.build(); + } + + private static OnlineQueryParamsComplete simpleParams() { + return params(null); + } + + /** Real Arrow IPC bytes for a one-column table with {@code rowCount} rows. */ + private byte[] scalarBytes(int rowCount) throws Exception { + List ids = new ArrayList<>(); + for (int i = 0; i < rowCount; i++) { + ids.add(i); + } + Map> inputs = new HashMap<>(); + inputs.put("user.id", ids); + return FeatherProcessor.inputsToArrowBytes(inputs, this.allocator); + } + + private OnlineQueryBulkResponse bulkResponse(int rowCount) throws Exception { + return OnlineQueryBulkResponse.newBuilder() + .setScalarsData(ByteString.copyFrom(scalarBytes(rowCount))) + .setResponseMeta(OnlineQueryMetadata.newBuilder().setQueryId("q-" + rowCount).build()) + .build(); + } + + private static GenericSingleResponse wrap(OnlineQueryBulkResponse bulk) { + return GenericSingleResponse.newBuilder().setBulkResponse(bulk).build(); + } + + private static ChalkError error(String message) { + return ChalkError.newBuilder() + .setCode(ErrorCode.ERROR_CODE_RESOLVER_FAILED) + .setCategory(ErrorCodeCategory.ERROR_CODE_CATEGORY_REQUEST) + .setMessage(message) + .build(); + } + + // ------------------------------------------------------- request building + + @Test + public void testBuildBulkRequestGolden() throws Exception { + var now = List.of(ZonedDateTime.now(), ZonedDateTime.now().minusDays(1)); + Map plannerOptions = new HashMap<>(); + plannerOptions.put("planner_version", "2"); + OnlineQueryParamsComplete p = OnlineQueryParams.builder() + .withInput("user.id", List.of(1, 2, 3)) + .withOutputs("user.socure_score", "user.email") + .withTags(List.of("tag1", "tag2")) + .withRequiredResolverTags(List.of("rrt1")) + .withMeta(Map.of("metaKey", "metaValue")) + .withCorrelationId("corr-1") + .withPreviewDeploymentId("preview-1") + .withEnvironmentId("env-1") + .withQueryName("my_query") + .withQueryNameVersion("2") + .withNow(now) + .withPlannerOptions(plannerOptions) + .withIncludeMeta(true) + .build(); + + byte[] feather = FeatherProcessor.inputsToArrowBytes(p.getInputs(), this.allocator); + var request = GRPCClient.buildBulkRequest(p, feather, null); + + assertEquals(FeatherBodyType.FEATHER_BODY_TYPE_TABLE, request.getBodyType()); + assertEquals(ByteString.copyFrom(feather), request.getInputsFeather()); + + assertEquals(2, request.getOutputsCount()); + assertEquals("user.socure_score", request.getOutputs(0).getFeatureFqn()); + assertEquals("user.email", request.getOutputs(1).getFeatureFqn()); + + assertEquals(2, request.getNowCount()); + assertEquals(now.get(0).toEpochSecond(), request.getNow(0).getSeconds()); + assertEquals(now.get(0).getNano(), request.getNow(0).getNanos()); + + var context = request.getContext(); + assertEquals(List.of("tag1", "tag2"), context.getTagsList()); + assertEquals(List.of("rrt1"), context.getRequiredResolverTagsList()); + assertEquals("corr-1", context.getCorrelationId()); + assertEquals("preview-1", context.getDeploymentId()); + assertEquals("env-1", context.getEnvironment()); + assertEquals("my_query", context.getQueryName()); + assertEquals("2", context.getQueryNameVersion()); + assertEquals("2", context.getOptionsOrThrow("planner_version").getStringValue()); + + var options = request.getResponseOptions(); + assertTrue(options.getIncludeMeta()); + assertTrue(options.getEncodingOptions().getEncodeStructsAsObjects()); + assertEquals("metaValue", options.getMetadataOrThrow("metaKey")); + } + + @Test + public void testBuildBulkRequestExplainForcesIncludeMeta() throws Exception { + OnlineQueryParamsComplete p = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withExplain(true) + .build(); + + byte[] feather = FeatherProcessor.inputsToArrowBytes(p.getInputs(), this.allocator); + var request = GRPCClient.buildBulkRequest(p, feather, null); + + assertTrue(request.getResponseOptions().getIncludeMeta()); + assertTrue(request.getResponseOptions().hasExplain()); + } + + @Test + public void testBuildBulkRequestBranchFallback() throws Exception { + byte[] feather = FeatherProcessor.inputsToArrowBytes( + simpleParams().getInputs(), this.allocator); + + OnlineQueryParamsComplete withBranch = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withBranch("params-branch") + .build(); + + // params branch wins over the client branch + assertEquals( + "params-branch", + GRPCClient.buildBulkRequest(withBranch, feather, "client-branch").getContext().getBranchId()); + // client branch is used when the params do not set one + assertEquals( + "client-branch", + GRPCClient.buildBulkRequest(simpleParams(), feather, "client-branch").getContext().getBranchId()); + // neither set + assertEquals( + "", + GRPCClient.buildBulkRequest(simpleParams(), feather, null).getContext().getBranchId()); + } + + @Test + public void testBuildMultiRequestPreservesOrderAndUsesBulk() throws Exception { + OnlineQueryParamsComplete a = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .build(); + OnlineQueryParamsComplete b = OnlineQueryParams.builder() + .withInput("merchant.id", List.of("m1", "m2")) + .withOutputs("merchant.risk_score") + .build(); + + var request = GRPCClient.buildMultiRequest(List.of(a, b), this.allocator, null); + + assertEquals(2, request.getQueriesCount()); + assertTrue(request.getQueries(0).hasBulkRequest()); + assertTrue(request.getQueries(1).hasBulkRequest()); + + // Order is preserved: each sub-request matches the params at the same index. + assertEquals( + "user.socure_score", + request.getQueries(0).getBulkRequest().getOutputs(0).getFeatureFqn()); + assertEquals( + "merchant.risk_score", + request.getQueries(1).getBulkRequest().getOutputs(0).getFeatureFqn()); + } + + // ------------------------------------------------------- settings resolvers + + @Test + public void testResolveDeadlineTakesTheLongest() { + assertTrue(GRPCClient.resolveDeadline(List.of(simpleParams(), simpleParams())).isEmpty()); + + OnlineQueryParamsComplete shortTimeout = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withTimeout(Duration.ofSeconds(1)) + .build(); + OnlineQueryParamsComplete longTimeout = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withTimeout(Duration.ofSeconds(30)) + .build(); + + assertEquals( + Duration.ofSeconds(30), + GRPCClient.resolveDeadline(List.of(shortTimeout, longTimeout, simpleParams())).orElseThrow()); + } + + @Test + public void testResolveQueryNameRequiresAgreement() throws Exception { + assertEquals("same", GRPCClient.resolveQueryName(List.of(params("same"), params("same")))); + assertNull(GRPCClient.resolveQueryName(List.of(simpleParams(), simpleParams()))); + + var conflict = assertThrows( + ClientException.class, + () -> GRPCClient.resolveQueryName(List.of(params("a"), params("b")))); + assertTrue(conflict.getMessage().contains("queryName")); + assertTrue(conflict.getMessage().contains("index 1")); + + // an unset name conflicts with a set one, because the header is per-request + assertThrows( + ClientException.class, + () -> GRPCClient.resolveQueryName(List.of(params("a"), simpleParams()))); + } + + /** + * When both the name and the version disagree, the name is the more useful + * thing to report: it is what the caller actually set. + */ + @Test + public void testResolveQueryNameReportsNameBeforeVersion() { + OnlineQueryParamsComplete a = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withQueryName("a") + .withQueryNameVersion("1") + .build(); + OnlineQueryParamsComplete b = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withQueryName("b") + .withQueryNameVersion("2") + .build(); + + var conflict = assertThrows( + ClientException.class, () -> GRPCClient.resolveQueryName(List.of(a, b))); + assertTrue(conflict.getMessage().contains("queryName,")); + assertTrue(!conflict.getMessage().contains("queryNameVersion")); + } + + @Test + public void testResolveEnvironmentRequiresAgreement() throws Exception { + OnlineQueryParamsComplete envA = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withEnvironmentId("env-a") + .build(); + OnlineQueryParamsComplete envB = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withEnvironmentId("env-b") + .build(); + + assertEquals("env-a", GRPCClient.resolveEnvironment(List.of(envA, envA))); + + var conflict = assertThrows( + ClientException.class, () -> GRPCClient.resolveEnvironment(List.of(envA, envB))); + assertTrue(conflict.getMessage().contains("environmentId")); + assertTrue(conflict.getMessage().contains("'env-a'")); + assertTrue(conflict.getMessage().contains("'env-b'")); + } + + @Test + public void testResolveBranchUsesEffectiveValue() throws Exception { + OnlineQueryParamsComplete branchA = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withBranch("a") + .build(); + + // the client-level branch fills in for the query that does not set one + assertEquals("a", GRPCClient.resolveBranch(List.of(branchA, simpleParams()), "a")); + // and a per-query branch that disagrees with it is rejected + assertThrows( + ClientException.class, + () -> GRPCClient.resolveBranch(List.of(branchA, simpleParams()), "other")); + } + + // -------------------------------------------------------------- validation + + @Test + public void testValidateMultiParams() { + assertThrows(ClientException.class, () -> GRPCClient.validateMultiParams(null)); + + List withNull = new ArrayList<>(); + withNull.add(simpleParams()); + withNull.add(null); + var nullError = assertThrows( + ClientException.class, () -> GRPCClient.validateMultiParams(withNull)); + assertTrue(nullError.getMessage().contains("index 1")); + + OnlineQueryParamsComplete versionWithoutName = OnlineQueryParams.builder() + .withInput("user.id", List.of(1)) + .withOutputs("user.socure_score") + .withQueryNameVersion("3") + .build(); + var versionError = assertThrows( + ClientException.class, + () -> GRPCClient.validateMultiParams(List.of(simpleParams(), versionWithoutName))); + assertTrue(versionError.getMessage().contains("queryNameVersion")); + assertTrue(versionError.getMessage().contains("index 1")); + } + + // ---------------------------------------------------------- response decode + + @Test + public void testDecodeMultiResponseHappyPath() throws Exception { + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(3))) + .addResponses(wrap(bulkResponse(5))) + .build(); + + try (OnlineQueryMultiResult multi = + GRPCClient.decodeMultiResponse(response, 2, TRACE_ID, this.allocator)) { + assertEquals(2, multi.getResults().size()); + assertEquals(0, multi.getGlobalErrors().length); + assertEquals(TRACE_ID, multi.getTraceId()); + + // positional: the 3-row response first, then the 5-row one + assertEquals(3, multi.getResults().get(0).getScalarsTable().getRowCount()); + assertEquals(5, multi.getResults().get(1).getScalarsTable().getRowCount()); + + // the trace id is copied onto every sub-result + assertEquals(TRACE_ID, multi.getResults().get(0).getMeta().getTraceId()); + assertEquals(TRACE_ID, multi.getResults().get(1).getMeta().getTraceId()); + + // each sub-result owns a distinct allocator, so closing one cannot + // release the buffers of another + assertNotSame( + multi.getResults().get(0).getAllocator(), + multi.getResults().get(1).getAllocator()); + } + + // closing the wrapper released everything + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + @Test + public void testDecodeMultiResponseCarriesPerQueryErrors() throws Exception { + var failing = OnlineQueryBulkResponse.newBuilder() + .addErrors(error("resolver exploded")) + .build(); + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(2))) + .addResponses(wrap(failing)) + .build(); + + try (OnlineQueryMultiResult multi = + GRPCClient.decodeMultiResponse(response, 2, TRACE_ID, this.allocator)) { + // the batch as a whole succeeded + assertEquals(0, multi.getGlobalErrors().length); + // query 0 is unaffected by query 1's failure + assertEquals(0, multi.getResults().get(0).getErrors().length); + assertEquals(2, multi.getResults().get(0).getScalarsTable().getRowCount()); + // query 1 carries its own error + assertEquals(1, multi.getResults().get(1).getErrors().length); + assertEquals("resolver exploded", multi.getResults().get(1).getErrors()[0].getMessage()); + } + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + @Test + public void testDecodeMultiResponseGlobalErrorsWithNoResponses() throws Exception { + var response = OnlineQueryMultiResponse.newBuilder() + .addErrors(error("planner blew up")) + .build(); + + try (OnlineQueryMultiResult multi = + GRPCClient.decodeMultiResponse(response, 2, TRACE_ID, this.allocator)) { + assertTrue(multi.getResults().isEmpty()); + assertEquals(1, multi.getGlobalErrors().length); + assertEquals("planner blew up", multi.getGlobalErrors()[0].getMessage()); + assertEquals(TRACE_ID, multi.getTraceId()); + } + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + /** + * The shape the engine produces when a sub-query fails to execute at all: every + * slot is present, the failed one is an empty bulk response carrying no errors of + * its own, and the reason is a global error naming the query by position. The + * empty result must not be mistaken for a successful one. + */ + @Test + public void testDecodeMultiResponseExecutionFailureIsGlobalWithEmptySlot() throws Exception { + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(2))) + .addResponses(wrap(OnlineQueryBulkResponse.newBuilder().build())) + .addErrors(error("Encountered an error occurred while executing query 2/2: rate limited")) + .build(); + + try (OnlineQueryMultiResult multi = + GRPCClient.decodeMultiResponse(response, 2, TRACE_ID, this.allocator)) { + // positional alignment is preserved + assertEquals(2, multi.getResults().size()); + + assertEquals(2, multi.getResults().get(0).getScalarsTable().getRowCount()); + assertEquals(0, multi.getResults().get(0).getErrors().length); + + // the failed sub-query's slot looks empty rather than failed ... + assertNull(multi.getResults().get(1).getScalarsTable()); + assertEquals(0, multi.getResults().get(1).getErrors().length); + + // ... and the reason lives in the global errors, naming the query + assertEquals(1, multi.getGlobalErrors().length); + assertTrue(multi.getGlobalErrors()[0].getMessage().contains("query 2/2")); + } + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + @Test + public void testDecodeMultiResponseCountMismatch() throws Exception { + var tooFew = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(1))) + .build(); + var tooFewError = assertThrows( + ClientException.class, + () -> GRPCClient.decodeMultiResponse(tooFew, 2, TRACE_ID, this.allocator)); + assertTrue(tooFewError.getMessage().contains("Expected 2 sub-responses")); + assertEquals(0, this.allocator.getAllocatedMemory()); + + var tooMany = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(1))) + .addResponses(wrap(bulkResponse(1))) + .addResponses(wrap(bulkResponse(1))) + .build(); + assertThrows( + ClientException.class, + () -> GRPCClient.decodeMultiResponse(tooMany, 2, TRACE_ID, this.allocator)); + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + @Test + public void testDecodeMultiResponseRejectsSingleResponse() throws Exception { + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(1))) + .addResponses(GenericSingleResponse.newBuilder() + .setSingleResponse(OnlineQueryResponse.newBuilder().build()) + .build()) + .build(); + + var error = assertThrows( + ClientException.class, + () -> GRPCClient.decodeMultiResponse(response, 2, TRACE_ID, this.allocator)); + assertTrue(error.getMessage().contains("index 1")); + assertTrue(error.getMessage().contains("SINGLE_RESPONSE")); + // the sub-result decoded at index 0 was released + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + @Test + public void testDecodeMultiResponseRejectsUnsetOneof() throws Exception { + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(GenericSingleResponse.newBuilder().build()) + .build(); + + var error = assertThrows( + ClientException.class, + () -> GRPCClient.decodeMultiResponse(response, 1, TRACE_ID, this.allocator)); + assertTrue(error.getMessage().contains("QUERY_NOT_SET")); + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + /** + * A sub-response whose Arrow payload cannot be decoded must fail the batch + * without leaking the results decoded before it. This is the case that the + * old non-rethrowing decode path got wrong. + */ + @Test + public void testDecodeMultiResponseDecodeFailureLeaksNothing() throws Exception { + var garbage = OnlineQueryBulkResponse.newBuilder() + .setScalarsData(ByteString.copyFromUtf8("this is not arrow data")) + .build(); + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(4))) + .addResponses(wrap(garbage)) + .build(); + + var error = assertThrows( + ClientException.class, + () -> GRPCClient.decodeMultiResponse(response, 2, TRACE_ID, this.allocator)); + assertTrue(error.getMessage().contains("index 1")); + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + @Test + public void testDecodeBulkResponseReleasesAllocatorOnFailure() throws Exception { + var garbage = OnlineQueryBulkResponse.newBuilder() + .setScalarsData(ByteString.copyFromUtf8("not arrow")) + .build(); + var childAllocator = this.allocator.newChildAllocator( + "test_decode_failure", 0, FeatherProcessor.ALLOCATOR_SIZE_RESPONSE); + + assertThrows( + ClientException.class, + () -> GRPCClient.decodeBulkResponse(garbage, TRACE_ID, childAllocator)); + // decodeBulkResponse owns the allocator it is handed and must release it + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + // ------------------------------------------------------------ result wrapper + + @Test + public void testMultiResultCloseIsIdempotent() throws Exception { + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(bulkResponse(2))) + .build(); + + OnlineQueryMultiResult multi = + GRPCClient.decodeMultiResponse(response, 1, TRACE_ID, this.allocator); + multi.close(); + multi.close(); + assertEquals(0, this.allocator.getAllocatedMemory()); + } + + @Test + public void testMultiResultCloseClosesEveryResultDespiteFailure() { + var failing = new OnlineQueryResult(null, Map.of(), null, null, null) { + @Override + public void close() { + throw new IllegalStateException("boom"); + } + }; + List closed = new ArrayList<>(); + var tracking = new OnlineQueryResult(null, Map.of(), null, null, null) { + @Override + public void close() { + closed.add(true); + } + }; + + var multi = new OnlineQueryMultiResult(List.of(failing, tracking), new ai.chalk.exceptions.ServerError[0], null); + + var thrown = assertThrows(IllegalStateException.class, multi::close); + assertEquals("boom", thrown.getMessage()); + // the failure at index 0 did not stop index 1 from being closed + assertEquals(1, closed.size()); + } +} From 0450ccb43aa06fc69f6d435408f4b2e6fd727069 Mon Sep 17 00:00:00 2001 From: sjmignot Date: Thu, 10 Sep 2026 11:05:16 -0700 Subject: [PATCH 2/2] Address review: deadline, environment and version resolution Fixes found reviewing the multi-query settings resolvers. - resolveDeadline dropped the client-level timeout as soon as any one query set its own. A query that sets no timeout is bounded by the client-level timeout, so that is its effective timeout and it has to take part in the maximum; otherwise a client built with a 5m timeout batching a 1s query alongside an unset one deadlined the whole request at 1s, killing a query that would have had 5m on its own. This contradicted the method's own stated invariant. A query that is effectively unbounded now yields no deadline at all. - resolveEnvironment compared raw per-query values, so naming the client's own environment on one query and leaving it unset on another was reported as a conflict even though both resolve to the same environment. It now compares the effective value, as resolveBranch already did. - Dropped the queryNameVersion agreement requirement. The justification does not hold for the version: there is no x-chalk-query-name-version header, and each sub-request carries its own version in its OnlineQueryContext, so queries differing only by version are satisfiable. - unmarshal() NPE'd on a null scalars table, which is exactly what a sub-query that failed to execute leaves behind, so the documented empty slot made the README's own pattern fail with an opaque error. It now explains where to find the reason. - Guard the varargs overload against a null array so it reports the same ClientException as the List overload rather than a raw NPE, and drop a repeated format argument that made the count-mismatch message tautological. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 14 +- .../java/ai/chalk/client/ChalkClient.java | 16 +- src/main/java/ai/chalk/client/GRPCClient.java | 57 ++++--- .../ai/chalk/internal/arrow/Unmarshaller.java | 11 ++ .../ai/chalk/client/TestGrpcMultiQuery.java | 141 +++++++++++++----- .../chalk/client/TestGrpcMultiQueryUser.java | 10 ++ 6 files changed, 183 insertions(+), 66 deletions(-) create mode 100644 src/test/java/ai/chalk/client/TestGrpcMultiQueryUser.java diff --git a/README.md b/README.md index f958562e..5552d045 100644 --- a/README.md +++ b/README.md @@ -418,6 +418,10 @@ try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(users, merch // this query failed; the others are unaffected continue; } + if (result.getScalarsTable() == null) { + // this query never ran -- see getGlobalErrors() above for the reason + continue; + } // do something with the result } } @@ -437,8 +441,10 @@ A few things to know: - **Close the wrapper, not the results.** Closing the `OnlineQueryMultiResult` closes every one of its results and releases all of their Arrow memory. - **Some settings apply to the whole request,** which carries one deadline and one set of headers. - Every query must agree on `environmentId`, `branch`, `queryName` and `queryNameVersion`, or the - call throws `ClientException` — a request resolves one branch to one deployment and routes named - queries by an exact match on the query name, so a mixed batch could not be honored. Run those as - separate queries. The deadline is the longest of the per-query timeouts, falling back to the + Every query must agree on `environmentId`, `branch` and `queryName`, or the call throws + `ClientException` — a request resolves one branch to one deployment and routes named queries by + an exact match on the query name, so a mixed batch could not be honored. Run those as separate + queries. Leaving one of these unset just takes the client's value, which is not a disagreement, + and `queryNameVersion` may differ freely since each query carries its own. The deadline is the + longest of the effective per-query timeouts, where a query that sets none is bounded by the client-level timeout, so per-query timeouts are not independently enforced. diff --git a/src/main/java/ai/chalk/client/ChalkClient.java b/src/main/java/ai/chalk/client/ChalkClient.java index f8af98ed..cf60165e 100644 --- a/src/main/java/ai/chalk/client/ChalkClient.java +++ b/src/main/java/ai/chalk/client/ChalkClient.java @@ -188,15 +188,17 @@ default int ping() throws ChalkException { *

Some settings apply to the whole request. A request carries one * deadline and one set of headers, so: *

    - *
  • Every query must agree on {@code environmentId}, {@code branch}, - * {@code queryName} and {@code queryNameVersion}. A disagreement throws + *
  • Every query must agree on {@code environmentId}, {@code branch} and + * {@code queryName}. A disagreement throws * {@link ai.chalk.exceptions.ClientException}, because a request resolves * one branch to one deployment and routes named queries by an exact match * on the query name, so a mixed batch could not be honored. Run those as - * separate queries.
    • - *
  • The deadline is the longest of the per-query timeouts, falling back to - * the client-level timeout. Per-query timeouts are not independently - * enforced.
    • + * separate queries. A query that leaves one of these unset simply takes the + * client's value, which is not a disagreement. {@code queryNameVersion} may + * differ freely: each query carries its own. + *
  • The deadline is the longest of the effective per-query timeouts, where a + * query that sets none is bounded by the client-level timeout. Per-query + * timeouts are not independently enforced.
    • *
  • Branch selection is currently inert on the gRPC path for both single and * multi queries: the client does not yet send the header the API server * routes branches on.
    • @@ -228,7 +230,7 @@ default OnlineQueryMultiResult onlineQueryMulti(List */ default OnlineQueryMultiResult onlineQueryMulti(OnlineQueryParamsComplete... params) throws ChalkException { - return onlineQueryMulti(Arrays.asList(params)); + return onlineQueryMulti(params == null ? null : Arrays.asList(params)); } /** diff --git a/src/main/java/ai/chalk/client/GRPCClient.java b/src/main/java/ai/chalk/client/GRPCClient.java index 5ea4519e..97857667 100644 --- a/src/main/java/ai/chalk/client/GRPCClient.java +++ b/src/main/java/ai/chalk/client/GRPCClient.java @@ -240,8 +240,8 @@ public OnlineQueryMultiResult onlineQueryMulti(List p // Resolve the request-level settings before serializing anything, so a batch that // cannot be satisfied fails before doing the work of encoding its inputs. - Duration deadline = resolveDeadline(paramsList).orElse(null); - String environmentId = resolveEnvironment(paramsList); + Duration deadline = resolveDeadline(paramsList, this.timeout).orElse(null); + String environmentId = resolveEnvironment(paramsList, this.resolvedEnvironmentId); String queryName = resolveQueryName(paramsList); // Validated for agreement but deliberately unused: the gRPC client does not send the // "x-chalk-branch-id" header that the API server routes branches on, so branch @@ -514,8 +514,7 @@ static OnlineQueryMultiResult decodeMultiResponse( // mismatch leaves every result unattributable. Fail rather than guess. if (response.getResponsesCount() != expectedCount) { throw new ClientException(String.format( - "Expected %d sub-responses for %d queries, but got %d%s", - expectedCount, + "Expected %d sub-responses, one per query, but got %d%s", expectedCount, response.getResponsesCount(), globalErrors.length == 0 ? "" : "; errors: " + describeErrors(globalErrors) @@ -587,18 +586,30 @@ static void validateMultiParams(List paramsList) thro } /** - * Resolves the deadline for the whole request as the longest of the per-query - * timeouts, so that no sub-query is cancelled before the time its caller - * allowed. Returns empty when no query sets one, which lets the client-level - * timeout apply. Per-query timeouts are not independently enforced: a single + * Resolves the deadline for the whole request as the longest of the effective + * per-query timeouts, so that no sub-query is cancelled before the time its + * caller allowed. Per-query timeouts are not independently enforced: a single * request carries a single deadline. + * + *

    A query that sets no timeout of its own is bounded by the client-level + * timeout, so that is its effective timeout and it takes part in the maximum. + * Leaving it out would let one query's short timeout cut short a query that was + * relying on a longer client-level default. If any query is effectively + * unbounded -- it sets no timeout and the client sets none either -- the request + * gets no deadline. */ - static Optional resolveDeadline(List paramsList) { + static Optional resolveDeadline( + List paramsList, + Optional clientLevelTimeout + ) { Duration longest = null; for (OnlineQueryParamsComplete params : paramsList) { - Duration timeout = params.getTimeout(); + Duration timeout = params.getTimeout() != null + ? params.getTimeout() + : clientLevelTimeout.orElse(null); if (timeout == null) { - continue; + // This query is unbounded, so the request cannot carry a deadline. + return Optional.empty(); } if (longest == null || timeout.compareTo(longest) > 0) { longest = timeout; @@ -607,9 +618,17 @@ static Optional resolveDeadline(List params return Optional.ofNullable(longest); } - static @Nullable String resolveEnvironment(List paramsList) - throws ClientException { - return requireAgreement(paramsList, OnlineQueryParamsComplete::getEnvironmentId, "environmentId"); + static @Nullable String resolveEnvironment( + List paramsList, + @Nullable String clientEnvironmentId + ) throws ClientException { + return requireAgreement( + paramsList, + params -> isBlank(params.getEnvironmentId()) + ? clientEnvironmentId + : params.getEnvironmentId(), + "environmentId" + ); } static @Nullable String resolveBranch( @@ -623,11 +642,15 @@ static Optional resolveDeadline(List params ); } + /* + * Only the name has to agree. Envoy route-matches named queries on the + * "x-chalk-query-name" header, of which the request has one. There is no + * corresponding header for the version -- each sub-request carries its own in its + * OnlineQueryContext -- so queries that differ only by version are satisfiable. + */ static @Nullable String resolveQueryName(List paramsList) throws ClientException { - String queryName = requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryName, "queryName"); - requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryNameVersion, "queryNameVersion"); - return queryName; + return requireAgreement(paramsList, OnlineQueryParamsComplete::getQueryName, "queryName"); } /* diff --git a/src/main/java/ai/chalk/internal/arrow/Unmarshaller.java b/src/main/java/ai/chalk/internal/arrow/Unmarshaller.java index 9e2d27ba..5a12fdd8 100644 --- a/src/main/java/ai/chalk/internal/arrow/Unmarshaller.java +++ b/src/main/java/ai/chalk/internal/arrow/Unmarshaller.java @@ -30,6 +30,17 @@ public class Unmarshaller { public static List prefixToSkip = List.of(Constants.chalkDunderPrefix); public static T[] unmarshalOnlineQueryResult(OnlineQueryResult result, Class target) throws ClientException { + if (result.getScalarsTable() == null) { + // No scalar data came back. For a multi-query this is what a sub-query that failed + // to execute leaves behind: an empty result whose reason sits in the multi-result's + // global errors rather than on the result itself. + throw new ClientException( + "Cannot unmarshal an online query result that has no scalar data. Check " + + "getErrors() on this result, and, if it came from onlineQueryMulti, " + + "getGlobalErrors() on the OnlineQueryMultiResult -- a query that " + + "failed to execute is reported there." + ); + } try { var rootFeatureClasses = unmarshalTable(result.getScalarsTable(), target); // This only unmarshals if has-many features gets returned in `result.groupsTables` diff --git a/src/test/java/ai/chalk/client/TestGrpcMultiQuery.java b/src/test/java/ai/chalk/client/TestGrpcMultiQuery.java index 1fb58e88..90bf1bdb 100644 --- a/src/test/java/ai/chalk/client/TestGrpcMultiQuery.java +++ b/src/test/java/ai/chalk/client/TestGrpcMultiQuery.java @@ -24,6 +24,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotSame; @@ -209,24 +210,57 @@ public void testBuildMultiRequestPreservesOrderAndUsesBulk() throws Exception { // ------------------------------------------------------- settings resolvers - @Test - public void testResolveDeadlineTakesTheLongest() { - assertTrue(GRPCClient.resolveDeadline(List.of(simpleParams(), simpleParams())).isEmpty()); - - OnlineQueryParamsComplete shortTimeout = OnlineQueryParams.builder() + private static OnlineQueryParamsComplete paramsWithTimeout(Duration timeout) { + return OnlineQueryParams.builder() .withInput("user.id", List.of(1)) .withOutputs("user.socure_score") - .withTimeout(Duration.ofSeconds(1)) - .build(); - OnlineQueryParamsComplete longTimeout = OnlineQueryParams.builder() - .withInput("user.id", List.of(1)) - .withOutputs("user.socure_score") - .withTimeout(Duration.ofSeconds(30)) + .withTimeout(timeout) .build(); + } + + @Test + public void testResolveDeadlineTakesTheLongest() { + var noClientTimeout = Optional.empty(); + + // nothing set anywhere -> no deadline + assertTrue(GRPCClient.resolveDeadline( + List.of(simpleParams(), simpleParams()), noClientTimeout).isEmpty()); + + var shortTimeout = paramsWithTimeout(Duration.ofSeconds(1)); + var longTimeout = paramsWithTimeout(Duration.ofSeconds(30)); + + assertEquals( + Duration.ofSeconds(30), + GRPCClient.resolveDeadline(List.of(shortTimeout, longTimeout), noClientTimeout) + .orElseThrow()); + } + + /** + * A query that sets no timeout is bounded by the client-level timeout, so that is + * its effective timeout. Ignoring it would let one query's short timeout cut short + * a query that was relying on a longer client default. + */ + @Test + public void testResolveDeadlineHonoursClientLevelTimeout() { + var fast = paramsWithTimeout(Duration.ofSeconds(1)); + var clientTimeout = Optional.of(Duration.ofMinutes(5)); + + // `fast` wants 1s, the unset one is entitled to the client's 5m -> 5m wins + assertEquals( + Duration.ofMinutes(5), + GRPCClient.resolveDeadline(List.of(fast, simpleParams()), clientTimeout).orElseThrow()); + // every query explicit -> the client default is irrelevant assertEquals( Duration.ofSeconds(30), - GRPCClient.resolveDeadline(List.of(shortTimeout, longTimeout, simpleParams())).orElseThrow()); + GRPCClient.resolveDeadline( + List.of(fast, paramsWithTimeout(Duration.ofSeconds(30))), clientTimeout) + .orElseThrow()); + + // an unset query with no client timeout is unbounded, so the request cannot + // carry a deadline at all + assertTrue(GRPCClient.resolveDeadline( + List.of(fast, simpleParams()), Optional.empty()).isEmpty()); } @Test @@ -246,51 +280,58 @@ public void testResolveQueryNameRequiresAgreement() throws Exception { () -> GRPCClient.resolveQueryName(List.of(params("a"), simpleParams()))); } - /** - * When both the name and the version disagree, the name is the more useful - * thing to report: it is what the caller actually set. - */ @Test - public void testResolveQueryNameReportsNameBeforeVersion() { - OnlineQueryParamsComplete a = OnlineQueryParams.builder() + public void testResolveEnvironmentRequiresAgreement() throws Exception { + OnlineQueryParamsComplete envA = OnlineQueryParams.builder() .withInput("user.id", List.of(1)) .withOutputs("user.socure_score") - .withQueryName("a") - .withQueryNameVersion("1") + .withEnvironmentId("env-a") .build(); - OnlineQueryParamsComplete b = OnlineQueryParams.builder() + OnlineQueryParamsComplete envB = OnlineQueryParams.builder() .withInput("user.id", List.of(1)) .withOutputs("user.socure_score") - .withQueryName("b") - .withQueryNameVersion("2") + .withEnvironmentId("env-b") .build(); + assertEquals("env-a", GRPCClient.resolveEnvironment(List.of(envA, envA), null)); + var conflict = assertThrows( - ClientException.class, () -> GRPCClient.resolveQueryName(List.of(a, b))); - assertTrue(conflict.getMessage().contains("queryName,")); - assertTrue(!conflict.getMessage().contains("queryNameVersion")); + ClientException.class, + () -> GRPCClient.resolveEnvironment(List.of(envA, envB), null)); + assertTrue(conflict.getMessage().contains("environmentId")); + assertTrue(conflict.getMessage().contains("'env-a'")); + assertTrue(conflict.getMessage().contains("'env-b'")); + + // An unset environmentId means the client's environment, so naming that same + // environment explicitly on one query is not a conflict. + assertEquals( + "env-a", + GRPCClient.resolveEnvironment(List.of(envA, simpleParams()), "env-a")); + assertThrows( + ClientException.class, + () -> GRPCClient.resolveEnvironment(List.of(envA, simpleParams()), "env-other")); } + /** + * The version rides along in each sub-request's own context and has no header of + * its own, so queries that differ only by version are satisfiable. + */ @Test - public void testResolveEnvironmentRequiresAgreement() throws Exception { - OnlineQueryParamsComplete envA = OnlineQueryParams.builder() + public void testResolveQueryNameAllowsDifferingVersions() throws Exception { + OnlineQueryParamsComplete v1 = OnlineQueryParams.builder() .withInput("user.id", List.of(1)) .withOutputs("user.socure_score") - .withEnvironmentId("env-a") + .withQueryName("same") + .withQueryNameVersion("1") .build(); - OnlineQueryParamsComplete envB = OnlineQueryParams.builder() + OnlineQueryParamsComplete v2 = OnlineQueryParams.builder() .withInput("user.id", List.of(1)) .withOutputs("user.socure_score") - .withEnvironmentId("env-b") + .withQueryName("same") + .withQueryNameVersion("2") .build(); - assertEquals("env-a", GRPCClient.resolveEnvironment(List.of(envA, envA))); - - var conflict = assertThrows( - ClientException.class, () -> GRPCClient.resolveEnvironment(List.of(envA, envB))); - assertTrue(conflict.getMessage().contains("environmentId")); - assertTrue(conflict.getMessage().contains("'env-a'")); - assertTrue(conflict.getMessage().contains("'env-b'")); + assertEquals("same", GRPCClient.resolveQueryName(List.of(v1, v2))); } @Test @@ -441,6 +482,30 @@ public void testDecodeMultiResponseExecutionFailureIsGlobalWithEmptySlot() throw assertEquals(0, this.allocator.getAllocatedMemory()); } + /** + * Unmarshalling the empty slot left by a query that failed to execute must explain + * itself rather than surfacing an opaque NPE from deep inside the Arrow reader. + */ + @Test + public void testUnmarshalEmptySlotExplainsItself() throws Exception { + var response = OnlineQueryMultiResponse.newBuilder() + .addResponses(wrap(OnlineQueryBulkResponse.newBuilder().build())) + .addErrors(error("Encountered an error occurred while executing query 1/1: rate limited")) + .build(); + + try (OnlineQueryMultiResult multi = + GRPCClient.decodeMultiResponse(response, 1, TRACE_ID, this.allocator)) { + OnlineQueryResult empty = multi.getResults().get(0); + assertNull(empty.getScalarsTable()); + + var thrown = assertThrows( + ClientException.class, () -> empty.unmarshal(TestGrpcMultiQueryUser.class)); + assertTrue(thrown.getMessage().contains("no scalar data")); + assertTrue(thrown.getMessage().contains("getGlobalErrors()")); + } + assertEquals(0, this.allocator.getAllocatedMemory()); + } + @Test public void testDecodeMultiResponseCountMismatch() throws Exception { var tooFew = OnlineQueryMultiResponse.newBuilder() diff --git a/src/test/java/ai/chalk/client/TestGrpcMultiQueryUser.java b/src/test/java/ai/chalk/client/TestGrpcMultiQueryUser.java new file mode 100644 index 00000000..8a1dcc13 --- /dev/null +++ b/src/test/java/ai/chalk/client/TestGrpcMultiQueryUser.java @@ -0,0 +1,10 @@ +package ai.chalk.client; + +import ai.chalk.features.Feature; +import ai.chalk.features.FeaturesClass; + +/** Minimal feature class used only to exercise the unmarshal guard. */ +public class TestGrpcMultiQueryUser extends FeaturesClass { + public Feature id; + public Feature socure_score; +}