diff --git a/app.py b/app.py index 31047b9..6177ba2 100644 --- a/app.py +++ b/app.py @@ -443,26 +443,9 @@ def init_db(): rollback_if_postgres(conn) ensure_kline_tables(conn) init_strategy_tables(conn) - conn.execute( - """CREATE TABLE IF NOT EXISTS account_risk_state ( - id INTEGER PRIMARY KEY CHECK (id = 1), - trading_day TEXT, - manual_close_count INTEGER DEFAULT 0, - cooloff_until_ms INTEGER, - cooloff_hours INTEGER, - daily_frozen INTEGER DEFAULT 0, - last_close_at_ms INTEGER, - updated_at TEXT - )""" - ) - if not conn.execute("SELECT id FROM account_risk_state WHERE id=1").fetchone(): - conn.execute( - "INSERT INTO account_risk_state (id, trading_day, manual_close_count, daily_frozen) " - "VALUES (1, '', 0, 0)" - ) - conn.commit() from risk.account_risk_lib import ensure_account_risk_schema from recommend_store import ensure_recommend_tables + ensure_account_risk_schema(conn) ensure_recommend_tables(conn) from ai_messages import ensure_ai_messages_table diff --git a/db_conn.py b/db_conn.py index d770b76..d6bc38e 100644 --- a/db_conn.py +++ b/db_conn.py @@ -120,6 +120,14 @@ def is_schema_migration_error(exc: BaseException) -> bool: return False +def is_missing_relation_error(exc: BaseException) -> bool: + """表/视图不存在。""" + if is_schema_migration_error(exc): + msg = str(exc).lower() + return any(x in msg for x in ("no such table", "does not exist", "undefined table")) + return False + + def rollback_if_postgres(conn: "DbConnection") -> None: if is_postgres(): try: diff --git a/risk/account_risk_lib.py b/risk/account_risk_lib.py index 05d6094..f60d094 100644 --- a/risk/account_risk_lib.py +++ b/risk/account_risk_lib.py @@ -12,7 +12,7 @@ from datetime import datetime from typing import Any, Callable, Optional, TypeVar from zoneinfo import ZoneInfo -from db_conn import OperationalError +from db_conn import OperationalError, is_missing_relation_error, rollback_if_postgres T = TypeVar("T") @@ -99,6 +99,36 @@ def trading_day_reset_hour() -> int: _SCHEMA_READY = False +ACCOUNT_RISK_STATE_SQL = """ +CREATE TABLE IF NOT EXISTS account_risk_state ( + id INTEGER NOT NULL PRIMARY KEY CHECK (id = 1), + trading_day TEXT, + manual_close_count INTEGER DEFAULT 0, + cooloff_until_ms INTEGER, + cooloff_hours INTEGER, + daily_frozen INTEGER DEFAULT 0, + last_close_at_ms INTEGER, + updated_at TEXT +) +""" + +SEED_ACCOUNT_RISK_SQL = """ +INSERT INTO account_risk_state (id, trading_day, manual_close_count, daily_frozen) +VALUES (1, '', 0, 0) +ON CONFLICT(id) DO NOTHING +""" + + +def _account_risk_table_exists(conn) -> bool: + try: + conn.execute("SELECT 1 FROM account_risk_state WHERE id=1") + return True + except Exception as exc: + if is_missing_relation_error(exc): + rollback_if_postgres(conn) + return False + raise + def _db_retry(action: Callable[[], T], *, retries: int = 8, base_delay: float = 0.03) -> T: last: OperationalError | None = None @@ -118,32 +148,11 @@ def _db_retry(action: Callable[[], T], *, retries: int = 8, base_delay: float = def ensure_account_risk_schema(conn) -> None: global _SCHEMA_READY - if _SCHEMA_READY: - try: - row = conn.execute( - "SELECT to_regclass('public.account_risk_state') AS reg" - ).fetchone() - if row and row["reg"]: - return - except Exception: - pass - _SCHEMA_READY = False - conn.execute( - """CREATE TABLE IF NOT EXISTS account_risk_state ( - id INTEGER PRIMARY KEY CHECK (id = 1), - trading_day TEXT, - manual_close_count INTEGER DEFAULT 0, - cooloff_until_ms INTEGER, - cooloff_hours INTEGER, - daily_frozen INTEGER DEFAULT 0, - last_close_at_ms INTEGER, - updated_at TEXT - )""" - ) - if not conn.execute("SELECT id FROM account_risk_state WHERE id=1").fetchone(): - conn.execute( - "INSERT INTO account_risk_state (id, trading_day, manual_close_count, daily_frozen) VALUES (1, '', 0, 0)" - ) + if _SCHEMA_READY and _account_risk_table_exists(conn): + return + _SCHEMA_READY = False + conn.execute(ACCOUNT_RISK_STATE_SQL) + conn.execute(SEED_ACCOUNT_RISK_SQL) conn.commit() _SCHEMA_READY = True @@ -330,7 +339,17 @@ def get_risk_status( ) -> dict: def _load() -> dict: ensure_account_risk_schema(conn) - row = conn.execute("SELECT * FROM account_risk_state WHERE id=1").fetchone() + try: + row = conn.execute("SELECT * FROM account_risk_state WHERE id=1").fetchone() + except Exception as exc: + if is_missing_relation_error(exc): + global _SCHEMA_READY + _SCHEMA_READY = False + rollback_if_postgres(conn) + ensure_account_risk_schema(conn) + row = conn.execute("SELECT * FROM account_risk_state WHERE id=1").fetchone() + else: + raise td = trading_day_label(now) stored = str(_row_get(row, "trading_day") or "") if stored != td: diff --git a/scripts/ensure_account_risk_pg.py b/scripts/ensure_account_risk_pg.py new file mode 100644 index 0000000..35113a7 --- /dev/null +++ b/scripts/ensure_account_risk_pg.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +"""Ensure account_risk_state exists (PostgreSQL hotfix).""" +from __future__ import annotations + +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from dotenv import load_dotenv + +load_dotenv(ROOT / ".env") + +from db_conn import connect_db, database_label # noqa: E402 +from risk.account_risk_lib import ( # noqa: E402 + _SCHEMA_READY, + ensure_account_risk_schema, +) + +import risk.account_risk_lib as risk_mod + +risk_mod._SCHEMA_READY = False +conn = connect_db() +try: + ensure_account_risk_schema(conn) + row = conn.execute("SELECT * FROM account_risk_state WHERE id=1").fetchone() + print("OK", database_label(), "row=", dict(row) if row else None) +finally: + conn.close()