feat: переезд на tortose

shared/datebase_base.py

@@ -1,19 +1,5 @@
-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)
+from tortoise import Tortoise
 
 
 class DatabaseConfig(pydantic.BaseModel):
@@ -25,141 +11,71 @@ class DatabaseConfig(pydantic.BaseModel):
 
 
 class DatabaseBase:
-    def __init__(self, config: DatabaseConfig, migrations_path: str = 'migration/versions') -> None:
+    def __init__(self, config: DatabaseConfig) -> 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
+        await Tortoise.init(
+            db_url=str(self.config.URL),
+            modules={'models': ['src.domain']},
+            use_tz=True,
+            timezone='UTC',
+        )
 
-            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}')
+        await self.ping()
+        await self._check_migrations()
 
     async def close(self) -> None:
-        async with self._connect_lock:
-            if self.engine is None:
-                logging.warning('Database already closed')
-                return
+        await Tortoise.close_connections()
 
-            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')
+    async def ping(self) -> None:
+        from tortoise import connections
 
+        conn = connections.get('default')
         try:
-            await session.commit()
-        finally:
-            await session.close()
-            _session_ctx.set(None)
+            await conn.execute_query('SELECT 1')
+        except Exception as e:
+            import logging
 
-    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()
+            logging.error(f'Database ping failed: {e}')
             raise
 
     async def _check_migrations(self) -> None:
-        if self.session_factory is None:
-            raise RuntimeError('Database not connected. Call connect() first.')
+        import logging
 
-        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
+        from tortoise import connections
 
-            result = await session.execute(text('SELECT version_num FROM alembic_version'))
-            current_version = result.scalar()
+        conn = connections.get('default')
 
-            if not current_version:
-                logging.error('No migration version found in database. Run: alembic upgrade head')
-                return
+        # Проверяем существование таблицы alembic_version
+        result = await conn.execute_query(
+            "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'alembic_version')"
+        )
+        if not result[1][0]['exists']:
+            logging.error('Migrations not applied. Run: alembic upgrade head')
+            return
 
-            try:
-                from alembic.config import Config
-                from alembic.script import ScriptDirectory
+        # Получаем текущую версию
+        result = await conn.execute_query('SELECT version_num FROM alembic_version')
+        if not result[1]:
+            logging.error('No migration version found in database. Run: alembic upgrade head')
+            return
 
-                alembic_cfg = Config('alembic.ini')
-                script = ScriptDirectory.from_config(alembic_cfg)
-                head_revision = script.get_current_head()
+        current_version = result[1][0]['version_num']
 
-                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}')
+        # Сравниваем с head версией
+        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}')