app/model/income_statement.py

from datetime import datetime
from typing import Any, Dict, List
from sqlalchemy import ForeignKey
from sqlalchemy.sql import expression as sql
from sqlalchemy.orm import relationship, Mapped, mapped_column
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.dialects.postgresql import JSON

from app.model.base import BaseModel
from app.engine.postgresdb import Base
from app.schema.index import IncomeStatement as IncomeStatementSchema


class IncomeStatement(Base, BaseModel):
    __tablename__ = "income_statement"

    date_from: Mapped[datetime]
    date_to: Mapped[datetime]
    income: Mapped[Dict] = mapped_column(JSON, nullable=True)
    expenses: Mapped[Dict] = mapped_column(JSON, nullable=True)

    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
    user = relationship("User", back_populates="income_statements")

    def __str__(self) -> str:
        return f"IncomeStatement(id={self.id}, user_id={self.user_id}, date_from={self.date_from}, date_to={self.date_to}, income={self.income}, expenses={self.expenses})"

    @classmethod
    async def create(cls: "type[IncomeStatement]", db: AsyncSession, **kwargs: IncomeStatementSchema) -> "IncomeStatement":
        income_statement = cls(**kwargs)
        db.add(income_statement)
        await db.commit()
        await db.refresh(income_statement)
        return income_statement

    @classmethod
    async def get_by_user(cls: "type[IncomeStatement]", db: AsyncSession, user_id: int) -> "List[IncomeStatement]":
        query = sql.select(cls).where(cls.user_id == user_id)
        income_statements = await db.scalars(query)
        return income_statements.all()

    @classmethod
    async def get(cls: "type[IncomeStatement]", db: AsyncSession, id: int) -> "IncomeStatement":
        query = sql.select(cls).where(cls.id == id)
        income_statement = await db.scalar(query)
        return income_statement