fix: fix query context precedence layer

Author: jtuglu1

server/src/main/java/org/apache/druid/server/QueryLifecycle.java

@@ -166,7 +166,21 @@ public <T> QueryResponse<T> 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 <T> QueryResponse<T> runSimple(
+      final Query<T> query,
+      final AuthenticationResult authenticationResult,
+      final AuthorizationResult authorizationResult,
+      @Nullable final Set<String> userProvidedContextKeys
+  )
+  {
+    initialize(query, userProvidedContextKeys);
 
     final Sequence<T> results;
 
@@ -212,43 +226,64 @@ 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<String> userProvidedContextKeys)
   {
     transition(State.NEW, State.INITIALIZED);
 
     userContextKeys = new HashSet<>(baseQuery.getContext().keySet());
+    final Set<String> 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<String, Object> contextWithDefaults = new HashMap<>(queryConfigProvider.getContext());
-    applyPerDatasourcePerSegmentTimeout(baseQuery, contextWithDefaults, queryId);
     Map<String, Object> finalContext = QueryContexts.override(contextWithDefaults, baseQuery.getContext());
+    applyPerDatasourcePerSegmentTimeout(baseQuery, finalContext, effectiveUserKeys, queryId);
     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.
+   * 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<String, Object> contextWithDefaults,
+      final Map<String, Object> finalContext,
+      final Set<String> userProvidedContextKeys,
       final String queryId
   )
   {
     if (perSegmentTimeoutConfig.isEmpty()) {
       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;
       }

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<String, PerSegmentTimeoutConfig> 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<String, PerSegmentTimeoutConfig> 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();

sql/src/main/java/org/apache/druid/sql/calcite/run/NativeQueryMaker.java

@@ -186,7 +186,10 @@ private <T> QueryResponse<Object[]> execute(
     final QueryResponse<T> results = queryLifecycle.runSimple(
         (Query<T>) 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(