src/modules/database/database.py - master - sam_single_1/ServiceVkAuth

src/modules/database/database.py
105 строк3 KB
sam_single_1
Before refactor
22 окт 2024, 19:17
22 окт 2024, 19:1783c90dd
from uuid import UUID
from typing import Any, List, Dict
from pydantic import BaseModel
import asyncpg
import os
from dotenv import load_dotenv
import datetime


load_dotenv()
db_name = os.getenv("DB_NAME")
db_host = os.getenv("DB_HOST")
db_port = os.getenv("DB_PORT")
db_password = os.getenv("DB_PASSWORD")
db_user = os.getenv("DB_USER")
if (
    db_name is None
    or db_host is None
    or db_port is None
    or db_password is None
    or db_user is None
):
    raise Exception


class DataBaseExec(BaseModel):
    SQL: str
    Values: List[Dict[str, str | int | str | datetime.datetime | bool | None | UUID]]


class DataBaseResults(BaseModel):
    has_error: bool
    error: Exception | None
    other_error_options: str | None
    Values: List[Dict[str, int | str | datetime.datetime | bool | None | UUID]]
    result_len: int
    row_len: int


async def database_common(exec_expression: DataBaseExec) -> DataBaseResults:
    conn_str: str = (
        f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
    )
    conn_query: str = exec_expression.SQL
    conn_params = exec_expression.Values
    try:
        conn: asyncpg.Connection = await asyncpg.connect(conn_str)
        fetch_result: List[asyncpg.Record] = await conn.fetchmany(
            conn_query, conn_params
        )
    except Exception as e:
        return DataBaseResults(
            has_error=True,
            error=e,
            Values=[],
            result_len=0,
            row_len=0,
            other_error_options=None,
        )
    else:
        if type(fetch_result) is List[asyncpg.Record]:

            result_length = len(fetch_result)
            if result_length == 0:
                return DataBaseResults(
                    has_error=False,
                    error=None,
                    Values=[],
                    row_len=0,
                    result_len=0,
                    other_error_options=None,
                )
            else:
                row_length = len(fetch_result[0])

                result = list()

                for record in fetch_result:
                    keys = record.keys()
                    new_line = dict()
                    for key in keys:
                        value = record.get(key)
                        new_line[key] = value
                    result.append(new_line)

                return DataBaseResults(
                    has_error=False,
                    error=None,
                    other_error_options=None,
                    Values=result,
                    result_len=result_length,
                    row_len=row_length,
                )
        else:
            return DataBaseResults(
                has_error=True,
                other_error_options="Wrong type from asyncpg",
                error=None,
                Values=[],
                result_len=0,
                row_len=0,
            )

    finally:
        await conn.close()