Initial code for job scheduler.

1 files changed, 59 insertions, 1 deletions

diff --git a/src/db.py b/src/db.py
index 012dfac..57f4003 100755
--- a/src/db.py
+++ b/src/db.py
@@ -10,8 +10,12 @@
 import os
 import sys
 import time
-import couch
+from contextlib import contextmanager
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
 
+import couch
 from index import CouchIindex
 
 
@@ -137,3 +141,57 @@ class DictDB():
             return None
 
         return key
+
+
+def get_conn_str():
+    try:
+        dialect = os.environ['SQL_DIALECT']
+        database = os.environ['SQL_DATABASE']
+    except KeyError:
+        print('The environment variables SQL_DIALECT and SQL_DATABASE must ' +
+              'be set.')
+        sys.exit(-1)
+
+    if dialect != 'sqlite':
+        try:
+            hostname = os.environ['SQL_HOSTNAME']
+            username = os.environ['SQL_USERNAME']
+            password = os.environ['SQL_PASSWORD']
+        except KeyError:
+            print('The environment variables SQL_DIALECT, SQL_NAME, ' +
+                  'SQL_HOSTNAME, SQL_USERNAME and SQL_PASSWORD must ' +
+                  'be set.')
+            sys.exit(-1)
+
+    if dialect == 'sqlite':
+        conn_str = f"{dialect}:///{database}.db"
+    else:
+        conn_str = f"{dialect}://{username}:{password}@{hostname}" + \
+            "/{database}"
+
+    return conn_str
+
+
+class SqlDB():
+    def get_session(conn_str):
+        if 'sqlite' in conn_str:
+            engine = create_engine(conn_str)
+        else:
+            engine = create_engine(conn_str, pool_size=50, max_overflow=0)
+        Session = sessionmaker(bind=engine)
+
+        return Session()
+
+    @classmethod
+    @contextmanager
+    def sql_session(cls, **kwargs):
+        session = cls.get_session(get_conn_str())
+
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()