models.py

from __future__ import annotations
from typing import Optional
from datetime import datetime

from sqlalchemy import CheckConstraint
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, session
from sqlalchemy.schema import PrimaryKeyConstraint
from sqlalchemy.dialects.sqlite import insert


class Base(DeclarativeBase):
    pass


class Stock(Base):
    __tablename__ = "stock_ohlcv"

    symbol: Mapped[str]
    open: Mapped[float]
    high: Mapped[float]
    low: Mapped[float]
    close: Mapped[Optional[float]]
    adjusted_close: Mapped[Optional[float]]
    volume: Mapped[int]
    dividend_amount: Mapped[Optional[float]]
    timestamp: Mapped[datetime]

    __table_args__ = (
        PrimaryKeyConstraint("symbol", "timestamp"),
        CheckConstraint(
            "close IS NOT NULL OR adjusted_close IS NOT NULL",
            name="check_close_or_adjusted_close_not_null",
        ),
    )

    @classmethod
    def insert_multiple(cls, rows: List[Stock], session: session):
        stock_dicts = []
        for row in rows:
            stock_dicts.append(
                {
                    key: value
                    for key, value in row.__dict__.items()
                    if key != "_sa_instance_state"
                }
            )

        for i in range(0, len(stock_dicts), 100):
            stmt = (
                insert(Stock).values(stock_dicts[i : i + 100]).on_conflict_do_nothing()
            )
            session.execute(stmt)