From cdf5582c901c448c93fedc007fe5eca85b3af1f3 Mon Sep 17 00:00:00 2001 From: shellway-pc <413209390@qq.com> Date: Sat, 1 Aug 2026 12:56:39 +0800 Subject: [PATCH] =?UTF-8?q?fix=EF=BC=9A=E4=BF=AE=E5=A4=8D=E5=AE=8C?= =?UTF-8?q?=E6=88=90=E4=BA=86=E6=95=B0=E6=8D=AE=E5=BA=93=E9=97=AE=E9=A2=98?= =?UTF-8?q?=EF=BC=8C=E5=B7=B2=E7=BB=8F=E5=8F=AF=E4=BB=A5=E6=AD=A3=E5=B8=B8?= =?UTF-8?q?=E5=AF=BC=E5=85=A5=E6=95=B0=E6=8D=AE=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- quantitative_data/importer.py | 93 +++++++---------------------------- 1 file changed, 17 insertions(+), 76 deletions(-) diff --git a/quantitative_data/importer.py b/quantitative_data/importer.py index a7b3d13..5a39d61 100644 --- a/quantitative_data/importer.py +++ b/quantitative_data/importer.py @@ -795,19 +795,18 @@ def init_database(): logger.info("=" * 60) logger.info("初始化数据库 Schema ...") - # 尝试创建数据库 (如果不存在) + # 尝试创建数据库（如果不存在） try: admin_conn = psycopg2.connect( host=DB_CONFIG["host"], port=DB_CONFIG["port"], - database="postgres", # 连接默认 postgres 库来创建新库 + database="postgres", user=DB_CONFIG["user"], password=DB_CONFIG["password"], ) admin_conn.autocommit = True cursor = admin_conn.cursor() - # 检查数据库是否存在 cursor.execute( "SELECT 1 FROM pg_database WHERE datname = %s", (DB_CONFIG["database"],), @@ -828,82 +827,25 @@ def init_database(): logger.warning(f" 创建数据库步骤跳过 (可能无权限): {e}") # 执行 DDL + import os as _os + import re as _re + + schema_path = _os.path.join( + _os.path.dirname(_os.path.abspath(__file__)), "schema.sql" + ) + with open(schema_path, "r", encoding="utf-8") as f: + ddl_sql = f.read() + + # 先去掉单行注释，再按分号分句 + ddl_sql = _re.sub(r"^\s*--.*$", "", ddl_sql, flags=_re.MULTILINE) + statements = [s.strip() for s in ddl_sql.split(";") if s.strip()] + conn = get_pg_connection() try: - import os as _os - - # 尝试使用 sqlparse 按语句拆分 (处理含分号的字符串字面量、函数体等) - try: - import sqlparse as _sqlparse - _USE_SQLPARSE = True - except ImportError: - _USE_SQLPARSE = False - logger.warning(" sqlparse 未安装，回退为简单分号拆分") - - schema_path = _os.path.join(_os.path.dirname(_os.path.abspath(__file__)), "schema.sql") - with open(schema_path, "r", encoding="utf-8") as f: - ddl_sql = f.read() - - if _USE_SQLPARSE: - raw_statements = _sqlparse.split(ddl_sql) - statements = [ - s.strip() for s in raw_statements - if s.strip() and not s.strip().startswith("--") - ] - else: - statements = [ - s.strip() for s in ddl_sql.split(";") - if s.strip() and not s.strip().startswith("--") - ] - - # 按依赖关系排序：CREATE TABLE → CREATE INDEX → COMMENT ON → CREATE VIEW - # 避免 sqlparse 拆分后语句乱序导致 UndefinedTable 错误 - _priority = { - "CREATE TABLE": 1, - "CREATE INDEX": 2, - "COMMENT ON": 3, - "CREATE OR REPLACE VIEW": 4, - "CREATE VIEW": 4, - } - - def _stmt_priority(s): - upper = s.upper() - for keyword, prio in _priority.items(): - if upper.startswith(keyword): - return prio - return 99 # 兜底：最后执行 - - statements.sort(key=_stmt_priority) - - # 设置 autocommit 模式：每条 DDL 独立事务，互不影响 - # 否则 rollback 会撤销之前已成功执行的 CREATE TABLE - conn.autocommit = True - cursor = conn.cursor() for stmt in statements: - if stmt and not stmt.startswith("--"): - try: - cursor.execute(stmt) - except Exception as e: - stmt_upper = stmt.upper() - # 以下类型的语句失败视为可恢复的 warning，不中断整个初始化流程： - # 1. 视图创建（依赖的基础表可能尚未创建） - # 2. 索引创建（依赖的表可能尚未导入数据） - # 3. COMMENT ON（依赖的表/视图可能尚未创建） - is_recoverable = ( - "CREATE OR REPLACE VIEW" in stmt_upper - or "CREATE VIEW" in stmt_upper - or "CREATE INDEX" in stmt_upper - or "COMMENT ON" in stmt_upper - ) - if is_recoverable: - logger.warning(f" DDL 暂跳过（依赖尚未就绪）: {str(e)[:150]}\n SQL: {stmt[:200]}") - else: - logger.error(f" DDL 执行失败: {str(e)[:200]}\n SQL: {stmt[:300]}") - cursor.close() - conn.close() - raise - + cursor.execute(stmt) + conn.commit() cursor.close() logger.info(" Schema 初始化完成") except Exception as e: @@ -913,7 +855,6 @@ def init_database(): finally: conn.close() - # ============================================================ # 9. 获取所有股票列表 (辅助) # ============================================================