From 5445b3a00ed8cb2bf73ba26fc2d9a81d6003aab6 Mon Sep 17 00:00:00 2001
From: jzonthemtn <jzemerick@opensourceconnections.com>
Date: Thu, 22 Feb 2024 15:12:45 -0500
Subject: [PATCH] Adding query-id header to the list of headers.

---
 src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java b/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java
index 84fd8bd..81aaf85 100644
--- a/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java
+++ b/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java
@@ -55,6 +55,7 @@ public class UserBehaviorLoggingPlugin extends Plugin implements ActionPlugin {
     public Collection<RestHeaderDefinition> getRestHeaders() {
         return List.of(
                 new RestHeaderDefinition(HeaderConstants.EVENT_STORE_HEADER.getHeader(), false),
+                new RestHeaderDefinition(HeaderConstants.QUERY_ID_HEADER.getHeader(), false),
                 new RestHeaderDefinition(HeaderConstants.SESSION_ID_HEADER.getHeader(), false),
                 new RestHeaderDefinition(HeaderConstants.USER_ID_HEADER.getHeader(), false)
         );
@@ -64,6 +65,7 @@ public Collection<RestHeaderDefinition> getRestHeaders() {
     public Collection<String> getTaskHeaders() {
         return List.of(
                 HeaderConstants.EVENT_STORE_HEADER.getHeader(),
+                HeaderConstants.QUERY_ID_HEADER.getHeader(),
                 HeaderConstants.SESSION_ID_HEADER.getHeader(),
                 HeaderConstants.USER_ID_HEADER.getHeader()
         );