From 93ca0e8b5901f46e22c4061d91a8a0bdb543b845 Mon Sep 17 00:00:00 2001
From: Rashi J <rashi.jaiswal@atlan.com>
Date: Tue, 29 Oct 2024 15:24:22 +0530
Subject: [PATCH] add headers to task

---
 .../org/apache/atlas/tasks/AtlasTaskService.java    | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/repository/src/main/java/org/apache/atlas/tasks/AtlasTaskService.java b/repository/src/main/java/org/apache/atlas/tasks/AtlasTaskService.java
index 21fd5f3ce1..2f769bcbd7 100644
--- a/repository/src/main/java/org/apache/atlas/tasks/AtlasTaskService.java
+++ b/repository/src/main/java/org/apache/atlas/tasks/AtlasTaskService.java
@@ -38,6 +38,8 @@ public class AtlasTaskService implements TaskService {
 
     private final List<String> retryAllowedStatuses;
 
+    private static final String ATLAN_HEADER_PREFIX_PATTERN = "x-atlan-";
+
     @Inject
     public AtlasTaskService(AtlasGraph graph) {
         this.graph = graph;
@@ -224,6 +226,7 @@ private void deleteTask(String taskGuid) throws AtlasBaseException {
     @Override
     public AtlasVertex createTaskVertex(AtlasTask task) {
         AtlasVertex ret = graph.addVertex();
+        Map<String, String> requestContextHeaders = RequestContext.get().getRequestContextHeaders();
 
         setEncodedProperty(ret, Constants.TASK_GUID, task.getGuid());
         setEncodedProperty(ret, Constants.TASK_TYPE_PROPERTY_KEY, Constants.TASK_TYPE_NAME);
@@ -253,6 +256,16 @@ public AtlasVertex createTaskVertex(AtlasTask task) {
             setEncodedProperty(ret, Constants.TASK_END_TIME, task.getEndTime().getTime());
         }
 
+        if (MapUtils.isNotEmpty(requestContextHeaders)) {
+            for (Map.Entry<String, String> entry : requestContextHeaders.entrySet()) {
+                String key = entry.getKey();
+                if (key.startsWith(ATLAN_HEADER_PREFIX_PATTERN)) {
+                    String val = entry.getValue();
+                    setEncodedProperty(ret, key, val);
+                }
+            }
+        }
+
         setEncodedProperty(ret, Constants.TASK_PARAMETERS, AtlasJson.toJson(task.getParameters()));
         setEncodedProperty(ret, Constants.TASK_ATTEMPT_COUNT, task.getAttemptCount());
         setEncodedProperty(ret, Constants.TASK_ERROR_MESSAGE, task.getErrorMessage());