FIX: Ensure exclusive access to dft database

Author: effigies

nibabel/dft.py

@@ -11,6 +11,7 @@
 """
 
 
+import contextlib
 import os
 from os.path import join as pjoin
 import tempfile
@@ -74,7 +75,7 @@ def __getattribute__(self, name):
         val = object.__getattribute__(self, name)
         if name == 'series' and val is None:
             val = []
-            with _db_nochange() as c:
+            with DB.readonly_cursor() as c:
                 c.execute("SELECT * FROM series WHERE study = ?", (self.uid, ))
                 cols = [el[0] for el in c.description]
                 for row in c:
@@ -106,7 +107,7 @@ def __getattribute__(self, name):
         val = object.__getattribute__(self, name)
         if name == 'storage_instances' and val is None:
             val = []
-            with _db_nochange() as c:
+            with DB.readonly_cursor() as c:
                 query = """SELECT *
                              FROM storage_instance
                             WHERE series = ?
@@ -227,7 +228,7 @@ def __init__(self, d):
     def __getattribute__(self, name):
         val = object.__getattribute__(self, name)
         if name == 'files' and val is None:
-            with _db_nochange() as c:
+            with DB.readonly_cursor() as c:
                 query = """SELECT directory, name
                              FROM file
                             WHERE storage_instance = ?
@@ -241,34 +242,6 @@ def dicom(self):
         return pydicom.read_file(self.files[0])
 
 
-class _db_nochange:
-    """context guard for read-only database access"""
-
-    def __enter__(self):
-        self.c = DB.cursor()
-        return self.c
-
-    def __exit__(self, type, value, traceback):
-        if type is None:
-            self.c.close()
-        DB.rollback()
-
-
-class _db_change:
-    """context guard for database access requiring a commit"""
-
-    def __enter__(self):
-        self.c = DB.cursor()
-        return self.c
-
-    def __exit__(self, type, value, traceback):
-        if type is None:
-            self.c.close()
-            DB.commit()
-        else:
-            DB.rollback()
-
-
 def _get_subdirs(base_dir, files_dict=None, followlinks=False):
     dirs = []
     for (dirpath, dirnames, filenames) in os.walk(base_dir, followlinks=followlinks):
@@ -288,7 +261,7 @@ def update_cache(base_dir, followlinks=False):
     for d in dirs:
         os.stat(d)
         mtimes[d] = os.stat(d).st_mtime
-    with _db_nochange() as c:
+    with DB.readwrite_cursor() as c:
         c.execute("SELECT path, mtime FROM directory")
         db_mtimes = dict(c)
         c.execute("SELECT uid FROM study")
@@ -297,7 +270,6 @@ def update_cache(base_dir, followlinks=False):
         series = [row[0] for row in c]
         c.execute("SELECT uid FROM storage_instance")
         storage_instances = [row[0] for row in c]
-    with _db_change() as c:
         for dir in sorted(mtimes.keys()):
             if dir in db_mtimes and mtimes[dir] <= db_mtimes[dir]:
                 continue
@@ -316,7 +288,7 @@ def get_studies(base_dir=None, followlinks=False):
     if base_dir is not None:
         update_cache(base_dir, followlinks)
     if base_dir is None:
-        with _db_nochange() as c:
+        with DB.readonly_cursor() as c:
             c.execute("SELECT * FROM study")
             studies = []
             cols = [el[0] for el in c.description]
@@ -331,7 +303,7 @@ def get_studies(base_dir=None, followlinks=False):
                                WHERE uid IN (SELECT storage_instance
                                                FROM file
                                               WHERE directory = ?))"""
-    with _db_nochange() as c:
+    with DB.readonly_cursor() as c:
         study_uids = {}
         for dir in _get_subdirs(base_dir, followlinks=followlinks):
             c.execute(query, (dir, ))
@@ -443,7 +415,7 @@ def _update_file(c, path, fname, studies, series, storage_instances):
 
 
 def clear_cache():
-    with _db_change() as c:
+    with DB.readwrite_cursor() as c:
         c.execute("DELETE FROM file")
         c.execute("DELETE FROM directory")
         c.execute("DELETE FROM storage_instance")
@@ -478,26 +450,56 @@ def clear_cache():
                           mtime INTEGER NOT NULL,
                           storage_instance TEXT DEFAULT NULL REFERENCES storage_instance,
                           PRIMARY KEY (directory, name))""")
-DB_FNAME = pjoin(tempfile.gettempdir(), f'dft.{getpass.getuser()}.sqlite')
-DB = None
 
 
-def _init_db(verbose=True):
-    """ Initialize database """
-    if verbose:
-        logger.info('db filename: ' + DB_FNAME)
-    global DB
-    DB = sqlite3.connect(DB_FNAME, check_same_thread=False)
-    with _db_change() as c:
-        c.execute("SELECT COUNT(*) FROM sqlite_master WHERE type = 'table'")
-        if c.fetchone()[0] == 0:
-            logger.debug('create')
-            for q in CREATE_QUERIES:
-                c.execute(q)
+class _DB:
+    def __init__(self, verbose=True):
+        self.fname = pjoin(tempfile.gettempdir(), f'dft.{getpass.getuser()}.sqlite')
+        if verbose:
+            logger.info('db filename: ' + self.fname)
+        self.session = sqlite3.connect(self.fname, isolation_level="EXCLUSIVE")
+        try:
+            with self.readwrite_cursor() as c:
+                c.execute("SELECT COUNT(*) FROM sqlite_master WHERE type = 'table'")
+                if c.fetchone()[0] == 0:
+                    logger.debug('create')
+                    for q in CREATE_QUERIES:
+                        c.execute(q)
+        except sqlite3.OperationalError as e:
+            # Race condition, but verify that tables actually got created
+            with self.readonly_cursor() as c:
+                c.execute("SELECT COUNT(*) FROM sqlite_master WHERE type = 'table'")
+                if c.fetchone()[0] != 5:
+                    raise DFTError("Could not construct DFT database") from e
+
+    def __repr__(self):
+        return f"<DFT {self.fname!r}>"
+
+    @contextlib.contextmanager
+    def readonly_cursor(self):
+        cursor = self.session.cursor()
+        try:
+            yield cursor
+        finally:
+            cursor.close()
+            self.session.rollback()
+
+    @contextlib.contextmanager
+    def readwrite_cursor(self):
+        cursor = self.session.cursor()
+        try:
+            yield cursor
+        except Exception:
+            self.session.rollback()
+            raise
+        finally:
+            cursor.close()
+        self.session.commit()
 
 
+DB = None
 if os.name == 'nt':
     warnings.warn('dft needs FUSE which is not available for windows')
 else:
-    _init_db()
+    DB = _DB()
 # eof