(pinot) branch master updated: feat: expose BrokerResponse to subclasses via overridable onQueryCompletion hook (#18247)

[xi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new 3eff094baa2 feat: expose BrokerResponse to subclasses via overridable onQueryCompletion hook (#18247)
3eff094baa2 is described below

commit 3eff094baa2164a2592462a81fce1ee7738d28be
Author: st-omarkhalid <om...@startree.ai>
AuthorDate: Sat Apr 18 21:19:59 2026 -0500

    feat: expose BrokerResponse to subclasses via overridable onQueryCompletion hook (#18247)
    
    * feat: expose BrokerResponse to subclasses via overridable onQueryCompletion hook
    
    Add a protected `onQueryCompletion(RequestContext, BrokerResponse)` method to
    `BaseBrokerRequestHandler` so subclasses can intercept the fully-populated response
    (serverStats, partialResult, pools, rlsFiltersApplied, etc.) without changing the
    `BrokerQueryEventListener` SPI.
    
    Add a protected `createSingleStageBrokerRequestHandler(...)` factory method to
    `BaseBrokerStarter` so subclasses (e.g. StarTreePinotBrokerStarter) can supply a
    custom `SingleConnectionBrokerRequestHandler` subclass that overrides the hook.
    
    The default behaviour is unchanged: the existing SPI listener is still invoked.
    
    Co-Authored-By: Claude Sonnet 4.6 <no...@anthropic.com>
    
    * test: add unit test for onQueryCompletion hook in BaseBrokerRequestHandler
    
    Verifies that the protected onQueryCompletion hook is invoked with a
    non-null BrokerResponse after every query, covering the new hook added
    alongside the factory method in BaseBrokerStarter.
    
    Co-Authored-By: Claude Sonnet 4.6 <no...@anthropic.com>
    
    * fix: sort BrokerResponse import alphabetically to satisfy Spotless
    
    Co-Authored-By: Claude Sonnet 4.6 <no...@anthropic.com>
    
    ---------
    
    Co-authored-by: Claude Sonnet 4.6 <no...@anthropic.com>
---
 .../broker/broker/helix/BaseBrokerStarter.java     | 20 ++++++++-
 .../requesthandler/BaseBrokerRequestHandler.java   | 13 +++++-
 .../BaseSingleStageBrokerRequestHandlerTest.java   | 50 ++++++++++++++++++++++
 3 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
index 790b632019b..1a30dea0ea7 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
@@ -47,6 +47,7 @@ import org.apache.pinot.broker.broker.AccessControlFactory;
 import org.apache.pinot.broker.broker.BrokerAdminApiApplication;
 import org.apache.pinot.broker.grpc.BrokerGrpcServer;
 import org.apache.pinot.broker.queryquota.HelixExternalViewBasedQueryQuotaManager;
+import org.apache.pinot.broker.queryquota.QueryQuotaManager;
 import org.apache.pinot.broker.requesthandler.BaseSingleStageBrokerRequestHandler;
 import org.apache.pinot.broker.requesthandler.BrokerRequestHandler;
 import org.apache.pinot.broker.requesthandler.BrokerRequestHandlerDelegate;
@@ -249,6 +250,23 @@ public abstract class BaseBrokerStarter implements ServiceStartable {
     return new WorkerManager(brokerId, hostname, port, routingManager);
   }
 
+  /**
+   * Override to supply a custom {@link SingleConnectionBrokerRequestHandler} subclass (e.g. one
+   * that overrides {@code onQueryCompletion(RequestContext, BrokerResponse)} for async query
+   * logging). The default implementation returns a plain {@link SingleConnectionBrokerRequestHandler}.
+   */
+  protected SingleConnectionBrokerRequestHandler createSingleStageBrokerRequestHandler(
+      PinotConfiguration config, String brokerId, BrokerRequestIdGenerator requestIdGenerator,
+      RoutingManager routingManager, AccessControlFactory accessControlFactory,
+      QueryQuotaManager queryQuotaManager, TableCache tableCache, NettyConfig nettyConfig,
+      TlsConfig tlsConfig, ServerRoutingStatsManager serverRoutingStatsManager,
+      FailureDetector failureDetector, ThreadAccountant threadAccountant,
+      MultiClusterRoutingContext multiClusterRoutingContext) {
+    return new SingleConnectionBrokerRequestHandler(config, brokerId, requestIdGenerator, routingManager,
+        accessControlFactory, queryQuotaManager, tableCache, nettyConfig, tlsConfig,
+        serverRoutingStatsManager, failureDetector, threadAccountant, multiClusterRoutingContext);
+  }
+
   private void setupHelixSystemProperties() {
     // NOTE: Helix will disconnect the manager and disable the instance if it detects flapping (too frequent disconnect
     // from ZooKeeper). Setting flapping time window to a small value can avoid this from happening. Helix ignores the
@@ -440,7 +458,7 @@ public abstract class BaseBrokerStarter implements ServiceStartable {
         brokerContext.setServerHttpsContext(sslContext);
       }
       singleStageBrokerRequestHandler =
-          new SingleConnectionBrokerRequestHandler(_brokerConf, brokerId, requestIdGenerator, _routingManager,
+          createSingleStageBrokerRequestHandler(_brokerConf, brokerId, requestIdGenerator, _routingManager,
               _accessControlFactory, _queryQuotaManager, _tableCache, nettyDefaults, tlsDefaults,
               _serverRoutingStatsManager, _failureDetector, _threadAccountant, multiClusterRoutingContext);
     }
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
index 30d09fa2360..1517b698a9d 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
@@ -219,11 +219,22 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
             accessControl);
     brokerResponse.setBrokerId(_brokerId);
     brokerResponse.setRequestId(Long.toString(requestId));
-    _brokerQueryEventListener.onQueryCompletion(requestContext);
+    onQueryCompletion(requestContext, brokerResponse);
 
     return brokerResponse;
   }
 
+  /**
+   * Called after every successfully executed query with the fully-populated {@link RequestContext}
+   * and {@link BrokerResponse}. The default implementation fires the configured
+   * {@link org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListener}. Subclasses may
+   * override to intercept the complete response (e.g. for async query-log pipelines) while still
+   * calling {@code super} to preserve the SPI listener behaviour.
+   */
+  protected void onQueryCompletion(RequestContext requestContext, BrokerResponse brokerResponse) {
+    _brokerQueryEventListener.onQueryCompletion(requestContext);
+  }
+
   protected abstract BrokerResponse handleRequest(long requestId, String query, SqlNodeAndOptions sqlNodeAndOptions,
       JsonNode request, @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext,
       @Nullable HttpHeaders httpHeaders, AccessControl accessControl)
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java
index ca216e1b145..aa31a46f82e 100644
--- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java
@@ -35,6 +35,7 @@ import org.apache.pinot.common.request.BrokerRequest;
 import org.apache.pinot.common.request.Expression;
 import org.apache.pinot.common.request.Function;
 import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.common.response.BrokerResponse;
 import org.apache.pinot.common.response.broker.BrokerResponseNative;
 import org.apache.pinot.core.routing.RoutingTable;
 import org.apache.pinot.core.routing.SegmentsToQuery;
@@ -95,6 +96,55 @@ public class BaseSingleStageBrokerRequestHandlerTest {
     }
   }
 
+  @Test
+  public void testOnQueryCompletionHookReceivesBrokerResponse() {
+    // Verify that the overridable onQueryCompletion(RequestContext, BrokerResponse) hook is invoked
+    // and receives the BrokerResponse that handleRequest() produced.
+    AtomicReference<BrokerResponse> capturedResponse = new AtomicReference<>();
+
+    PinotConfiguration config = new PinotConfiguration();
+    BrokerQueryEventListenerFactory.init(config);
+    BrokerMetrics.register(mock(BrokerMetrics.class));
+    QueryQuotaManager queryQuotaManager = mock(QueryQuotaManager.class);
+    when(queryQuotaManager.acquire(anyString())).thenReturn(true);
+    when(queryQuotaManager.acquireDatabase(anyString())).thenReturn(true);
+    when(queryQuotaManager.acquireApplication(anyString())).thenReturn(true);
+    TableCache tableCache = mock(TableCache.class);
+
+    BaseSingleStageBrokerRequestHandler handler =
+        new BaseSingleStageBrokerRequestHandler(config, "testBrokerId", new BrokerRequestIdGenerator(),
+            mock(org.apache.pinot.core.routing.RoutingManager.class), new AllowAllAccessControlFactory(),
+            queryQuotaManager, tableCache, ThreadAccountantUtils.getNoOpAccountant(), null) {
+          @Override
+          public void start() {
+          }
+
+          @Override
+          public void shutDown() {
+          }
+
+          @Override
+          protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
+              BrokerRequest serverBrokerRequest, TableRouteInfo route, long timeoutMs, ServerStats serverStats,
+              RequestContext requestContext) {
+            return new BrokerResponseNative();
+          }
+
+          @Override
+          protected void onQueryCompletion(RequestContext requestContext, BrokerResponse brokerResponse) {
+            capturedResponse.set(brokerResponse);
+          }
+        };
+
+    try {
+      handler.handleRequest("SELECT 1");
+    } catch (Exception ignored) {
+      // routing may fail — we only care that the hook was called with a non-null response
+    }
+    Assert.assertNotNull(capturedResponse.get(),
+        "onQueryCompletion hook must be called with the BrokerResponse from handleRequest");
+  }
+
   @Test
   public void testGetActualColumnNameCaseSensitive() {
     Map<String, String> columnNameMap = new HashMap<>();


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org