From 738df559cb0efeb77c787e6f1a508c03091124db Mon Sep 17 00:00:00 2001 From: Tianyu Liu Date: Sat, 13 Sep 2025 18:46:16 +0200 Subject: [PATCH] add db models --- backend/ruff.toml | 4 +- backend/trading_journal/db.py | 6 ++- backend/trading_journal/db_migration.py | 32 +++--------- backend/trading_journal/models.py | 9 +++- backend/trading_journal/models_v1.py | 68 +++++++++++++++++++++++++ 5 files changed, 90 insertions(+), 29 deletions(-) create mode 100644 backend/trading_journal/models_v1.py diff --git a/backend/ruff.toml b/backend/ruff.toml index a4b5c62..4bcadc2 100644 --- a/backend/ruff.toml +++ b/backend/ruff.toml @@ -4,7 +4,7 @@ line-length = 144 [lint] select = ["ALL"] fixable = ["UP034", "I001"] -ignore = ["T201", "D", "ANN101", "TD002", "TD003"] +ignore = ["T201", "D", "ANN101", "TD002", "TD003", "TRY003", "EM102"] [lint.extend-per-file-ignores] -"test*.py" = ["S101"] \ No newline at end of file +"test*.py" = ["S101"] diff --git a/backend/trading_journal/db.py b/backend/trading_journal/db.py index c147df7..229284a 100644 --- a/backend/trading_journal/db.py +++ b/backend/trading_journal/db.py @@ -7,7 +7,7 @@ from sqlalchemy import event from sqlalchemy.pool import StaticPool from sqlmodel import Session, create_engine -import trading_journal.db_migration +from trading_journal import db_migration if TYPE_CHECKING: from collections.abc import Generator @@ -33,6 +33,8 @@ class Database: def _enable_sqlite_pragmas(dbapi_conn: DBAPIConnection, _connection_record: object) -> None: try: cur = dbapi_conn.cursor() + cur.execute("PRAGMA journal_mode=WAL;") + cur.execute("PRAGMA synchronous=NORMAL;") cur.execute("PRAGMA foreign_keys=ON;") cur.execute("PRAGMA busy_timeout=30000;") cur.close() @@ -43,7 +45,7 @@ class Database: event.listen(self._engine, "connect", _enable_sqlite_pragmas) def init_db(self) -> None: - pass + db_migration.run_migrations(self._engine) def get_session(self) -> Generator[Session, None, None]: session = Session(self._engine) diff --git a/backend/trading_journal/db_migration.py b/backend/trading_journal/db_migration.py index 678634f..fe9b0ba 100644 --- a/backend/trading_journal/db_migration.py +++ b/backend/trading_journal/db_migration.py @@ -1,12 +1,13 @@ from __future__ import annotations -from typing import Callable +from typing import TYPE_CHECKING, Callable from sqlalchemy import text -from sqlalchemy.engine import Connection, Engine from sqlmodel import SQLModel -# 最新 schema 版本号 +if TYPE_CHECKING: + from sqlalchemy.engine import Connection, Engine + LATEST_VERSION = 1 @@ -17,6 +18,8 @@ def _mig_0_1(engine: Engine) -> None: """ # Ensure all models are imported before this is called (import side-effect registers tables) # e.g. trading_journal.models is imported in the caller / app startup. + from trading_journal import models_v1 # noqa: PLC0415, F401 + SQLModel.metadata.create_all(bind=engine) @@ -31,7 +34,7 @@ def _get_sqlite_user_version(conn: Connection) -> int: return int(row[0]) if row and row[0] is not None else 0 -def _set_sqlite_user_version(conn, v: int) -> None: +def _set_sqlite_user_version(conn: Connection, v: int) -> None: conn.execute(text(f"PRAGMA user_version = {int(v)}")) @@ -54,23 +57,4 @@ def run_migrations(engine: Engine, target_version: int | None = None) -> int: _set_sqlite_user_version(conn, cur_version + 1) cur_version += 1 return cur_version - else: - # generic migrations table for non-sqlite - conn.execute( - text(""" - CREATE TABLE IF NOT EXISTS migrations ( - version INTEGER PRIMARY KEY, - applied_at TEXT DEFAULT CURRENT_TIMESTAMP - ) - """) - ) - row = conn.execute(text("SELECT MAX(version) FROM migrations")).fetchone() - cur_version = int(row[0]) if row and row[0] is not None else 0 - while cur_version < target: - fn = MIGRATIONS.get(cur_version) - if fn is None: - raise RuntimeError(f"No migration from {cur_version} -> {cur_version + 1}") - fn(engine) - conn.execute(text("INSERT INTO migrations(version) VALUES (:v)"), {"v": cur_version + 1}) - cur_version += 1 - return cur_version + return -1 # unknown / unsupported driver; no-op diff --git a/backend/trading_journal/models.py b/backend/trading_journal/models.py index 051c876..7382741 100644 --- a/backend/trading_journal/models.py +++ b/backend/trading_journal/models.py @@ -2,7 +2,6 @@ from __future__ import annotations from datetime import date, datetime # noqa: TC003 from enum import Enum -from typing import TYPE_CHECKING from sqlmodel import Column, DateTime, Field, Relationship, SQLModel from sqlmodel import Enum as SQLEnum @@ -15,6 +14,13 @@ class TradeType(str, Enum): EXERCISE_CALL = "EXERCISE_CALL" +class TradeStrategy(str, Enum): + WHEELS = "WHEEL" + FX = "FX" + SPOT = "SPOT" + OTHER = "OTHER" + + class CycleStatus(str, Enum): OPEN = "OPEN" CLOSED = "CLOSED" @@ -33,6 +39,7 @@ class Trades(SQLModel, table=True): symbol: str underlying_currency: str trade_type: TradeType = Field(sa_column=Column(SQLEnum(TradeType, name="trade_type_enum"), nullable=False)) + trade_strategy: TradeStrategy = Field(sa_column=Column(SQLEnum(TradeStrategy, name="trade_strategy_enum"), nullable=False)) trade_time_utc: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False)) expiry_date: date | None = Field(default=None, nullable=True) strike_price_cents: int | None = Field(default=None, nullable=True) diff --git a/backend/trading_journal/models_v1.py b/backend/trading_journal/models_v1.py new file mode 100644 index 0000000..7382741 --- /dev/null +++ b/backend/trading_journal/models_v1.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from datetime import date, datetime # noqa: TC003 +from enum import Enum + +from sqlmodel import Column, DateTime, Field, Relationship, SQLModel +from sqlmodel import Enum as SQLEnum + + +class TradeType(str, Enum): + SELL_PUT = "SELL_PUT" + ASSIGNMENT = "ASSIGNMENT" + SELL_CALL = "SELL_CALL" + EXERCISE_CALL = "EXERCISE_CALL" + + +class TradeStrategy(str, Enum): + WHEELS = "WHEEL" + FX = "FX" + SPOT = "SPOT" + OTHER = "OTHER" + + +class CycleStatus(str, Enum): + OPEN = "OPEN" + CLOSED = "CLOSED" + + +class FundingSource(str, Enum): + CASH = "CASH" + MARGIN = "MARGIN" + MIXED = "MIXED" + + +class Trades(SQLModel, table=True): + __tablename__ = "trades" + id: str | None = Field(default=None, primary_key=True) + user_id: str + symbol: str + underlying_currency: str + trade_type: TradeType = Field(sa_column=Column(SQLEnum(TradeType, name="trade_type_enum"), nullable=False)) + trade_strategy: TradeStrategy = Field(sa_column=Column(SQLEnum(TradeStrategy, name="trade_strategy_enum"), nullable=False)) + trade_time_utc: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False)) + expiry_date: date | None = Field(default=None, nullable=True) + strike_price_cents: int | None = Field(default=None, nullable=True) + quantity: int + price_cents: int + gross_cash_flow_cents: int + commission_cents: int + net_cash_flow_cents: int + cycle_id: str | None = Field(default=None, foreign_key="cycles.id", nullable=True) + cycle: Cycles | None = Relationship(back_populates="trades") + + +class Cycles(SQLModel, table=True): + __tablename__ = "cycles" + id: str | None = Field(default=None, primary_key=True) + user_id: str + symbol: str + underlying_currency: str + start_date: date + end_date: date | None = Field(default=None, nullable=True) + status: CycleStatus = Field(sa_column=Column(SQLEnum(CycleStatus, name="cycle_status_enum"), nullable=False)) + funding_source: FundingSource = Field(sa_column=Column(SQLEnum(FundingSource, name="funding_source_enum"), nullable=False)) + capital_exposure_cents: int + loan_amount_cents: int | None = Field(default=None, nullable=True) + loan_interest_rate_bps: int | None = Field(default=None, nullable=True) + trades: list[Trades] = Relationship(back_populates="cycle")