From 6abe06d9357b08d3ca40f0787f3454ef1c13fe8f Mon Sep 17 00:00:00 2001
From: dekun <dekun@local>
Date: Wed, 1 Jul 2026 08:22:54 +0800
Subject: [PATCH] Ensure all PG tables on init; fix migration commits per
 table.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 app.py                                |  9 +++++++++
 scripts/migrate_sqlite_to_postgres.py | 20 +++++++++-----------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/app.py b/app.py
index 5d09d20..3caa530 100644
--- a/app.py
+++ b/app.py
@@ -301,6 +301,12 @@ def get_stats_data() -> dict:
 
 
 def init_db():
+    import strategy.strategy_db as strategy_db
+    import risk.account_risk_lib as account_risk_lib
+
+    strategy_db._TABLES_READY = False
+    account_risk_lib._SCHEMA_READY = False
+
     conn = get_db()
     c = conn.cursor()
     c.execute("CREATE TABLE IF NOT EXISTS settings (key TEXT PRIMARY KEY, value TEXT)")
@@ -441,6 +447,9 @@ def init_db():
     from recommend_store import ensure_recommend_tables
     ensure_account_risk_schema(conn)
     ensure_recommend_tables(conn)
+    from ai_messages import ensure_ai_messages_table
+
+    ensure_ai_messages_table(conn)
     conn.commit()
     conn.close()
 
diff --git a/scripts/migrate_sqlite_to_postgres.py b/scripts/migrate_sqlite_to_postgres.py
index 66aec55..cfdd579 100644
--- a/scripts/migrate_sqlite_to_postgres.py
+++ b/scripts/migrate_sqlite_to_postgres.py
@@ -38,14 +38,10 @@ def _pg_columns(pg_conn, table: str) -> list[str]:
 
 
 def _reset_sequences(pg_conn, table: str, pk: str = "id") -> None:
-    try:
-        pg_conn.execute(
-            f"SELECT setval(pg_get_serial_sequence('{table}', '{pk}'), "
-            f"COALESCE((SELECT MAX({pk}) FROM {table}), 1), true)"
-        )
-        pg_conn.commit()
-    except Exception:
-        rollback_if_postgres(pg_conn)
+    pg_conn.execute(
+        f"SELECT setval(pg_get_serial_sequence('{table}', '{pk}'), "
+        f"COALESCE((SELECT MAX({pk}) FROM {table}), 1), true)"
+    )
 
 
 def migrate(*, sqlite_path: str | None = None, dry_run: bool = False) -> dict:
@@ -94,9 +90,11 @@ def migrate(*, sqlite_path: str | None = None, dry_run: bool = False) -> dict:
                 for row in rows:
                     dst.execute(insert_sql, tuple(row[c] for c in cols))
                 if "id" in cols:
-                    _reset_sequences(dst, table, "id")
-                else:
-                    dst.commit()
+                    try:
+                        _reset_sequences(dst, table, "id")
+                    except Exception:
+                        rollback_if_postgres(dst)
+                dst.commit()
                 stats[table] = len(rows)
                 print(f"  {table}: {len(rows)} 行")
             except Exception as exc: