Adds support for Streaming Query

pom.xml

@@ -51,7 +51,7 @@
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
 
-        <axon.version>4.5.9</axon.version>
+        <axon.version>4.6.0-SNAPSHOT</axon.version>
 
         <spring.boot.version>2.6.7</spring.boot.version>
 

tracing/src/main/java/org/axonframework/extensions/tracing/TracingQueryGateway.java

@@ -26,13 +26,16 @@
 import org.axonframework.messaging.responsetypes.ResponseType;
 import org.axonframework.queryhandling.DefaultQueryGateway;
 import org.axonframework.queryhandling.GenericQueryMessage;
+import org.axonframework.queryhandling.GenericStreamingQueryMessage;
 import org.axonframework.queryhandling.GenericSubscriptionQueryMessage;
 import org.axonframework.queryhandling.QueryBus;
 import org.axonframework.queryhandling.QueryGateway;
 import org.axonframework.queryhandling.QueryMessage;
 import org.axonframework.queryhandling.SubscriptionQueryBackpressure;
 import org.axonframework.queryhandling.SubscriptionQueryMessage;
 import org.axonframework.queryhandling.SubscriptionQueryResult;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
@@ -41,6 +44,7 @@
 import static org.axonframework.common.BuilderUtils.assertNonNull;
 import static org.axonframework.common.ObjectUtils.nullSafeTypeOf;
 import static org.axonframework.messaging.GenericMessage.asMessage;
+import static org.axonframework.queryhandling.QueryMessage.queryName;
 
 /**
  * A tracing {@link QueryGateway} which activates a calling {@link Span}, when the {@link CompletableFuture} completes.
@@ -120,6 +124,29 @@ public <R, Q> Stream<R> scatterGather(String queryName,
         );
     }
 
+    @Override
+    public <R, Q> Publisher<R> streamingQuery(Q query, Class<R> responseType) {
+        return streamingQuery(queryName(query), query, responseType);
+    }
+
+    @Override
+    public <R, Q> Publisher<R> streamingQuery(String queryName, Q query, Class<R> responseType) {
+        GenericStreamingQueryMessage<Q, R> queryMessagesMessage = new GenericStreamingQueryMessage<>(query,
+                                                                                                               queryName,
+                                                                                                               responseType);
+        return getWithSpan(
+                "streamingQuery_" + SpanUtils.messageName(nullSafeTypeOf(query), queryName),
+                queryMessagesMessage,
+                (childSpan) -> Flux.from(delegate.streamingQuery(queryName, queryMessagesMessage, responseType))
+                        .doOnSubscribe(unused ->  childSpan.log("subscriptionStarted"))
+                        .doOnNext(unused -> childSpan.log("answerReceived"))
+                        .doFinally(unused -> {
+                            childSpan.log("subscriptionTerminated");
+                            childSpan.finish();
+                        })
+        );
+    }
+
     @Override
     public <Q, I, U> SubscriptionQueryResult<I, U> subscriptionQuery(String queryName,
                                                                      Q query,

tracing/src/test/java/org/axonframework/extensions/tracing/TracingQueryGatewayTest.java

@@ -31,6 +31,7 @@
 import org.axonframework.queryhandling.QueryResponseMessage;
 import org.axonframework.queryhandling.SubscriptionQueryResult;
 import org.junit.jupiter.api.*;
+import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
@@ -272,6 +273,41 @@ private <I, U> SubscriptionQueryResult createSubscriptionQueryResult(I initial,
                 () -> true);
     }
 
+    @Test
+    void testStreamingQuery() {
+
+        Publisher<QueryResponseMessage<Object>> updates = Flux.just(new GenericQueryResponseMessage<>("answer1"),
+                                                                    new GenericQueryResponseMessage<>("answer2"));
+        when(mockQueryBus.streamingQuery(any()))
+                .thenReturn(updates);
+
+        MockSpan span = mockTracer.buildSpan("testStreamingQueryResults").start();
+        ScopeManager scopeManager = mockTracer.scopeManager();
+        try (final Scope ignored = scopeManager.activate(span)) {
+            Publisher<String> stringPublisher = testSubject.streamingQuery(new MyQuery(),
+                                                                           String.class);
+            // check the following results are there
+            StepVerifier.create(Flux.from(stringPublisher))
+                        .expectNext("answer1")
+                        .expectNext("answer2")
+                        .expectComplete()
+                        .verify();
+
+            // Verify the parent span is restored, and that a child span was finished.
+            Span activeSpan = mockTracer.activeSpan();
+            assertEquals(span, activeSpan);
+
+            List<MockSpan> mockSpans = mockTracer.finishedSpans();
+            assertEquals(1, mockSpans.size());
+            assertEquals("streamingQuery_MyQuery", mockSpans.get(0).operationName());
+            assertNotNull(mockSpans.get(0).logEntries());
+            assertFalse(mockSpans.get(0).logEntries().isEmpty());
+            assertNotNull(mockSpans.get(0).tags());
+            assertFalse(mockSpans.get(0).tags().isEmpty());
+        }
+        assertNull(scopeManager.activeSpan(), "There should be no activeSpan");
+    }
+
     private static class MyQuery {
 
     }