ATLAS-3256 Modify export API to process with relationshipAttributes

repository/src/main/java/org/apache/atlas/repository/impexp/EntitiesExtractor.java

+/**
+ * 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.atlas.repository.impexp;
+
+import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.typedef.AtlasEntityDef;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.type.AtlasTypeRegistry;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class EntitiesExtractor {
+    static final String PROPERTY_GUID = "__guid";
+    private static final String VERTEX_BASED_EXTRACT = "default";
+    private static final String INCREMENTAL_EXTRACT = "incremental";
+    private static final String RELATION_BASED_EXTRACT = "relationship";
+
+    private Map<String, ExtractStrategy> extractors = new HashMap<>();
+    private ExtractStrategy extractor;
+
+    public EntitiesExtractor(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry) {
+        extractors.put(VERTEX_BASED_EXTRACT, new VertexExtractor(atlasGraph, typeRegistry));
+        extractors.put(INCREMENTAL_EXTRACT, new IncrementalExportEntityProvider(atlasGraph));
+        extractors.put(RELATION_BASED_EXTRACT, new RelationshipAttributesExtractor(typeRegistry));
+    }
+
+    public void get(AtlasEntity entity, ExportService.ExportContext context) {
+        if(extractor == null) {
+            extractor = extractors.get(VERTEX_BASED_EXTRACT);
+        }
+
+        switch (context.fetchType) {
+            case CONNECTED:
+                extractor.connectedFetch(entity, context);
+                break;
+
+            case INCREMENTAL:
+                if (context.isHiveDBIncrementalSkipLineage()) {
+                    extractors.get(INCREMENTAL_EXTRACT).fullFetch(entity, context);
+                    break;
+                }
+
+            case FULL:
+            default:
+                extractor.fullFetch(entity, context);
+        }
+    }
+
+    public void setExtractor(AtlasEntityDef atlasEntityDef) {
+        extractor = extractUsing(atlasEntityDef);
+    }
+
+    public void close() {
+        for (ExtractStrategy es : extractors.values()) {
+            es.close();
+        }
+    }
+
+    private ExtractStrategy extractUsing(AtlasEntityDef atlasEntityDef) {
+        return (atlasEntityDef == null || atlasEntityDef.getRelationshipAttributeDefs().size() == 0)
+                ? extractors.get(VERTEX_BASED_EXTRACT)
+                : extractors.get(RELATION_BASED_EXTRACT);
+    }
+}

repository/src/main/java/org/apache/atlas/repository/impexp/ExtractStrategy.java

+/**
+ * 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.atlas.repository.impexp;
+
+import org.apache.atlas.model.instance.AtlasEntity;
+
+public interface ExtractStrategy {
+
+    void connectedFetch(AtlasEntity entity, ExportService.ExportContext context);
+    void fullFetch(AtlasEntity entity, ExportService.ExportContext context);
+    void close();
+}

repository/src/main/java/org/apache/atlas/repository/impexp/IncrementalExportEntityProvider.java

@@ -18,6 +18,8 @@
 
 package org.apache.atlas.repository.impexp;
 
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.util.UniqueList;
 import org.slf4j.Logger;
@@ -28,11 +30,10 @@ import javax.script.ScriptEngine;
 import javax.script.ScriptException;
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 
-public class IncrementalExportEntityProvider {
+public class IncrementalExportEntityProvider implements ExtractStrategy {
     private static final Logger LOG = LoggerFactory.getLogger(IncrementalExportEntityProvider.class);
 
     private static final String QUERY_PARAMETER_START_GUID = "startGuid";
@@ -50,9 +51,23 @@ public class IncrementalExportEntityProvider {
     private ScriptEngine scriptEngine;
 
     @Inject
-    public IncrementalExportEntityProvider(AtlasGraph atlasGraph, ScriptEngine scriptEngine) {
+    public IncrementalExportEntityProvider(AtlasGraph atlasGraph) {
         this.atlasGraph = atlasGraph;
-        this.scriptEngine = scriptEngine;
+        try {
+            this.scriptEngine = atlasGraph.getGremlinScriptEngine();
+        } catch (AtlasBaseException e) {
+            LOG.error("Error instantiating script engine.", e);
+        }
+    }
+
+    @Override
+    public void fullFetch(AtlasEntity entity, ExportService.ExportContext context) {
+        populate(entity.getGuid(), context.changeMarker, context.guidsToProcess);
+    }
+
+    @Override
+    public void connectedFetch(AtlasEntity entity, ExportService.ExportContext context) {
+
     }
 
     public void populate(String dbEntityGuid, long timeStamp, UniqueList<String> guidsToProcess) {
@@ -63,6 +78,13 @@ public class IncrementalExportEntityProvider {
         }
     }
 
+    @Override
+    public void close() {
+        if (scriptEngine != null) {
+            atlasGraph.releaseGremlinScriptEngine(scriptEngine);
+        }
+    }
+
     private void partial(String dbEntityGuid, long timeStamp, UniqueList<String> guidsToProcess) {
         guidsToProcess.addAll(fetchGuids(dbEntityGuid, QUERY_TABLE, timeStamp));
         guidsToProcess.addAll(fetchGuids(dbEntityGuid, QUERY_SD, timeStamp));
@@ -98,7 +120,7 @@ public class IncrementalExportEntityProvider {
             }
 
             for (Map<String, Object> item : result) {
-                guids.add((String) item.get(ExportService.PROPERTY_GUID));
+                guids.add((String) item.get(EntitiesExtractor.PROPERTY_GUID));
             }
 
             return guids;

repository/src/main/java/org/apache/atlas/repository/impexp/RelationshipAttributesExtractor.java

+/**
+ * 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.atlas.repository.impexp;
+
+import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasRelatedObjectId;
+import org.apache.atlas.model.typedef.AtlasEntityDef;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.type.AtlasTypeUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+public class RelationshipAttributesExtractor implements ExtractStrategy {
+
+    private static final Logger LOG = LoggerFactory.getLogger(RelationshipAttributesExtractor.class);
+
+    private final AtlasTypeRegistry typeRegistry;
+
+    public RelationshipAttributesExtractor(AtlasTypeRegistry typeRegistry) {
+        this.typeRegistry = typeRegistry;
+    }
+
+    @Override
+    public void fullFetch(AtlasEntity entity, ExportService.ExportContext context) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> fullFetch({}): guidsToProcess {}", AtlasTypeUtil.getAtlasObjectId(entity), context.guidsToProcess.size());
+        }
+
+        List<AtlasRelatedObjectId> atlasRelatedObjectIdList = getRelatedObjectIds(entity);
+
+        for (AtlasRelatedObjectId ar : atlasRelatedObjectIdList) {
+            boolean isLineage = isLineageType(ar.getTypeName());
+
+            if (context.skipLineage && isLineage) {
+                continue;
+            }
+            context.addToBeProcessed(isLineage, ar.getGuid(), ExportService.TraversalDirection.BOTH);
+        }
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== fullFetch({}): guidsToProcess {}", entity.getGuid(), context.guidsToProcess.size());
+        }
+    }
+
+    @Override
+    public void connectedFetch(AtlasEntity entity, ExportService.ExportContext context) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> connectedFetch({}): guidsToProcess {} isSkipConnectedFetch :{}", AtlasTypeUtil.getAtlasObjectId(entity), context.guidsToProcess.size(), context.isSkipConnectedFetch);
+        }
+
+        List<AtlasRelatedObjectId> atlasRelatedObjectIdList = getRelatedObjectIds(entity);
+        for (AtlasRelatedObjectId ar : atlasRelatedObjectIdList) {
+            boolean isLineage = isLineageType(ar.getTypeName());
+
+            if (context.skipLineage && isLineage) {
+                continue;
+            }
+            if (!context.isSkipConnectedFetch || isLineage) {
+                context.addToBeProcessed(isLineage, ar.getGuid(), ExportService.TraversalDirection.BOTH);
+            }
+        }
+
+        if(isLineageType(entity.getTypeName())){
+            context.isSkipConnectedFetch = false;
+        }else{
+            context.isSkipConnectedFetch = true;
+        }
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> connectedFetch({}): guidsToProcess {}, isSkipConnectedFetch :{}", AtlasTypeUtil.getAtlasObjectId(entity), context.guidsToProcess.size(), context.isSkipConnectedFetch);
+        }
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private boolean isLineageType(String typeName) {
+        AtlasEntityDef entityDef = typeRegistry.getEntityDefByName(typeName);
+        return entityDef.getSuperTypes().contains("Process");
+    }
+
+    private List<AtlasRelatedObjectId> getRelatedObjectIds(AtlasEntity entity) {
+        List<AtlasRelatedObjectId> relatedObjectIds = new ArrayList<>();
+
+        for (Object o : entity.getRelationshipAttributes().values()) {
+            if (o instanceof AtlasRelatedObjectId) {
+                relatedObjectIds.add((AtlasRelatedObjectId) o);
+            } else if (o instanceof Collection) {
+                relatedObjectIds.addAll((List) o);
+            }
+        }
+
+        return relatedObjectIds;
+    }
+}

repository/src/main/java/org/apache/atlas/repository/impexp/VertexExtractor.java

+/**
+ * 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.atlas.repository.impexp;
+
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.type.AtlasEntityType;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.type.AtlasTypeUtil;
+import org.apache.atlas.util.AtlasGremlinQueryProvider;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.script.ScriptEngine;
+import javax.script.ScriptException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.atlas.repository.impexp.EntitiesExtractor.PROPERTY_GUID;
+
+public class VertexExtractor implements ExtractStrategy {
+    private static final Logger LOG = LoggerFactory.getLogger(VertexExtractor.class);
+
+    private static final String PROPERTY_IS_PROCESS = "isProcess";
+    private static final String QUERY_BINDING_START_GUID = "startGuid";
+
+    private final AtlasGremlinQueryProvider gremlinQueryProvider;
+
+    private final Map<String, Object> bindings;
+    private AtlasGraph atlasGraph;
+    private AtlasTypeRegistry typeRegistry;
+    private ScriptEngine scriptEngine;
+
+    public VertexExtractor(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry) {
+        this.atlasGraph = atlasGraph;
+        this.typeRegistry = typeRegistry;
+        try {
+            this.scriptEngine = atlasGraph.getGremlinScriptEngine();
+        } catch (AtlasBaseException e) {
+            LOG.error("Script Engine: Instantiation failed!");
+        }
+        this.gremlinQueryProvider = AtlasGremlinQueryProvider.INSTANCE;
+        this.bindings = new HashMap<>();
+    }
+
+    @Override
+    public void fullFetch(AtlasEntity entity, ExportService.ExportContext context) {
+        if (LOG.isDebugEnabled()){
+            LOG.debug("==> fullFetch({}): guidsToProcess {}", AtlasTypeUtil.getAtlasObjectId(entity), context.guidsToProcess.size());
+        }
+
+        String query = this.gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_BY_GUID_FULL);
+
+        bindings.clear();
+        bindings.put(QUERY_BINDING_START_GUID, entity.getGuid());
+
+        List<Map<String, Object>> result = executeGremlinQuery(query, context);
+
+        if (CollectionUtils.isEmpty(result)) {
+            return;
+        }
+
+        for (Map<String, Object> hashMap : result) {
+            String guid = (String) hashMap.get(PROPERTY_GUID);
+            boolean isLineage = (boolean) hashMap.get(PROPERTY_IS_PROCESS);
+
+            if (context.getSkipLineage() && isLineage) continue;
+
+            if (!context.guidsProcessed.contains(guid)) {
+                context.addToBeProcessed(isLineage, guid, ExportService.TraversalDirection.BOTH);
+            }
+        }
+    }
+
+    @Override
+    public void connectedFetch(AtlasEntity entity, ExportService.ExportContext context) {
+        if (LOG.isDebugEnabled()){
+            LOG.debug("==> connectedFetch({}): guidsToProcess {}", AtlasTypeUtil.getAtlasObjectId(entity), context.guidsToProcess.size());
+        }
+
+        ExportService.TraversalDirection direction = context.guidDirection.get(entity.getGuid());
+
+        if (direction == null || direction == ExportService.TraversalDirection.UNKNOWN) {
+            getConnectedEntityGuids(entity, context, ExportService.TraversalDirection.OUTWARD, ExportService.TraversalDirection.INWARD);
+        } else {
+            if (isProcessEntity(entity)) {
+                direction = ExportService.TraversalDirection.OUTWARD;
+            }
+
+            getConnectedEntityGuids(entity, context, direction);
+        }
+    }
+
+    @Override
+    public void close() {
+        if (scriptEngine != null) {
+            atlasGraph.releaseGremlinScriptEngine(scriptEngine);
+        }
+    }
+
+    private void getConnectedEntityGuids(AtlasEntity entity, ExportService.ExportContext context, ExportService.TraversalDirection... directions) {
+        if (directions == null) {
+            return;
+        }
+
+        for (ExportService.TraversalDirection direction : directions) {
+            String query = getQueryForTraversalDirection(direction);
+
+            bindings.clear();
+            bindings.put(QUERY_BINDING_START_GUID, entity.getGuid());
+
+            List<Map<String, Object>> result = executeGremlinQuery(query, context);
+
+            if (CollectionUtils.isEmpty(result)) {
+                continue;
+            }
+
+            for (Map<String, Object> hashMap : result) {
+                String guid = (String) hashMap.get(PROPERTY_GUID);
+                ExportService.TraversalDirection currentDirection = context.guidDirection.get(guid);
+                boolean isLineage = (boolean) hashMap.get(PROPERTY_IS_PROCESS);
+
+                if (context.skipLineage && isLineage) continue;
+
+                if (currentDirection == null) {
+                    context.addToBeProcessed(isLineage, guid, direction);
+
+                } else if (currentDirection == ExportService.TraversalDirection.OUTWARD && direction == ExportService.TraversalDirection.INWARD) {
+                    // the entity should be reprocessed to get inward entities
+                    context.guidsProcessed.remove(guid);
+                    context.addToBeProcessed(isLineage, guid, direction);
+                }
+            }
+        }
+    }
+
+    private boolean isProcessEntity(AtlasEntity entity) {
+        String typeName = entity.getTypeName();
+        AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
+
+        return entityType.isSubTypeOf(AtlasBaseTypeDef.ATLAS_TYPE_PROCESS);
+    }
+
+    private String getQueryForTraversalDirection(ExportService.TraversalDirection direction) {
+        switch (direction) {
+            case INWARD:
+                return this.gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_BY_GUID_CONNECTED_IN_EDGE);
+
+            default:
+            case OUTWARD:
+                return this.gremlinQueryProvider.getQuery(AtlasGremlinQueryProvider.AtlasGremlinQuery.EXPORT_BY_GUID_CONNECTED_OUT_EDGE);
+        }
+    }
+
+    private List<Map<String, Object>> executeGremlinQuery(String query, ExportService.ExportContext context) {
+        try {
+            return (List<Map<String, Object>>) atlasGraph.executeGremlinScript(scriptEngine, bindings, query, false);
+        } catch (ScriptException e) {
+            LOG.error("Script execution failed for query: ", query, e);
+            return null;
+        }
+    }
+}

repository/src/test/java/org/apache/atlas/repository/impexp/IncrementalExportEntityProviderTest.java

@@ -63,7 +63,7 @@ public class IncrementalExportEntityProviderTest extends ExportImportTestBase {
         verifyCreatedEntities(entityStore, entityGuids, 2);
 
         gremlinScriptEngine = atlasGraph.getGremlinScriptEngine();
-        incrementalExportEntityProvider = new IncrementalExportEntityProvider(atlasGraph, gremlinScriptEngine);
+        incrementalExportEntityProvider = new IncrementalExportEntityProvider(atlasGraph);
     }
 
     @AfterClass