ATLAS-3091: AtlasGraphProvider Support for Graph with Batch-loading Enable

graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/GraphDatabase.java

@@ -37,6 +37,11 @@ public interface GraphDatabase<V, E> {
      */
     AtlasGraph<V, E> getGraph();
 
+    /*
+     * Get graph initialized for bulk loading. This instance is used for high-performance ingest.
+     */
+    AtlasGraph<V, E> getGraphBulkLoading();
+
     /**
      * Sets things up so that getGraph() will return a graph that can be used for running
      * tests.

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusGraph.java

@@ -78,6 +78,7 @@ import java.util.stream.StreamSupport;
 
 import static org.apache.atlas.repository.Constants.INDEX_SEARCH_VERTEX_PREFIX_DEFAULT;
 import static org.apache.atlas.repository.Constants.INDEX_SEARCH_VERTEX_PREFIX_PROPERTY;
+import static org.apache.atlas.repository.graphdb.janus.AtlasJanusGraphDatabase.getGraphInstance;
 
 /**
  * Janus implementation of AtlasGraph.
@@ -91,11 +92,15 @@ public class AtlasJanusGraph implements AtlasGraph<AtlasJanusVertex, AtlasJanusE
     private final StandardJanusGraph          janusGraph;
 
     public AtlasJanusGraph() {
+        this(getGraphInstance());
+    }
+
+    public AtlasJanusGraph(JanusGraph graphInstance) {
         //determine multi-properties once at startup
         JanusGraphManagement mgmt = null;
 
         try {
-            mgmt = AtlasJanusGraphDatabase.getGraphInstance().openManagement();
+            mgmt = graphInstance.openManagement();
 
             Iterable<PropertyKey> keys = mgmt.getRelationTypes(PropertyKey.class);
 
@@ -278,7 +283,7 @@ public class AtlasJanusGraph implements AtlasGraph<AtlasJanusVertex, AtlasJanusE
     }
 
     private JanusGraph getGraph() {
-        return AtlasJanusGraphDatabase.getGraphInstance();
+        return getGraphInstance();
     }
 
     @Override

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusGraphDatabase.java

@@ -264,6 +264,11 @@ public class AtlasJanusGraphDatabase implements GraphDatabase<AtlasJanusVertex,
         return atlasGraphInstance;
     }
 
+    @Override
+    public AtlasGraph<AtlasJanusVertex, AtlasJanusEdge> getGraphBulkLoading() {
+        return new AtlasJanusGraph(getBulkLoadingGraphInstance());
+    }
+
     private static void startLocalSolr() {
         if (isEmbeddedSolr()) {
             try {

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

@@ -72,6 +72,20 @@ public class AtlasGraphProvider implements IAtlasGraphProvider {
         }
     }
 
+    public AtlasGraph getBulkLoading() {
+        try {
+            GraphDatabase<?, ?> graphDB = null;
+            synchronized (AtlasGraphProvider.class) {
+                Class implClass = AtlasRepositoryConfiguration.getGraphDatabaseImpl();
+                graphDB = (GraphDatabase<?, ?>) implClass.newInstance();
+            }
+
+            return graphDB.getGraphBulkLoading();
+        } catch (IllegalAccessException | InstantiationException e) {
+            throw new RuntimeException("Error initializing graph database", e);
+        }
+    }
+
     @VisibleForTesting
     public static void cleanup() {
         getGraphDatabase().cleanup();

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

@@ -29,4 +29,6 @@ import org.apache.atlas.repository.graphdb.AtlasGraph;
 public interface IAtlasGraphProvider {
     
     AtlasGraph get() throws RepositoryException;
+
+    AtlasGraph getBulkLoading() throws RepositoryException;
 }

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/BulkLoadingTest.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.store.graph.v2;
+
+import org.apache.atlas.TestModules;
+import org.apache.atlas.repository.RepositoryException;
+import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.testng.annotations.Guice;
+import org.testng.annotations.Test;
+
+import javax.inject.Inject;
+
+import static org.testng.Assert.assertNotNull;
+
+@Guice(modules = TestModules.TestOnlyModule.class)
+public class BulkLoadingTest {
+
+    @Inject
+    AtlasGraphProvider graphProvider;
+
+    @Test(expectedExceptions = RuntimeException.class)
+    public void instantiateBulkLoadingForBerkeleyDB() throws RepositoryException {
+        assertNotNull(graphProvider.get());
+        assertNotNull(graphProvider.getBulkLoading());
+    }
+}