Implemented metadatadb as a context object instead of a class.

lark

@@ -91,7 +91,7 @@ def _kwargs_parse(kwargs_list):
 action_object = sys.argv[1]
 action = sys.argv[2]
 
-db = metadatadb.MetadataDB(os.path.join(data_path, SQLITE_FILENAME))
+metadatadb.configure(os.path.join(data_path, SQLITE_FILENAME))
 
 # TODO: Use a real UI library. This mess is just intended for development.
 if action_object == 'platform':
@@ -102,21 +102,25 @@ if action_object == 'platform':
 
         platform_data = metadatadb.Platform(shortcode=platform_shortcode,
                                         fullname=platform_name)
-        db.add_platform(platform_data)
+        with metadatadb.get_db_session() as session:
+            session.add(platform_data)
 
     elif action == 'list':
         # TODO: Filter support is exclusively limited to SQLAlchemy's filter.ilike function. Figure
         # out a good way to include other filters.
         filters = _kwargs_parse(sys.argv[3:])
-        platform_results = db.search_platforms(**filters)
-        print(platform_results)
+        with metadatadb.get_db_session() as session:
+            print(metadatadb.search(session, metadatadb.Platform, **filters))
 
     elif action == 'remove':
         constraints = sys.argv[3:]
-        platforms = db.search_platforms(constraints)
-        for platform in platforms:
-            print('Removing %s.' % platform.fullname)
-            db.remove_platform(platform)
+        filters = _kwargs_parse(sys.argv[3:])
+
+        with metadatadb.get_db_session() as session:
+            platforms = metadatadb.search(session, metadatadb.Platform, **filters)
+            for platform in platforms:
+                print('Removing %s.' % platform.fullname)
+                session.delete(platform)
 
     elif action == 'test':
         # TODO: Delete this action before merging into dev. It's just for ugly testing.
@@ -126,7 +130,7 @@ if action_object == 'platform':
         platform_data = metadatadb.Platform(shortcode=platform_shortcode,
                                             fullname=platform_name)
 
-        #db.add_platform(platform_data)
-        print(db.search_platforms())
+        with metadatadb.get_db_session() as session:
+            print(metadatadb.search(session, metadatadb.Platform))
 else:
     print('Unknown object.')

metadatadb.py

@@ -1,10 +1,9 @@
 '''
 metadatadb.py
 '''
+import contextlib
 import collections
-import dataclasses
 import hashlib
-import sqlite3
 import uuid
 
 import sqlalchemy
@@ -14,6 +13,7 @@ import sqlalchemy.orm
 # TODO: l o g g i n g
 HASH_CHUNK_SIZE = 10485760 # 10mb
 _db_session_maker = sqlalchemy.orm.sessionmaker()
+_engine = None
 
 ImageData = collections.namedtuple('ImageData', 'UUID, sha1sum, filename, release_group')
 ReleaseGroupData = collections.namedtuple('ReleaseGroupData', 'UUID, name, platform')
@@ -49,72 +49,34 @@ def get_file_sha1sum(filename):
 
     return sha1sum.hexdigest()
 
-class MetadataDB:
-    def __init__(self, db_path):
-        '''
-        If db file does not exist, create it and create necessary tables.
-        Either way, create a connection object.
-        '''
-        # TODO: This process needs real error handling.
-        # TODO: Add DAT import/credit support.
-        self._engine = sqlalchemy.create_engine('sqlite:///%s' % db_path)
+def configure(db_path):
+    _engine = sqlalchemy.create_engine('sqlite:///%s' % db_path)
 
-        _SQLBase.metadata.create_all(self._engine)
-        _db_session_maker.configure(bind=self._engine)
-        # TODO: Using One Big Session may have unintended consequences in other, less linear
-        # applications. For now, it works.
-        self._session = _db_session_maker()
+    _SQLBase.metadata.create_all(_engine)
+    _db_session_maker.configure(bind=_engine)
 
-    def add_image(self, image_data):
-        pass
+def search(session, table_object, inclusive=True, **constraints):
+    query = session.query(table_object)
 
-    def update_image(self, image_data):
-        pass
+    for key, value in constraints.items():
+        query = query.filter(getattr(table_object, key).ilike('%%%s%%' % value))
 
-    def remove_image(self, image_data):
-        pass
+    item_list = []
+    for item in query:
+        item_list.append(item)
 
-    def add_release_group(self, release_group_data):
-        pass
+    return item_list
 
-    def update_release_group(self, release_group_data):
-        pass
+@contextlib.contextmanager
+def get_db_session():
+    session = _db_session_maker()
+    try:
+        yield session
 
-    def remove_release_group(self, release_group_data):
-        pass
-
-    def add_platform(self, platform):
-        ''' Add a platform shortcode to the database. '''
-        self._session.add(platform)
-        self._session.commit()
-
-    def update_platform(self, platform):
-        pass
-
-    def remove_platform(self, platform):
-        '''Remove a specific platform from the database. '''
-        self._session.delete(platform)
-        self._session.commit()
-
-    def search_platforms(self, inclusive=True, **constraints):
-        '''Search for platforms, given the parameters. '''
-
-        query = self._session.query(Platform)
-
-        for key, value in constraints.items():
-            query = query.filter(getattr(Platform, key).ilike('%%%s%%' % value))
-
-        platform_list = []
-        for platform in query:
-            platform_list.append(platform)
-
-        return platform_list
-
-    def add_dat(self, dat_data):
-        pass
-
-    def update_dat(self, dat_data):
-        pass
-
-    def remove_dat(self, dat_data):
-        pass
+    except:
+        session.rollback()
+        raise
+    else:
+        session.commit()
+    finally:
+        session.close()