api/utils/postgres_client.py - master - saba07022001/intelclinic

api/utils/postgres_client.py
88 строк3 KB
saba070201
Initial commit
04 янв 2025, 01:31
04 янв 2025, 01:31b450ee4
from __future__ import annotations

import asyncpg


from enum import Enum
from typing import Iterator


from core.settings import Settings


settings = Settings()


class IsolationMode(str, Enum):
    SERIALIZABLE = "serializable"
    REPEATABLE_READ = "repeatable_read"
    READ_COMMITTED = "read_committed"


class PostgresClient(object):
    pool = None
    instance = None

    def __new__(cls) -> PostgresClient:
        if not cls.instance:
            cls.instance = super().__new__(cls)
        return cls.instance

    @classmethod
    async def start(cls, dsn: str = settings.POSTGRES_DSN) -> None:
        cls.pool = await asyncpg.create_pool(
            dsn=dsn,
            min_size=settings.POSTGRES_MIN_CONN_SIZE,
            max_size=settings.POSTGRES_MAX_CONN_SIZE,
            command_timeout=settings.POSTGRES_CONN_TIMEOUT,
        )

    @classmethod
    async def stop(cls) -> None:
        await cls.pool.close()

    async def fetch_one(self, query: str) -> dict | None:
        async with self.pool.acquire() as conn:
            try:
                return [dict(row) for row in await conn.fetch(query)][0]
            except IndexError:
                return None

    async def fetch_all(self, query: str) -> list[dict | None]:
        async with self.pool.acquire() as conn:
            return [dict(row) for row in await conn.fetch(query)]

    async def iterator(self, query: str) -> Iterator[dict]:
        async with self.pool.acquire() as conn:
            async with conn.transaction():
                async for row in conn.cursor(query):
                    yield row

    async def execute(self, query: str) -> None:
        async with self.pool.acquire() as conn:
            return await conn.execute(query)

    async def execute_in_transaction(
        self,
        query: str,
        isolation: IsolationMode | None = None,
        readonly: bool = False,
    ) -> None:
        async with self.pool.acquire() as conn:
            async with conn.transaction(isolation=isolation, readonly=readonly):
                return await conn.execute(query)

    async def executemany(self, query: str, values: list[tuple]) -> None:
        async with self.pool.acquire() as conn:
            return await conn.executemany(query, values)

    async def executemany_in_transaction(
        self,
        query: str,
        values: list[tuple],
        isolation: IsolationMode | None = None,
    ) -> None:
        async with self.pool.acquire() as conn:
            conn: asyncpg.Connection
            async with conn.transaction(isolation=isolation):
                return await conn.executemany(query, values)