From af877edffce1a2968f087c238109368c130ed93b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jesse=20Tu=C4=9Flu?= Date: Thu, 2 Jul 2026 17:01:32 -0700 Subject: [PATCH] fix: fix query context precedence layer --- .../apache/druid/server/QueryLifecycle.java | 55 +++++++++++++++---- .../PerSegmentTimeoutInjectionTest.java | 42 ++++++++++++++ .../sql/calcite/run/NativeQueryMaker.java | 5 +- .../query-context-completions.ts | 4 ++ 4 files changed, 95 insertions(+), 11 deletions(-) diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java index 9a514b85e2d5..ed2a6a89b903 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -166,7 +166,21 @@ public QueryResponse runSimple( final AuthorizationResult authorizationResult ) { - initialize(query); + return runSimple(query, authenticationResult, authorizationResult, null); + } + + /** + * As {@link #runSimple(Query, AuthenticationResult, AuthorizationResult)}, but takes the user-provided context keys. + * See {@link #initialize(Query, Set)}. + */ + public QueryResponse runSimple( + final Query query, + final AuthenticationResult authenticationResult, + final AuthorizationResult authorizationResult, + @Nullable final Set userProvidedContextKeys + ) + { + initialize(query, userProvidedContextKeys); final Sequence results; @@ -212,19 +226,34 @@ public void after(final boolean isDone, final Throwable thrown) * @throws DruidException if the current state is not NEW, which indicates a bug */ public void initialize(final Query baseQuery) + { + initialize(baseQuery, null); + } + + /** + * As {@link #initialize(Query)}, but takes the keys the caller set in the context. Pass {@code null} to treat the + * whole context as caller-set (native queries). The SQL layer merges static defaults into the context, so it must + * pass the real user-provided keys. + * + * @throws DruidException if the current state is not NEW, which indicates a bug + */ + public void initialize(final Query baseQuery, @Nullable final Set userProvidedContextKeys) { transition(State.NEW, State.INITIALIZED); userContextKeys = new HashSet<>(baseQuery.getContext().keySet()); + final Set effectiveUserKeys = + userProvidedContextKeys != null ? userProvidedContextKeys : baseQuery.getContext().keySet(); + String queryId = baseQuery.getId(); if (Strings.isNullOrEmpty(queryId)) { queryId = UUID.randomUUID().toString(); } - // Start with system defaults, apply per-datasource override, then user context wins + // Precedence, high to low: user context > per-datasource per-segment timeout > static defaults > code default. Map contextWithDefaults = new HashMap<>(queryConfigProvider.getContext()); - applyPerDatasourcePerSegmentTimeout(baseQuery, contextWithDefaults, queryId); Map finalContext = QueryContexts.override(contextWithDefaults, baseQuery.getContext()); + applyPerDatasourcePerSegmentTimeout(baseQuery, finalContext, effectiveUserKeys, queryId); finalContext.put(BaseQuery.QUERY_ID, queryId); this.baseQuery = baseQuery.withOverriddenContext(finalContext); @@ -232,16 +261,17 @@ public void initialize(final Query baseQuery) } /** - * If a per-datasource per-segment timeout is configured, injects it into the context defaults. - * User context (applied later via {@link QueryContexts#override}) will override this if set explicitly. - * In monitorOnly mode, logs the configured timeout but does not inject it. + * If a per-datasource per-segment timeout is configured, injects it into {@code finalContext}, overriding a value + * merged in from static defaults. A perSegmentTimeout the caller set (per {@code userProvidedContextKeys}) is left + * untouched. In monitorOnly mode, logs the configured timeout but does not inject it. * - * For queries involving multiple datasources (e.g., joins or unions), the timeout from the first matching datasource is applied - * since getTableNames() returns a Set, the match order is non-deterministic. + * For multi-datasource queries (joins/unions), the first matching datasource wins; getTableNames() returns a Set, so + * the match order is non-deterministic. */ private void applyPerDatasourcePerSegmentTimeout( final Query query, - final Map contextWithDefaults, + final Map finalContext, + final Set userProvidedContextKeys, final String queryId ) { @@ -249,6 +279,11 @@ private void applyPerDatasourcePerSegmentTimeout( return; } + // Caller set perSegmentTimeout: respect it. + if (userProvidedContextKeys.contains(QueryContexts.PER_SEGMENT_TIMEOUT_KEY)) { + return; + } + for (String tableName : query.getDataSource().getTableNames()) { PerSegmentTimeoutConfig dsConfig = perSegmentTimeoutConfig.get(tableName); if (dsConfig != null) { @@ -260,7 +295,7 @@ private void applyPerDatasourcePerSegmentTimeout( queryId ); } else { - contextWithDefaults.put(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, dsConfig.getPerSegmentTimeoutMs()); + finalContext.put(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, dsConfig.getPerSegmentTimeoutMs()); } return; } diff --git a/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java b/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java index 440f065ba485..d99cfff2a514 100644 --- a/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java +++ b/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java @@ -44,6 +44,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Set; public class PerSegmentTimeoutInjectionTest { @@ -195,6 +196,47 @@ DATASOURCE, new PerSegmentTimeoutConfig(5000, false) Assert.assertEquals(2000L, lifecycle.getQuery().context().getPerSegmentTimeout()); } + @Test + public void testPerDatasourceTimeout_staticDefaultInContextNotUserProvided_dynamicWins() + { + // SQL path: a static default is merged into the context but the caller did not set perSegmentTimeout, so the + // dynamic config overrides it. + Map config = Map.of( + DATASOURCE, new PerSegmentTimeoutConfig(5000, false) + ); + + expectDefaults(); + + TimeseriesQuery queryWithMergedDefault = baseQuery.withOverriddenContext( + Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, "0") + ); + + QueryLifecycle lifecycle = createLifecycle(config); + lifecycle.initialize(queryWithMergedDefault, Collections.emptySet()); + + Assert.assertEquals(5000L, lifecycle.getQuery().context().getPerSegmentTimeout()); + } + + @Test + public void testPerDatasourceTimeout_userProvidedViaProvenance_userWins() + { + // Same context value, but the caller set perSegmentTimeout, so it wins over the dynamic config. + Map config = Map.of( + DATASOURCE, new PerSegmentTimeoutConfig(5000, false) + ); + + expectDefaults(); + + TimeseriesQuery queryWithUserTimeout = baseQuery.withOverriddenContext( + Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, 2000L) + ); + + QueryLifecycle lifecycle = createLifecycle(config); + lifecycle.initialize(queryWithUserTimeout, Set.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY)); + + Assert.assertEquals(2000L, lifecycle.getQuery().context().getPerSegmentTimeout()); + } + private void expectDefaults() { EasyMock.expect(queryConfig.getContext()).andReturn(Map.of()).anyTimes(); diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeQueryMaker.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeQueryMaker.java index 8020b50f6c2e..aa0cca9832a0 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeQueryMaker.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeQueryMaker.java @@ -186,7 +186,10 @@ private QueryResponse execute( final QueryResponse results = queryLifecycle.runSimple( (Query) query, authenticationResult, - authorizationResult + authorizationResult, + // SQL merges static defaults into the context; pass the user-set keys so dynamic config can override a + // merged-in default but not a value the caller set. + plannerContext.authContextKeys() ); return mapResultSequence( diff --git a/web-console/src/dialogs/edit-context-dialog/query-context-completions.ts b/web-console/src/dialogs/edit-context-dialog/query-context-completions.ts index d7e67b7f492d..68392e6a6dd8 100644 --- a/web-console/src/dialogs/edit-context-dialog/query-context-completions.ts +++ b/web-console/src/dialogs/edit-context-dialog/query-context-completions.ts @@ -25,6 +25,10 @@ export const QUERY_CONTEXT_COMPLETIONS: JsonCompletionRule[] = [ isObject: true, completions: [ { value: 'timeout', documentation: 'Query timeout in milliseconds' }, + { + value: 'perSegmentTimeout', + documentation: 'Per-segment processing timeout in milliseconds', + }, { value: 'priority', documentation: 'Query priority (higher = more important)',