Skip to content

Commit

Permalink
Merge pull request #3697 from atlanhq/amk-49-metastore-master
Browse files Browse the repository at this point in the history 
AMK-49 store headers in task
  • Loading branch information
@rashi-j
rashi-j authored Nov 14, 2024
2 parents 7c0472a + 9c230fb commit 4160f80
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 3 deletions.
20 changes: 19 additions & 1 deletion common/src/main/java/org/apache/atlas/repository/Constants.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -358,7 +358,25 @@ public final class Constants {
public static final String TASK_ENTITY_GUID = encodePropertyKey(TASK_PREFIX + "entityGuid");
public static final String TASK_CLASSIFICATION_TYPENAME = encodePropertyKey(TASK_PREFIX + "classificationTypeName");
public static final String ACTIVE_STATE_VALUE = "ACTIVE";

public static final String TASK_HEADER_ATLAN_AGENT = "x-atlan-agent";
public static final String TASK_HEADER_ATLAN_AGENT_ID = "x-atlan-agent-id";
public static final String TASK_HEADER_ATLAN_PKG_NAME = "x-atlan-package-name";
public static final String TASK_HEADER_ATLAN_AGENT_WORKFLOW_ID = "x-atlan-agent-workflow-id";
public static final String TASK_HEADER_ATLAN_VIA_UI = "x-atlan-via-ui";
public static final String TASK_HEADER_ATLAN_REQUEST_ID = "x-atlan-request-id";
public static final String TASK_HEADER_ATLAN_GOOGLE_SHEETS_ID = "x-atlan-google-sheets-id";
public static final String TASK_HEADER_ATLAN_MS_EXCEL_ID = "x-atlan-microsoft-excel-id";
public static final Set<String> TASK_HEADER_SET = new HashSet<String>() {{
add(TASK_HEADER_ATLAN_AGENT);
add(TASK_HEADER_ATLAN_AGENT_ID);
add(TASK_HEADER_ATLAN_VIA_UI);
add(TASK_HEADER_ATLAN_PKG_NAME);
add(TASK_HEADER_ATLAN_AGENT_WORKFLOW_ID);
add(TASK_HEADER_ATLAN_REQUEST_ID);
add(TASK_HEADER_ATLAN_GOOGLE_SHEETS_ID);
add(TASK_HEADER_ATLAN_MS_EXCEL_ID);
}};

/**
* Index Recovery vertex property keys.
*/
Expand Down
17 changes: 15 additions & 2 deletions repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -330,6 +330,11 @@ private void initialize(AtlasGraph graph) throws RepositoryException, IndexExcep
HashMap<String, HashMap<String, Object>> KEYWORD_MULTIFIELD = new HashMap<>();
KEYWORD_MULTIFIELD.put("keyword", ES_KEYWORD_FIELD);

HashMap<String, Object> ES_KEYWORD_WO_NORMALIZER = new HashMap<>();
ES_KEYWORD_WO_NORMALIZER.put("type", "keyword");
HashMap<String, HashMap<String, Object>> KEYWORD_FIELD = new HashMap<>();
KEYWORD_FIELD.put("keyword", ES_KEYWORD_WO_NORMALIZER);

HashMap<String, Object> ES_GLOSSARY_ANALYZER_TEXT_FIELD = new HashMap<>();
ES_GLOSSARY_ANALYZER_TEXT_FIELD.put("type", "text");
ES_GLOSSARY_ANALYZER_TEXT_FIELD.put("analyzer", "atlan_glossary_analyzer");
Expand Down Expand Up @@ -393,7 +398,7 @@ private void initialize(AtlasGraph graph) throws RepositoryException, IndexExcep
createCommonVertexIndex(management, TASK_GUID, UniqueKind.GLOBAL_UNIQUE, String.class, SINGLE, true, false);
createCommonVertexIndex(management, TASK_TYPE_PROPERTY_KEY, UniqueKind.NONE, String.class, SINGLE, true, false);
createCommonVertexIndex(management, TASK_CREATED_TIME, UniqueKind.NONE, Long.class, SINGLE, true, false);
createCommonVertexIndex(management, TASK_STATUS, UniqueKind.NONE, String.class, SINGLE, true, false);
createCommonVertexIndex(management, TASK_STATUS, UniqueKind.NONE, String.class, SINGLE, true, false, false, new HashMap<>(), KEYWORD_FIELD);

createCommonVertexIndex(management, TASK_TYPE, UniqueKind.NONE, String.class, SINGLE, true, false, true);
createCommonVertexIndex(management, TASK_CREATED_BY, UniqueKind.NONE, String.class, SINGLE, false, false, true);
Expand All @@ -407,7 +412,15 @@ private void initialize(AtlasGraph graph) throws RepositoryException, IndexExcep
createCommonVertexIndex(management, TASK_TIME_TAKEN_IN_SECONDS, UniqueKind.NONE, Long.class, SINGLE, false, false);
createCommonVertexIndex(management, TASK_START_TIME, UniqueKind.NONE, Long.class, SINGLE, false, false);
createCommonVertexIndex(management, TASK_END_TIME, UniqueKind.NONE, Long.class, SINGLE, false, false);

createCommonVertexIndex(management, TASK_HEADER_ATLAN_AGENT, UniqueKind.NONE, String.class, SINGLE, false, false, true);
createCommonVertexIndex(management, TASK_HEADER_ATLAN_AGENT_ID, UniqueKind.NONE, String.class, SINGLE, false, false, true);
createCommonVertexIndex(management, TASK_HEADER_ATLAN_PKG_NAME, UniqueKind.NONE, String.class, SINGLE, false, false, true);
createCommonVertexIndex(management, TASK_HEADER_ATLAN_AGENT_WORKFLOW_ID, UniqueKind.NONE, String.class, SINGLE, false, false, true);
createCommonVertexIndex(management, TASK_HEADER_ATLAN_VIA_UI, UniqueKind.NONE, String.class, SINGLE, false, false, true);
createCommonVertexIndex(management, TASK_HEADER_ATLAN_REQUEST_ID, UniqueKind.NONE, String.class, SINGLE, false, false, true);
createCommonVertexIndex(management, TASK_HEADER_ATLAN_GOOGLE_SHEETS_ID, UniqueKind.NONE, String.class, SINGLE, false, false, true);
createCommonVertexIndex(management, TASK_HEADER_ATLAN_MS_EXCEL_ID, UniqueKind.NONE, String.class, SINGLE, false, false, true);

// index recovery
createCommonVertexIndex(management, PROPERTY_KEY_INDEX_RECOVERY_NAME, UniqueKind.GLOBAL_UNIQUE, String.class, SINGLE, true, false);

Expand Down
12 changes: 12 additions & 0 deletions repository/src/main/java/org/apache/atlas/tasks/AtlasTaskService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import org.apache.atlas.utils.AtlasPerfMetrics;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.collections.MapUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
Expand Down Expand Up @@ -225,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);
Expand Down Expand Up @@ -254,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().toLowerCase().trim();
if (Constants.TASK_HEADER_SET.contains(key)) {
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());
Expand Down

0 comments on commit 4160f80

Please sign in to comment.