dotorm 2.0.8__py3-none-any.whl

dotorm/databases/mysql/pool.py

@@ -0,0 +1,69 @@
+"""MySQL connection pool management."""
+
+import asyncio
+import logging
+import time
+
+try:
+    import aiomysql
+except ImportError:
+    ...
+from ..abstract.types import ContainerSettings, MysqlPoolSettings
+
+
+log = logging.getLogger("dotorm")
+
+
+class ContainerMysql:
+    """
+    MySQL connection pool container.
+
+    Manages pool lifecycle.
+    """
+
+    def __init__(
+        self,
+        pool_settings: MysqlPoolSettings,
+        container_settings: ContainerSettings,
+    ):
+        self.pool_settings = pool_settings
+        self.container_settings = container_settings
+        self.pool: "aiomysql.Pool | None" = None
+
+    async def create_pool(self) -> "aiomysql.Pool":
+        """Create connection pool with retry on failure."""
+        try:
+            start_time = time.time()
+            self.pool = await aiomysql.create_pool(
+                **self.pool_settings.model_dump(),
+                minsize=5,
+                maxsize=15,
+                autocommit=True,
+                # 15 minutes
+                pool_recycle=60 * 15,
+            )
+
+            log.debug(
+                "Connection MySQL db: %s, created time: [%0.3fs]",
+                self.pool_settings.db,
+                time.time() - start_time,
+            )
+            return self.pool
+
+        except (ConnectionError, TimeoutError, aiomysql.OperationalError):
+            log.exception(
+                "MySQL create pool connection lost, reconnect after %d seconds",
+                self.container_settings.reconnect_timeout,
+            )
+            await asyncio.sleep(self.container_settings.reconnect_timeout)
+            return await self.create_pool()
+        except Exception as e:
+            log.exception("MySQL create pool error:")
+            raise e
+
+    async def close_pool(self):
+        """Close connection pool."""
+        if self.pool:
+            self.pool.close()
+            await self.pool.wait_closed()
+            self.pool = None

dotorm/databases/mysql/session.py

@@ -0,0 +1,128 @@
+"""MySQL session implementations."""
+
+from typing import Any, Callable, TYPE_CHECKING
+
+from ..abstract.session import SessionAbstract
+from ..abstract.dialect import MySQLDialect, CursorType
+
+
+if TYPE_CHECKING:
+    import aiomysql
+
+
+# Shared dialect instance
+_dialect = MySQLDialect()
+
+
+class MysqlSession(SessionAbstract):
+    """Base MySQL session."""
+
+    @staticmethod
+    async def _do_execute(
+        cursor: "aiomysql.Cursor",
+        stmt: str,
+        values: Any,
+        cursor_type: CursorType,
+    ) -> Any:
+        """
+        Execute query on cursor (shared logic).
+
+        Args:
+            cursor: aiomysql cursor
+            stmt: SQL with %s placeholders
+            values: Query values
+            cursor_type: Cursor type
+
+        Returns:
+            Raw result from aiomysql
+        """
+        # executemany
+        if cursor_type == "executemany":
+            if not values:
+                raise ValueError("executemany requires values")
+            await cursor.executemany(stmt, values)
+            return None
+
+        # Execute query
+        if values:
+            await cursor.execute(stmt, values)
+        else:
+            await cursor.execute(stmt)
+
+        # void - execute only (INSERT/UPDATE/DELETE without return)
+        if cursor_type == "void":
+            return None
+
+        # lastrowid special case
+        if cursor_type == "lastrowid":
+            return cursor.lastrowid
+
+        # fetch operations
+        method = getattr(cursor, _dialect.get_cursor_method(cursor_type))
+        return await method()
+
+
+class TransactionSession(MysqlSession):
+    """
+    Session for transactional queries.
+    Uses single connection within transaction context.
+    """
+
+    def __init__(
+        self, connection: "aiomysql.Connection", cursor: "aiomysql.Cursor"
+    ) -> None:
+        self.connection = connection
+        self.cursor = cursor
+
+    async def execute(
+        self,
+        stmt: str,
+        values: Any = None,
+        *,
+        prepare: Callable | None = None,
+        cursor: CursorType = "fetchall",
+    ) -> Any:
+        stmt = _dialect.convert_placeholders(stmt)
+        result = await self._do_execute(self.cursor, stmt, values, cursor)
+        result = _dialect.convert_result(result, cursor)
+
+        if prepare and result:
+            return prepare(result)
+        return result
+
+
+class NoTransactionSession(MysqlSession):
+    """
+    Session for non-transactional queries.
+    Acquires connection from pool per query.
+    """
+
+    default_pool: "aiomysql.Pool | None" = None
+
+    def __init__(self, pool: "aiomysql.Pool | None" = None) -> None:
+        if pool is None:
+            assert self.default_pool is not None
+            self.pool = self.default_pool
+        else:
+            self.pool = pool
+
+    async def execute(
+        self,
+        stmt: str,
+        values: Any = None,
+        *,
+        prepare: Callable | None = None,
+        cursor: CursorType = "fetchall",
+    ) -> Any:
+        import aiomysql
+        
+        stmt = _dialect.convert_placeholders(stmt)
+
+        async with self.pool.acquire() as conn:
+            async with conn.cursor(aiomysql.DictCursor) as cur:
+                result = await self._do_execute(cur, stmt, values, cursor)
+                result = _dialect.convert_result(result, cursor)
+
+                if prepare and result:
+                    return prepare(result)
+                return result

dotorm/databases/mysql/transaction.py

@@ -0,0 +1,39 @@
+"""MySQL transaction management."""
+
+try:
+    import aiomysql
+except ImportError:
+    aiomysql = None  # type: ignore
+
+from .session import TransactionSession
+
+
+class ContainerTransaction:
+    """
+    Transaction context manager for MySQL.
+
+    Acquires connection, executes queries,
+    commits on success, rollbacks on exception.
+    """
+
+    def __init__(self, pool: "aiomysql.Pool"):
+        self.pool = pool
+
+    async def __aenter__(self):
+        connection: "aiomysql.Connection" = await self.pool.acquire()
+        cursor: "aiomysql.Cursor" = await connection.cursor(
+            aiomysql.DictCursor
+        )
+        self.session = TransactionSession(connection, cursor)
+        return self.session
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is not None:
+            # Выпало исключение вызвать ролбек
+            await self.session.connection.rollback()
+        else:
+            # Не выпало исключение вызвать комит
+            await self.session.connection.commit()
+        await self.session.cursor.close()
+        # В любом случае закрыть соединение и курсор
+        self.pool.release(self.session.connection)

dotorm/databases/postgres/__init__.py

@@ -0,0 +1,23 @@
+"""PostgreSQL database support."""
+
+from .pool import ContainerPostgres
+from .session import (
+    PostgresSession,
+    TransactionSession,
+    NoTransactionSession,
+    NoTransactionNoPoolSession,
+)
+from .transaction import ContainerTransaction, get_current_session
+from ..abstract.dialect import CursorType, PostgresDialect
+
+__all__ = [
+    "ContainerPostgres",
+    "PostgresSession",
+    "TransactionSession",
+    "NoTransactionSession",
+    "NoTransactionNoPoolSession",
+    "ContainerTransaction",
+    "get_current_session",
+    "CursorType",
+    "PostgresDialect",
+]

dotorm/databases/postgres/pool.py

@@ -0,0 +1,133 @@
+"""PostgreSQL connection pool management."""
+
+import logging
+import asyncio
+import time
+
+try:
+    import asyncpg
+except ImportError:
+    ...
+
+from .transaction import ContainerTransaction
+from ..abstract.types import ContainerSettings, PostgresPoolSettings
+from .session import NoTransactionNoPoolSession
+
+
+log = logging.getLogger("dotorm")
+
+
+class ContainerPostgres:
+    """
+    PostgreSQL connection pool container.
+
+    Manages pool lifecycle and provides utilities
+    for database and table creation.
+    """
+
+    def __init__(
+        self,
+        pool_settings: PostgresPoolSettings,
+        container_settings: ContainerSettings,
+    ):
+        self.pool_settings = pool_settings
+        self.container_settings = container_settings
+        self.pool: "asyncpg.Pool | None" = None
+
+    async def create_pool(self) -> "asyncpg.Pool":
+        """Create connection pool with retry on failure."""
+        try:
+            start_time = time.time()
+            pool = await asyncpg.create_pool(
+                **self.pool_settings.model_dump(),
+                min_size=5,
+                max_size=15,
+                command_timeout=60,
+                # 15 minutes
+                # max_inactive_connection_lifetime
+                # pool_recycle=60 * 15,
+            )
+            assert pool is not None
+            self.pool = pool
+
+            log.debug(
+                "Connection PostgreSQL db: %s, created time: [%0.3fs]",
+                self.pool_settings.database,
+                time.time() - start_time,
+            )
+            return self.pool
+        except (
+            asyncpg.InvalidCatalogNameError,
+            asyncpg.ConnectionDoesNotExistError,
+        ):
+            # БД не существует — пробрасываем наверх для создания
+            raise
+        except (
+            ConnectionError,
+            TimeoutError,
+            asyncpg.exceptions.ConnectionFailureError,
+        ):
+            log.exception(
+                "Postgres create pool connection lost, reconnect after %d seconds",
+                self.container_settings.reconnect_timeout,
+            )
+            await asyncio.sleep(self.container_settings.reconnect_timeout)
+            return await self.create_pool()
+        except Exception as e:
+            log.exception("Postgres create pool error:")
+            raise e
+
+    async def close_pool(self):
+        """Close connection pool."""
+        if self.pool:
+            # await self.pool.close()
+            self.pool.terminate()
+            self.pool = None
+
+    async def create_database(self):
+        """Create database if it doesn't exist."""
+        db_to_create = self.pool_settings.database
+        # Connect to default postgres database
+        temp_settings = self.pool_settings.model_copy()
+        temp_settings.database = "postgres"
+
+        conn = await NoTransactionNoPoolSession.get_connection(temp_settings)
+        sql = f'CREATE DATABASE "{db_to_create}"'
+        try:
+            await conn.execute(sql)
+        finally:
+            await conn.close()
+
+    async def create_and_update_tables(self, models: list):
+        """
+        Create/update tables for given models.
+
+        Args:
+            models: List of DotModel classes
+        """
+        stmt_foreign_keys: list[tuple[str, str]] = []
+
+        async with ContainerTransaction(self.pool) as session:
+            # создаем модели в БД, без FK
+            for model in models:
+                foreign_keys = await model.__create_table__(session)
+                stmt_foreign_keys += foreign_keys
+
+            if not stmt_foreign_keys:
+                return
+
+            # Дедупликация по имени FK (M2M таблицы могут дублироваться)
+            unique_fks = {
+                fk_name: fk_sql for fk_name, fk_sql in stmt_foreign_keys
+            }
+
+            # Получаем существующие FK одним запросом
+            existing_fk_result = await session.execute(
+                "SELECT conname FROM pg_constraint WHERE contype = 'f'"
+            )
+            existing_fk_names = {row["conname"] for row in existing_fk_result}
+
+            # Создаём только отсутствующие FK
+            for fk_name, fk_sql in unique_fks.items():
+                if fk_name not in existing_fk_names:
+                    await session.execute(fk_sql)

dotorm/databases/postgres/session.py

@@ -0,0 +1,174 @@
+"""PostgreSQL session implementations."""
+
+from typing import Any, Callable, TYPE_CHECKING
+
+from ..abstract.types import PostgresPoolSettings
+from ..abstract.session import SessionAbstract
+from ..abstract.dialect import PostgresDialect, CursorType
+
+
+if TYPE_CHECKING:
+    import asyncpg
+    from asyncpg.transaction import Transaction
+
+
+# Shared dialect instance
+_dialect = PostgresDialect()
+
+
+class PostgresSession(SessionAbstract):
+    """Base PostgreSQL session."""
+
+    @staticmethod
+    async def _do_execute(
+        conn: "asyncpg.Connection",
+        stmt: str,
+        values: Any,
+        cursor: CursorType,
+    ) -> Any:
+        """
+        Execute query on connection (shared logic).
+
+        Args:
+            conn: asyncpg connection
+            stmt: SQL with $1, $2... placeholders (already converted)
+            values: Query values
+            cursor: Cursor type
+
+        Returns:
+            Raw result from asyncpg
+        """
+        # executemany
+        if cursor == "executemany":
+            if not values:
+                raise ValueError("executemany requires values")
+            # Handle [[v1, v2], [v3, v4]] or [[[v1, v2], [v3, v4]]] format
+            rows = (
+                values[0]
+                if isinstance(values[0], list)
+                and values[0]
+                and isinstance(values[0][0], (list, tuple))
+                else values
+            )
+            for row in rows:
+                await conn.execute(stmt, *row)
+            return None
+
+        # void - execute only (INSERT/UPDATE/DELETE without return)
+        if cursor == "void":
+            if values:
+                await conn.execute(stmt, *values)
+            else:
+                await conn.execute(stmt)
+            return None
+
+        # fetch operations
+        method = getattr(conn, _dialect.get_cursor_method(cursor))
+        if values:
+            return await method(stmt, *values)
+        return await method(stmt)
+
+
+class TransactionSession(PostgresSession):
+    """
+    Session for transactional queries.
+    Uses single connection within transaction context.
+    """
+
+    def __init__(
+        self, connection: "asyncpg.Connection", transaction: "Transaction"
+    ) -> None:
+        self.connection = connection
+        self.transaction = transaction
+
+    async def execute(
+        self,
+        stmt: str,
+        values: Any = None,
+        *,
+        prepare: Callable | None = None,
+        cursor: CursorType = "fetchall",
+    ) -> Any:
+        stmt = _dialect.convert_placeholders(stmt)
+        result = await self._do_execute(self.connection, stmt, values, cursor)
+        result = _dialect.convert_result(result, cursor)
+
+        if prepare and result:
+            return prepare(result)
+        return result
+
+
+class NoTransactionSession(PostgresSession):
+    """
+    Session for non-transactional queries.
+    Acquires connection from pool per query.
+    """
+
+    default_pool: "asyncpg.Pool | None" = None
+
+    def __init__(self, pool: "asyncpg.Pool | None" = None) -> None:
+        if pool is None:
+            assert self.default_pool is not None
+            self.pool = self.default_pool
+        else:
+            self.pool = pool
+
+    async def execute(
+        self,
+        stmt: str,
+        values: Any = None,
+        *,
+        prepare: Callable | None = None,
+        cursor: CursorType = "fetchall",
+    ) -> Any:
+        stmt = _dialect.convert_placeholders(stmt)
+
+        async with self.pool.acquire() as conn:
+            result = await self._do_execute(conn, stmt, values, cursor)
+            result = _dialect.convert_result(result, cursor)
+
+            if prepare and result:
+                return prepare(result)
+            return result
+
+
+class NoTransactionNoPoolSession(PostgresSession):
+    """
+    Session without pool.
+    Opens connection, executes, closes. For admin tasks.
+    """
+
+    @classmethod
+    async def get_connection(
+        cls, settings: PostgresPoolSettings
+    ) -> "asyncpg.Connection":
+        """Create new connection without pool."""
+        import asyncpg
+
+        return await asyncpg.connect(**settings.model_dump())
+
+    @classmethod
+    async def execute(
+        cls,
+        settings: PostgresPoolSettings,
+        stmt: str,
+        values: Any = None,
+        *,
+        prepare: Callable | None = None,
+        cursor: str = "execute",
+    ) -> Any:
+        conn = await cls.get_connection(settings)
+
+        try:
+            if values:
+                await conn.execute(stmt, values)
+            else:
+                await conn.execute(stmt)
+
+            rows = await getattr(conn, cursor)()
+
+            if prepare:
+                return prepare(rows)
+            return rows
+        finally:
+            await conn.close()

dotorm/databases/postgres/transaction.py

@@ -0,0 +1,82 @@
+"""PostgreSQL transaction management."""
+
+from contextvars import ContextVar
+
+try:
+    import asyncpg
+    from asyncpg.transaction import Transaction
+except ImportError:
+    asyncpg = None  # type: ignore
+    Transaction = None  # type: ignore
+
+from .session import TransactionSession
+
+
+# Context variable для хранения текущей сессии транзакции
+_current_session: ContextVar["TransactionSession | None"] = ContextVar(
+    "current_session", default=None
+)
+
+
+def get_current_session() -> "TransactionSession | None":
+    """Получить текущую сессию из контекста (если есть активная транзакция)."""
+    return _current_session.get()
+
+
+class ContainerTransaction:
+    """
+    Transaction context manager for PostgreSQL.
+
+    Acquires connection, starts transaction, executes queries,
+    commits on success, rollbacks on exception.
+
+    Автоматически устанавливает текущую сессию в contextvars,
+    так что методы ORM могут использовать её без явной передачи.
+
+    Example:
+        async with ContainerTransaction(pool) as session:
+            await session.execute("INSERT INTO users ...")
+            # Или без явной передачи session:
+            await User.create(payload=user)  # session подставится из контекста
+            # Commits on exit
+    """
+
+    default_pool: "asyncpg.Pool | None" = None
+
+    def __init__(self, pool: "asyncpg.Pool | None" = None):
+        self.session_factory = TransactionSession
+        if pool is None:
+            assert self.default_pool is not None
+            self.pool = self.default_pool
+        else:
+            self.pool = pool
+        self._token = None
+
+    async def __aenter__(self):
+        connection: "asyncpg.Connection" = await self.pool.acquire()
+        transaction = connection.transaction()
+
+        assert isinstance(transaction, Transaction)
+        assert isinstance(connection, asyncpg.Connection)
+
+        await transaction.start()
+        self.session = self.session_factory(connection, transaction)
+
+        # Устанавливаем текущую сессию в контекст
+        self._token = _current_session.set(self.session)
+
+        return self.session
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        # Сбрасываем контекст
+        if self._token is not None:
+            _current_session.reset(self._token)
+
+        if exc_type is not None:
+            # Выпало исключение вызвать ролбек
+            await self.session.transaction.rollback()
+        else:
+            # Не выпало исключение вызвать комит
+            await self.session.transaction.commit()
+        # В любом случае вернуть соединение в пул
+        await self.pool.release(self.session.connection)