Added support for streaming initial results of subscription queries

src/main/java/io/axoniq/axonserver/connector/query/QueryChannel.java

@@ -52,16 +52,33 @@ public interface QueryChannel {
      * an initial query result followed by a stream of updates.
      *
      * @param query      the subscription {@link QueryRequest} to send
-     * @param updateType the type of updates expected from this subscription query
      * @param bufferSize the number of updates to be buffered by the update result stream
      * @param fetchSize  the number of updates to be consumed from the update stream prior to refilling it
      * @return the {@link SubscriptionQueryResult} containing the initial result and update stream
      */
     SubscriptionQueryResult subscriptionQuery(QueryRequest query,
-                                              SerializedObject updateType,
                                               int bufferSize,
                                               int fetchSize);
 
+    /**
+     * Sends out a subscription {@code query} to AxonServer for routing to the appropriate handler. Allows for receiving
+     * an initial query result followed by a stream of updates.
+     *
+     * @param query      the subscription {@link QueryRequest} to send
+     * @param updateType the type of updates expected from this subscription query
+     * @param bufferSize the number of updates to be buffered by the update result stream
+     * @param fetchSize  the number of updates to be consumed from the update stream prior to refilling it
+     * @return the {@link SubscriptionQueryResult} containing the initial result and update stream
+     * @deprecated in favor of {@link #subscriptionQuery(QueryRequest, int, int)}. UpdateType is not used.
+     */
+    @Deprecated(since = "2025.2.0")
+    default SubscriptionQueryResult subscriptionQuery(QueryRequest query,
+                                                      @SuppressWarnings("unused") SerializedObject updateType,
+                                                      int bufferSize,
+                                                      int fetchSize) {
+        return subscriptionQuery(query, bufferSize, fetchSize);
+    }
+
     /**
      * Prepares this {@link QueryChannel} to disconnect, by unsubscribing all registered query handlers. Will wait with
      * a certain cut off until all acknowledgments of unsubscribing have been received.

src/main/java/io/axoniq/axonserver/connector/query/SubscriptionQueryResult.java

@@ -35,11 +35,26 @@ public interface SubscriptionQueryResult {
      * <p>
      * Invoking this method will send a request for the initial result, if that hasn't been requested before. Subsequent
      * invocations of this method will return the same CompletableFuture instance.
+     * <p>
+     *     Note that calling this method will not prevent the {@link #initialResults()} call from retrieving a new set
+     *     of initial results, or vice versa.
      *
      * @return a CompletableFuture that completes with the initial result of the subscription query
+     * @deprecated in favor of {@link #initialResults()}, which returns a stream of results.
      */
+    @Deprecated(since = "2025.2.0", forRemoval = true)
     CompletableFuture<QueryResponse> initialResult();
 
+    /**
+     * Returns a ResultStream that represent the initial result of the subscription query.
+     * <p>
+     * Invoking this method will send a request for the initial result if that hasn't been requested before. Further
+     * invocations of this method will return the same ResultStream instance.
+     *
+     * @return a ResultStream that provides the initial result of the subscription query
+     */
+    ResultStream<QueryResponse> initialResults();
+
     /**
      * Returns the stream of updates to the queried projection. The stream can be read in a blocking and non-blocking
      * fashion, as desired by downstream processes.

src/main/java/io/axoniq/axonserver/connector/query/impl/QueryChannelImpl.java

@@ -43,7 +43,6 @@
 import io.axoniq.axonserver.grpc.InstructionAck;
 import io.axoniq.axonserver.grpc.ProcessingInstruction;
 import io.axoniq.axonserver.grpc.ProcessingKey;
-import io.axoniq.axonserver.grpc.SerializedObject;
 import io.axoniq.axonserver.grpc.control.ClientIdentification;
 import io.axoniq.axonserver.grpc.query.QueryComplete;
 import io.axoniq.axonserver.grpc.query.QueryProviderInbound;
@@ -367,7 +366,7 @@ public ResultStream<QueryResponse> query(QueryRequest query) {
         if (query.getMessageIdentifier().isEmpty()) {
             query = query.toBuilder().setMessageIdentifier(UUID.randomUUID().toString()).build();
         }
-        AbstractBufferedStream<QueryResponse, QueryRequest> results = new AbstractBufferedStream<QueryResponse, QueryRequest>(
+        AbstractBufferedStream<QueryResponse, QueryRequest> results = new AbstractBufferedStream<>(
                 clientIdentification.getClientId(), 32, 8
         ) {
 
@@ -386,7 +385,7 @@ public void close() {
                 outboundStream().cancel("Client cancelled the stream.", null);
             }
         };
-        if ("".equals(query.getMessageIdentifier())) {
+        if (query.getMessageIdentifier().isEmpty()) {
             logger.debug("No message identifier has been set on the query. Adding a random identifier now.");
             query = query.toBuilder().setMessageIdentifier(UUID.randomUUID().toString()).build();
         }
@@ -396,7 +395,6 @@ public void close() {
 
     @Override
     public SubscriptionQueryResult subscriptionQuery(QueryRequest query,
-                                                     SerializedObject updateResponseType,
                                                      int bufferSize,
                                                      int fetchSize) {
         logger.trace("Sending subscription query over QueryChannel for context '{}'.", context);
@@ -409,15 +407,16 @@ public SubscriptionQueryResult subscriptionQuery(QueryRequest query,
         String subscriptionId = finalQuery.getMessageIdentifier();
         CompletableFuture<QueryResponse> initialResultFuture = new CompletableFuture<>();
         SubscriptionQueryStream subscriptionStream = new SubscriptionQueryStream(
-                subscriptionId, initialResultFuture, QueryChannelImpl.this.clientIdentification.getClientId(),
-                bufferSize, fetchSize
+                subscriptionId, initialResultFuture,
+                clientIdentification.getClientId(), () -> QueryChannelImpl.this.query(finalQuery),
+                bufferSize,
+                fetchSize
         );
         StreamObserver<SubscriptionQueryRequest> upstream = queryServiceStub.subscription(subscriptionStream);
         subscriptionStream.enableFlowControl();
         SubscriptionQuery subscriptionQuery = SubscriptionQuery.newBuilder()
                                                                .setQueryRequest(finalQuery)
                                                                .setSubscriptionIdentifier(subscriptionId)
-                                                               .setUpdateResponseType(updateResponseType)
                                                                .build();
         upstream.onNext(SubscriptionQueryRequest.newBuilder().setSubscribe(subscriptionQuery).build());
         return new SubscriptionQueryResult() {
@@ -438,9 +437,14 @@ public CompletableFuture<QueryResponse> initialResult() {
                 return initialResultFuture;
             }
 
+            @Override
+            public ResultStream<QueryResponse> initialResults() {
+                return subscriptionStream.initialStream();
+            }
+
             @Override
             public ResultStream<QueryUpdate> updates() {
-                return subscriptionStream.buffer();
+                return subscriptionStream.updatesBuffer();
             }
         };
     }
@@ -588,7 +592,7 @@ private boolean booleanProcessingValue(List<ProcessingInstruction> instructions,
     }
 
     private void handleQuery(QueryProviderInbound inbound, ReplyChannel<QueryProviderOutbound> result) {
-        doHandleQuery(inbound, new ReplyChannel<QueryResponse>() {
+        doHandleQuery(inbound, new ReplyChannel<>() {
             @Override
             public void send(QueryResponse response) {
                 result.send(QueryProviderOutbound.newBuilder().setQueryResponse(response).build());
@@ -615,15 +619,14 @@ public void completeWithError(ErrorMessage errorMessage) {
             public void completeWithError(ErrorCategory errorCategory, String message) {
                 result.completeWithError(errorCategory, message);
             }
-
         });
     }
 
     private void getInitialResult(QueryProviderInbound query, ReplyChannel<QueryProviderOutbound> result) {
         String subscriptionId = query.getSubscriptionQueryRequest().getGetInitialResult().getSubscriptionIdentifier();
         doHandleQuery(
                 query.getSubscriptionQueryRequest().getGetInitialResult().getQueryRequest(),
-                new ReplyChannel<QueryResponse>() {
+                new ReplyChannel<>() {
 
                     @Override
                     public void send(QueryResponse response) {

src/main/java/io/axoniq/axonserver/connector/query/impl/SubscriptionQueryStream.java

@@ -33,6 +33,10 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.concurrent.CompletableFuture;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 /**
  * A {@link FlowControlledStream} implementation to return the results of a {@link SubscriptionQuery}.
@@ -45,23 +49,29 @@ public class SubscriptionQueryStream extends FlowControlledStream<SubscriptionQu
     private final CompletableFuture<QueryResponse> initialResultFuture;
     private final AbstractBufferedStream<QueryUpdate, SubscriptionQueryRequest> updateBuffer;
     private final String clientId;
+    private final Supplier<ResultStream<QueryResponse>> initialResultSupplier;
+
+    private final AtomicReference<ResultStream<QueryResponse>> initialResult = new AtomicReference<>();
 
     /**
      * Instantiates a {@link SubscriptionQueryStream} to stream {@link SubscriptionQuery} results.
      *
-     * @param subscriptionQueryId the identifier of this subscription query
-     * @param initialResultFuture the initial result of the subscription query
-     * @param clientId            the identifier of the client initiating the subscription query
-     * @param bufferSize          the size of the update {@link #buffer()}
-     * @param fetchSize           the number of updates to be consumed prior to refilling the update {@link #buffer()}
+     * @param subscriptionQueryId   the identifier of this subscription query
+     * @param clientId              the identifier of the client initiating the subscription query
+     * @param initialResultSupplier a {@link Supplier} providing the initial {@link ResultStream} of the subscription query
+     * @param bufferSize            the size of the update {@link #updatesBuffer()}
+     * @param fetchSize             the number of updates to be consumed prior to refilling the update
+     *                              {@link #updatesBuffer()}
      */
     public SubscriptionQueryStream(String subscriptionQueryId,
                                    CompletableFuture<QueryResponse> initialResultFuture,
                                    String clientId,
+                                   Supplier<ResultStream<QueryResponse>> initialResultSupplier,
                                    int bufferSize,
                                    int fetchSize) {
         super(clientId, bufferSize, fetchSize);
         this.clientId = clientId;
+        this.initialResultSupplier = initialResultSupplier;
         this.subscriptionQueryId = subscriptionQueryId;
         this.initialResultFuture = initialResultFuture;
         this.updateBuffer = new SubscriptionQueryUpdateBuffer(clientId, subscriptionQueryId, bufferSize, fetchSize);
@@ -72,10 +82,78 @@ public SubscriptionQueryStream(String subscriptionQueryId,
      *
      * @return the {@link ResultStream} buffering the {@link QueryUpdate}s to this subscription query
      */
-    public ResultStream<QueryUpdate> buffer() {
+    public ResultStream<QueryUpdate> updatesBuffer() {
         return updateBuffer;
     }
 
+    /**
+     *
+     * Returns the {@link ResultStream} buffering the initial results to this subscription query.
+     *
+     * @return the {@link ResultStream} buffering the initial results to this subscription query
+     */
+    public ResultStream<QueryResponse> initialStream() {
+        return initialResult.updateAndGet(i -> i == null ? buildInitialResultStream() : i);
+    }
+
+    private ResultStream<QueryResponse> buildInitialResultStream() {
+        ResultStream<QueryResponse> delegate = initialResultSupplier.get();
+        delegate.onAvailable(() -> {
+            if (delegate.getError().isPresent()) {
+                // close the update stream, since initial results completed with an error
+                updatesBuffer().close();
+            }
+        });
+        return new ResultStream<>() {
+            @Override
+            public QueryResponse peek() {
+                return delegate.peek();
+            }
+
+            @Override
+            public QueryResponse nextIfAvailable() {
+                return delegate.nextIfAvailable();
+            }
+
+            @Override
+            public QueryResponse nextIfAvailable(long timeout, TimeUnit unit) throws InterruptedException {
+                return delegate.nextIfAvailable(timeout, unit);
+            }
+
+            @Override
+            public QueryResponse next() throws InterruptedException {
+                return delegate.next();
+            }
+
+            @Override
+            public void onAvailable(Runnable callback) {
+                delegate.onAvailable(() -> {
+                    if (delegate.getError().isPresent()) {
+                        // close the update stream, since initial results completed with an error
+                        updatesBuffer().close();
+                    }
+                    callback.run();
+                });
+            }
+
+            @Override
+            public void close() {
+                delegate.close();
+                SubscriptionQueryStream.this.updatesBuffer().close();
+            }
+
+            @Override
+            public boolean isClosed() {
+                return delegate.isClosed();
+            }
+
+            @Override
+            public Optional<Throwable> getError() {
+                return delegate.getError();
+            }
+        };
+    }
+
     @Override
     public void onNext(SubscriptionQueryResponse value) {
         switch (value.getResponseCase()) {
@@ -111,14 +189,19 @@ public void onNext(SubscriptionQueryResponse value) {
                 break;
             default:
                 logger.info("Received unsupported message from SubscriptionQuery. "
-                                    + "It doesn't declare one of the expected types");
+                                    + "It doesn't declare one of the expected types: {}",
+                            value.getResponseCase());
                 break;
         }
     }
 
     @Override
     public void onError(Throwable t) {
         initialResultFuture.completeExceptionally(t);
+        ResultStream<QueryResponse> initialResult = this.initialResult.get();
+        if (initialResult != null && !initialResult.isClosed()) {
+            initialResult.close();
+        }
         updateBuffer.onError(t);
 
         try {

src/test/java/io/axoniq/axonserver/connector/demo/ConnectorRunner_SubscriptionQuery.java

@@ -60,8 +60,18 @@ public static void main(String[] args) throws Exception {
                                                                                        .addProcessingInstructions(ProcessingInstruction.newBuilder().setKey(ProcessingKey.NR_OF_RESULTS).setValue(MetaDataValue.newBuilder().setNumberValue(20).build()).build())
                                                                                        .setQuery("java.lang.String")
                                                                                        .setPayload(SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Hello world")).build()).build(),
-                                                                           SerializedObject.newBuilder().setType("java.lang.String").setData(ByteString.copyFromUtf8("Hello")).build(), 1000, 100);
-                result.initialResult().whenComplete((r, e) -> System.out.println("Initial result: " + r.getPayload().getData().toStringUtf8()));
+                                                                           1000, 100);
+                ResultStream<QueryResponse> initialResult = result.initialResults();
+                initialResult.onAvailable(() -> {
+                    QueryResponse r = initialResult.nextIfAvailable();
+                    if (r != null) {
+                        System.out.println("Initial result: " + r.getPayload().getData().toStringUtf8());
+                    } else if (initialResult.getError().isPresent()) {
+                        System.out.println("Initial result error: " + initialResult.getError().get().getMessage());
+                    } else if (initialResult.isClosed()) {
+                        System.out.println("Initial result stream closed");
+                    }
+                });
                 ResultStream<QueryUpdate> updates = result.updates();
                 long deadline = System.currentTimeMillis() + 15000;
                 while (!updates.isClosed()) {
@@ -73,7 +83,6 @@ public static void main(String[] args) throws Exception {
                         updates.close();
                     }
                 }
-                result.initialResult().get(5, TimeUnit.SECONDS);
             } finally {
                 contextConnection.disconnect();
                 testSubject.shutdown();
@@ -95,6 +104,7 @@ public static void main(String[] args) {
                 @Override
                 public void handle(QueryRequest q, ReplyChannel<QueryResponse> r) {
                     System.out.println("Handled query");
+                    r.send(QueryResponse.newBuilder().setRequestIdentifier(q.getMessageIdentifier()).setPayload(q.getPayload()).build());
                     r.sendLast(QueryResponse.newBuilder().setRequestIdentifier(q.getMessageIdentifier()).setPayload(q.getPayload()).build());
                 }