feat: creatives and external channels

shared/datebase_base.py

@@ -0,0 +1,165 @@
+import asyncio
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+
+import pydantic
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+
+_session_ctx: ContextVar[AsyncSession | None] = ContextVar('session', default=None)
+
+
+class DatabaseConfig(pydantic.BaseModel):
+    URL: pydantic.PostgresDsn
+    ECHO: bool = False
+    ECHO_POOL: bool = False
+    POOL_SIZE: int = 50
+    MAX_OVERFLOW: int = 10
+
+
+class DatabaseBase:
+    def __init__(self, config: DatabaseConfig, migrations_path: str = 'migration/versions') -> None:
+        self.config = config
+        self.migrations_path = migrations_path
+        self.engine: AsyncEngine | None = None
+        self.session_factory: async_sessionmaker[AsyncSession] | None = None
+        self._connect_lock = asyncio.Lock()
+
+    async def connect(self) -> None:
+        async with self._connect_lock:
+            if self.engine is not None:
+                logging.warning('Database already connected')
+                return
+
+            self.engine = create_async_engine(
+                url=str(self.config.URL),
+                echo=self.config.ECHO,
+                echo_pool=self.config.ECHO_POOL,
+                pool_size=self.config.POOL_SIZE,
+                max_overflow=self.config.MAX_OVERFLOW,
+            )
+
+            self.session_factory = async_sessionmaker(
+                bind=self.engine,
+                autoflush=False,
+                autocommit=False,
+                expire_on_commit=False,
+            )
+
+            await self.ping()
+            await self._check_migrations()
+
+            logging.info('Database initialized')
+
+    async def ping(self) -> None:
+        if self.session_factory is None:
+            raise RuntimeError('Database not connected. Call connect() first.')
+
+        try:
+            async with self.session_factory() as session:
+                await session.execute(text('SELECT 1'))
+        except Exception as e:
+            logging.error(f'Database ping failed: {e}')
+
+    async def close(self) -> None:
+        async with self._connect_lock:
+            if self.engine is None:
+                logging.warning('Database already closed')
+                return
+
+            await self.engine.dispose()
+            self.engine = None
+            self.session_factory = None
+
+            logging.info('Database closed')
+
+    @property
+    def session(self) -> AsyncSession:
+        session = _session_ctx.get()
+        if session is None:
+            raise RuntimeError('No active transaction. Use transaction() or begin() first.')
+        return session
+
+    async def begin(self) -> None:
+        if _session_ctx.get() is not None:
+            raise RuntimeError('Transaction already started')
+
+        if self.session_factory is None:
+            raise RuntimeError('Database not connected. Call connect() first.')
+
+        session = self.session_factory()
+        _session_ctx.set(session)
+
+    async def commit(self) -> None:
+        session = _session_ctx.get()
+        if session is None:
+            raise RuntimeError('No active transaction')
+
+        try:
+            await session.commit()
+        finally:
+            await session.close()
+            _session_ctx.set(None)
+
+    async def rollback(self) -> None:
+        session = _session_ctx.get()
+        if session is None:
+            raise RuntimeError('No active transaction')
+
+        try:
+            await session.rollback()
+        finally:
+            await session.close()
+            _session_ctx.set(None)
+
+    @asynccontextmanager
+    async def transaction(self) -> AsyncGenerator[None]:
+        await self.begin()
+        try:
+            yield
+            await self.commit()
+        except Exception:
+            await self.rollback()
+            raise
+
+    async def _check_migrations(self) -> None:
+        if self.session_factory is None:
+            raise RuntimeError('Database not connected. Call connect() first.')
+
+        async with self.session_factory() as session:
+            result = await session.execute(
+                text("SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'alembic_version')")
+            )
+            if not result.scalar():
+                logging.error('Migrations not applied. Run: alembic upgrade head')
+                return
+
+            result = await session.execute(text('SELECT version_num FROM alembic_version'))
+            current_version = result.scalar()
+
+            if not current_version:
+                logging.error('No migration version found in database. Run: alembic upgrade head')
+                return
+
+            try:
+                from alembic.config import Config
+                from alembic.script import ScriptDirectory
+
+                alembic_cfg = Config('alembic.ini')
+                script = ScriptDirectory.from_config(alembic_cfg)
+                head_revision = script.get_current_head()
+
+                if head_revision and current_version != head_revision:
+                    logging.warning(
+                        f'Database version {current_version} is outdated (latest: {head_revision}). '
+                        f'Run: alembic upgrade head'
+                    )
+            except Exception as e:
+                logging.debug(f'Could not check migration head: {e}')