Ongoing transition to DB backend (SQLite)

nanocompore/DataStore.py

@@ -1,16 +1,25 @@
 # -*- coding: utf-8 -*-
 
-from collections import *
+from enum import Enum
 import datetime
 import os
-import sqlite3 as lite
+import sqlite3
+import contextlib
 
 # Third party
 from loguru import logger
-import nanocompore as pkg
+from nanocompore.common import NanocomporeError
+
+
+class DBCreateMode(Enum):
+    """Options for handling (non-) existence of the SQLite database file"""
+    MUST_EXIST = "r" # open for reading, error if file doesn't exist
+    CREATE_MAYBE = "a" # use an existing database, otherwise create one
+    OVERWRITE = "w" # always create a new database, overwrite if it exists
+
 
 class DataStore(object):
-    """ Init analysis and check args"""
+    """Store Nanocompore data in an SQLite database"""
 
     create_reads_query = ("CREATE TABLE IF NOT EXISTS reads ("
                           "id INTEGER NOT NULL PRIMARY KEY,"
@@ -22,9 +31,13 @@ class DataStore(object):
                           "numevents INT NOT NULL,"
                           "numsignals INT NOT NULL,"
                           "dwelltime REAL NOT NULL,"
+                          "kmers INT NOT NULL,"
+                          "missing_kmers INT NOT NULL,"
+                          "NNNNN_kmers INT NOT NULL,"
+                          "mismatch_kmers INT NOT NULL,"
+                          "valid_kmers INT NOT NULL,"
                           "FOREIGN KEY(sampleid) REFERENCES samples(id)"
-                          "FOREIGN KEY(transcriptid) REFERENCES transcripts(id),"
-                          "UNIQUE(id, name)"
+                          "FOREIGN KEY(transcriptid) REFERENCES transcripts(id)"
                           ")"
                           )
 
@@ -44,67 +57,69 @@ class DataStore(object):
                           "FOREIGN KEY(readid) REFERENCES reads(id)"
                           ")"
                           )
+    # TODO: 'sequence' is stored redundantly - move it to a separate table
+    # TODO: encode 'status' as int to save space (foreign key referencing a table with all possible statuses)
 
     create_samples_query = ("CREATE TABLE IF NOT EXISTS samples ("
                             "id INTEGER NOT NULL PRIMARY KEY,"
                             "name VARCHAR NOT NULL UNIQUE"
                             ")"
                             )
-
+    # TODO: add 'condition' column
 
     create_transcripts_query = ("CREATE TABLE IF NOT EXISTS transcripts ("
                                 "id INTEGER NOT NULL PRIMARY KEY,"
                                 "name VARCHAR NOT NULL UNIQUE"
                                 ")"
                                 )
 
-
-
-    def __init__(self, db_path:str):
-        self.__db_path=db_path
-        db_is_new = not os.path.exists(self.__db_path)
-        logger.debug(f"DB file doesn't exist: {db_is_new}")
-        if db_is_new: self.__init_db()
+    def __init__(self,
+                 db_path:str,
+                 create_mode=DBCreateMode.MUST_EXIST):
+        self.__db_path = db_path
+        self.__create_mode = create_mode
+        self.__connection = None
+        self.__cursor = None
 
     def __enter__(self):
-        self.__open_db_connection()
-        return self
-
-    def __exit__(self,exc_type, exc_value, traceback):
-        self.__connection.commit()
-        self.__close_db_connection()
-
-    def __open_db_connection(self):
+        if self.__create_mode == DBCreateMode.MUST_EXIST and not os.path.exists(self.__db_path):
+            raise NanocomporeError(f"Database file '{self.__db_path}' does not exist")
+        if self.__create_mode == DBCreateMode.OVERWRITE:
+            with contextlib.suppress(FileNotFoundError): # file may not exist
+                os.remove(self.__db_path)
+                logger.debug(f"Removed existing database file '{self.__db_path}'")
         try:
-            logger.debug("Connecting to DB")
-            self.__connection = lite.connect(self.__db_path);
+            logger.debug("Connecting to database")
+            self.__connection = sqlite3.connect(self.__db_path)
+            self.__connection.row_factory = sqlite3.Row
             self.__cursor = self.__connection.cursor()
         except:
             logger.error("Error connecting to database")
             raise
+        if self.__create_mode == DBCreateMode.OVERWRITE or \
+           (self.__create_mode == DBCreateMode.CREATE_MAYBE and not os.path.exists(self.__db_path)):
+            self.__init_db()
+        return self
 
-    def __close_db_connection(self):
+    def __exit__(self, exc_type, exc_value, traceback):
         if self.__connection:
-            logger.debug("Closing connection to DB")
+            logger.debug("Closing database connection")
             self.__connection.commit()
             self.__connection.close()
             self.__connection = None
             self.__cursor = None
 
     def __init_db(self):
-        logger.debug("Setting up DB tables")
-        self.__open_db_connection()
+        logger.debug("Setting up database tables")
         try:
             self.__cursor.execute(self.create_reads_query)
             self.__cursor.execute(self.create_kmers_query)
             self.__cursor.execute(self.create_samples_query)
             self.__cursor.execute(self.create_transcripts_query)
+            self.__connection.commit()
         except:
-            self.__close_db_connection()
-            logger.error("Error creating tables")
+            logger.error("Error creating database tables")
             raise
-        self.__connection.commit()
-        self.__close_db_connection()
 
     def store_read(self, read):
         """
@@ -116,13 +131,14 @@ def store_read(self, read):
         """
         tx_id = self.get_transcript_id_by_name(read.ref_id, create_if_not_exists=True)
         sample_id = self.get_sample_id_by_name(read.sample_name, create_if_not_exists=True)
+        values = (read.read_id, sample_id, tx_id, read.ref_start, read.ref_end,
+                  read.n_events, read.n_signals, read.dwell_time) + tuple(read.kmers_status.values())
         try:
-            self.__cursor.execute("INSERT INTO reads VALUES(NULL, ?, ?, ?, ?, ?, ?, ?, ?)",
-                                    (read.read_id, sample_id, tx_id, read.ref_start, read.ref_end,
-                                     read.n_events, read.n_signals, read.dwell_time))
+            self.__cursor.execute("INSERT INTO reads VALUES(NULL" + ", ?" * len(values) + ")",
+                                  values)
             read_id = self.__cursor.lastrowid
         except Exception:
-            logger.error("Error inserting read into DB")
+            logger.error("Error inserting read into database")
             raise Exception
 
         for kmer in read.kmer_l:
@@ -144,7 +160,7 @@ def __store_kmer(self, kmer, read_id):
                                res["num_signals"], res["status"], res["dwell_time"],
                                res["NNNNN_dwell_time"], res["mismatch_dwell_time"], res["median"], res["mad"]))
         except Exception:
-            logger.error("Error inserting kmer into DB")
+            logger.error("Error inserting kmer into database")
             raise Exception
 
     def get_transcript_id_by_name(self, tx_name, create_if_not_exists=False):
@@ -161,7 +177,7 @@ def get_transcript_id_by_name(self, tx_name, create_if_not_exists=False):
             try:
                 self.__cursor.execute(query)
             except Exception:
-                logger.error("There was an error while inserting a new transcript in the DB")
+                logger.error("Error while inserting transcript into the database")
                 raise Exception
 
         query = f"SELECT id from transcripts WHERE name = '{tx_name}'"
@@ -170,7 +186,7 @@ def get_transcript_id_by_name(self, tx_name, create_if_not_exists=False):
             record = self.__cursor.fetchone()
             self.__connection.commit()
         except Exception:
-            logger.error("There was an error while selecting the transcript_id from the DB")
+            logger.error("Error while selecting transcript ID from the database")
             raise Exception
         if record is not None:
             return record[0]
@@ -191,7 +207,7 @@ def get_sample_id_by_name(self, sample_name, create_if_not_exists=False):
             try:
                 self.__cursor.execute(query)
             except Exception:
-                logger.error("There was an error while inserting a new sample in the DB")
+                logger.error("Error while inserting sample into the database")
                 raise Exception
 
         query = f"SELECT id from samples WHERE name = '{sample_name}'"
@@ -200,9 +216,38 @@ def get_sample_id_by_name(self, sample_name, create_if_not_exists=False):
             record = self.__cursor.fetchone()
             self.__connection.commit()
         except Exception:
-            logger.error("There was an error while selecting the sample_id from the DB")
+            logger.error("Error while selecting sample ID from the database")
             raise Exception
         if record is not None:
             return record[0]
         else:
             return None
+
+    @property
+    def cursor(self):
+        return self.__cursor
+
+    def get_samples(self, sample_dict=None):
+        if not self.__connection:
+            raise NanocomporeError("Database connection not yet opened")
+        expected_samples = []
+        if sample_dict: # query only relevant samples
+            for samples in sample_dict.values():
+                expected_samples += samples
+            if not expected_samples:
+                raise NanocomporeError("No sample names in 'sample_dict'")
+            where = " WHERE name IN ('%s')" % "', '".join(expected_samples)
+        else:
+            where = ""
+        db_samples = {}
+        try:
+            self.__cursor.execute("SELECT * FROM samples" + where)
+            for row in self.__cursor:
+                db_samples[row["id"]] = row["name"]
+        except Exception:
+            logger.error("Error reading sample names from database")
+            raise Exception
+        for sample in expected_samples: # check that requested samples are in DB
+            if sample not in db_samples.values():
+                raise NanocomporeError(f"Sample '{sample}' not present in database")
+        return db_samples