ATLAS-1388: Cache entities that are created/updated

release-log.txt

@@ -9,6 +9,7 @@ ATLAS-1060 Add composite indexes for exact match performance improvements for al
 ATLAS-1127 Modify creation and modification timestamps to Date instead of Long(sumasai)
 
 ALL CHANGES:
+ATLAS-1388 Cache entities that are created/updated (jnhagelb)
 ATLAS-1369 Optimize Gremlin queries generated by DSL translator (jnhagelb)
 ATLAS-1517: updated hive_model to include schema related attributes (sarath.kum4r@gmail.com via mneethiraj)
 ATLAS-1514 Remove duplicates from class array attribute when target is deleted (dkantor)

repository/src/main/java/org/apache/atlas/discovery/graph/DefaultGraphPersistenceStrategy.java

@@ -18,9 +18,12 @@
 
 package org.apache.atlas.discovery.graph;
 
-import com.google.common.collect.ImmutableCollection;
-import com.google.common.collect.ImmutableList;
+import java.util.List;
+
+import javax.inject.Inject;
+
 import org.apache.atlas.AtlasException;
+import org.apache.atlas.RequestContext;
 import org.apache.atlas.groovy.GroovyExpression;
 import org.apache.atlas.query.GraphPersistenceStrategies;
 import org.apache.atlas.query.GraphPersistenceStrategies$class;
@@ -48,8 +51,8 @@ import org.apache.atlas.typesystem.types.TypeSystem;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import javax.inject.Inject;
-import java.util.List;
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableList;
 
 /**
  * Default implementation of GraphPersistenceStrategy.
@@ -178,9 +181,12 @@ public class DefaultGraphPersistenceStrategy implements GraphPersistenceStrategi
 
             case CLASS:
                 AtlasVertex classVertex = (AtlasVertex) value;
-                ITypedReferenceableInstance classInstance = metadataRepository.getGraphToInstanceMapper()
-                    .mapGraphToTypedInstance(GraphHelper.getSingleValuedProperty(classVertex, Constants.GUID_PROPERTY_KEY, String.class),
-                        classVertex);
+                String guid = classVertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class);
+                // Check if the instance we need was previously loaded.
+                ITypedReferenceableInstance classInstance = RequestContext.get().getInstance(guid);
+                if (classInstance == null) {
+                    classInstance = metadataRepository.getGraphToInstanceMapper().mapGraphToTypedInstance(guid, classVertex);
+                }
                 return dataType.convert(classInstance, Multiplicity.OPTIONAL);
 
             default:

repository/src/main/java/org/apache/atlas/repository/graph/FullTextMapper.java

@@ -17,11 +17,11 @@
  */
 package org.apache.atlas.repository.graph;
 
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 import org.apache.atlas.AtlasException;
+import org.apache.atlas.RequestContext;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.typesystem.ITypedInstance;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
@@ -44,20 +44,19 @@ public class FullTextMapper {
     private static final GraphHelper graphHelper = GraphHelper.getInstance();
 
     private static final String FULL_TEXT_DELIMITER = " ";
-    private final Map<String, ITypedReferenceableInstance> instanceCache;
 
     FullTextMapper(TypedInstanceToGraphMapper typedInstanceToGraphMapper,
             GraphToTypedInstanceMapper graphToTypedInstanceMapper) {
         this.graphToTypedInstanceMapper = graphToTypedInstanceMapper;
         this.typedInstanceToGraphMapper = typedInstanceToGraphMapper;
-        instanceCache = new HashMap<>();
     }
 
     public String mapRecursive(AtlasVertex instanceVertex, boolean followReferences) throws AtlasException {
         String guid = GraphHelper.getGuid(instanceVertex);
         ITypedReferenceableInstance typedReference;
-        if (instanceCache.containsKey(guid)) {
-            typedReference = instanceCache.get(guid);
+        RequestContext context = RequestContext.get();
+        typedReference = context.getInstance(guid);
+        if (typedReference != null) {
 
             if (LOG.isDebugEnabled()) {
                 LOG.debug("Cache hit: guid = {}, entityId = {}", guid, typedReference.getId()._getId());
@@ -65,7 +64,7 @@ public class FullTextMapper {
         } else {
             typedReference =
                     graphToTypedInstanceMapper.mapGraphToTypedInstance(guid, instanceVertex);
-            instanceCache.put(guid, typedReference);
+            context.cache(typedReference);
 
             if (LOG.isDebugEnabled()) {
                 LOG.debug("Cache miss: guid = {}, entityId = {}", guid, typedReference.getId().getId());

repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java

@@ -191,6 +191,10 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
                 Constants.STATE_PROPERTY_KEY, Id.EntityState.ACTIVE.name());
 
         String guid = GraphHelper.getGuid(instanceVertex);
+        ITypedReferenceableInstance cached = RequestContext.get().getInstance(guid);
+        if(cached != null) {
+            return cached;
+        }
         return graphToInstanceMapper.mapGraphToTypedInstance(guid, instanceVertex);
     }
 

repository/src/main/java/org/apache/atlas/repository/graph/GraphToTypedInstanceMapper.java

@@ -17,8 +17,18 @@
  */
 package org.apache.atlas.repository.graph;
 
-import com.google.inject.Singleton;
+import static org.apache.atlas.repository.graph.GraphHelper.string;
+
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 import org.apache.atlas.AtlasException;
+import org.apache.atlas.RequestContext;
 import org.apache.atlas.repository.RepositoryException;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
@@ -28,10 +38,8 @@ import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.typesystem.ITypedInstance;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedStruct;
-import org.apache.atlas.typesystem.Referenceable;
 import org.apache.atlas.typesystem.persistence.AtlasSystemAttributes;
 import org.apache.atlas.typesystem.persistence.Id;
-import org.apache.atlas.typesystem.persistence.ReferenceableInstance;
 import org.apache.atlas.typesystem.types.AttributeInfo;
 import org.apache.atlas.typesystem.types.ClassType;
 import org.apache.atlas.typesystem.types.DataTypes;
@@ -43,15 +51,7 @@ import org.apache.atlas.typesystem.types.TypeSystem;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.math.BigDecimal;
-import java.math.BigInteger;
-import java.util.ArrayList;
-import java.util.Date;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
-import static org.apache.atlas.repository.graph.GraphHelper.string;
+import com.google.inject.Singleton;
 
 @Singleton
 public final class GraphToTypedInstanceMapper {
@@ -69,7 +69,14 @@ public final class GraphToTypedInstanceMapper {
     public ITypedReferenceableInstance mapGraphToTypedInstance(String guid, AtlasVertex instanceVertex)
         throws AtlasException {
 
-        if (LOG.isDebugEnabled()) {
+        if(LOG.isDebugEnabled()) {
+            //We don't do a cache check here since we want that to be at a higher level
+            //where the vertex lookup can also be avoided.  However, this is a convenient
+            //place to add a check to see if there are any places that were missed.
+            if(RequestContext.get().getInstance(guid) != null) {
+                LOG.warn("Looking up previously cached guid at: ", new Exception());
+            }
+
             LOG.debug("Mapping graph root vertex {} to typed instance for guid {}", instanceVertex, guid);
         }
 
@@ -99,7 +106,7 @@ public final class GraphToTypedInstanceMapper {
 
         mapVertexToInstance(instanceVertex, typedInstance, classType.fieldMapping().fields);
         mapVertexToInstanceTraits(instanceVertex, typedInstance, traits);
-
+        RequestContext.get().cache(typedInstance);
         return typedInstance;
     }
 
@@ -209,6 +216,10 @@ public final class GraphToTypedInstanceMapper {
             if (attributeInfo.isComposite) {
                 //Also, when you retrieve a type's instance, you get the complete object graph of the composites
                 LOG.debug("Found composite, mapping vertex to instance");
+                ITypedReferenceableInstance cached = RequestContext.get().getInstance(guid);
+                if(cached != null) {
+                    return cached;
+                }
                 return mapGraphToTypedInstance(guid, referenceVertex);
             } else {
                 String state = GraphHelper.getStateAsString(referenceVertex);

repository/src/main/java/org/apache/atlas/repository/graph/TypedInstanceToGraphMapper.java

@@ -125,6 +125,9 @@ public final class TypedInstanceToGraphMapper {
             throw new UnsupportedOperationException("Not handled - " + operation);
         }
 
+        for(ITypedReferenceableInstance instance : typedInstances) {
+            addToEntityCache(requestContext, instance);
+        }
     }
 
     private Collection<IReferenceableInstance> walkClassInstances(ITypedReferenceableInstance typedInstance)
@@ -825,4 +828,23 @@ public final class TypedInstanceToGraphMapper {
     public AtlasVertex lookupVertex(Id refId) {
         return idToVertexMap.get(refId);
     }
+
+    private void addToEntityCache(RequestContext context, ITypedReferenceableInstance instance)
+            throws EntityNotFoundException {
+
+        Id instanceId = instance.getId();
+        if(instanceId.isUnassigned()) {
+            if(instance instanceof ReferenceableInstance) {
+                //When the id is unassigned, we can only cache the instance of it is
+                //an instance of ReferenceableInstance, since replaceWithNewId is not
+                //currently in the ITypedReferenceableInstance interface.
+                Id id = getId(instance);
+                ((ReferenceableInstance)instance).replaceWithNewId(id);
+                context.cache(instance);
+            }
+        }
+        else {
+            context.cache(instance);
+        }
+    }
 }

repository/src/test/java/org/apache/atlas/TestUtils.java

@@ -18,14 +18,33 @@
 
 package org.apache.atlas;
 
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
-import com.google.inject.Provider;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createClassTypeDef;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createOptionalAttrDef;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createRequiredAttrDef;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createStructTypeDef;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createTraitTypeDef;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createUniqueRequiredAttrDef;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Date;
+import java.util.List;
 
 import org.apache.atlas.listener.EntityChangeListener;
 import org.apache.atlas.listener.TypesChangeListener;
 import org.apache.atlas.repository.MetadataRepository;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.apache.atlas.repository.graph.GraphBackedMetadataRepository;
 import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
@@ -59,23 +78,9 @@ import org.apache.commons.lang.RandomStringUtils;
 import org.codehaus.jettison.json.JSONArray;
 import org.testng.Assert;
 
-import java.io.File;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.math.BigDecimal;
-import java.math.BigInteger;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.Date;
-import java.util.List;
-
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createClassTypeDef;
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createOptionalAttrDef;
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createRequiredAttrDef;
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createStructTypeDef;
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createTraitTypeDef;
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createUniqueRequiredAttrDef;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.inject.Provider;
 
 /**
  * Test utility class.
@@ -505,11 +510,14 @@ public final class TestUtils {
         }
         return null;
     }
-    
-    public static void resetRequestContext() {      
+
+    public static void resetRequestContext() {
+        //reset the context while preserving the user
+        String user = RequestContext.get().getUser();
         RequestContext.createContext();
+        RequestContext.get().setUser(user);
     }
-    
+
     public static void setupGraphProvider(MetadataRepository repo) throws AtlasException {
         TypeCache typeCache = null;
         try {
@@ -538,10 +546,92 @@ public final class TestUtils {
         getGraph().commit();
 
     }
-    
+
     public static AtlasGraph getGraph() {
 
         return AtlasGraphProvider.getGraphInstance();
-       
+
+    }
+
+    /**
+     * Adds a proxy wrapper around the specified MetadataService that automatically
+     * resets the request context before every call.
+     *
+     * @param delegate
+     * @return
+     */
+    public static MetadataService addSessionCleanupWrapper(final MetadataService delegate) {
+
+        return (MetadataService)Proxy.newProxyInstance(Thread.currentThread().getContextClassLoader(),
+                new Class[]{MetadataService.class}, new InvocationHandler() {
+
+            @Override
+            public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+
+                try {
+                    resetRequestContext();
+                    Object result = method.invoke(delegate, args);
+
+                    return result;
+                }
+                catch(InvocationTargetException e) {
+                    e.getCause().printStackTrace();
+                    throw e.getCause();
+                }
+                catch(Throwable t) {
+                    t.printStackTrace();
+                    throw t;
+                }
+            }
+
+        });
+    }
+
+    /**
+     * Adds a proxy wrapper around the specified MetadataRepository that automatically
+     * resets the request context before every call and either commits or rolls
+     * back the graph transaction after every call.
+     *
+     * @param delegate
+     * @return
+     */
+    public static MetadataRepository addTransactionWrapper(final MetadataRepository delegate) {
+        return (MetadataRepository)Proxy.newProxyInstance(Thread.currentThread().getContextClassLoader(),
+                new Class[]{MetadataRepository.class}, new InvocationHandler() {
+
+            @Override
+            public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+                boolean useTransaction = GraphBackedMetadataRepository.class.getMethod(
+                        method.getName(), method.getParameterTypes())
+                        .isAnnotationPresent(GraphTransaction.class);
+                try {
+                    resetRequestContext();
+                    Object result = method.invoke(delegate, args);
+                    if(useTransaction) {
+                        System.out.println("Committing changes");
+                        getGraph().commit();
+                        System.out.println("Commit succeeded.");
+                    }
+                    return result;
+                }
+                catch(InvocationTargetException e) {
+                    e.getCause().printStackTrace();
+                    if(useTransaction) {
+                        System.out.println("Rolling back changes due to exception.");
+                        getGraph().rollback();
+                    }
+                    throw e.getCause();
+                }
+                catch(Throwable t) {
+                    t.printStackTrace();
+                    if(useTransaction) {
+                        System.out.println("Rolling back changes due to exception.");
+                        getGraph().rollback();
+                    }
+                    throw t;
+                }
+            }
+
+        });
     }
 }

repository/src/test/java/org/apache/atlas/discovery/GraphBackedDiscoveryServiceTest.java

@@ -87,6 +87,7 @@ public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
     public void setUp() throws Exception {
         super.setUp();
 
+        repositoryService = TestUtils.addTransactionWrapper(repositoryService);
         final TypeSystem typeSystem = TypeSystem.getInstance();
         Collection<String> oldTypeNames = new HashSet<>();
         oldTypeNames.addAll(typeSystem.getTypeNames());

repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryDeleteTestBase.java

@@ -114,43 +114,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         new GraphBackedSearchIndexer(new AtlasTypeRegistry());
         final GraphBackedMetadataRepository delegate = new GraphBackedMetadataRepository(getDeleteHandler(typeSystem));
 
-        repositoryService = (MetadataRepository)Proxy.newProxyInstance(Thread.currentThread().getContextClassLoader(),
-                new Class[]{MetadataRepository.class}, new InvocationHandler() {
-
-            @Override
-            public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
-                boolean useTransaction = GraphBackedMetadataRepository.class.getMethod(
-                        method.getName(), method.getParameterTypes())
-                        .isAnnotationPresent(GraphTransaction.class);
-                try {
-
-                    Object result = method.invoke(delegate, args);
-                    if(useTransaction) {
-                        System.out.println("Committing changes");
-                        TestUtils.getGraph().commit();
-                        System.out.println("Commit succeeded.");
-                    }
-                    return result;
-                }
-                catch(InvocationTargetException e) {
-                    e.getCause().printStackTrace();
-                    if(useTransaction) {
-                        System.out.println("Rolling back changes due to exception.");
-                        TestUtils.getGraph().rollback();
-                    }
-                    throw e.getCause();
-                }
-                catch(Throwable t) {
-                    t.printStackTrace();
-                    if(useTransaction) {
-                        System.out.println("Rolling back changes due to exception.");
-                        TestUtils.getGraph().rollback();
-                    }
-                    throw t;
-                }
-            }
-
-        });
+        repositoryService = TestUtils.addTransactionWrapper(delegate);
 
         TestUtils.defineDeptEmployeeTypes(typeSystem);
         TestUtils.createHiveTypes(typeSystem);
@@ -531,7 +495,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
     /**
      * Verify deleting an entity which is contained by another
      * entity through a bi-directional composite reference.
-     * 
+     *
      * @throws Exception
      */
     @Test
@@ -633,21 +597,21 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
     @Test
     public void testDisconnectUnidirectionalArrayReferenceFromStructAndTraitTypes() throws Exception {
         // Define class types.
-        HierarchicalTypeDefinition<ClassType> structTargetDef = TypesUtil.createClassTypeDef("StructTarget", 
+        HierarchicalTypeDefinition<ClassType> structTargetDef = TypesUtil.createClassTypeDef("StructTarget",
             ImmutableSet.<String>of(), TypesUtil.createOptionalAttrDef("attr1", DataTypes.STRING_TYPE));
-        HierarchicalTypeDefinition<ClassType> traitTargetDef = TypesUtil.createClassTypeDef("TraitTarget", 
+        HierarchicalTypeDefinition<ClassType> traitTargetDef = TypesUtil.createClassTypeDef("TraitTarget",
             ImmutableSet.<String>of(), TypesUtil.createOptionalAttrDef("attr1", DataTypes.STRING_TYPE));
-        HierarchicalTypeDefinition<ClassType> structContainerDef = TypesUtil.createClassTypeDef("StructContainer", 
+        HierarchicalTypeDefinition<ClassType> structContainerDef = TypesUtil.createClassTypeDef("StructContainer",
             ImmutableSet.<String>of(), TypesUtil.createOptionalAttrDef("struct", "TestStruct"));
 
         // Define struct and trait types which have a unidirectional array reference
         // to a class type.
-        StructTypeDefinition structDef = TypesUtil.createStructTypeDef("TestStruct", 
+        StructTypeDefinition structDef = TypesUtil.createStructTypeDef("TestStruct",
             new AttributeDefinition("target", DataTypes.arrayTypeName("StructTarget"), Multiplicity.OPTIONAL, false, null),
             new AttributeDefinition("nestedStructs", DataTypes.arrayTypeName("NestedStruct"), Multiplicity.OPTIONAL, false, null));
-        StructTypeDefinition nestedStructDef = TypesUtil.createStructTypeDef("NestedStruct", 
+        StructTypeDefinition nestedStructDef = TypesUtil.createStructTypeDef("NestedStruct",
             TypesUtil.createOptionalAttrDef("attr1", DataTypes.STRING_TYPE));
-        HierarchicalTypeDefinition<TraitType> traitDef = TypesUtil.createTraitTypeDef("TestTrait", ImmutableSet.<String>of(), 
+        HierarchicalTypeDefinition<TraitType> traitDef = TypesUtil.createTraitTypeDef("TestTrait", ImmutableSet.<String>of(),
             new AttributeDefinition("target", DataTypes.arrayTypeName("TraitTarget"), Multiplicity.OPTIONAL, false, null));
 
         TypesDef typesDef = TypesUtil.getTypesDef(ImmutableList.<EnumTypeDefinition>of(), ImmutableList.of(structDef, nestedStructDef),
@@ -669,9 +633,9 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         ClassType traitTargetType = typeSystem.getDataType(ClassType.class, "TraitTarget");
         ClassType structContainerType = typeSystem.getDataType(ClassType.class, "StructContainer");
 
-        ITypedReferenceableInstance structTargetConvertedEntity = 
+        ITypedReferenceableInstance structTargetConvertedEntity =
             structTargetType.convert(structTargetEntity, Multiplicity.REQUIRED);
-        ITypedReferenceableInstance traitTargetConvertedEntity = 
+        ITypedReferenceableInstance traitTargetConvertedEntity =
             traitTargetType.convert(traitTargetEntity, Multiplicity.REQUIRED);
         ITypedReferenceableInstance structContainerConvertedEntity =
             structContainerType.convert(structContainerEntity, Multiplicity.REQUIRED);
@@ -755,13 +719,13 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
     @Test
     public void testDisconnectMapReferenceFromClassType() throws Exception {
         // Define type for map value.
-        HierarchicalTypeDefinition<ClassType> mapValueDef = TypesUtil.createClassTypeDef("MapValue", 
+        HierarchicalTypeDefinition<ClassType> mapValueDef = TypesUtil.createClassTypeDef("MapValue",
             ImmutableSet.<String>of(),
             new AttributeDefinition("biMapOwner", "MapOwner", Multiplicity.OPTIONAL, false, "biMap"));
 
         // Define type with unidirectional and bidirectional map references,
         // where the map value is a class reference to MapValue.
-        HierarchicalTypeDefinition<ClassType> mapOwnerDef = TypesUtil.createClassTypeDef("MapOwner", 
+        HierarchicalTypeDefinition<ClassType> mapOwnerDef = TypesUtil.createClassTypeDef("MapOwner",
             ImmutableSet.<String>of(),
             new AttributeDefinition("map", DataTypes.mapTypeName(DataTypes.STRING_TYPE.getName(),
                         "MapValue"), Multiplicity.OPTIONAL, false, null),
@@ -811,7 +775,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         }
 
         // Delete the map value instance.
-        // This should disconnect the references from the map owner instance. 
+        // This should disconnect the references from the map owner instance.
         deleteEntities(mapValueGuid);
         assertEntityDeleted(mapValueGuid);
         assertTestDisconnectMapReferenceFromClassType(mapOwnerGuid);

repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryTest.java

@@ -28,6 +28,7 @@ import org.apache.atlas.TestUtils;
 import org.apache.atlas.discovery.graph.GraphBackedDiscoveryService;
 import org.apache.atlas.query.QueryParams;
 import org.apache.atlas.repository.Constants;
+import org.apache.atlas.repository.MetadataRepository;
 import org.apache.atlas.repository.RepositoryException;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
@@ -76,7 +77,7 @@ import java.util.concurrent.Future;
 
 import javax.inject.Inject;
 
-import scala.actors.threadpool.Arrays;
+import java.util.Arrays;
 
 import static org.apache.atlas.typesystem.types.utils.TypesUtil.createClassTypeDef;
 import static org.apache.atlas.typesystem.types.utils.TypesUtil.createUniqueRequiredAttrDef;
@@ -95,7 +96,7 @@ import static org.testng.Assert.assertTrue;
 public class GraphBackedMetadataRepositoryTest {
 
     @Inject
-    private GraphBackedMetadataRepository repositoryService;
+    private MetadataRepository repositoryService;
 
     @Inject
     private GraphBackedDiscoveryService discoveryService;
@@ -109,6 +110,8 @@ public class GraphBackedMetadataRepositoryTest {
         typeSystem = TypeSystem.getInstance();
         typeSystem.reset();
 
+        assertTrue(repositoryService instanceof GraphBackedMetadataRepository);
+        repositoryService = TestUtils.addTransactionWrapper(repositoryService);
         new GraphBackedSearchIndexer(new AtlasTypeRegistry());
 
         TestUtils.defineDeptEmployeeTypes(typeSystem);

repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1Test.java

@@ -97,6 +97,7 @@ public class AtlasEntityStoreV1Test {
 
     @BeforeClass
     public void setUp() throws Exception {
+        metadataService = TestUtils.addSessionCleanupWrapper(metadataService);
         new GraphBackedSearchIndexer(typeRegistry);
         final AtlasTypesDef deptTypesDef = TestUtilsV2.defineDeptEmployeeTypes();
         typeDefStore.createTypesDef(deptTypesDef);
@@ -112,6 +113,7 @@ public class AtlasEntityStoreV1Test {
     @AfterClass
     public void clear() {
         AtlasGraphProvider.cleanup();
+        TestUtils.resetRequestContext();
     }
 
     @BeforeTest

repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java

@@ -18,9 +18,30 @@
 
 package org.apache.atlas.service;
 
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
-import com.google.inject.Inject;
+import static org.apache.atlas.TestUtils.COLUMNS_ATTR_NAME;
+import static org.apache.atlas.TestUtils.COLUMN_TYPE;
+import static org.apache.atlas.TestUtils.PII;
+import static org.apache.atlas.TestUtils.TABLE_TYPE;
+import static org.apache.atlas.TestUtils.createColumnEntity;
+import static org.apache.atlas.TestUtils.createDBEntity;
+import static org.apache.atlas.TestUtils.createInstance;
+import static org.apache.atlas.TestUtils.createTableEntity;
+import static org.apache.atlas.TestUtils.randomString;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createClassTypeDef;
+import static org.apache.atlas.typesystem.types.utils.TypesUtil.createOptionalAttrDef;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
@@ -32,6 +53,7 @@ import org.apache.atlas.discovery.graph.GraphBackedDiscoveryService;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.listener.ChangedTypeDefs;
 import org.apache.atlas.listener.EntityChangeListener;
+import org.apache.atlas.listener.TypeDefChangeListener;
 import org.apache.atlas.query.QueryParams;
 import org.apache.atlas.repository.audit.EntityAuditRepository;
 import org.apache.atlas.repository.audit.HBaseBasedAuditRepository;
@@ -72,37 +94,17 @@ import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Collections;
-
-import static org.apache.atlas.TestUtils.COLUMNS_ATTR_NAME;
-import static org.apache.atlas.TestUtils.COLUMN_TYPE;
-import static org.apache.atlas.TestUtils.DATABASE_TYPE;
-import static org.apache.atlas.TestUtils.PII;
-import static org.apache.atlas.TestUtils.TABLE_TYPE;
-import static org.apache.atlas.TestUtils.createColumnEntity;
-import static org.apache.atlas.TestUtils.createDBEntity;
-import static org.apache.atlas.TestUtils.createInstance;
-import static org.apache.atlas.TestUtils.createTableEntity;
-import static org.apache.atlas.TestUtils.randomString;
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createClassTypeDef;
-import static org.apache.atlas.typesystem.types.utils.TypesUtil.createOptionalAttrDef;
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertNotNull;
-import static org.testng.Assert.assertNull;
-import static org.testng.Assert.assertTrue;
-import static org.testng.Assert.fail;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.inject.Inject;
 
 @Guice(modules = RepositoryMetadataModule.class)
 public class DefaultMetadataServiceTest {
     @Inject
     private MetadataService metadataService;
 
+    private TypeDefChangeListener typeDefChangeListener;
+
     @Inject
     private EntityAuditRepository auditRepository;
 
@@ -114,12 +116,16 @@ public class DefaultMetadataServiceTest {
     private Referenceable table;
 
     private Id tableId;
-    
+
     private final String NAME = "name";
 
 
     @BeforeTest
     public void setUp() throws Exception {
+
+        typeDefChangeListener = (DefaultMetadataService)metadataService;
+        metadataService = TestUtils.addSessionCleanupWrapper(metadataService);
+
         if (auditRepository instanceof HBaseBasedAuditRepository) {
             HBaseTestUtils.startCluster();
             ((HBaseBasedAuditRepository) auditRepository).start();
@@ -1218,7 +1224,7 @@ public class DefaultMetadataServiceTest {
             List<String> beforeChangeTypeNames = new ArrayList<>();
             beforeChangeTypeNames.addAll(metadataService.getTypeNames(new HashMap<TypeCache.TYPE_FILTER, String>()));
 
-            ((DefaultMetadataService)metadataService).onChange(new ChangedTypeDefs());
+            typeDefChangeListener.onChange(new ChangedTypeDefs());
 
             List<String> afterChangeTypeNames = new ArrayList<>();
             afterChangeTypeNames.addAll(metadataService.getTypeNames(new HashMap<TypeCache.TYPE_FILTER, String>()));
@@ -1269,7 +1275,7 @@ public class DefaultMetadataServiceTest {
                 deletedEntities.add(entity.getId()._getId());
             }
         }
-        
+
         public List<String> getDeletedEntities() {
             return deletedEntities;
         }

server-api/src/main/java/org/apache/atlas/RequestContext.java

@@ -18,6 +18,14 @@
 
 package org.apache.atlas;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
 import org.apache.atlas.metrics.Metrics;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.persistence.Id;
@@ -26,12 +34,6 @@ import org.apache.atlas.typesystem.types.TypeSystem;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Set;
-
 public class RequestContext {
     private static final Logger LOG = LoggerFactory.getLogger(RequestContext.class);
 
@@ -41,6 +43,7 @@ public class RequestContext {
     private Set<String> updatedEntityIds = new LinkedHashSet<>();
     private Set<String> deletedEntityIds = new LinkedHashSet<>();
     private List<ITypedReferenceableInstance> deletedEntities = new ArrayList<>();
+    private Map<String,ITypedReferenceableInstance> entityCache = new HashMap<>();
 
     private String user;
     private long requestTime;
@@ -71,7 +74,27 @@ public class RequestContext {
         return context;
     }
 
+    /**
+     * Adds the specified instance to the cache
+     *
+     */
+    public void cache(ITypedReferenceableInstance instance) {
+        entityCache.put(instance.getId()._getId(), instance);
+    }
+
+    /**
+     * Checks if an instance with the given guid is in the cache for this request.  Either returns the instance
+     * or null if it is not in the cache.
+     *
+     * @param guid the guid to find
+     * @return Either the instance or null if it is not in the cache.
+     */
+    public ITypedReferenceableInstance getInstance(String guid) {
+        return entityCache.get(guid);
+    }
+
     public static void clear() {
+        CURRENT_CONTEXT.get().entityCache.clear();
         CURRENT_CONTEXT.remove();
     }
 
@@ -122,7 +145,7 @@ public class RequestContext {
     public long getRequestTime() {
         return requestTime;
     }
-    
+
     public boolean isDeletedEntity(String entityGuid) {
         return deletedEntityIds.contains(entityGuid);
     }