Optimized db.py for future use

moonstreamdb-v3/moonstreamdbv3/db.py

@@ -13,46 +13,49 @@ from sqlalchemy.orm import Session, sessionmaker
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# try:
-#     MOONSTREAM_DB_URI = os.environ.get("MOONSTREAM_DB_URI")
-#     if MOONSTREAM_DB_URI is None:
-#         raise Warning("MOONSTREAM_DB_URI environment variable must be set")
+try:
+    MOONSTREAM_DB_URI = os.environ.get("MOONSTREAM_DB_URI")
+    if MOONSTREAM_DB_URI is None:
+        raise Warning("MOONSTREAM_DB_URI environment variable must be set")
 
-#     MOONSTREAM_DB_URI_READ_ONLY = os.environ.get("MOONSTREAM_DB_URI_READ_ONLY")
-#     if MOONSTREAM_DB_URI_READ_ONLY is None:
-#         raise Warning("MOONSTREAM_DB_URI_READ_ONLY environment variable must be set")
+    MOONSTREAM_DB_URI_READ_ONLY = os.environ.get("MOONSTREAM_DB_URI_READ_ONLY")
+    if MOONSTREAM_DB_URI_READ_ONLY is None:
+        raise Warning("MOONSTREAM_DB_URI_READ_ONLY environment variable must be set")
 
-#     MOONSTREAM_POOL_SIZE_RAW = os.environ.get("MOONSTREAM_POOL_SIZE")
-#     MOONSTREAM_POOL_SIZE = 1
-#     try:
-#         if MOONSTREAM_POOL_SIZE_RAW is not None:
-#             MOONSTREAM_POOL_SIZE = int(MOONSTREAM_POOL_SIZE_RAW)
-#     except:
-#         raise ValueError(
-#             f"Could not parse MOONSTREAM_POOL_SIZE as int: {MOONSTREAM_POOL_SIZE_RAW}"
-#         )
+    MOONSTREAM_POOL_SIZE_RAW = os.environ.get("MOONSTREAM_POOL_SIZE")
+    MOONSTREAM_POOL_SIZE = 1
+    try:
+        if MOONSTREAM_POOL_SIZE_RAW is not None:
+            MOONSTREAM_POOL_SIZE = int(MOONSTREAM_POOL_SIZE_RAW)
+    except:
+        raise ValueError(
+            f"Could not parse MOONSTREAM_POOL_SIZE as int: {MOONSTREAM_POOL_SIZE_RAW}"
+        )
 
-#     MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW = os.environ.get(
-#         "MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS"
-#     )
-#     MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS = 30000
-#     try:
-#         if MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW is not None:
-#             MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS = int(
-#                 MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW
-#             )
-#     except:
-#         raise ValueError(
-#             f"MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS must be an integer: {MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW}"
-#         )
-# except ValueError as e:
-#     raise ValueError(e)
-# except Warning:
-#     logger.warning("Database variables not set")
+    MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW = os.environ.get(
+        "MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS"
+    )
+    MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS = 30000
+    try:
+        if MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW is not None:
+            MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS = int(
+                MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW
+            )
+    except:
+        raise ValueError(
+            f"MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS must be an integer: {MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS_RAW}"
+        )
+except ValueError as e:
+    raise ValueError(e)
+except Warning:
+    logger.warning("Database variables not set")
 
 
 def create_moonstream_engine(
-    url: str, pool_size: int, statement_timeout: int, pool_pre_ping: bool = False
+    url: str,
+    pool_size: int,
+    statement_timeout: int,
+    pool_pre_ping: bool = False,
 ):
     # Pooling: https://docs.sqlalchemy.org/en/14/core/pooling.html#sqlalchemy.pool.QueuePool
     # Statement timeout: https://stackoverflow.com/a/44936982
@@ -63,49 +66,76 @@ def create_moonstream_engine(
         connect_args={"options": f"-c statement_timeout={statement_timeout}"},
     )
 
-# engine = create_moonstream_engine(
-#     url=MOONSTREAM_DB_URI,  # type: ignore
-#     pool_size=MOONSTREAM_POOL_SIZE,
-#     statement_timeout=MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS,
-# )
-# SessionLocal = sessionmaker(bind=engine)
+
+class MoonstreamDBEngine:
+    def __init__(self) -> None:
+        self._engine = create_moonstream_engine(
+            url=MOONSTREAM_DB_URI,  # type: ignore
+            pool_size=MOONSTREAM_POOL_SIZE,
+            statement_timeout=MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS,
+        )
+        self._session_local = sessionmaker(bind=self.engine)
+
+        self._yield_db_session_ctx = contextmanager(self.yield_db_session)
+
+    @property
+    def engine(self):
+        return self._engine
+
+    @property
+    def session_local(self):
+        return self._session_local
+
+    @property
+    def yield_db_session_ctx(self):
+        return self._yield_db_session_ctx
+
+    def yield_db_session(
+        self,
+    ) -> Generator[Session, None, None]:
+        """
+        Yields a database connection (created using environment variables).
+        As per FastAPI docs:
+        https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-dependency
+        """
+        session = self._session_local()
+        try:
+            yield session
+        finally:
+            session.close()
 
 
-# def yield_db_session() -> Generator[Session, None, None]:
-#     """
-#     Yields a database connection (created using environment variables).
-#     As per FastAPI docs:
-#     https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-dependency
-#     """
-#     session = SessionLocal()
-#     try:
-#         yield session
-#     finally:
-#         session.close()
+class MoonstreamDBEngineRO:
+    def __init__(self) -> None:
+        self._RO_engine = create_moonstream_engine(
+            url=MOONSTREAM_DB_URI_READ_ONLY,  # type: ignore
+            pool_size=MOONSTREAM_POOL_SIZE,
+            statement_timeout=MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS,
+        )
+        self._RO_session_local = sessionmaker(bind=self.RO_engine)
 
+        self._RO_yield_db_session_ctx = contextmanager(self.yield_db_read_only_session)
 
-# yield_db_session_ctx = contextmanager(yield_db_session)
+    @property
+    def RO_engine(self):
+        return self._RO_engine
 
-# # Read only
-# RO_engine = create_moonstream_engine(
-#     url=MOONSTREAM_DB_URI_READ_ONLY,  # type: ignore
-#     pool_size=MOONSTREAM_POOL_SIZE,
-#     statement_timeout=MOONSTREAM_DB_STATEMENT_TIMEOUT_MILLIS,
-# )
-# RO_SessionLocal = sessionmaker(bind=RO_engine)
+    @property
+    def RO_session_local(self):
+        return self._RO_session_local
 
+    @property
+    def RO_yield_db_session_ctx(self):
+        return self._RO_yield_db_session_ctx
 
-# def yield_db_read_only_session() -> Generator[Session, None, None]:
-#     """
-#     Yields a database connection (created using environment variables).
-#     As per FastAPI docs:
-#     https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-dependency
-#     """
-#     session = RO_SessionLocal()
-#     try:
-#         yield session
-#     finally:
-#         session.close()
-
-
-# yield_db_read_only_session_ctx = contextmanager(yield_db_read_only_session)
+    def yield_db_read_only_session(self) -> Generator[Session, None, None]:
+        """
+        Yields read-only database connection (created using environment variables).
+        As per FastAPI docs:
+        https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-dependency
+        """
+        session = self._RO_session_local()
+        try:
+            yield session
+        finally:
+            session.close()