modules/DB/database.py - main - andreyTkach/petProject

modules/DB/database.py
207 строк6 KB
TkachevAV
new chance in logic
18 авг 2025, 19:09
18 авг 2025, 19:094c53161
import uuid
import json
from datetime import datetime, timezone
from typing import Any
from loger import logger

from config.dbSettings import settings
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy.dialects.postgresql import UUID, JSON, TIMESTAMP
from sqlalchemy import select, update
from sqlalchemy.exc import IntegrityError


from asyncpg.exceptions import UniqueViolationError

from modules.DB.tools import toHash
from models.shemas import reqAddContract, ContractDB, reqUpdateContract

INFINITY = datetime.max.replace(tzinfo=timezone.utc)


engine = create_async_engine(settings.linkForConnection)
session = async_sessionmaker(engine, expire_on_commit=False)


class Model(DeclarativeBase):
    pass


class ContractsORM(Model):
    __tablename__ = "contracts"

    id: Mapped[int] = mapped_column(primary_key=True)
    contract_version: Mapped[str]
    uuid_obj: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True))
    contract_obj: Mapped[dict[str, Any]] = mapped_column(JSON)
    hash_contract: Mapped[str] = mapped_column(unique=True)
    date_from: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True))
    date_to: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True))
    is_active: Mapped[bool] = mapped_column(default=False)


async def create_tables():
    async with engine.begin() as conn:
        await conn.run_sync(Model.metadata.create_all)


async def delete_tables():
    async with engine.begin() as conn:
        await conn.run_sync(Model.metadata.drop_all)


def connection(method):
    async def wrapper(*args, **kwargs):
        async with session() as sn:
            try:
                result = await method(*args, session=sn, **kwargs)
                await sn.commit()
                return result
            except Exception as e:
                await sn.rollback()
                logger.error(f"Error in method {method}")
                raise

    return wrapper


class ErrorWithUniqueHash(Exception):
    pass


class ErrorWithDate(Exception):
    pass


# Можно вынести в отдельный файл
@connection
async def addContract(
    db_contract: reqAddContract, session: AsyncSession
) -> tuple[UUID, int]:
    obj_str = json.dumps(
        db_contract.contract_obj, sort_keys=True, separators=(",", ":")
    )
    hashObj = toHash(obj_str)
    uuidObj = uuid.uuid4()

    date_to = db_contract.date_to
    updated_model = db_contract.model_dump()
    if date_to is None:
        updated_model["date_to"] = INFINITY

    contract_to_DB = ContractDB(
        **updated_model, uuid_obj=uuidObj, hash_contract=hashObj
    )
    contract = ContractsORM(**contract_to_DB.model_dump())

    session.add(contract)

    try:
        await session.flush()
    except IntegrityError as e:
        if isinstance(
            e.orig.__cause__, UniqueViolationError
        ):  # еще ошибка от тригеррах на датах
            raise ErrorWithUniqueHash("Hash contract is not unique")
        raise

    await session.refresh(contract)
    return contract.uuid_obj, contract.id


@connection
async def getActualContract(uuid_obj: UUID, session: AsyncSession) -> ContractsORM:
    stmt = select(ContractsORM).where(ContractsORM.uuid_obj == uuid_obj)
    result = await session.execute(stmt)
    tempResult = result.scalars().all()
    if not tempResult:
        return None
    tempResultTime = tempResult[-1].date_to.replace(tzinfo=timezone.utc)
    local_time = datetime.now(timezone.utc)
    if not (tempResultTime is None):
        if tempResultTime < local_time:
            return None
    return tempResult


@connection
async def getLastVersion(uuid_obj: UUID, session: AsyncSession) -> ContractsORM:
    stmt = select(ContractsORM).where(ContractsORM.uuid_obj == uuid_obj)
    result = await session.execute(stmt)

    tempResult = result.scalars().all()
    if not tempResult:
        return None
    
    return tempResult[-1]


@connection
async def getHistoryContract(
    uuid_obj: UUID, session: AsyncSession
) -> list[ContractsORM]:
    stmt = select(ContractsORM).where(ContractsORM.uuid_obj == uuid_obj)
    result = await session.execute(stmt)
    return result.scalars().all()


@connection
async def updateVersionContract(
    oldVersion: ContractsORM, update: reqUpdateContract, session: AsyncSession
) -> int:
    newStrObj = json.dumps(update.contract_obj, sort_keys=True, separators=(",", ":"))
    newStrObjHash = toHash(newStrObj)

    if oldVersion.hash_contract == newStrObjHash:
        raise ErrorWithUniqueHash("Hash contract is not unique")

    if update.date_to is None:  # если не имеет даты окончания, то ставим в infinity
        update.date_to = INFINITY

    update_old = False
    changeDate = False
    oldVersionDate = oldVersion.date_to
    if oldVersionDate.tzinfo is None:
        oldVersionDate = oldVersionDate.replace(tzinfo=timezone.utc)

    if oldVersionDate == INFINITY:
        if update.date_from is None:
            update.date_from = datetime.now(timezone.utc)
            update_old = True
    else:
        if update.date_from is None or update.date_from < oldVersionDate:
            update.date_from = oldVersionDate
            changeDate = True

    upDatingContract = ContractsORM(
        **update.model_dump(),
        hash_contract=newStrObjHash,
    )

    session.add(upDatingContract)
    await session.flush()

    await session.refresh(upDatingContract)
    if not (upDatingContract.id is None):
        await updateOldDate(oldVersion, update.date_from, session, update_old)

    return upDatingContract.id, changeDate


async def updateOldDate(
    oldVersionContract: ContractsORM,
    newDateTo: datetime,
    session: AsyncSession,
    flag: bool,
):
    if flag:
        return

    stmt = (
        update(ContractsORM)
        .where(ContractsORM.id == oldVersionContract.id)
        .values(date_to=newDateTo)
    )
    result = await session.execute(stmt)
    return result