import enum import itertools import json from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession from src.db.models.form_type import FormType class DirectionEnum(str, enum.Enum): SUPPORT = "Support" DEVELOPMENT = "Development" ACCEPTABLE_SHEETS = set(itertools.chain(*FormType.SHEETS_BY_FORM_TYPE.values())) ACCEPTABLE_SECTIONS = set(itertools.chain(*FormType.SECTIONS_BY_FORM_TYPE.values())) class SheetValidationError(Exception): def __init__(self, message, *args): super().__init__(*args) self.message = message class SheetRepository: def __init__(self, db: AsyncSession): self.db = db async def get( self, form_id: int, sheet: str, direction: str | None, sections: list[str] | None = None, ) -> list[tuple]: if any( [ not isinstance(form_id, int), sheet not in ACCEPTABLE_SHEETS, direction is not None and direction not in [s.value for s in DirectionEnum], sections and (set(sections) - ACCEPTABLE_SECTIONS), ] ): raise SheetValidationError("Некорректные данные") if direction: func_query = "v3.v_form_view(:form_id, :sheet, :sections, :direction)" else: func_query = "v3.v_form_view(:form_id, :sheet, :sections)" return ( await self.db.execute( text( f""" SELECT row_type, depth, sort_order, data FROM {func_query} """ ), { "form_id": form_id, "sheet": sheet, "direction": direction, "sections": sections, } ) ).all() async def update_cell( self, form_id: int, sheet: str, direction: str | None, sections: list[str] | None, line_id: int, column: str, value: str | int | float | None = None, user_id: int = None, ) -> list[tuple]: if any( [ not isinstance(form_id, int), sheet not in ACCEPTABLE_SHEETS, direction is not None and direction not in [s.value for s in DirectionEnum], sections and (set(sections) - ACCEPTABLE_SECTIONS), ] ): raise SheetValidationError("Некорректные данные") if user_id: func_query = "v3.upd_form_cell(:form_id, :sheet, :line_id, :column, :value, :user_id, :direction, :sections)" else: func_query = "v3.upd_form_cell(:form_id, :sheet, :line_id, :column, :value, :direction, :sections)" return [ tuple(r) for r in ( await self.db.execute( text( f""" SELECT row_type, depth, sort_order, data FROM {func_query} """ ), { "form_id": form_id, "sheet": sheet, "direction": direction, "sections": sections, "line_id": line_id, "column": column, "value": json.dumps(value), "user_id": user_id, } ) ).all() ] async def update_cells( self, form_id: int, sheet: str, direction: str | None, sections: list[str] | None, changes: list[dict], user_id: int = None, ) -> list[tuple]: if any( [ not isinstance(form_id, int), sheet not in ACCEPTABLE_SHEETS, direction is not None and direction not in [s.value for s in DirectionEnum], sections and (set(sections) - ACCEPTABLE_SECTIONS), ] ): raise SheetValidationError("Некорректные данные") changes = json.dumps(changes) if user_id: func_query = "v3.upd_form_cells(:form_id, :sheet, :changes, :user_id, :direction, :sections)" else: func_query = "v3.upd_form_cells(:form_id, :sheet, :changes, :direction, :sections)" return [ tuple(r) for r in ( await self.db.execute( text( f""" SELECT row_type, depth, sort_order, data FROM {func_query} """ ), { "form_id": form_id, "sheet": sheet, "direction": direction, "sections": sections, "changes": changes, "user_id": user_id, } ) ).all() ] async def add_line( self, form_id: int, sheet: str, expense_item_id: int | None = None, item_id: str | None = None, section_code: str | None = None, direction: str | None = None, name: str | None = None, internal_order: str | None = None, vsp_id: int | None = None, project_id: int | None = None, justification: str | None = None, contract_number: str | None = None, contract_end_date: str | None = None, user_id: int | None = None, ) -> list[tuple]: return [ tuple(r) for r in ( await self.db.execute( text( """ SELECT row_type, depth, sort_order, data FROM v3.add_budget_line( :p_form_id, :p_expense_item_id, :p_sheet, :p_item_id, :p_section_code, :p_direction, :p_name, :p_internal_order, :p_vsp_id, :p_project_id, :p_justification, :p_contract_number, :p_contract_end_date ) """ ), { "p_form_id": form_id, "p_expense_item_id": expense_item_id, "p_sheet": sheet, "p_item_id": item_id, "p_section_code": section_code, "p_direction": direction, "p_name": name, "p_internal_order": internal_order, "p_vsp_id": vsp_id, "p_project_id": project_id, "p_justification": justification, "p_contract_number": contract_number, "p_contract_end_date": contract_end_date, } ) ).all() ] async def delete_line( self, sheet: str, row_id: int, direction: str | None = None, ) -> list[tuple]: return [ tuple(r) for r in ( await self.db.execute( text( "SELECT row_type, depth, sort_order, data " "FROM v3.del_budget_line(:p_row_id, :p_direction, :p_sheet)" ), { "p_row_id": row_id, "p_direction": direction, "p_sheet": sheet, } ) ).all() ]