🎨 Reorganize back-end code

2023-09-20 22:08:32 +02:00
parent bdbc72ab39
commit 3434802b31
28 changed files with 29 additions and 36 deletions
--- a/backend/api/db/db.py
+++ b/backend/api/db/db.py
@@ -0,0 +1,76 @@
+from asyncio import sleep
+from logging import getLogger
+from typing import Annotated, AsyncIterator
+
+from fastapi import Depends
+from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError, SQLAlchemyError
+from sqlalchemy.ext.asyncio import (
+    async_sessionmaker,
+    AsyncEngine,
+    AsyncSession,
+    create_async_engine,
+)
+
+from .base_class import Base
+from settings import DatabaseSettings
+
+logger = getLogger(__name__)
+
+
+class Database:
+    def __init__(self) -> None:
+        self._async_engine: AsyncEngine | None = None
+        self._async_session_local: async_sessionmaker[AsyncSession] | None = None
+
+    async def get_session(self) -> AsyncSession | None:
+        try:
+            return self._async_session_local()  # type: ignore
+
+        except (SQLAlchemyError, AttributeError) as e:
+            logger.exception(e)
+
+        return None
+
+    # TODO: Preserve UserLastStopSearchResults table from drop.
+    async def connect(
+        self, settings: DatabaseSettings, clear_static_data: bool = False
+    ) -> bool:
+        password = settings.password
+        path = (
+            f"{settings.driver}://{settings.user}:"
+            f"{password.get_secret_value() if password is not None else ''}"
+            f"@{settings.host}:{settings.port}/{settings.name}"
+        )
+        self._async_engine = create_async_engine(
+            path, pool_pre_ping=True, pool_size=10, max_overflow=20
+        )
+
+        if self._async_engine is not None:
+            SQLAlchemyInstrumentor().instrument(engine=self._async_engine.sync_engine)
+
+            self._async_session_local = async_sessionmaker(
+                bind=self._async_engine,
+                # autoflush=False,
+                expire_on_commit=False,
+                class_=AsyncSession,
+            )
+
+            ret = False
+            while not ret:
+                try:
+                    async with self._async_engine.begin() as session:
+                        if clear_static_data:
+                            await session.run_sync(Base.metadata.drop_all)
+                            await session.run_sync(Base.metadata.create_all)
+                        ret = True
+                except OperationalError as err:
+                    logger.error(err)
+                    await sleep(1)
+
+        return True
+
+    async def disconnect(self) -> None:
+        if self._async_engine is not None:
+            await self._async_engine.dispose()