feat: Propagate Request Headers

docs/configuration/index.md

@@ -348,6 +348,24 @@ Coordinator and Overlord log changes to lookups, segment load/drop rules, and dy
 |`druid.audit.manager.maxPayloadSizeBytes`|The maximum size of audit payload to store in Druid's metadata store audit table. If the size of audit payload exceeds this value, the audit log would be stored with a message indicating that the payload was omitted instead. Setting `maxPayloadSizeBytes` to -1 (default value) disables this check, meaning Druid will always store audit payload regardless of it's size. Setting to any negative number other than `-1` is invalid. Human-readable format is supported, see [here](human-readable-byte.md).  |-1|
 |`druid.audit.manager.skipNullField`|If true, the audit payload stored in metadata store will exclude any field with null value. |false|
 
+### Request header propagation
+
+Druid can capture configured inbound HTTP headers and propagate their values through the query context and into audit events. This is useful for correlating Druid queries and configuration changes with an external distributed-trace or request-tracking system.
+
+|Property|Description|Default|
+|--------|-----------|-------|
+|`druid.audit.requestHeaders.headerToContextKey`|JSON map of inbound HTTP header name to the [query context](../querying/query-context.md) key the header value is bound to. For each configured header present on a request, the value is captured by a servlet filter and injected into the query context, from where Druid's native sub-query context propagation carries it to data servers (for example, broker to historical). The captured header values are also recorded, keyed by their context key (for example `traceId`), in the `requestMetadata` map of the audit entry for any configuration change made via an audited API. Mapping a header to a reserved Druid context key (`queryId`, `subQueryId`, or `sqlQueryId`) is rejected at startup, because it would let a client overwrite the server-assigned value. Set to an empty map (`{}`) to disable header propagation.|`{"X-Druid-Trace-Id":"traceId"}`|
+
+For example, to also capture a tenant identifier:
+
+```properties
+druid.audit.requestHeaders.headerToContextKey={"X-Druid-Trace-Id":"traceId","X-Tenant-Id":"tenantId"}
+```
+
+A client-supplied value for a configured context key in the query body is ignored (stripped); only a value captured from the corresponding inbound header takes effect. This prevents a client from spoofing the trace identifier by setting it in the query JSON.
+
+When query-context authorization is enabled (`druid.auth.authorizeQueryContextParams=true`), a value captured from a header is subject to the same `QUERY_CONTEXT` `WRITE` authorization as a value supplied in the query body. To allow the default `traceId` (or any other propagated key) without an explicit grant, add it to `druid.auth.unsecuredContextKeys`.
+
 ### Metadata storage
 
 These properties specify the JDBC connection and other configuration around the metadata storage. The only services that connect to the metadata storage with these properties are the [Coordinator](../design/coordinator.md) and [Overlord](../design/overlord.md).

extensions-contrib/rabbit-stream-indexing-service/src/main/java/org/apache/druid/indexing/rabbitstream/RabbitStreamIndexTaskClientFactory.java

@@ -21,19 +21,30 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.inject.Inject;
+import org.apache.druid.audit.RequestHeaderContextConfig;
 import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.guice.annotations.EscalatedGlobal;
 import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory;
 import org.apache.druid.java.util.http.client.HttpClient;
 
+import javax.annotation.Nullable;
+
 @LazySingleton
 public class RabbitStreamIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory<String, Long>
 {
   @Inject
   public RabbitStreamIndexTaskClientFactory(
       @EscalatedGlobal HttpClient httpClient,
-      @Json ObjectMapper mapper)
+      @Json ObjectMapper mapper,
+      @Nullable RequestHeaderContextConfig requestHeaderContextConfig)
+  {
+    super(httpClient, mapper, requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig());
+  }
+
+  public RabbitStreamIndexTaskClientFactory(
+      HttpClient httpClient,
+      ObjectMapper mapper)
   {
     super(httpClient, mapper);
   }

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTaskClientFactory.java

@@ -21,20 +21,32 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.inject.Inject;
+import org.apache.druid.audit.RequestHeaderContextConfig;
 import org.apache.druid.data.input.kafka.KafkaTopicPartition;
 import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.guice.annotations.EscalatedGlobal;
 import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory;
 import org.apache.druid.java.util.http.client.HttpClient;
 
+import javax.annotation.Nullable;
+
 @LazySingleton
 public class KafkaIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory<KafkaTopicPartition, Long>
 {
   @Inject
   public KafkaIndexTaskClientFactory(
       @EscalatedGlobal HttpClient httpClient,
-      @Json ObjectMapper mapper
+      @Json ObjectMapper mapper,
+      @Nullable RequestHeaderContextConfig requestHeaderContextConfig
+  )
+  {
+    super(httpClient, mapper, requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig());
+  }
+
+  public KafkaIndexTaskClientFactory(
+      HttpClient httpClient,
+      ObjectMapper mapper
   )
   {
     super(httpClient, mapper);

extensions-core/kinesis-indexing-service/src/main/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskClientFactory.java

@@ -21,19 +21,31 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.inject.Inject;
+import org.apache.druid.audit.RequestHeaderContextConfig;
 import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.guice.annotations.EscalatedGlobal;
 import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskClientFactory;
 import org.apache.druid.java.util.http.client.HttpClient;
 
+import javax.annotation.Nullable;
+
 @LazySingleton
 public class KinesisIndexTaskClientFactory extends SeekableStreamIndexTaskClientFactory<String, String>
 {
   @Inject
   public KinesisIndexTaskClientFactory(
       @EscalatedGlobal HttpClient httpClient,
-      @Json ObjectMapper mapper
+      @Json ObjectMapper mapper,
+      @Nullable RequestHeaderContextConfig requestHeaderContextConfig
+  )
+  {
+    super(httpClient, mapper, requestHeaderContextConfig != null ? requestHeaderContextConfig : new RequestHeaderContextConfig());
+  }
+
+  public KinesisIndexTaskClientFactory(
+      HttpClient httpClient,
+      ObjectMapper mapper
   )
   {
     super(httpClient, mapper);

indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskClientFactory.java

@@ -20,6 +20,7 @@
 package org.apache.druid.indexing.seekablestream;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.audit.RequestHeaderContextConfig;
 import org.apache.druid.indexing.common.TaskInfoProvider;
 import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor;
 import org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorTuningConfig;
@@ -36,14 +37,25 @@ public abstract class SeekableStreamIndexTaskClientFactory<PartitionIdType, Sequ
 
   private final HttpClient httpClient;
   private final ObjectMapper jsonMapper;
+  private final RequestHeaderContextConfig requestHeaderContextConfig;
 
   protected SeekableStreamIndexTaskClientFactory(
       final HttpClient httpClient,
       final ObjectMapper jsonMapper
   )
+  {
+    this(httpClient, jsonMapper, new RequestHeaderContextConfig());
+  }
+
+  protected SeekableStreamIndexTaskClientFactory(
+      final HttpClient httpClient,
+      final ObjectMapper jsonMapper,
+      final RequestHeaderContextConfig requestHeaderContextConfig
+  )
   {
     this.httpClient = httpClient;
     this.jsonMapper = jsonMapper;
+    this.requestHeaderContextConfig = requestHeaderContextConfig;
   }
 
   /**
@@ -72,7 +84,7 @@ public SeekableStreamIndexTaskClient<PartitionIdType, SequenceOffsetType> build(
 
     return new SeekableStreamIndexTaskClientAsyncImpl<>(
         dataSource,
-        new ServiceClientFactoryImpl(httpClient, connectExec),
+        new ServiceClientFactoryImpl(httpClient, connectExec, requestHeaderContextConfig),
         taskInfoProvider,
         jsonMapper,
         tuningConfig.getHttpTimeout(),

processing/src/main/java/org/apache/druid/audit/AuditManager.java

@@ -31,6 +31,7 @@ public interface AuditManager
 
   String X_DRUID_AUTHOR = "X-Druid-Author";
   String X_DRUID_COMMENT = "X-Druid-Comment";
+  String X_DRUID_TRACE_ID = "X-Druid-Trace-Id";
 
   void doAudit(AuditEntry event);
 

processing/src/main/java/org/apache/druid/audit/RequestHeaderContextConfig.java

@@ -0,0 +1,159 @@
+/*
+ * 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.audit;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.query.BaseQuery;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.function.BiConsumer;
+
+/**
+ * Configures which inbound HTTP headers are captured and where they're stored in the
+ * {@link org.apache.druid.query.Query} context. The map is keyed by header name (case as
+ * sent on the wire) and the value is the reserved context key.
+ *
+ * <p>Defaults: {@code X-Druid-Trace-Id → traceId} so trace propagation works out of the
+ * box. Operators can extend by setting:
+ * <pre>
+ *   druid.audit.requestHeaders.headerToContextKey={"X-Druid-Trace-Id": "traceId",
+ *                                                  "X-Custom-Foo": "foo"}
+ * </pre>
+ *
+ * <p>The values in this map are reserved context keys: user-supplied values for these keys
+ * are stripped from the query context before the filter-captured value is merged
+ * (anti-spoof). Setting overlapping keys to existing query-context keys is therefore
+ * discouraged — pick context-key names that don't collide with standard Druid context.
+ */
+public class RequestHeaderContextConfig
+{
+  private static final Map<String, String> DEFAULT_HEADER_TO_CONTEXT_KEY = ImmutableMap.of(
+      AuditManager.X_DRUID_TRACE_ID, "traceId"
+  );
+
+  /**
+   * Context keys that are server-managed by Druid and must not be configurable as targets
+   * for header propagation. Mapping a header to one of these would let a client overwrite
+   * the server-assigned value (e.g. the random per-request {@code queryId}), breaking
+   * query identification, cancellation, and metric correlation.
+   */
+  private static final Set<String> RESERVED_DRUID_CONTEXT_KEYS = ImmutableSet.of(
+      BaseQuery.QUERY_ID,
+      BaseQuery.SUB_QUERY_ID,
+      BaseQuery.SQL_QUERY_ID
+  );
+
+  // @JsonProperty on the field is REQUIRED: JsonConfigurator.verifyClazzIsConfigurable()
+  // rejects config classes whose bean-property fields lack a Jackson field annotation, which
+  // would fail server startup (JettyServerModule binds this via JsonConfigProvider).
+  // Declared as ImmutableMap (not Map) so getHeaderToContextKey() can return the field
+  // reference directly without exposing a mutable internal representation (CodeQL).
+  @JsonProperty
+  private final ImmutableMap<String, String> headerToContextKey;
+
+  public RequestHeaderContextConfig()
+  {
+    this(DEFAULT_HEADER_TO_CONTEXT_KEY);
+  }
+
+  /**
+   * @param headerToContextKey explicit mapping; null means "use defaults" (the typical case
+   *     when {@code JsonConfigProvider} deserializes from an empty config block);
+   *     {@link java.util.Collections#emptyMap()} explicitly disables propagation.
+   */
+  @JsonCreator
+  public RequestHeaderContextConfig(
+      @JsonProperty("headerToContextKey") Map<String, String> headerToContextKey
+  )
+  {
+    final Map<String, String> raw = headerToContextKey == null
+                                    ? DEFAULT_HEADER_TO_CONTEXT_KEY
+                                    : headerToContextKey;
+    for (String contextKey : raw.values()) {
+      if (RESERVED_DRUID_CONTEXT_KEYS.contains(contextKey)) {
+        throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+                            .ofCategory(DruidException.Category.INVALID_INPUT)
+                            .build(
+                                "druid.audit.requestHeaders.headerToContextKey maps a header to "
+                                + "reserved Druid context key[%s]; this would let a client overwrite "
+                                + "the server-assigned value. Reserved keys are %s.",
+                                contextKey, RESERVED_DRUID_CONTEXT_KEYS
+                            );
+      }
+    }
+    this.headerToContextKey = ImmutableMap.copyOf(raw);
+  }
+
+  /**
+   * Returns the configured map of {@code header → context-key}. Immutable. Empty means
+   * no header propagation is performed (the feature is effectively disabled).
+   */
+  @JsonProperty
+  public Map<String, String> getHeaderToContextKey()
+  {
+    return headerToContextKey;
+  }
+
+  /**
+   * For each configured {@code header → contextKey}, if the supplied query context
+   * carries a value for {@code contextKey}, invokes {@code headerSetter} with
+   * {@code (headerName, stringifiedValue)}. Used by internal RPC clients
+   * (broker → historical, etc.) to attach propagated headers to outbound requests so
+   * the receiving node's filter captures them just as it would for a client-originated
+   * request.
+   */
+  public void applyToOutboundRequest(Map<String, Object> queryContext, BiConsumer<String, String> headerSetter)
+  {
+    if (queryContext == null || queryContext.isEmpty() || headerToContextKey.isEmpty()) {
+      return;
+    }
+    for (Map.Entry<String, String> entry : headerToContextKey.entrySet()) {
+      final Object value = queryContext.get(entry.getValue());
+      if (value != null) {
+        headerSetter.accept(entry.getKey(), value.toString());
+      }
+    }
+  }
+
+  /**
+   * For each configured {@code header → contextKey}, if the supplied captured-header map (keyed
+   * by context key, as held by {@code RequestHeaderContext}) carries a value for {@code contextKey},
+   * invokes {@code headerSetter} with {@code (headerName, value)}. Used by the shared outbound
+   * HTTP client to forward configured headers on ALL inter-service calls made on the request
+   * thread (not just queries), so the chain of internal calls carries the same header values.
+   */
+  public void applyCapturedHeaders(Map<String, String> capturedByContextKey, BiConsumer<String, String> headerSetter)
+  {
+    if (capturedByContextKey == null || capturedByContextKey.isEmpty() || headerToContextKey.isEmpty()) {
+      return;
+    }
+    for (Map.Entry<String, String> entry : headerToContextKey.entrySet()) {
+      final String value = capturedByContextKey.get(entry.getValue());
+      if (value != null) {
+        headerSetter.accept(entry.getKey(), value);
+      }
+    }
+  }
+}