fix: fix query context precedence layer

Author: jtuglu1

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

@@ -166,7 +166,24 @@ public <T> QueryResponse<T> runSimple(
       final AuthorizationResult authorizationResult
   )
   {
-    initialize(query);
+    return runSimple(query, authenticationResult, authorizationResult, null);
+  }
+
+  /**
+   * Variant of {@link #runSimple(Query, AuthenticationResult, AuthorizationResult)} that also accepts the set of
+   * context keys the caller genuinely supplied (as opposed to values merged in from static/system defaults). The SQL
+   * layer merges static defaults into the query context before it reaches this class, so it must pass the real
+   * user-provided keys; otherwise a merged-in default is indistinguishable from a user-supplied value. 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 +229,76 @@ 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);
+  }
+
+  /**
+   * Initializes this object to execute a specific query. Does not actually execute the query.
+   * <p>
+   * The {@code state} transitions from NEW, to INITIALIZED.
+   *
+   * @param baseQuery               the query
+   * @param userProvidedContextKeys the context keys the caller genuinely supplied, or {@code null} to treat the whole
+   *                                query context as caller-supplied. Native queries can pass {@code null} since their
+   *                                context is exactly the caller's context. The SQL layer merges static/system defaults
+   *                                into the query context upstream, so it must pass the real user-provided keys so a
+   *                                merged-in default is not mistaken for an explicit user value.
+   * @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());
+    // For native queries the query context is exactly the caller's context, so its key set is authoritative. For SQL
+    // queries, static/system defaults are merged into the query context upstream, so the caller passes the real
+    // user-provided keys (otherwise a merged-in default looks user-supplied).
+    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, highest to lowest: user-supplied context > per-datasource per-segment timeout dynamic config >
+    // static/system defaults > code default. User context is merged last so it wins; the per-datasource timeout is then
+    // applied over the merged result, but only when the user did not supply perSegmentTimeout themselves.
     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 any value
+   * that was merged in from static/system defaults. A perSegmentTimeout the caller supplied explicitly wins and is left
+   * untouched (identified via {@code userProvidedContextKeys}). 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<String, Object> contextWithDefaults,
+      final Map<String, Object> finalContext,
+      final Set<String> userProvidedContextKeys,
       final String queryId
   )
   {
     if (perSegmentTimeoutConfig.isEmpty()) {
       return;
     }
 
+    // The caller explicitly set perSegmentTimeout: respect it (highest precedence).
+    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 +310,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,48 @@ DATASOURCE, new PerSegmentTimeoutConfig(5000, false)
     Assert.assertEquals(2000L, lifecycle.getQuery().context().getPerSegmentTimeout());
   }
 
+  @Test
+  public void testPerDatasourceTimeout_staticDefaultInContextNotUserProvided_dynamicWins()
+  {
+    // Simulates the SQL path: a static/system default was merged into the query context upstream, but the user did not
+    // actually set perSegmentTimeout. The per-datasource dynamic config must override the merged-in default.
+    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);
+    // Empty user-provided key set: perSegmentTimeout in the context came from a default, not the caller.
+    lifecycle.initialize(queryWithMergedDefault, Collections.emptySet());
+
+    Assert.assertEquals(5000L, lifecycle.getQuery().context().getPerSegmentTimeout());
+  }
+
+  @Test
+  public void testPerDatasourceTimeout_userProvidedViaProvenance_userWins()
+  {
+    // Same context value as above, but the caller genuinely set perSegmentTimeout; it must win 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,11 @@ private <T> QueryResponse<Object[]> execute(
     final QueryResponse<T> results = queryLifecycle.runSimple(
         (Query<T>) query,
         authenticationResult,
-        authorizationResult
+        authorizationResult,
+        // The SQL layer merges static/system defaults into the query context; pass the genuinely user-provided keys so
+        // per-datasource per-segment timeout config can override a merged-in default while still yielding to an
+        // explicit user value.
+        plannerContext.authContextKeys()
     );
 
     return mapResultSequence(