New environment db

light/data_model/db/base.py

@@ -7,9 +7,30 @@
 
 from abc import ABC, abstractmethod
 from omegaconf import MISSING, DictConfig
-from contextlib import contextmanager
+from sqlalchemy import create_engine
 from enum import Enum
 from typing import Optional, Union, Dict, Any
+from dataclasses import dataclass
+from tempfile import mkdtemp
+import shutil
+import os
+import json
+
+from hydra.core.config_store import ConfigStore
+
+DEFAULT_LOG_PATH = "".join(
+    [os.path.abspath(os.path.dirname(__file__)), "/../../../logs"]
+)
+
+
+@dataclass
+class LightDBConfig:
+    backend: str = "test"
+    file_root: Optional[str] = DEFAULT_LOG_PATH
+
+
+cs = ConfigStore.instance()
+cs.store(name="config1", node=LightDBConfig)
 
 
 class DBStatus(Enum):
@@ -48,6 +69,18 @@ def __init__(self, config: "DictConfig"):
         Create this database, either connecting to a remote host or local
         files and instances.
         """
+        # TODO replace with a swappable engine that persists the data
+        self.backend = config.backend
+        if config.backend == "test":
+            self.engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
+            self.made_temp_dir = config.file_root is None
+            if self.made_temp_dir:
+                self.file_root = mkdtemp()
+            else:
+                self.file_root = config.file_root
+        else:
+            raise NotImplementedError()
+        self._complete_init(config)
 
     @abstractmethod
     def _complete_init(self, config: "DictConfig"):
@@ -61,20 +94,60 @@ def _validate_init(self):
         Ensure that this database is initialized correctly
         """
 
+    def _enforce_get_first(self, session, stmt, error_text) -> Any:
+        """
+        Enforce getting the first element using stmt, raise a key_error
+        with error_text if fails to find
+        """
+        result = session.scalars(stmt).first()
+        if result is None:
+            raise KeyError(error_text)
+        return result
+
+    def file_path_exists(self, file_path: str) -> bool:
+        """
+        Determine if the given file path exists on this storage
+        """
+        if self.backend in ["test", "local"]:
+            full_path = os.path.join(self.file_root, file_path)
+            return os.path.exists(full_path)
+        else:
+            raise NotImplementedError
+
     def write_data_to_file(
         self, data: Union[str, Dict[str, Any]], filename: str, json_encode: bool = False
-    ):
+    ) -> None:
         """
         Write the given data to the provided filename
         in the correct storage location (local or remote)
         """
-        # Expects data to be a string, unless json_encode is True
+        if self.backend in ["test", "local"]:
+            full_path = os.path.join(self.file_root, filename)
+            os.makedirs(os.path.dirname(full_path), exist_ok=True)
+            with open(full_path, "w+") as target_file:
+                if json_encode:
+                    json.dump(data, target_file)
+                else:
+                    target_file.write(data)
+        else:
+            raise NotImplementedError()
 
-    def read_data_from_file(self, filename: str, json_encoded: bool = False):
+    def read_data_from_file(
+        self, filename: str, json_encoded: bool = False
+    ) -> Union[str, Dict[str, Any]]:
         """
         Read the data from the given filename from wherever it
         is currently stored (local or remote)
         """
+        if self.backend in ["test", "local"]:
+            full_path = os.path.join(self.file_root, filename)
+            with open(full_path, "r") as target_file:
+                if json_encoded:
+                    return json.load(target_file)
+                else:
+                    return target_file.read()
+        else:
+            raise NotImplementedError()
 
     def open_file(self):
         try:
@@ -83,13 +156,7 @@ def open_file(self):
         finally:
             file.close()
 
-    @contextmanager
-    def get_database_connection(self):
-        """Get a connection to the database that can be used for a transaction"""
-        try:
-            # Get DB connection
-            # yield db connection
-            raise NotImplementedError()
-        finally:
-            # close DB connection, rollback if there's an issue
-            pass
+    def shutdown(self):
+        if self.backend == "test":
+            if self.made_temp_dir:
+                shutil.rmtree(self.file_root)