fix:修复了财务数据导入过程中出现缺失列的bug。
This commit is contained in:
@@ -98,6 +98,26 @@ def safe_float(val):
|
||||
return None
|
||||
|
||||
|
||||
def _get_table_columns(table_name: str, conn) -> set:
|
||||
"""
|
||||
查询数据库表的实际列名集合
|
||||
"""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_name = %s
|
||||
""",
|
||||
(table_name,),
|
||||
)
|
||||
cols = {row[0] for row in cursor.fetchall()}
|
||||
finally:
|
||||
cursor.close()
|
||||
return cols
|
||||
|
||||
|
||||
def batch_insert(table_name: str, df: pd.DataFrame, conn, conflict_columns: List[str]):
|
||||
"""
|
||||
使用 execute_values 批量 UPSERT (INSERT ... ON CONFLICT)
|
||||
@@ -105,15 +125,44 @@ def batch_insert(table_name: str, df: pd.DataFrame, conn, conflict_columns: List
|
||||
- df: 待导入 DataFrame
|
||||
- conn: psycopg2 连接
|
||||
- conflict_columns: 冲突列 (唯一约束列),冲突时更新其他列
|
||||
|
||||
自动过滤 DataFrame 中数据库表不存在的列,避免
|
||||
'column "xxx" of relation "yyy" does not exist' 错误。
|
||||
"""
|
||||
if df.empty:
|
||||
logger.warning(f" {table_name}: 空数据,跳过")
|
||||
return 0
|
||||
|
||||
# ---- 动态过滤:只保留数据库表中存在的列 ----
|
||||
db_columns = _get_table_columns(table_name, conn)
|
||||
if not db_columns:
|
||||
# 表可能不存在,回退到不过滤(后续会抛出真实错误)
|
||||
logger.warning(f" {table_name}: 未查询到表结构,使用原始列")
|
||||
else:
|
||||
extra_cols = [c for c in df.columns if c not in db_columns]
|
||||
if extra_cols:
|
||||
logger.warning(
|
||||
f" {table_name}: 过滤掉表中不存在的列 {extra_cols}"
|
||||
)
|
||||
df = df[[c for c in df.columns if c in db_columns]]
|
||||
|
||||
if df.empty:
|
||||
logger.warning(f" {table_name}: 过滤后无可用列,跳过")
|
||||
return 0
|
||||
|
||||
# pd.NaT / pd.NaT / numpy NaN 等无法被 psycopg2 识别,统一替换为 Python None
|
||||
df = df.where(pd.notna(df), None)
|
||||
|
||||
columns = list(df.columns)
|
||||
|
||||
# 过滤后验证冲突列仍存在
|
||||
missing_conflict = [c for c in conflict_columns if c not in columns]
|
||||
if missing_conflict:
|
||||
logger.error(
|
||||
f" {table_name}: 冲突列 {missing_conflict} 不在可用列中,跳过"
|
||||
)
|
||||
return 0
|
||||
|
||||
rows = [tuple(row) for row in df.itertuples(index=False)]
|
||||
|
||||
# 构建 ON CONFLICT 子句(使用 sql.Identifier 防止注入/语法错误)
|
||||
@@ -886,6 +935,62 @@ def import_index_daily(
|
||||
# 8. 初始化数据库 Schema
|
||||
# ============================================================
|
||||
|
||||
# ---- 表结构迁移映射: 为已存在的旧表补充新列 ----
|
||||
# key: 表名, value: 需要确保存在的列 -> (列定义类型, 注释)
|
||||
_SCHEMA_MIGRATIONS = {
|
||||
"income": {
|
||||
"fv_value_chg_gain": "NUMERIC(20,4)",
|
||||
},
|
||||
"balancesheet": {
|
||||
"total_share": "NUMERIC(20,4)",
|
||||
},
|
||||
"cashflow": {
|
||||
"finan_exp": "NUMERIC(20,4)",
|
||||
},
|
||||
"fina_indicator": {
|
||||
"ca_turn": "NUMERIC(16,4)",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _migrate_schema(conn):
|
||||
"""
|
||||
迁移已存在的旧表: 使用 ALTER TABLE ... ADD COLUMN IF NOT EXISTS
|
||||
补齐 Tushare API 返回但旧 schema 缺失的列。
|
||||
|
||||
- 对已存在的表生效 (CREATE TABLE IF NOT EXISTS 不会修改旧表)
|
||||
- ADD COLUMN IF NOT EXISTS 幂等,可安全重复执行
|
||||
- 与其他存储引擎不同,PostgreSQL 的 ADD COLUMN 是 O(1) 元数据操作
|
||||
"""
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
for table_name, columns in _SCHEMA_MIGRATIONS.items():
|
||||
# 先检查表是否存在
|
||||
cursor.execute(
|
||||
"SELECT to_regclass(%s)",
|
||||
(table_name,),
|
||||
)
|
||||
if cursor.fetchone()[0] is None:
|
||||
continue
|
||||
|
||||
for col_name, col_type in columns.items():
|
||||
cursor.execute(
|
||||
sql.SQL("ALTER TABLE {} ADD COLUMN IF NOT EXISTS {} {}").format(
|
||||
sql.Identifier(table_name),
|
||||
sql.Identifier(col_name),
|
||||
sql.SQL(col_type),
|
||||
)
|
||||
)
|
||||
logger.info(f" 迁移: {table_name}.{col_name} 列已确认 ({col_type})")
|
||||
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logger.warning(f" 表结构迁移失败 (可忽略,导入时自动过滤): {e}")
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
|
||||
def init_database():
|
||||
"""
|
||||
执行 DDL,创建所有表结构
|
||||
@@ -946,6 +1051,9 @@ def init_database():
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
logger.info(" Schema 初始化完成")
|
||||
|
||||
# 迁移已存在的旧表: 补齐 Tushare API 新增列
|
||||
_migrate_schema(conn)
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
logger.error(f" Schema 初始化失败: {e}")
|
||||
|
||||
Reference in New Issue
Block a user