Fix account_risk_state missing on PostgreSQL: probe table before cache skip.

Co-authored-by: Cursor <cursoragent@cursor.com>

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: