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


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")

    @classmethod
    def create(cls: "type[IncomeStatement]", db: AsyncSession, **kwargs) -> "IncomeStatement":
        query = sql.insert(cls).values(**kwargs)
        income_statements = db.execute(query)
        income_statement = income_statements.first()
        db.commit()
        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

    @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