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}")
|
||||
|
||||
@@ -216,7 +216,8 @@ CREATE TABLE IF NOT EXISTS income (
|
||||
n_sec_uw_income NUMERIC(20,4), -- 证券承销业务净收入
|
||||
n_asset_mg_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), -- 加:投资净收益
|
||||
ass_invest_income NUMERIC(20,4), -- 其中:对联营企业和合营企业的投资收益
|
||||
forex_gain NUMERIC(20,4), -- 加:汇兑净收益
|
||||
@@ -325,6 +326,9 @@ CREATE TABLE IF NOT EXISTS balancesheet (
|
||||
decr_in_disbur NUMERIC(20,4), -- 待处理流动资产损益
|
||||
oth_nca NUMERIC(20,4), -- 其他非流动资产
|
||||
|
||||
-- 股东权益相关
|
||||
total_share NUMERIC(20,4), -- 总股本 (Tushare balancesheet 返回字段)
|
||||
|
||||
-- 负债
|
||||
total_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), -- 净利润
|
||||
fin_exp NUMERIC(20,4), -- 财务费用
|
||||
fin_exp NUMERIC(20,4), -- 财务费用 (兼容旧版字段)
|
||||
finan_exp NUMERIC(20,4), -- 财务费用 (Tushare cashflow 实际字段)
|
||||
c_fr_sale_sg NUMERIC(20,4), -- 销售商品、提供劳务收到的现金
|
||||
recp_tax_rends 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), -- 应收账款周转天数
|
||||
inv_turn NUMERIC(16,4), -- 存货周转率
|
||||
ar_turn NUMERIC(16,4), -- 应收账款周转率
|
||||
ca_turn NUMERIC(16,4), -- 流动资产周转率 (Tushare fina_indicator 实际字段)
|
||||
assets_turn NUMERIC(16,4), -- 总资产周转率
|
||||
-- 盈利能力
|
||||
roe NUMERIC(16,4), -- 净资产收益率
|
||||
|
||||
Reference in New Issue
Block a user