add headers to task

atlanhq · Oct 29, 2024 · 93ca0e8 · 93ca0e8
1 parent 96d2048
commit 93ca0e8
Showing 1 changed file with 13 additions and 0 deletions.
diff --git 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());