diff --git a/processing/src/main/java/org/apache/druid/query/QueryConfigProvider.java b/processing/src/main/java/org/apache/druid/query/QueryConfigProvider.java index 6b533ab6a37c..1559beb07328 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryConfigProvider.java +++ b/processing/src/main/java/org/apache/druid/query/QueryConfigProvider.java @@ -22,7 +22,7 @@ import java.util.Map; /** - * Provides the default query config applied to all incoming queries before per-query overrides are merged in. + * Provides the default query context applied to all incoming queries before per-query overrides are merged in. * *

On non-broker nodes this is backed by static runtime properties ({@link DefaultQueryConfig}). * On brokers, it is backed by {@code BrokerViewOfBrokerConfig}, which merges the static defaults with diff --git a/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java b/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java index b67b4bc46dab..2fbcfb2fbdfe 100644 --- a/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java +++ b/server/src/main/java/org/apache/druid/client/BrokerViewOfBrokerConfig.java @@ -36,6 +36,7 @@ import org.apache.druid.rpc.ServiceLocator; import org.apache.druid.rpc.StandardRetryPolicy; import org.apache.druid.server.broker.BrokerDynamicConfig; +import org.apache.druid.server.broker.QueryConfigSnapshot; import javax.validation.constraints.NotNull; import java.util.Map; @@ -127,4 +128,13 @@ public Map getContext() { return resolvedDefaultQueryContext.asMap(); } + + /** + * Captures the resolved default context and the dynamic config atomically, so a single query resolves its context + * and blocklist against one consistent snapshot rather than re-reading the live config. + */ + public synchronized QueryConfigSnapshot snapshotForQuery() + { + return new QueryConfigSnapshot(getContext(), getDynamicConfig()); + } } 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..e94b30cdd5ee 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycle.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycle.java @@ -37,9 +37,7 @@ import org.apache.druid.query.DruidMetrics; import org.apache.druid.query.GenericQueryMetricsFactory; import org.apache.druid.query.Query; -import org.apache.druid.query.QueryConfigProvider; import org.apache.druid.query.QueryContext; -import org.apache.druid.query.QueryContexts; import org.apache.druid.query.QueryInterruptedException; import org.apache.druid.query.QueryMetrics; import org.apache.druid.query.QueryPlus; @@ -49,7 +47,7 @@ import org.apache.druid.query.QueryToolChest; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.query.policy.PolicyEnforcer; -import org.apache.druid.server.broker.PerSegmentTimeoutConfig; +import org.apache.druid.server.broker.QueryConfigSnapshot; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.Action; import org.apache.druid.server.security.AuthConfig; @@ -64,7 +62,6 @@ import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; -import java.util.HashMap; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; @@ -99,11 +96,9 @@ public class QueryLifecycle private final ServiceEmitter emitter; private final RequestLogger requestLogger; private final AuthorizerMapper authorizerMapper; - private final QueryConfigProvider queryConfigProvider; private final AuthConfig authConfig; private final PolicyEnforcer policyEnforcer; - private final List queryBlocklist; - private final Map perSegmentTimeoutConfig; + private final QueryConfigSnapshot configSnapshot; private final long startMs; private final long startNs; @@ -113,8 +108,11 @@ public class QueryLifecycle @MonotonicNonNull private Query baseQuery; + /** + * Keys present on the query context as received; the candidate set fed to context-key authorization. + */ @MonotonicNonNull - private Set userContextKeys; + private Set authorizationContextKeys; public QueryLifecycle( final QueryRunnerFactoryConglomerate conglomerate, @@ -123,11 +121,9 @@ public QueryLifecycle( final ServiceEmitter emitter, final RequestLogger requestLogger, final AuthorizerMapper authorizerMapper, - final QueryConfigProvider queryConfigProvider, final AuthConfig authConfig, final PolicyEnforcer policyEnforcer, - final List queryBlocklist, - final Map perSegmentTimeoutConfig, + final QueryConfigSnapshot configSnapshot, final long startMs, final long startNs ) @@ -138,11 +134,9 @@ public QueryLifecycle( this.emitter = emitter; this.requestLogger = requestLogger; this.authorizerMapper = authorizerMapper; - this.queryConfigProvider = queryConfigProvider; this.authConfig = authConfig; this.policyEnforcer = policyEnforcer; - this.queryBlocklist = queryBlocklist; - this.perSegmentTimeoutConfig = perSegmentTimeoutConfig; + this.configSnapshot = configSnapshot; this.startMs = startMs; this.startNs = startNs; } @@ -166,7 +160,21 @@ public QueryResponse runSimple( final AuthorizationResult authorizationResult ) { - initialize(query); + return runSimple(query, authenticationResult, authorizationResult, null); + } + + /** + * As {@link #runSimple(Query, AuthenticationResult, AuthorizationResult)}, but takes the context keys the client + * actually set. See {@link #initialize(Query, Set)}. + */ + public QueryResponse runSimple( + final Query query, + final AuthenticationResult authenticationResult, + final AuthorizationResult authorizationResult, + @Nullable final Set clientProvidedQueryContextKeys + ) + { + initialize(query, clientProvidedQueryContextKeys); final Sequence results; @@ -212,61 +220,41 @@ 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 context keys the client actually set. Pass {@code null} to treat the + * whole context as client-set (native queries). The SQL layer merges static defaults into the context, so it must + * pass the real client-set keys so dynamic overrides can beat a merged-in default without overriding the client. + * + * @throws DruidException if the current state is not NEW, which indicates a bug + */ + public void initialize(final Query baseQuery, @Nullable final Set clientProvidedQueryContextKeys) { transition(State.NEW, State.INITIALIZED); - userContextKeys = new HashSet<>(baseQuery.getContext().keySet()); + final Map baseContext = baseQuery.getContext(); + authorizationContextKeys = new HashSet<>(baseContext.keySet()); + + // Keys the client actually set (native queries pass null, so the whole context is client-set). + final Set effectiveClientProvidedQueryContextKeys = + clientProvidedQueryContextKeys != null ? clientProvidedQueryContextKeys : baseContext.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 - Map contextWithDefaults = new HashMap<>(queryConfigProvider.getContext()); - applyPerDatasourcePerSegmentTimeout(baseQuery, contextWithDefaults, queryId); - Map finalContext = QueryContexts.override(contextWithDefaults, baseQuery.getContext()); + final Map finalContext = + configSnapshot.resolveContext(baseQuery, effectiveClientProvidedQueryContextKeys); finalContext.put(BaseQuery.QUERY_ID, queryId); this.baseQuery = baseQuery.withOverriddenContext(finalContext); this.toolChest = conglomerate.getToolChest(this.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. - * - * 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. - */ - private void applyPerDatasourcePerSegmentTimeout( - final Query query, - final Map contextWithDefaults, - final String queryId - ) - { - if (perSegmentTimeoutConfig.isEmpty()) { - return; - } - - for (String tableName : query.getDataSource().getTableNames()) { - PerSegmentTimeoutConfig dsConfig = perSegmentTimeoutConfig.get(tableName); - if (dsConfig != null) { - if (dsConfig.isMonitorOnly()) { - log.debug( - "Per-segment timeout [%d ms] configured for datasource [%s] in monitorOnly mode (not enforced) for query [%s].", - dsConfig.getPerSegmentTimeoutMs(), - tableName, - queryId - ); - } else { - contextWithDefaults.put(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, dsConfig.getPerSegmentTimeoutMs()); - } - return; - } - } - } - /** * Returns {@link AuthorizationResult} based on {@code DRUID_AUTHENTICATION_RESULT} in the given request, base query * would be transformed with restrictions on the AuthorizationResult. @@ -290,7 +278,7 @@ public AuthorizationResult authorize(HttpServletRequest req) AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR ), Iterables.transform( - authConfig.contextKeysToAuthorize(userContextKeys), + authConfig.contextKeysToAuthorize(authorizationContextKeys), contextParam -> new ResourceAction(new Resource(contextParam, ResourceType.QUERY_CONTEXT), Action.WRITE) ) ); @@ -328,7 +316,7 @@ public AuthorizationResult authorize(AuthenticationResult authenticationResult) AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR ), Iterables.transform( - authConfig.contextKeysToAuthorize(userContextKeys), + authConfig.contextKeysToAuthorize(authorizationContextKeys), contextParam -> new ResourceAction(new Resource(contextParam, ResourceType.QUERY_CONTEXT), Action.WRITE) ) ); @@ -364,7 +352,9 @@ private void preAuthorized( */ private void checkQueryBlocklist() { - if (queryBlocklist == null || queryBlocklist.isEmpty()) { + final List queryBlocklist = configSnapshot.getQueryBlocklist(); + + if (queryBlocklist.isEmpty()) { return; } diff --git a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java index 57fcd1b47e91..4b936d822e7d 100644 --- a/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java +++ b/server/src/main/java/org/apache/druid/server/QueryLifecycleFactory.java @@ -28,15 +28,12 @@ import org.apache.druid.query.QueryRunnerFactoryConglomerate; import org.apache.druid.query.QuerySegmentWalker; import org.apache.druid.query.policy.PolicyEnforcer; -import org.apache.druid.server.broker.PerSegmentTimeoutConfig; +import org.apache.druid.server.broker.QueryConfigSnapshot; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.AuthConfig; import org.apache.druid.server.security.AuthorizerMapper; import javax.annotation.Nullable; -import java.util.Collections; -import java.util.List; -import java.util.Map; @LazySingleton public class QueryLifecycleFactory @@ -80,15 +77,11 @@ public QueryLifecycleFactory( public QueryLifecycle factorize() { - final List queryBlocklist; - final Map perSegmentTimeoutConfig; - if (brokerViewOfBrokerConfig != null && brokerViewOfBrokerConfig.getDynamicConfig() != null) { - queryBlocklist = brokerViewOfBrokerConfig.getDynamicConfig().getQueryBlocklist(); - perSegmentTimeoutConfig = brokerViewOfBrokerConfig.getDynamicConfig().getPerSegmentTimeoutConfig(); - } else { - queryBlocklist = Collections.emptyList(); - perSegmentTimeoutConfig = Collections.emptyMap(); - } + // Capture one snapshot per query so the default context, per-query overrides, and blocklist are consistent. + final QueryConfigSnapshot configSnapshot = + brokerViewOfBrokerConfig == null + ? new QueryConfigSnapshot(queryConfigProvider.getContext(), null) + : brokerViewOfBrokerConfig.snapshotForQuery(); return new QueryLifecycle( conglomerate, @@ -97,11 +90,9 @@ public QueryLifecycle factorize() emitter, requestLogger, authorizerMapper, - queryConfigProvider, authConfig, policyEnforcer, - queryBlocklist, - perSegmentTimeoutConfig, + configSnapshot, System.currentTimeMillis(), System.nanoTime() ); diff --git a/server/src/main/java/org/apache/druid/server/broker/BrokerDynamicConfig.java b/server/src/main/java/org/apache/druid/server/broker/BrokerDynamicConfig.java index 57200e729d3a..919f8487a157 100644 --- a/server/src/main/java/org/apache/druid/server/broker/BrokerDynamicConfig.java +++ b/server/src/main/java/org/apache/druid/server/broker/BrokerDynamicConfig.java @@ -22,7 +22,10 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import org.apache.druid.common.config.Configs; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.query.Query; import org.apache.druid.query.QueryContext; +import org.apache.druid.query.QueryContexts; import org.apache.druid.server.QueryBlocklistRule; import javax.annotation.Nullable; @@ -40,6 +43,8 @@ */ public class BrokerDynamicConfig { + private static final Logger log = new Logger(BrokerDynamicConfig.class); + public static final String CONFIG_KEY = "broker.config"; /** @@ -91,6 +96,33 @@ public Map getPerSegmentTimeoutConfig() return perSegmentTimeoutConfig; } + /** + * Query-specific broker dynamic config query context overrides (e.g. per segment timeout). + */ + public QueryContext getQuerySpecificContextOverrides(Query query) + { + if (perSegmentTimeoutConfig.isEmpty()) { + return QueryContext.empty(); + } + + for (String tableName : query.getDataSource().getTableNames()) { + PerSegmentTimeoutConfig dsConfig = perSegmentTimeoutConfig.get(tableName); + if (dsConfig != null) { + if (dsConfig.isMonitorOnly()) { + log.debug( + "Per-segment timeout [%d ms] configured for datasource [%s] in monitorOnly mode (not enforced) for query [%s].", + dsConfig.getPerSegmentTimeoutMs(), + tableName, + query.getId() + ); + return QueryContext.empty(); + } + return QueryContext.of(Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, dsConfig.getPerSegmentTimeoutMs())); + } + } + return QueryContext.empty(); + } + @Override public boolean equals(Object o) { diff --git a/server/src/main/java/org/apache/druid/server/broker/QueryConfigSnapshot.java b/server/src/main/java/org/apache/druid/server/broker/QueryConfigSnapshot.java new file mode 100644 index 000000000000..461437d495c4 --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/broker/QueryConfigSnapshot.java @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.server.broker; + +import org.apache.druid.query.Query; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.server.QueryBlocklistRule; + +import javax.annotation.Nullable; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * A snapshot of the broker's resolved default query context and dynamic config, captured once for a single query so + * that context resolution and the blocklist read one consistent view (rather than re-reading the live config at + * different points in the query lifecycle). The context-resolution logic lives here, keeping it out of + * {@code QueryLifecycle}. On non-broker nodes {@code dynamicConfig} is null (defaults only, no per-query overrides). + */ +public class QueryConfigSnapshot +{ + private final Map defaultContext; + @Nullable + private final BrokerDynamicConfig dynamicConfig; + + public QueryConfigSnapshot(Map defaultContext, @Nullable BrokerDynamicConfig dynamicConfig) + { + this.defaultContext = defaultContext; + this.dynamicConfig = dynamicConfig; + } + + /** + * The final query context. Precedence high to low: a key the client set > per-query dynamic override > the query's + * own context > defaults. Dynamic overrides are applied over the merged context but skipped for keys the client set, + * so a client-provided value always wins. + */ + public Map resolveContext(Query query, Set clientProvidedQueryContextKeys) + { + Map result = QueryContexts.override(defaultContext, query.getContext()); + if (dynamicConfig != null) { + for (Map.Entry override : dynamicConfig.getQuerySpecificContextOverrides(query).asMap().entrySet()) { + if (!clientProvidedQueryContextKeys.contains(override.getKey())) { + result.put(override.getKey(), override.getValue()); + } + } + } + return result; + } + + public List getQueryBlocklist() + { + return dynamicConfig == null ? Collections.emptyList() : dynamicConfig.getQueryBlocklist(); + } +} 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..f95e38568aac 100644 --- a/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java +++ b/server/src/test/java/org/apache/druid/server/PerSegmentTimeoutInjectionTest.java @@ -23,7 +23,6 @@ import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.query.Druids; import org.apache.druid.query.GenericQueryMetricsFactory; -import org.apache.druid.query.QueryConfigProvider; import org.apache.druid.query.QueryContexts; import org.apache.druid.query.QueryRunnerFactoryConglomerate; import org.apache.druid.query.QuerySegmentWalker; @@ -31,7 +30,9 @@ import org.apache.druid.query.aggregation.CountAggregatorFactory; import org.apache.druid.query.policy.NoopPolicyEnforcer; import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.server.broker.BrokerDynamicConfig; import org.apache.druid.server.broker.PerSegmentTimeoutConfig; +import org.apache.druid.server.broker.QueryConfigSnapshot; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.AuthConfig; import org.apache.druid.server.security.AuthorizerMapper; @@ -44,10 +45,17 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Set; +/** + * Precedence of the per-query dynamic overrides at the {@link QueryLifecycle} level: an override beats a + * non-client-provided value in the context, but a value the client set wins. The datasource-matching logic itself is + * tested in {@code org.apache.druid.server.broker.BrokerDynamicConfigTest}. + */ public class PerSegmentTimeoutInjectionTest { private static final String DATASOURCE = "my_datasource"; + private static final String KEY = QueryContexts.PER_SEGMENT_TIMEOUT_KEY; private QueryRunnerFactoryConglomerate conglomerate; private QuerySegmentWalker texasRanger; @@ -55,7 +63,6 @@ public class PerSegmentTimeoutInjectionTest private ServiceEmitter emitter; private RequestLogger requestLogger; private AuthorizerMapper authzMapper; - private QueryConfigProvider queryConfig; private QueryToolChest toolChest; private final TimeseriesQuery baseQuery = Druids.newTimeseriesQueryBuilder() @@ -68,142 +75,73 @@ public class PerSegmentTimeoutInjectionTest public void setUp() { conglomerate = EasyMock.createMock(QueryRunnerFactoryConglomerate.class); - texasRanger = EasyMock.createMock(QuerySegmentWalker.class); - metricsFactory = EasyMock.createMock(GenericQueryMetricsFactory.class); - emitter = EasyMock.createMock(ServiceEmitter.class); + texasRanger = EasyMock.createNiceMock(QuerySegmentWalker.class); + metricsFactory = EasyMock.createNiceMock(GenericQueryMetricsFactory.class); + emitter = EasyMock.createNiceMock(ServiceEmitter.class); requestLogger = EasyMock.createNiceMock(RequestLogger.class); authzMapper = EasyMock.createNiceMock(AuthorizerMapper.class); - queryConfig = EasyMock.createMock(QueryConfigProvider.class); toolChest = EasyMock.createNiceMock(QueryToolChest.class); } @After public void tearDown() { - EasyMock.verify(conglomerate, queryConfig); + EasyMock.verify(conglomerate); } @Test - public void testPerDatasourceTimeout_applied() + public void testDynamicOverrideAppliedWhenClientDidNotSet() { - Map config = Map.of( - DATASOURCE, new PerSegmentTimeoutConfig(5000, false) - ); - - expectDefaults(); - - QueryLifecycle lifecycle = createLifecycle(config); + QueryLifecycle lifecycle = createLifecycle(perSegmentTimeout(5000)); lifecycle.initialize(baseQuery); Assert.assertEquals(5000L, lifecycle.getQuery().context().getPerSegmentTimeout()); } @Test - public void testPerDatasourceTimeout_userOverrideWins() + public void testDynamicOverridesNonClientValueInContext() { - Map config = Map.of( - DATASOURCE, new PerSegmentTimeoutConfig(5000, false) - ); + // SQL path: a default was merged into the context but the client did not set it, so the dynamic override wins. + TimeseriesQuery query = baseQuery.withOverriddenContext(Map.of(KEY, "0")); - expectDefaults(); + QueryLifecycle lifecycle = createLifecycle(perSegmentTimeout(5000)); + lifecycle.initialize(query, Collections.emptySet()); - TimeseriesQuery queryWithUserTimeout = baseQuery.withOverriddenContext( - Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, 2000L) - ); - - QueryLifecycle lifecycle = createLifecycle(config); - lifecycle.initialize(queryWithUserTimeout); - - Assert.assertEquals(2000L, lifecycle.getQuery().context().getPerSegmentTimeout()); + Assert.assertEquals(5000L, lifecycle.getQuery().context().getPerSegmentTimeout()); } @Test - public void testPerDatasourceTimeout_monitorOnlyDoesNotInject() + public void testClientProvidedValueWins() { - // monitorOnly=true: config exists but should NOT be enforced - Map config = Map.of( - DATASOURCE, new PerSegmentTimeoutConfig(5000, true) - ); - - expectDefaults(); + TimeseriesQuery query = baseQuery.withOverriddenContext(Map.of(KEY, 2000L)); - QueryLifecycle lifecycle = createLifecycle(config); - lifecycle.initialize(baseQuery); + QueryLifecycle lifecycle = createLifecycle(perSegmentTimeout(5000)); + lifecycle.initialize(query, Set.of(KEY)); - Assert.assertFalse( - "monitorOnly should not inject perSegmentTimeout", - lifecycle.getQuery().context().usePerSegmentTimeout() - ); + Assert.assertEquals(2000L, lifecycle.getQuery().context().getPerSegmentTimeout()); } @Test - public void testPerDatasourceTimeout_noMatchingDatasource() + public void testNoDynamicConfigMeansNoInjection() { - Map config = Map.of( - "other_datasource", new PerSegmentTimeoutConfig(5000, false) - ); - - expectDefaults(); - - QueryLifecycle lifecycle = createLifecycle(config); + QueryLifecycle lifecycle = createLifecycle(null); lifecycle.initialize(baseQuery); Assert.assertFalse(lifecycle.getQuery().context().usePerSegmentTimeout()); } - @Test - public void testPerDatasourceTimeout_overridesSystemDefault() - { - Map config = Map.of( - DATASOURCE, new PerSegmentTimeoutConfig(5000, false) - ); - - // System default sets perSegmentTimeout to 10000 - EasyMock.expect(queryConfig.getContext()) - .andReturn(Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, 10000L)) - .anyTimes(); - EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())).andReturn(toolChest).once(); - EasyMock.replay(conglomerate, queryConfig); - - QueryLifecycle lifecycle = createLifecycle(config); - lifecycle.initialize(baseQuery); - - Assert.assertEquals(5000L, lifecycle.getQuery().context().getPerSegmentTimeout()); - } - - @Test - public void testPrecedence_userOverridesPerDatasourceOverridesSystemDefault() + private static BrokerDynamicConfig perSegmentTimeout(long timeoutMs) { - // System default: 10000, per-datasource: 5000, user: 2000 — user should win - Map config = Map.of( - DATASOURCE, new PerSegmentTimeoutConfig(5000, false) - ); - - EasyMock.expect(queryConfig.getContext()) - .andReturn(Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, 10000L)) - .anyTimes(); - EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())).andReturn(toolChest).once(); - EasyMock.replay(conglomerate, queryConfig); - - TimeseriesQuery queryWithUserTimeout = baseQuery.withOverriddenContext( - Map.of(QueryContexts.PER_SEGMENT_TIMEOUT_KEY, 2000L) - ); - - QueryLifecycle lifecycle = createLifecycle(config); - lifecycle.initialize(queryWithUserTimeout); - - Assert.assertEquals(2000L, lifecycle.getQuery().context().getPerSegmentTimeout()); + return BrokerDynamicConfig.builder() + .withPerSegmentTimeoutConfig(Map.of(DATASOURCE, new PerSegmentTimeoutConfig(timeoutMs, false))) + .build(); } - private void expectDefaults() + private QueryLifecycle createLifecycle(BrokerDynamicConfig dynamicConfig) { - EasyMock.expect(queryConfig.getContext()).andReturn(Map.of()).anyTimes(); EasyMock.expect(conglomerate.getToolChest(EasyMock.anyObject())).andReturn(toolChest).once(); - EasyMock.replay(conglomerate, queryConfig); - } + EasyMock.replay(conglomerate); - private QueryLifecycle createLifecycle(Map perSegmentTimeoutConfig) - { return new QueryLifecycle( conglomerate, texasRanger, @@ -211,11 +149,9 @@ private QueryLifecycle createLifecycle(Map perS emitter, requestLogger, authzMapper, - queryConfig, new AuthConfig(), NoopPolicyEnforcer.instance(), - Collections.emptyList(), - perSegmentTimeoutConfig, + new QueryConfigSnapshot(Collections.emptyMap(), dynamicConfig), System.currentTimeMillis(), System.nanoTime() ); diff --git a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java index 5f8da0d21a39..3691f80ce276 100644 --- a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java @@ -59,6 +59,8 @@ import org.apache.druid.query.policy.RestrictAllTablesPolicyEnforcer; import org.apache.druid.query.policy.RowFilterPolicy; import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.server.broker.BrokerDynamicConfig; +import org.apache.druid.server.broker.QueryConfigSnapshot; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.security.Access; import org.apache.druid.server.security.Action; @@ -80,9 +82,7 @@ import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; -import java.util.Collections; import java.util.HashMap; -import java.util.List; import java.util.Map; import java.util.Optional; @@ -848,7 +848,6 @@ public void testRunSimple_queryBlocklisted() replayAll(); // Create lifecycle with blocklist - List queryBlocklist = ImmutableList.of(rule); QueryLifecycle lifecycle = new QueryLifecycle( conglomerate, texasRanger, @@ -856,11 +855,12 @@ public void testRunSimple_queryBlocklisted() emitter, requestLogger, authzMapper, - queryConfig, authConfig, policyEnforcer, - queryBlocklist, - Collections.emptyMap(), + new QueryConfigSnapshot( + ImmutableMap.of(), + BrokerDynamicConfig.builder().withQueryBlocklist(ImmutableList.of(rule)).build() + ), System.currentTimeMillis(), System.nanoTime() ); @@ -900,7 +900,6 @@ public void testRunSimple_queryNotBlocklisted() replayAll(); // Create lifecycle with blocklist - List queryBlocklist = ImmutableList.of(rule); QueryLifecycle lifecycle = new QueryLifecycle( conglomerate, texasRanger, @@ -908,11 +907,12 @@ public void testRunSimple_queryNotBlocklisted() emitter, requestLogger, authzMapper, - queryConfig, authConfig, policyEnforcer, - queryBlocklist, - Collections.emptyMap(), + new QueryConfigSnapshot( + ImmutableMap.of(), + BrokerDynamicConfig.builder().withQueryBlocklist(ImmutableList.of(rule)).build() + ), System.currentTimeMillis(), System.nanoTime() ); diff --git a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java index d7a611bfc320..28a5c244b28a 100644 --- a/server/src/test/java/org/apache/druid/server/QueryResourceTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryResourceTest.java @@ -74,6 +74,7 @@ import org.apache.druid.query.policy.NoopPolicyEnforcer; import org.apache.druid.query.policy.RowFilterPolicy; import org.apache.druid.query.timeboundary.TimeBoundaryResultValue; +import org.apache.druid.server.broker.QueryConfigSnapshot; import org.apache.druid.server.initialization.ServerConfig; import org.apache.druid.server.log.TestRequestLogger; import org.apache.druid.server.metrics.NoopServiceEmitter; @@ -772,11 +773,9 @@ public QueryLifecycle factorize() emitter, testRequestLogger, AuthTestUtils.TEST_AUTHORIZER_MAPPER, - overrideConfig, new AuthConfig(), NoopPolicyEnforcer.instance(), - null, - Collections.emptyMap(), + new QueryConfigSnapshot(overrideConfig.getContext(), null), System.currentTimeMillis(), System.nanoTime() ) diff --git a/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java b/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java index 6a466bcd269e..8a0d28dab2f5 100644 --- a/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java +++ b/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java @@ -24,7 +24,11 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import nl.jqno.equalsverifier.EqualsVerifier; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.query.Druids; import org.apache.druid.query.QueryContext; +import org.apache.druid.query.aggregation.CountAggregatorFactory; +import org.apache.druid.query.timeseries.TimeseriesQuery; import org.apache.druid.segment.TestHelper; import org.apache.druid.server.DefaultQueryBlocklistRule; import org.apache.druid.server.QueryBlocklistRule; @@ -204,4 +208,45 @@ public void testEquals() .usingGetClass() .verify(); } + + @Test + public void testContextOverridesInjectsPerSegmentTimeoutForMatchingDatasource() + { + BrokerDynamicConfig config = perSegmentTimeout("ds", new PerSegmentTimeoutConfig(5000, false)); + Assert.assertEquals(5000L, config.getQuerySpecificContextOverrides(query("ds")).getPerSegmentTimeout()); + } + + @Test + public void testContextOverridesEmptyForMonitorOnly() + { + BrokerDynamicConfig config = perSegmentTimeout("ds", new PerSegmentTimeoutConfig(5000, true)); + Assert.assertTrue(config.getQuerySpecificContextOverrides(query("ds")).isEmpty()); + } + + @Test + public void testContextOverridesEmptyForNonMatchingDatasource() + { + BrokerDynamicConfig config = perSegmentTimeout("other", new PerSegmentTimeoutConfig(5000, false)); + Assert.assertTrue(config.getQuerySpecificContextOverrides(query("ds")).isEmpty()); + } + + @Test + public void testContextOverridesEmptyWhenNoPerSegmentTimeoutConfigured() + { + Assert.assertTrue(BrokerDynamicConfig.builder().build().getQuerySpecificContextOverrides(query("ds")).isEmpty()); + } + + private static BrokerDynamicConfig perSegmentTimeout(String datasource, PerSegmentTimeoutConfig timeoutConfig) + { + return BrokerDynamicConfig.builder().withPerSegmentTimeoutConfig(Map.of(datasource, timeoutConfig)).build(); + } + + private static TimeseriesQuery query(String datasource) + { + return Druids.newTimeseriesQueryBuilder() + .dataSource(datasource) + .intervals(List.of(Intervals.ETERNITY)) + .aggregators(new CountAggregatorFactory("count")) + .build(); + } } 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)',