Chore/dat 288 optimize esbuild

esbuild/gdc_elasticsearch.py

@@ -19,6 +19,7 @@
 from gdcdatamodel.models import File
 from progressbar import ProgressBar, Percentage, Bar, ETA
 from psqlgraph import PsqlGraphDriver
+from esbuild.graph.common.cache import CachedGraph
 
 # TODO this could probably be bumped now that the number of bulk
 # threads in the config is higher, c.f.
@@ -44,6 +45,7 @@ def shouldnt_delete(node):
     delete them.
 
     """
+
     if isinstance(node, File) and node.derived_files:
         return True
     else:
@@ -55,8 +57,7 @@ class GDCElasticsearch(object):
     """
     """
 
-    def __init__(self, converter_class, es=None,
-                 index_base="gdc_from_graph"):
+    def __init__(self, converter_class, es=None, index_base="gdc_from_graph"):
         """Walks the graph to produce elasticsearch json documents.
 
         :param es: An instance of Elasticsearch class
@@ -75,22 +76,28 @@ def __init__(self, converter_class, es=None,
                            os.environ.get("ES_PASSWORD", "")),
                 timeout=9999)
 
-        self.graph = PsqlGraphDriver(
+        self.psqlgraph_args = (
             os.environ["PG_HOST"],
             os.environ["PG_USER"],
             os.environ["PG_PASS"],
             os.environ["PG_NAME"],
         )
 
-        self.converter = converter_class(self.graph)
+        self.graph = PsqlGraphDriver(self.psqlgraph_args)
+
+        caching_options = converter_class.get_caching_options()
+        cache = CachedGraph(self.graph, caching_options)
+        converter = converter_class(cache)
+
+        self.converter = converter
 
     def go(self, roll_alias=True):
         self.log.info("Caching database")
         # having a transation out here is important, since it ensures
         # that the cached database and which nodes get deleted is
         # consistent
         with self.graph.session_scope() as session:
-            self.converter.cache_database()
+            self.converter.cache.cache_database()
             self.log.info("Querying for old nodes to delete")
             to_delete = self.graph.nodes().sysan({"to_delete": True}).all()
             to_delete = [n for n in to_delete if not shouldnt_delete(n)]

esbuild/graph/active/builder.py

@@ -214,7 +214,7 @@ def get_parent_with_category(self, node, category):
             if l['dst_type']._dictionary['category'] == category
         ]
 
-        return self.neighbors_labeled(node, labels)
+        return self.cache.neighbors_labeled(node.node_id, labels)
 
     def get_child_with_category(self, node, category):
         """returns iterable of neighors from inbound edges with category"""
@@ -224,7 +224,7 @@ def get_child_with_category(self, node, category):
             if l['src_type']._dictionary['category'] == category
         ]
 
-        return self.neighbors_labeled(node, labels)
+        return self.cache.neighbors_labeled(node.node_id, labels)
 
     def add_file_analysis(self, node, doc):
         """Add the 'analysis' that produced the current file"""
@@ -315,7 +315,8 @@ def get_read_group_qc_docs(self, read_group):
         """Returns a list of documents for Read Group QCs"""
 
         read_group_qc_docs = []
-        rg_qcs = self.neighbors_labeled(read_group, 'read_group_qc')
+        rg_qcs = self.cache.neighbors_labeled(
+            read_group.node_id, 'read_group_qc')
         for read_group_qc in rg_qcs:
             read_group_qc_docs.append(self._get_base_doc(read_group_qc))
 
@@ -329,7 +330,7 @@ def get_file_read_groups(self, node):
         """
 
         paths = self.file_to_read_group_paths.get(node.label, [])
-        return set(self.walk_paths(node, paths))
+        return set(self.cache.walk_paths(node.node_id, paths))
 
     def get_analysis_read_groups(self, node):
         """Given a analysis node, traverse up the tree to read_groups:
@@ -356,7 +357,7 @@ def get_simple_file_doc(self, node):
 
         doc['data_format'] = self.get_data_format(node)
 
-        for dst in self.neighbors_labeled(node, 'data_subtype'):
+        for dst in self.cache.neighbors_labeled(node.node_id, 'data_subtype'):
             doc['data_type'] = dst['name']
 
         return doc
@@ -372,15 +373,17 @@ def get_file_associated_entities(self, node):
             entity
             for rg in self.get_file_read_groups(node)
             for entity in
-            self.neighbors_labeled(rg, self.possible_associated_entites)
+            self.cache.neighbors_labeled(
+                rg.node_id, self.possible_associated_entites)
         ]
 
         # Add entities with one step through a data_file
         entities += [
             entity
             for parent in self.get_parent_with_category(node, 'data_file')
             for entity in
-            self.neighbors_labeled(parent, self.possible_associated_entites)
+            self.cache.neighbors_labeled(
+                parent.node_id, self.possible_associated_entites)
         ]
 
         # Copy number paths
@@ -389,7 +392,8 @@ def get_file_associated_entities(self, node):
             list_product([['aliquot']], self.aliquot_to_copy_number_paths)
         ]
         entities += [
-            entity for entity in self.walk_paths(node, cnv_paths)
+            entity for entity in self.cache.walk_paths(
+                node.node_id, cnv_paths)
         ]
 
         return list(set(entities))