2025-04-18

[MtkN1]

Today I Learned ...

[MtkN1]

CSV ファイルを PostgreSQL テーブルに UPSERT する

import csv
from datetime import datetime as datetime_
from decimal import Decimal

from sqlalchemy import DateTime, Float, Numeric, create_engine
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker


class Base(DeclarativeBase): ...


class Stock(Base):
    __tablename__ = "stock_data"

    timestamp: Mapped[float] = mapped_column(Float, primary_key=True)
    open: Mapped[Decimal] = mapped_column(Numeric(10, 2))
    high: Mapped[Decimal] = mapped_column(Numeric(10, 2))
    low: Mapped[Decimal] = mapped_column(Numeric(10, 2))
    close: Mapped[Decimal] = mapped_column(Numeric(10, 2))
    volume: Mapped[Decimal] = mapped_column(Numeric(10, 2))
    datetime: Mapped[datetime_] = mapped_column(DateTime)

    def __repr__(self) -> str:
        return (
            f"<Stock(timestamp={self.timestamp}, open={self.open}, "
            f"high={self.high}, low={self.low}, close={self.close}, "
            f"volume={self.volume}, datetime={self.datetime})>"
        )


if __name__ == "__main__":
    before = datetime_.now()

    engine = create_engine("postgresql+psycopg://postgres@localhost:46234")
    Session = sessionmaker(engine)

    with (
        open("./btcusd_1-min_data.csv", "r", encoding="utf-8", newline="") as csvfile,
        Session.begin() as session,
    ):
        reader = csv.DictReader(csvfile)
        for row in reader:
            insert_stmt = insert(Stock).values(
                timestamp=float(row["Timestamp"]),
                open=Decimal(row["Open"]),
                high=Decimal(row["High"]),
                low=Decimal(row["Low"]),
                close=Decimal(row["Close"]),
                volume=Decimal(row["Volume"]),
                datetime=datetime_.fromisoformat(row["datetime"]),
            )
            upsert_stmt = insert_stmt.on_conflict_do_update(
                index_elements=[Stock.timestamp],
                set_={
                    "open": insert_stmt.excluded.open,
                    "high": insert_stmt.excluded.high,
                    "low": insert_stmt.excluded.low,
                    "close": insert_stmt.excluded.close,
                    "volume": insert_stmt.excluded.volume,
                    "datetime": insert_stmt.excluded.datetime,
                },
            )
            session.execute(upsert_stmt)

    after = datetime_.now()
    print(f"Execution time: {after - before}")