fix:修复了财务数据导入过程中出现缺失列的bug。

This commit is contained in:
2026-08-05 22:07:30 +08:00
parent 45a554585b
commit b6035a7edd
2 changed files with 116 additions and 2 deletions
+108
View File
@@ -98,6 +98,26 @@ def safe_float(val):
return None 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]): def batch_insert(table_name: str, df: pd.DataFrame, conn, conflict_columns: List[str]):
""" """
使用 execute_values 批量 UPSERT (INSERT ... ON CONFLICT) 使用 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 - df: 待导入 DataFrame
- conn: psycopg2 连接 - conn: psycopg2 连接
- conflict_columns: 冲突列 (唯一约束列),冲突时更新其他列 - conflict_columns: 冲突列 (唯一约束列),冲突时更新其他列
自动过滤 DataFrame 中数据库表不存在的列,避免
'column "xxx" of relation "yyy" does not exist' 错误。
""" """
if df.empty: if df.empty:
logger.warning(f" {table_name}: 空数据,跳过") logger.warning(f" {table_name}: 空数据,跳过")
return 0 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 # pd.NaT / pd.NaT / numpy NaN 等无法被 psycopg2 识别,统一替换为 Python None
df = df.where(pd.notna(df), None) df = df.where(pd.notna(df), None)
columns = list(df.columns) 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)] rows = [tuple(row) for row in df.itertuples(index=False)]
# 构建 ON CONFLICT 子句(使用 sql.Identifier 防止注入/语法错误) # 构建 ON CONFLICT 子句(使用 sql.Identifier 防止注入/语法错误)
@@ -886,6 +935,62 @@ def import_index_daily(
# 8. 初始化数据库 Schema # 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(): def init_database():
""" """
执行 DDL,创建所有表结构 执行 DDL,创建所有表结构
@@ -946,6 +1051,9 @@ def init_database():
conn.commit() conn.commit()
cursor.close() cursor.close()
logger.info(" Schema 初始化完成") logger.info(" Schema 初始化完成")
# 迁移已存在的旧表: 补齐 Tushare API 新增列
_migrate_schema(conn)
except Exception as e: except Exception as e:
conn.rollback() conn.rollback()
logger.error(f" Schema 初始化失败: {e}") logger.error(f" Schema 初始化失败: {e}")
+8 -2
View File
@@ -216,7 +216,8 @@ CREATE TABLE IF NOT EXISTS income (
n_sec_uw_income NUMERIC(20,4), -- 证券承销业务净收入 n_sec_uw_income NUMERIC(20,4), -- 证券承销业务净收入
n_asset_mg_income NUMERIC(20,4), -- 受托客户资产管理业务净收入 n_asset_mg_income NUMERIC(20,4), -- 受托客户资产管理业务净收入
oth_b_income NUMERIC(20,4), -- 其他业务收入 oth_b_income NUMERIC(20,4), -- 其他业务收入
fv_value_chg NUMERIC(20,4), -- 加:公允价值变动净收益 fv_value_chg NUMERIC(20,4), -- 加:公允价值变动净收益 (兼容旧版字段)
fv_value_chg_gain NUMERIC(20,4), -- 加:公允价值变动净收益 (Tushare income 实际字段)
invest_income NUMERIC(20,4), -- 加:投资净收益 invest_income NUMERIC(20,4), -- 加:投资净收益
ass_invest_income NUMERIC(20,4), -- 其中:对联营企业和合营企业的投资收益 ass_invest_income NUMERIC(20,4), -- 其中:对联营企业和合营企业的投资收益
forex_gain NUMERIC(20,4), -- 加:汇兑净收益 forex_gain NUMERIC(20,4), -- 加:汇兑净收益
@@ -325,6 +326,9 @@ CREATE TABLE IF NOT EXISTS balancesheet (
decr_in_disbur NUMERIC(20,4), -- 待处理流动资产损益 decr_in_disbur NUMERIC(20,4), -- 待处理流动资产损益
oth_nca NUMERIC(20,4), -- 其他非流动资产 oth_nca NUMERIC(20,4), -- 其他非流动资产
-- 股东权益相关
total_share NUMERIC(20,4), -- 总股本 (Tushare balancesheet 返回字段)
-- 负债 -- 负债
total_liab NUMERIC(20,4), -- 负债合计 total_liab NUMERIC(20,4), -- 负债合计
total_cur_liab NUMERIC(20,4), -- 流动负债合计 total_cur_liab NUMERIC(20,4), -- 流动负债合计
@@ -363,7 +367,8 @@ CREATE TABLE IF NOT EXISTS cashflow (
-- 经营活动 -- 经营活动
net_profit NUMERIC(20,4), -- 净利润 net_profit NUMERIC(20,4), -- 净利润
fin_exp NUMERIC(20,4), -- 财务费用 fin_exp NUMERIC(20,4), -- 财务费用 (兼容旧版字段)
finan_exp NUMERIC(20,4), -- 财务费用 (Tushare cashflow 实际字段)
c_fr_sale_sg NUMERIC(20,4), -- 销售商品、提供劳务收到的现金 c_fr_sale_sg NUMERIC(20,4), -- 销售商品、提供劳务收到的现金
recp_tax_rends NUMERIC(20,4), -- 收到的税费返还 recp_tax_rends NUMERIC(20,4), -- 收到的税费返还
n_depos_incr_fi NUMERIC(20,4), -- 客户存款和同业存放款项净增加额 n_depos_incr_fi NUMERIC(20,4), -- 客户存款和同业存放款项净增加额
@@ -429,6 +434,7 @@ CREATE TABLE IF NOT EXISTS fina_indicator (
arturn_days NUMERIC(16,4), -- 应收账款周转天数 arturn_days NUMERIC(16,4), -- 应收账款周转天数
inv_turn NUMERIC(16,4), -- 存货周转率 inv_turn NUMERIC(16,4), -- 存货周转率
ar_turn NUMERIC(16,4), -- 应收账款周转率 ar_turn NUMERIC(16,4), -- 应收账款周转率
ca_turn NUMERIC(16,4), -- 流动资产周转率 (Tushare fina_indicator 实际字段)
assets_turn NUMERIC(16,4), -- 总资产周转率 assets_turn NUMERIC(16,4), -- 总资产周转率
-- 盈利能力 -- 盈利能力
roe NUMERIC(16,4), -- 净资产收益率 roe NUMERIC(16,4), -- 净资产收益率