separate subgraphs into different neo4js

dataload/scripts/dataload.py

@@ -13,12 +13,8 @@
 for subgraph in config['subgraphs']:
     print(f"===== LOADING SUBGRAPH: {subgraph} =====")
     os.environ['GREBI_SUBGRAPH'] = subgraph
-    res = os.system(f'NXF_WORK=work_{subgraph} nextflow {GREBI_DATALOAD_HOME}/nextflow/01_create_subgraph.nf -c {GREBI_NEXTFLOW_CONFIG} -resume')
+    res = os.system(f'NXF_WORK=work_{subgraph} nextflow {GREBI_DATALOAD_HOME}/nextflow/load_subgraph.nf -c {GREBI_NEXTFLOW_CONFIG} -resume')
     if res != 0:
         exit(res)
     print(f"===== FINISHED LOADING SUBGRAPH: {subgraph} =====")
 
-res = os.system(f'NXF_WORK=work_combined nextflow {GREBI_DATALOAD_HOME}/nextflow/02_create_dbs.nf -c {GREBI_NEXTFLOW_CONFIG} -resume')
-if res != 0:
-    exit(res)
-

webapp/grebi_api/src/main/java/uk/ac/ebi/grebi/GrebiApi.java

@@ -34,17 +34,25 @@ public static void main(String[] args) throws ParseException, org.apache.commons
         Set<String> rocksDbSubgraphs = null;
         Set<String> solrSubgraphs = null;
         Set<String> summarySubgraphs = null;
+        Set<String> neoSubgraphs = null;
 
         while(true) {
             try {
                 neo = new GrebiNeoRepo();
                 solr = new GrebiSolrRepo();
                 summary = new GrebiSummaryRepo();
+                neo = new GrebiNeoRepo();
                 rocksDbSubgraphs = (new ResolverClient()).getSubgraphs();
                 solrSubgraphs = solr.getSubgraphs();
                 summarySubgraphs = summary.getSubgraphs();
-                if(new HashSet<>(List.of(rocksDbSubgraphs, solrSubgraphs, summarySubgraphs)).size() != 1) {
-                    throw new RuntimeException("RocksDB/Solr/the summary jsons do not seem to contain the same subgraphs. Found: " + String.join(",", rocksDbSubgraphs) + " for RocksDB (from resolver service) and " + String.join(",", solrSubgraphs) + " for Solr (from list of solr cores) and " + String.join(",", summarySubgraphs) + " for the summary jsons (from summary server)");
+                neoSubgraphs = neo.getSubgraphs();
+                if(new HashSet<>(List.of(rocksDbSubgraphs, solrSubgraphs, summarySubgraphs, neoSubgraphs)).size() != 1) {
+                    throw new RuntimeException("RocksDB/Solr/the summary jsons/Neo4j do not seem to contain the same subgraphs. Found: "
+                            + String.join(",", rocksDbSubgraphs) + " for RocksDB (from resolver service) and "
+                            + String.join(",", solrSubgraphs) + " for Solr (from list of solr cores) and "
+                            + String.join(",", summarySubgraphs) + " for the summary jsons (from summary server) and"
+                            + String.join(",", neoSubgraphs) + " for Neo4j (from all neo4j hosts)"
+                    );
                 }
                 break;
             } catch(Exception e) {

webapp/grebi_api/src/main/java/uk/ac/ebi/grebi/db/Neo4jClient.java

@@ -18,26 +18,15 @@
 
 public class Neo4jClient {
 
-    static final String NEO4J_HOST = System.getenv("GREBI_NEO4J_HOST");
-
-    public static String getNeo4jHost() {
-        if(NEO4J_HOST != null)
-            return NEO4J_HOST;
-        return "bolt://localhost:7687/";
-    }
-
+    private final Driver driver;
     private Gson gson = new Gson();
 
-    private Driver driver;
+    public Neo4jClient(String host) {
+        this.driver = GraphDatabase.driver(host);
+    }
 
     public Driver getDriver() {
-
-        if(driver == null) {
-            driver = GraphDatabase.driver(getNeo4jHost());
-        }
-
         return driver;
-
     }
 
     public Session getSession() {
@@ -52,7 +41,6 @@ public List<Map<String,Object>> rawQuery(String query) {
         Session session = getSession();
 
         Result result = session.run(query);
-
         List<Map<String,Object>> list = result.stream().map(r -> r.asMap()).collect(Collectors.toList());
 
         session.close();

webapp/grebi_api/src/main/java/uk/ac/ebi/grebi/repo/GrebiNeoRepo.java

@@ -12,31 +12,60 @@
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.TreeMap;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 
 public class GrebiNeoRepo {
 
-    Neo4jClient neo4jClient = new Neo4jClient();
+    static final String[] NEO4J_HOSTS =
+            System.getenv("GREBI_NEO4J_HOSTS").split(";");
+
+    public static String[] getNeo4jHosts() {
+        if(NEO4J_HOSTS != null)
+            return NEO4J_HOSTS;
+        return List.of("bolt://localhost:7687/").toArray(new String[0]);
+    }
+
+    Map<String, Neo4jClient> subgraphToClient = new HashMap<>();
+
     ResolverClient resolver = new ResolverClient();
     Gson gson = new Gson();
 
-    public GrebiNeoRepo() throws IOException {}
+    public GrebiNeoRepo() throws IOException {
+
+        for(String host : getNeo4jHosts()) {
+            Neo4jClient client = new Neo4jClient(host);
+
+            String subgraph = (String)
+                    client.rawQuery("MATCH (n:GraphNode) RETURN n.`grebi:subgraph` AS subgraph LIMIT 1")
+                        .get(0).get("subgraph");
+
+            subgraphToClient.put(subgraph, client);
+        }
+    }
+
+    private Neo4jClient getClient(String subgraph) {
+        var client = subgraphToClient.get(subgraph);
+        if(client != null)
+            return client;
+        throw new IllegalArgumentException("subgraph " + subgraph + " not found");
+    }
+
+    public Set<String> getSubgraphs() {
+        return subgraphToClient.keySet();
+    }
 
     final String STATS_QUERY = new String(GrebiApi.class.getResourceAsStream("/cypher/stats.cypher").readAllBytes(), StandardCharsets.UTF_8);
     final String INCOMING_EDGES_QUERY = new String(GrebiApi.class.getResourceAsStream("/cypher/incoming_edges.cypher").readAllBytes(), StandardCharsets.UTF_8);
 
-    public Map<String,Object> getStats() {
-        EagerResult props_res = neo4jClient.getDriver().executableQuery(STATS_QUERY).withConfig(QueryConfig.builder().withDatabase("neo4j").build()).execute();
-        return props_res.records().get(0).values().get(0).asMap();
-    }
-    public List<Record> cypher(String query, String resVar) {
-        EagerResult res = neo4jClient.getDriver().executableQuery(query).withConfig(QueryConfig.builder().withDatabase("neo4j").build()).execute();
-        return List.of();
+    public Map<String, Map<String,Object>> getStats() {
+        Map<String, Map<String,Object>> subgraphToStats = new HashMap<>();
+        for(var subgraph : subgraphToClient.keySet()) {
+            EagerResult props_res = getClient(subgraph).getDriver().executableQuery(STATS_QUERY).withConfig(QueryConfig.builder().withDatabase("neo4j").build()).execute();
+            subgraphToStats.put(subgraph, props_res.records().get(0).values().get(0).asMap());
+        }
+        return subgraphToStats;
     }
 
     public class EdgeAndNode {
@@ -48,7 +77,7 @@ public EdgeAndNode(Map<String,Object> edge, Map<String,Object> node) {
     }
 
     public List<EdgeAndNode> getIncomingEdges(String subgraph, String nodeId, Pageable pageable) {
-        EagerResult res = neo4jClient.getDriver().executableQuery(INCOMING_EDGES_QUERY)
+        EagerResult res = getClient(subgraph).getDriver().executableQuery(INCOMING_EDGES_QUERY)
             .withParameters(Map.of(
                     "nodeId", subgraph + ":" + nodeId,
                     "offset", pageable.getOffset(),