Added db caching to greatly improve workflow generation time.

chipathlon/db.py

@@ -30,6 +30,7 @@ class MongoDB(object):
         self.password = password
         self.client = MongoClient(host)
         self.db = self.client.chipseq
+        self.cache = {}
         try:
             self.db.authenticate(username, password, mechanism="SCRAM-SHA-1")
         except:
@@ -39,6 +40,32 @@ class MongoDB(object):
         self.gfs = gridfs.GridFS(self.db)
         return
 
+    def add_cache(self, function, key, data):
+        """
+        :param function: The function name
+        :type function: str
+        :param key: The key to add the cache entry under.
+        :type key: Any hashable
+        :param data: The data to add to the cache.
+        :type data: Object
+        """
+        if function not in self.cache:
+            self.cache[function] = {}
+        self.cache[function][key] = data
+        return
+
+    def get_cache(self, function, key):
+        """
+        :param function: The function name
+        :type function: str
+        :param key: The key to get from the cache.
+        :type key: Any hashable
+        """
+        if function in self.cache:
+            if key in self.cache[function]:
+                return self.cache[function][key]
+        return None
+
     def delete_result(self, result, genome):
         """
         :param result: The result to delete
@@ -279,19 +306,25 @@ class MongoDB(object):
         valid = True
         msg = ""
         data = {}
-        cursor = self.db.samples.find({
-            "accession": accession,
-            "file_type": file_type
-        })
-        if cursor.count() == 1:
-            data = cursor.next()
+        check_cache = self.get_cache("get_sample", accession)
+        if check_cache is not None:
+            msg = "Retrieved data from cache."
+            data = check_cache
         else:
-            valid = False
-            msg = "Found %s files with accession: %s, file_type: %s. Should only be 1." % (
-                cursor.count(),
-                accession,
-                file_type
-            )
+            cursor = self.db.samples.find({
+                "accession": accession,
+                "file_type": file_type
+            })
+            if cursor.count() == 1:
+                data = cursor.next()
+                self.add_cache("get_sample", accession, data)
+            else:
+                valid = False
+                msg = "Found %s files with accession: %s, file_type: %s. Should only be 1." % (
+                    cursor.count(),
+                    accession,
+                    file_type
+                )
         return (valid, msg, data)
 
     def get_samples(self, experiment_accession, file_type):
@@ -307,7 +340,7 @@ class MongoDB(object):
         valid -- Whether or not the accession / file_type combo is a valid exp
         msg -- Why it is or is not valid
         data -- A dictionary containing a list of all control / sample documents.
-        
+
         The data dictionary has two keys, "control" and "signal", each one containing
         a list of all metadata related to the experiment samples.  The sample metadata
         is taken directly from Mongo.