Skip to content

feat: expose BrokerResponse to subclasses via overridable onQueryCompletion hook#18247

Open
st-omarkhalid wants to merge 2 commits intoapache:masterfrom
st-omarkhalid:feature/broker-query-event-listener-response-oss
Open

feat: expose BrokerResponse to subclasses via overridable onQueryCompletion hook#18247
st-omarkhalid wants to merge 2 commits intoapache:masterfrom
st-omarkhalid:feature/broker-query-event-listener-response-oss

Conversation

@st-omarkhalid
Copy link
Copy Markdown

@st-omarkhalid st-omarkhalid commented Apr 17, 2026

Fixes #18248

Summary

This change makes it easier for subclasses to intercept the fully-populated BrokerResponse after a query executes successfully, without requiring changes to the BrokerQueryEventListener SPI.

Changes

BaseBrokerRequestHandler — adds a protected onQueryCompletion(RequestContext, BrokerResponse) hook:

  • The success-path call site now calls this.onQueryCompletion(rc, response) instead of directly invoking _brokerQueryEventListener.onQueryCompletion(rc)
  • The default implementation delegates to the existing SPI listener — no behaviour change for existing code
  • Error-path call sites (ACCESS_DENIED, SQL_PARSING) are intentionally unchanged — no BrokerResponse exists there

BaseBrokerStarter — adds a protected createSingleStageBrokerRequestHandler(...) factory method:

  • start() now calls this method instead of new SingleConnectionBrokerRequestHandler(...) directly
  • The default implementation returns a plain SingleConnectionBrokerRequestHandler — no behaviour change
  • Subclasses can override to inject a custom handler subclass that overrides the onQueryCompletion hook

Motivation

The BrokerQueryEventListener SPI only receives RequestContext, which is missing several fields that are only available on BrokerResponse after execution completes: isPartialResult(), getPools(), getRLSFiltersApplied(), getTablesQueried(), isNumGroupsWarningLimitReached(), etc.

This pattern (protected factory + protected hook) is already used in this codebase for createWorkerManager(...) and follows standard Java extension patterns. It allows subclasses to capture the complete response for use cases like async query logging pipelines without modifying the SPI contract.

Tests

  • BaseSingleStageBrokerRequestHandlerTest (13 tests) — all pass
  • LiteralOnlyBrokerRequestTest (8 tests) — all pass
  • Existing behaviour is unchanged — the default onQueryCompletion still calls the SPI listener

🤖 Generated with Claude Code

st-omarkhalid and others added 2 commits April 17, 2026 16:33
…letion 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 <noreply@anthropic.com>
…ndler

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 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Allow subclasses to intercept BrokerResponse via overridable onQueryCompletion hook

1 participant