1420 lines
44 KiB
Python
1420 lines
44 KiB
Python
"""
|
||
量化数据导入核心模块
|
||
连接 Docker PostgreSQL (192.168.27.11:5438)
|
||
从 Tushare 拉取数据并批量导入
|
||
"""
|
||
import time
|
||
import logging
|
||
from datetime import datetime, timedelta
|
||
from typing import Optional, List, Dict
|
||
|
||
import tushare as ts
|
||
import pandas as pd
|
||
import psycopg2
|
||
from psycopg2 import sql
|
||
from psycopg2.extras import execute_values
|
||
from sqlalchemy import create_engine
|
||
|
||
from config import DB_CONFIG, TUSHARE_TOKEN, BATCH_SIZE, START_DATE, END_DATE, PASSWORD_ENCODED
|
||
|
||
# ============================================================
|
||
# 日志配置
|
||
# ============================================================
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
handlers=[
|
||
logging.FileHandler("import_data.log", encoding="utf-8"),
|
||
logging.StreamHandler(),
|
||
],
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ============================================================
|
||
# 初始化连接
|
||
# ============================================================
|
||
|
||
# psycopg2 原生连接 (用于执行 DDL / 精细控制)
|
||
def get_pg_connection():
|
||
"""获取 psycopg2 原生连接"""
|
||
return psycopg2.connect(
|
||
host=DB_CONFIG["host"],
|
||
port=DB_CONFIG["port"],
|
||
database=DB_CONFIG["database"],
|
||
user=DB_CONFIG["user"],
|
||
password=DB_CONFIG["password"],
|
||
)
|
||
|
||
|
||
# SQLAlchemy 引擎 (用于 DataFrame.to_sql)
|
||
_sqlalchemy_engine = None
|
||
|
||
|
||
def get_sqlalchemy_engine():
|
||
"""获取 SQLAlchemy 引擎 (单例)"""
|
||
global _sqlalchemy_engine
|
||
if _sqlalchemy_engine is None:
|
||
db_url = (
|
||
f"postgresql://{DB_CONFIG['user']}:{PASSWORD_ENCODED}"
|
||
f"@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
|
||
)
|
||
_sqlalchemy_engine = create_engine(db_url, pool_size=5, max_overflow=10)
|
||
return _sqlalchemy_engine
|
||
|
||
|
||
# Tushare Pro API (单例)
|
||
_ts_pro = None
|
||
|
||
|
||
def get_ts_pro():
|
||
"""获取 Tushare Pro API 实例 (单例)"""
|
||
global _ts_pro
|
||
if _ts_pro is None:
|
||
ts.set_token(TUSHARE_TOKEN)
|
||
_ts_pro = ts.pro_api()
|
||
logger.info("Tushare Pro API 初始化完成")
|
||
return _ts_pro
|
||
|
||
|
||
# ============================================================
|
||
# 通用导入工具函数
|
||
# ============================================================
|
||
|
||
def normalize_columns(df: pd.DataFrame) -> pd.DataFrame:
|
||
"""
|
||
标准化列名：Tushare 返回的列名可能有大小写差异，统一转小写
|
||
"""
|
||
df.columns = [c.lower() for c in df.columns]
|
||
return df
|
||
|
||
|
||
def safe_float(val):
|
||
"""安全转换为浮点数，NaN -> None"""
|
||
try:
|
||
if pd.isna(val):
|
||
return None
|
||
return float(val)
|
||
except (ValueError, TypeError):
|
||
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)
|
||
- table_name: 目标表名
|
||
- 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 防止注入/语法错误）
|
||
conflict_identifiers = sql.SQL(", ").join(map(sql.Identifier, conflict_columns))
|
||
# 构建 UPDATE SET 子句 (排除冲突列)
|
||
update_cols = [c for c in columns if c not in conflict_columns]
|
||
if not update_cols:
|
||
# 没有需要更新的列，使用 DO NOTHING
|
||
upsert_sql = sql.SQL(
|
||
"INSERT INTO {table} ({cols}) VALUES %s "
|
||
"ON CONFLICT ({conflict}) DO NOTHING"
|
||
).format(
|
||
table=sql.Identifier(table_name),
|
||
cols=sql.SQL(", ").join(map(sql.Identifier, columns)),
|
||
conflict=conflict_identifiers,
|
||
)
|
||
else:
|
||
update_set = sql.SQL(", ").join(
|
||
sql.SQL("{col} = EXCLUDED.{col}").format(col=sql.Identifier(c))
|
||
for c in update_cols
|
||
)
|
||
upsert_sql = sql.SQL(
|
||
"INSERT INTO {table} ({cols}) VALUES %s "
|
||
"ON CONFLICT ({conflict}) DO UPDATE SET {update_set}"
|
||
).format(
|
||
table=sql.Identifier(table_name),
|
||
cols=sql.SQL(", ").join(map(sql.Identifier, columns)),
|
||
conflict=conflict_identifiers,
|
||
update_set=update_set,
|
||
)
|
||
|
||
cursor = conn.cursor()
|
||
try:
|
||
execute_values(cursor, upsert_sql.as_string(cursor), rows, page_size=BATCH_SIZE)
|
||
conn.commit()
|
||
logger.info(f" {table_name}: 成功导入 {len(rows)} 条记录")
|
||
return len(rows)
|
||
except Exception as e:
|
||
conn.rollback()
|
||
logger.error(f" {table_name}: 批量导入失败 - {e}")
|
||
raise
|
||
finally:
|
||
cursor.close()
|
||
|
||
|
||
def fetch_with_retry(fetch_func, max_retries: int = 3, delay: float = 2.0):
|
||
"""
|
||
带重试的数据获取装饰器
|
||
- fetch_func: 数据获取函数 (返回 DataFrame)
|
||
- max_retries: 最大重试次数
|
||
- delay: 重试间隔(秒)
|
||
"""
|
||
for attempt in range(max_retries):
|
||
try:
|
||
result = fetch_func()
|
||
if result is not None and not result.empty:
|
||
return result
|
||
logger.warning(f" 第 {attempt+1} 次获取返回空数据，重试...")
|
||
except Exception as e:
|
||
logger.warning(f" 第 {attempt+1} 次获取失败: {e}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(delay * (attempt + 1)) # 递增延迟
|
||
return pd.DataFrame()
|
||
|
||
|
||
# ============================================================
|
||
# 1. 导入股票基本信息
|
||
# ============================================================
|
||
|
||
def import_stock_basic():
|
||
"""
|
||
导入股票基本信息 (stock_basic)
|
||
Tushare: stock_basic
|
||
"""
|
||
# 确保数据库表结构已初始化
|
||
init_database()
|
||
|
||
logger.info("=" * 60)
|
||
logger.info("[1/7] 导入股票基本信息 (stock_basic) ...")
|
||
|
||
pro = get_ts_pro()
|
||
conn = get_pg_connection()
|
||
|
||
try:
|
||
# 获取全量股票基本信息
|
||
df = pro.stock_basic(
|
||
exchange="",
|
||
list_status="L",
|
||
fields="ts_code,symbol,name,area,industry,market,list_status,list_date,is_hs,act_name,act_ent_type",
|
||
)
|
||
if df is None or df.empty:
|
||
logger.warning("未获取到股票基本信息")
|
||
return
|
||
|
||
df = normalize_columns(df)
|
||
|
||
# 转换日期格式
|
||
if "list_date" in df.columns:
|
||
df["list_date"] = pd.to_datetime(df["list_date"], format="%Y%m%d", errors="coerce")
|
||
|
||
logger.info(f" 获取到 {len(df)} 条股票基本信息")
|
||
|
||
# 批量导入
|
||
conflict_cols = ["ts_code"]
|
||
batch_insert("stock_basic", df, conn, conflict_cols)
|
||
|
||
# 也尝试获取退市的股票
|
||
try:
|
||
df_d = pro.stock_basic(
|
||
exchange="",
|
||
list_status="D",
|
||
fields="ts_code,symbol,name,area,industry,market,list_status,list_date,is_hs",
|
||
)
|
||
if df_d is not None and not df_d.empty:
|
||
df_d = normalize_columns(df_d)
|
||
if "list_date" in df_d.columns:
|
||
df_d["list_date"] = pd.to_datetime(
|
||
df_d["list_date"], format="%Y%m%d", errors="coerce"
|
||
)
|
||
batch_insert("stock_basic", df_d, conn, conflict_cols)
|
||
logger.info(f" 额外导入 {len(df_d)} 条退市股票信息")
|
||
except Exception as e:
|
||
logger.warning(f" 获取退市股票信息失败: {e}")
|
||
|
||
except Exception as e:
|
||
logger.error(f" 导入股票基本信息失败: {e}")
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
# ============================================================
|
||
# 2. 导入交易日历
|
||
# ============================================================
|
||
|
||
def import_trade_cal(start_date: Optional[str] = None, end_date: Optional[str] = None):
|
||
"""
|
||
导入交易日历 (trade_cal)
|
||
Tushare: trade_cal
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
logger.info("=" * 60)
|
||
logger.info(f"[2/7] 导入交易日历 (trade_cal): {start_date} ~ {end_date}")
|
||
|
||
pro = get_ts_pro()
|
||
conn = get_pg_connection()
|
||
|
||
try:
|
||
# 上交所
|
||
for exchange, ex_name in [("SSE", "上交所"), ("SZSE", "深交所")]:
|
||
df = pro.trade_cal(
|
||
exchange=exchange,
|
||
start_date=start_date.replace("-", ""),
|
||
end_date=end_date.replace("-", ""),
|
||
)
|
||
if df is None or df.empty:
|
||
logger.warning(f" {ex_name} 交易日历为空")
|
||
continue
|
||
|
||
df = normalize_columns(df)
|
||
|
||
# 转换日期
|
||
for col in ["cal_date", "pretrade_date"]:
|
||
if col in df.columns:
|
||
df[col] = pd.to_datetime(df[col], format="%Y%m%d", errors="coerce")
|
||
|
||
# 重命名 is_open (Tushare 返回 0/1 整数)
|
||
if "is_open" in df.columns:
|
||
df["is_open"] = df["is_open"].astype(int)
|
||
|
||
conflict_cols = ["exchange", "cal_date"]
|
||
batch_insert("trade_cal", df, conn, conflict_cols)
|
||
logger.info(f" {ex_name}: {len(df)} 条交易日历")
|
||
|
||
except Exception as e:
|
||
logger.error(f" 导入交易日历失败: {e}")
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
time.sleep(0.5) # API 频率限制
|
||
|
||
|
||
# ============================================================
|
||
# 3. 导入日线行情 (核心表)
|
||
# ============================================================
|
||
|
||
def import_daily_for_stock(ts_code: str, start_date: str, end_date: str, conn) -> int:
|
||
"""
|
||
导入单只股票的日线行情（保留用于单只股票补充/重试）
|
||
返回导入的记录数
|
||
"""
|
||
pro = get_ts_pro()
|
||
|
||
def fetch():
|
||
return pro.daily_vip(
|
||
ts_code=ts_code,
|
||
start_date=start_date.replace("-", ""),
|
||
end_date=end_date.replace("-", ""),
|
||
)
|
||
|
||
df = fetch_with_retry(fetch, max_retries=2)
|
||
if df is None or df.empty:
|
||
return 0
|
||
|
||
df = normalize_columns(df)
|
||
|
||
if "trade_date" in df.columns:
|
||
df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")
|
||
|
||
numeric_cols = [
|
||
"open", "high", "low", "close", "pre_close", "change", "pct_chg",
|
||
"vol", "amount",
|
||
]
|
||
for col in numeric_cols:
|
||
if col in df.columns:
|
||
df[col] = pd.to_numeric(df[col], errors="coerce")
|
||
|
||
for col in ["turnover_rate", "volume_ratio", "ma5", "ma10", "ma20", "ma_v_5", "ma_v_10", "ma_v_20"]:
|
||
if col not in df.columns:
|
||
df[col] = None
|
||
|
||
conflict_cols = ["ts_code", "trade_date"]
|
||
return batch_insert("daily", df, conn, conflict_cols)
|
||
|
||
|
||
def _normalize_daily_df(df: pd.DataFrame) -> pd.DataFrame:
|
||
"""标准化 daily DataFrame 的列和类型 (供 import_daily_by_date 复用)"""
|
||
df = normalize_columns(df)
|
||
if "trade_date" in df.columns:
|
||
df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")
|
||
numeric_cols = [
|
||
"open", "high", "low", "close", "pre_close", "change", "pct_chg",
|
||
"vol", "amount",
|
||
]
|
||
for col in numeric_cols:
|
||
if col in df.columns:
|
||
df[col] = pd.to_numeric(df[col], errors="coerce")
|
||
for col in ["turnover_rate", "volume_ratio", "ma5", "ma10", "ma20", "ma_v_5", "ma_v_10", "ma_v_20"]:
|
||
if col not in df.columns:
|
||
df[col] = None
|
||
return df
|
||
|
||
|
||
def import_daily_by_date(
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
conn=None,
|
||
sleep_interval: float = 0.3,
|
||
skip_existing: bool = True,
|
||
):
|
||
"""
|
||
按交易日批量导入日线行情 (高效模式)
|
||
使用 pro.daily_vip(trade_date='YYYYMMDD') 一次性拉取全市场当日数据
|
||
大幅减少 API 调用次数: 约250交易日/年 × 16年 ≈ 4000次 (原来需要 5000股票 × 16年 = 80000次)
|
||
|
||
参数:
|
||
- skip_existing: 是否跳过数据库中已有数据的交易日 (默认 True，避免重复导入)
|
||
|
||
返回: 失败的交易日列表
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
own_conn = conn is None
|
||
if own_conn:
|
||
conn = get_pg_connection()
|
||
|
||
# 从 trade_cal 获取交易日列表
|
||
try:
|
||
cursor = conn.cursor()
|
||
cursor.execute(
|
||
"""
|
||
SELECT DISTINCT cal_date FROM trade_cal
|
||
WHERE is_open = 1
|
||
AND cal_date >= %s AND cal_date <= %s
|
||
ORDER BY cal_date
|
||
""",
|
||
(start_date, end_date),
|
||
)
|
||
trade_dates = [row[0] for row in cursor.fetchall()]
|
||
cursor.close()
|
||
finally:
|
||
if own_conn:
|
||
conn.close()
|
||
conn = get_pg_connection()
|
||
|
||
total_cal = len(trade_dates)
|
||
if total_cal == 0:
|
||
logger.warning(f" 日期范围 {start_date} ~ {end_date} 内无交易日")
|
||
if own_conn:
|
||
conn.close()
|
||
return []
|
||
|
||
# --- 跳过已有数据的交易日 ---
|
||
skipped_count = 0
|
||
if skip_existing:
|
||
try:
|
||
cursor = conn.cursor()
|
||
cursor.execute(
|
||
"""
|
||
SELECT DISTINCT trade_date FROM daily
|
||
WHERE trade_date >= %s AND trade_date <= %s
|
||
""",
|
||
(start_date, end_date),
|
||
)
|
||
existing_dates = set(row[0] for row in cursor.fetchall())
|
||
cursor.close()
|
||
if existing_dates:
|
||
# 统一格式为 date 类型再比较
|
||
trade_dates_filtered = [td for td in trade_dates if td not in existing_dates]
|
||
skipped_count = len(trade_dates) - len(trade_dates_filtered)
|
||
trade_dates = trade_dates_filtered
|
||
except Exception as e:
|
||
logger.warning(f" 查询已有交易日失败，将导入全部: {e}")
|
||
# 查询失败时回退到全量导入（安全策略）
|
||
|
||
total = len(trade_dates)
|
||
if total == 0:
|
||
logger.info(f" 日期范围 {start_date} ~ {end_date}: 所有 {total_cal} 个交易日均已有数据，无需导入")
|
||
if own_conn:
|
||
conn.close()
|
||
return []
|
||
|
||
if skipped_count > 0:
|
||
logger.info(
|
||
f" 日期范围 {start_date} ~ {end_date}: "
|
||
f"共 {total_cal} 个交易日，跳过 {skipped_count} 个已有数据，"
|
||
f"待导入 {total} 个"
|
||
)
|
||
else:
|
||
logger.info(f" 日期范围 {start_date} ~ {end_date}: 共 {total} 个交易日")
|
||
|
||
pro = get_ts_pro()
|
||
success_count = 0
|
||
fail_list = []
|
||
|
||
for i, td in enumerate(trade_dates, 1):
|
||
try:
|
||
td_str = td.strftime("%Y%m%d") if hasattr(td, "strftime") else str(td).replace("-", "")
|
||
|
||
def fetch():
|
||
return pro.daily_vip(trade_date=td_str)
|
||
|
||
df = fetch_with_retry(fetch, max_retries=3)
|
||
if df is None or df.empty:
|
||
logger.warning(f" [{td_str}] 返回空数据 (可能非交易日或API限制)")
|
||
continue
|
||
|
||
df = _normalize_daily_df(df)
|
||
conflict_cols = ["ts_code", "trade_date"]
|
||
batch_insert("daily", df, conn, conflict_cols)
|
||
success_count += 1
|
||
|
||
except Exception as e:
|
||
logger.error(f" [{td}] 导入失败: {e}")
|
||
fail_list.append(str(td))
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
|
||
if i % 50 == 0 or i == total:
|
||
logger.info(f" 进度: {i}/{total} 成功={success_count} 失败={len(fail_list)}")
|
||
|
||
time.sleep(sleep_interval)
|
||
|
||
if own_conn:
|
||
conn.close()
|
||
|
||
logger.info(
|
||
f" 日线行情按日期导入完成: 成功 {success_count}/{total} 个交易日"
|
||
)
|
||
if fail_list:
|
||
logger.warning(f" 失败日期({len(fail_list)}): {fail_list[:20]}...")
|
||
return fail_list
|
||
|
||
|
||
def import_daily_by_year(
|
||
start_year: int = 2010,
|
||
end_year: int = 2025,
|
||
sleep_interval: float = 0.3,
|
||
):
|
||
"""
|
||
按年份逐批导入日线行情 (按交易日循环拉取全市场数据)
|
||
不再需要 stock_list 参数 — 每次 API 调用拉取当日全市场数据
|
||
"""
|
||
logger.info("=" * 60)
|
||
logger.info(
|
||
f"[3/7] 按年导入日线行情 (按交易日): {start_year} ~ {end_year}"
|
||
)
|
||
|
||
for year in range(start_year, end_year + 1):
|
||
year_start = f"{year}-01-01"
|
||
year_end = f"{year}-12-31"
|
||
logger.info(f"--- 导入 {year} 年日线行情 ---")
|
||
import_daily_by_date(
|
||
start_date=year_start,
|
||
end_date=year_end,
|
||
sleep_interval=sleep_interval,
|
||
)
|
||
logger.info(f" {year} 年完成\n")
|
||
|
||
logger.info(" 所有年份日线行情导入完成!")
|
||
|
||
|
||
# ============================================================
|
||
# 4. 导入每日指标 (daily_basic) - 估值/基本面
|
||
# ============================================================
|
||
|
||
def import_daily_basic(
|
||
ts_code: Optional[str] = None,
|
||
trade_date: Optional[str] = None,
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
conn=None,
|
||
):
|
||
"""
|
||
导入每日指标 (daily_basic)
|
||
Tushare: daily_basic
|
||
- 支持按单只股票导入
|
||
- 支持按日期范围全市场导入
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
pro = get_ts_pro()
|
||
own_conn = conn is None
|
||
if own_conn:
|
||
conn = get_pg_connection()
|
||
|
||
try:
|
||
# 构建参数
|
||
kwargs = {
|
||
"start_date": start_date.replace("-", ""),
|
||
"end_date": end_date.replace("-", ""),
|
||
}
|
||
if ts_code:
|
||
kwargs["ts_code"] = ts_code
|
||
if trade_date:
|
||
kwargs["trade_date"] = trade_date.replace("-", "")
|
||
|
||
def fetch():
|
||
return pro.daily_basic_vip(**kwargs)
|
||
|
||
df = fetch_with_retry(fetch, max_retries=2)
|
||
if df is None or df.empty:
|
||
return 0
|
||
|
||
df = normalize_columns(df)
|
||
|
||
if "trade_date" in df.columns:
|
||
df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")
|
||
|
||
numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
|
||
for col in numeric_cols:
|
||
df[col] = pd.to_numeric(df[col], errors="coerce")
|
||
|
||
conflict_cols = ["ts_code", "trade_date"]
|
||
return batch_insert("daily_basic", df, conn, conflict_cols)
|
||
|
||
finally:
|
||
if own_conn:
|
||
conn.close()
|
||
|
||
|
||
def import_daily_basic_by_date(
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
):
|
||
"""
|
||
按日期批量导入每日指标 (全市场)
|
||
Tushare daily_basic 接口可按交易日获取全市场数据，比较高效
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
logger.info("=" * 60)
|
||
logger.info(f"[4/7] 导入每日指标 (daily_basic): {start_date} ~ {end_date}")
|
||
|
||
# 获取交易日列表
|
||
conn = get_pg_connection()
|
||
try:
|
||
cursor = conn.cursor()
|
||
cursor.execute(
|
||
"""
|
||
SELECT DISTINCT cal_date FROM trade_cal
|
||
WHERE is_open = 1
|
||
AND cal_date >= %s AND cal_date <= %s
|
||
ORDER BY cal_date
|
||
""",
|
||
(start_date, end_date),
|
||
)
|
||
trade_dates = [row[0].strftime("%Y%m%d") for row in cursor.fetchall()]
|
||
cursor.close()
|
||
finally:
|
||
conn.close()
|
||
|
||
total = len(trade_dates)
|
||
logger.info(f" 共 {total} 个交易日")
|
||
|
||
conn = get_pg_connection()
|
||
for i, td in enumerate(trade_dates, 1):
|
||
try:
|
||
pro = get_ts_pro()
|
||
|
||
def fetch_daily_basic():
|
||
return pro.daily_basic_vip(trade_date=td)
|
||
|
||
df = fetch_with_retry(fetch_daily_basic, max_retries=3)
|
||
if df is not None and not df.empty:
|
||
df = normalize_columns(df)
|
||
if "trade_date" in df.columns:
|
||
df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")
|
||
conflict_cols = ["ts_code", "trade_date"]
|
||
batch_insert("daily_basic", df, conn, conflict_cols)
|
||
except Exception as e:
|
||
logger.warning(f" [{td}] 导入失败: {e}")
|
||
conn.rollback()
|
||
|
||
if i % 20 == 0 or i == total:
|
||
logger.info(f" 进度: {i}/{total}")
|
||
time.sleep(0.3)
|
||
|
||
conn.close()
|
||
logger.info(" 每日指标导入完成")
|
||
|
||
|
||
# ============================================================
|
||
# 5. 导入复权因子
|
||
# ============================================================
|
||
|
||
def import_adj_factor(
|
||
ts_code: Optional[str] = None,
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
conn=None,
|
||
):
|
||
"""
|
||
导入复权因子 (adj_factor)
|
||
Tushare: adj_factor
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
pro = get_ts_pro()
|
||
own_conn = conn is None
|
||
if own_conn:
|
||
conn = get_pg_connection()
|
||
|
||
try:
|
||
kwargs = {
|
||
"start_date": start_date.replace("-", ""),
|
||
"end_date": end_date.replace("-", ""),
|
||
}
|
||
if ts_code:
|
||
kwargs["ts_code"] = ts_code
|
||
|
||
def fetch():
|
||
return pro.adj_factor_vip(**kwargs)
|
||
|
||
df = fetch_with_retry(fetch, max_retries=2)
|
||
if df is None or df.empty:
|
||
return 0
|
||
|
||
df = normalize_columns(df)
|
||
|
||
if "trade_date" in df.columns:
|
||
df["trade_date"] = pd.to_datetime(df["trade_date"], format="%Y%m%d", errors="coerce")
|
||
|
||
conflict_cols = ["ts_code", "trade_date"]
|
||
return batch_insert("adj_factor", df, conn, conflict_cols)
|
||
|
||
finally:
|
||
if own_conn:
|
||
conn.close()
|
||
|
||
|
||
def import_adj_factor_batch(
|
||
stock_list: List[str],
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
):
|
||
"""批量导入复权因子"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
total = len(stock_list)
|
||
logger.info("=" * 60)
|
||
logger.info(f"[5/7] 导入复权因子 (adj_factor): {total} 只股票")
|
||
|
||
conn = get_pg_connection()
|
||
for i, ts_code in enumerate(stock_list, 1):
|
||
try:
|
||
import_adj_factor(ts_code=ts_code, start_date=start_date, end_date=end_date, conn=conn)
|
||
except Exception as e:
|
||
logger.warning(f" [{ts_code}] 复权因子导入失败: {e}")
|
||
conn.rollback()
|
||
|
||
if i % 100 == 0 or i == total:
|
||
logger.info(f" 进度: {i}/{total}")
|
||
time.sleep(0.25)
|
||
|
||
conn.close()
|
||
logger.info(" 复权因子导入完成")
|
||
|
||
|
||
# ============================================================
|
||
# 6. 导入财务数据 (利润表、资产负债表、现金流量表、财务指标)
|
||
# ============================================================
|
||
|
||
def import_financial_statements(
|
||
stock_list: List[str],
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
):
|
||
"""
|
||
按股票批量导入三大报表 + 财务指标
|
||
- income: 利润表
|
||
- balancesheet: 资产负债表
|
||
- cashflow: 现金流量表
|
||
- fina_indicator: 财务指标
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
total = len(stock_list)
|
||
logger.info("=" * 60)
|
||
logger.info(
|
||
f"[6/7] 导入财务数据: {start_date} ~ {end_date}, "
|
||
f"共 {total} 只股票"
|
||
)
|
||
|
||
period_start = start_date.replace("-", "")
|
||
period_end = end_date.replace("-", "")
|
||
pro = get_ts_pro()
|
||
conn = get_pg_connection()
|
||
|
||
for i, ts_code in enumerate(stock_list, 1):
|
||
for table_name, fetch_method in [
|
||
("income", pro.income_vip),
|
||
("balancesheet", pro.balancesheet_vip),
|
||
("cashflow", pro.cashflow_vip),
|
||
("fina_indicator", pro.fina_indicator_vip),
|
||
]:
|
||
try:
|
||
if table_name == "fina_indicator":
|
||
# fina_indicator 参数略有不同
|
||
df = fetch_method(
|
||
ts_code=ts_code,
|
||
start_date=period_start,
|
||
end_date=period_end,
|
||
)
|
||
else:
|
||
df = fetch_method(
|
||
ts_code=ts_code,
|
||
start_date=period_start,
|
||
end_date=period_end,
|
||
)
|
||
|
||
if df is None or df.empty:
|
||
continue
|
||
|
||
df = normalize_columns(df)
|
||
|
||
# 转换日期列
|
||
for col in ["ann_date", "f_ann_date", "end_date"]:
|
||
if col in df.columns:
|
||
df[col] = pd.to_datetime(df[col], format="%Y%m%d", errors="coerce")
|
||
|
||
if table_name == "fina_indicator":
|
||
conflict_cols = ["ts_code", "end_date"]
|
||
else:
|
||
conflict_cols = ["ts_code", "end_date", "report_type"]
|
||
|
||
batch_insert(table_name, df, conn, conflict_cols)
|
||
|
||
except Exception as e:
|
||
logger.warning(f" [{ts_code}] {table_name}: {e}")
|
||
conn.rollback()
|
||
|
||
if i % 50 == 0 or i == total:
|
||
logger.info(f" 财务数据进度: {i}/{total}")
|
||
time.sleep(0.3)
|
||
|
||
conn.close()
|
||
logger.info(" 财务数据导入完成")
|
||
|
||
|
||
# ============================================================
|
||
# 7. 导入指数日线行情
|
||
# ============================================================
|
||
|
||
def import_index_daily(
|
||
index_codes: Optional[List[str]] = None,
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
):
|
||
"""
|
||
导入指数日线行情 (index_daily)
|
||
默认导入主要指数：上证指数、深证成指、沪深300、中证500、创业板指、科创50
|
||
"""
|
||
if index_codes is None:
|
||
index_codes = [
|
||
"000001.SH", # 上证指数
|
||
"399001.SZ", # 深证成指
|
||
"000300.SH", # 沪深300
|
||
"000905.SH", # 中证500
|
||
"399006.SZ", # 创业板指
|
||
"000688.SH", # 科创50
|
||
]
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
logger.info("=" * 60)
|
||
logger.info(f"[7/7] 导入指数日线行情: {len(index_codes)} 个指数")
|
||
|
||
pro = get_ts_pro()
|
||
conn = get_pg_connection()
|
||
|
||
for idx_code in index_codes:
|
||
try:
|
||
def fetch():
|
||
return pro.index_daily_vip(
|
||
ts_code=idx_code,
|
||
start_date=start_date.replace("-", ""),
|
||
end_date=end_date.replace("-", ""),
|
||
)
|
||
|
||
df = fetch_with_retry(fetch, max_retries=2)
|
||
if df is not None and not df.empty:
|
||
df = normalize_columns(df)
|
||
if "trade_date" in df.columns:
|
||
df["trade_date"] = pd.to_datetime(
|
||
df["trade_date"], format="%Y%m%d", errors="coerce"
|
||
)
|
||
conflict_cols = ["ts_code", "trade_date"]
|
||
n = batch_insert("index_daily", df, conn, conflict_cols)
|
||
logger.info(f" {idx_code}: {n} 条")
|
||
else:
|
||
logger.warning(f" {idx_code}: 无数据")
|
||
except Exception as e:
|
||
logger.error(f" {idx_code}: {e}")
|
||
conn.rollback()
|
||
time.sleep(0.3)
|
||
|
||
conn.close()
|
||
logger.info(" 指数日线行情导入完成")
|
||
|
||
|
||
# ============================================================
|
||
# 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，创建所有表结构
|
||
"""
|
||
logger.info("=" * 60)
|
||
logger.info("初始化数据库 Schema ...")
|
||
|
||
# 尝试创建数据库（如果不存在）
|
||
try:
|
||
admin_conn = psycopg2.connect(
|
||
host=DB_CONFIG["host"],
|
||
port=DB_CONFIG["port"],
|
||
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"],),
|
||
)
|
||
if cursor.fetchone() is None:
|
||
cursor.execute(
|
||
sql.SQL("CREATE DATABASE {} ENCODING 'UTF8'").format(
|
||
sql.Identifier(DB_CONFIG["database"])
|
||
)
|
||
)
|
||
logger.info(f" 数据库 {DB_CONFIG['database']} 创建成功")
|
||
else:
|
||
logger.info(f" 数据库 {DB_CONFIG['database']} 已存在")
|
||
|
||
cursor.close()
|
||
admin_conn.close()
|
||
except Exception as e:
|
||
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:
|
||
cursor = conn.cursor()
|
||
for stmt in statements:
|
||
cursor.execute(stmt)
|
||
conn.commit()
|
||
cursor.close()
|
||
logger.info(" Schema 初始化完成")
|
||
|
||
# 迁移已存在的旧表: 补齐 Tushare API 新增列
|
||
_migrate_schema(conn)
|
||
except Exception as e:
|
||
conn.rollback()
|
||
logger.error(f" Schema 初始化失败: {e}")
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
# ============================================================
|
||
# 9. 获取所有股票列表 (辅助)
|
||
# ============================================================
|
||
|
||
def get_all_stock_codes(include_delisted: bool = False) -> List[str]:
|
||
"""
|
||
从 Tushare 获取所有 A 股股票代码列表
|
||
"""
|
||
pro = get_ts_pro()
|
||
codes = []
|
||
|
||
for status, label in [("L", "上市"), ("D", "退市"), ("P", "暂停")]:
|
||
if status != "L" and not include_delisted:
|
||
continue
|
||
try:
|
||
df = pro.stock_basic(
|
||
exchange="",
|
||
list_status=status,
|
||
fields="ts_code",
|
||
)
|
||
if df is not None and not df.empty:
|
||
codes.extend(df["ts_code"].tolist())
|
||
except Exception as e:
|
||
logger.warning(f" 获取 {label} 股票列表失败: {e}")
|
||
|
||
logger.info(f" 获取到 {len(codes)} 只股票代码")
|
||
return codes
|
||
|
||
|
||
def get_stock_codes_from_db(conn=None) -> List[str]:
|
||
"""
|
||
从已导入的 stock_basic 表获取股票代码列表
|
||
"""
|
||
own_conn = conn is None
|
||
if own_conn:
|
||
conn = get_pg_connection()
|
||
|
||
try:
|
||
cursor = conn.cursor()
|
||
cursor.execute("SELECT ts_code FROM stock_basic WHERE list_status = 'L' ORDER BY ts_code")
|
||
codes = [row[0] for row in cursor.fetchall()]
|
||
cursor.close()
|
||
return codes
|
||
finally:
|
||
if own_conn:
|
||
conn.close()
|
||
|
||
|
||
# ============================================================
|
||
# 10. 一键全量导入
|
||
# ============================================================
|
||
|
||
def full_import(
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
import_financials: bool = True,
|
||
stock_codes: Optional[List[str]] = None,
|
||
):
|
||
"""
|
||
一键全量导入：
|
||
1. 初始化 Schema
|
||
2. 股票基本信息
|
||
3. 交易日历
|
||
4. 日线行情
|
||
5. 每日指标(估值)
|
||
6. 复权因子
|
||
7. 财务数据 (可选)
|
||
8. 指数日线行情
|
||
|
||
参数:
|
||
- start_date, end_date: 数据范围
|
||
- import_financials: 是否导入财务数据 (耗时较长)
|
||
- stock_codes: 指定股票列表，不传则全量导入
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
start_time = datetime.now()
|
||
logger.info("=" * 70)
|
||
logger.info(f" 开始全量数据导入: {start_date} ~ {end_date}")
|
||
logger.info(f" PostgreSQL: {DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}")
|
||
logger.info("=" * 70)
|
||
|
||
# Step 0: 初始化 Schema
|
||
init_database()
|
||
|
||
# Step 1: 股票基本信息
|
||
import_stock_basic()
|
||
|
||
# Step 2: 交易日历
|
||
import_trade_cal(start_date, end_date)
|
||
|
||
# 获取股票列表
|
||
if stock_codes is None:
|
||
stock_codes = get_stock_codes_from_db()
|
||
if not stock_codes:
|
||
logger.error("无法获取股票列表，请先导入 stock_basic")
|
||
return
|
||
|
||
# Step 3: 日线行情 (按年导入，按交易日循环拉取全市场数据)
|
||
import_daily_by_year(
|
||
start_year=int(start_date[:4]),
|
||
end_year=int(end_date[:4]),
|
||
)
|
||
|
||
# Step 4: 每日指标 (按日期导入)
|
||
import_daily_basic_by_date(start_date, end_date)
|
||
|
||
# Step 5: 复权因子
|
||
import_adj_factor_batch(stock_codes, start_date, end_date)
|
||
|
||
# Step 6: 财务数据
|
||
if import_financials:
|
||
import_financial_statements(stock_codes, start_date, end_date)
|
||
|
||
# Step 7: 指数日线行情
|
||
import_index_daily(start_date=start_date, end_date=end_date)
|
||
|
||
elapsed = datetime.now() - start_time
|
||
logger.info("=" * 70)
|
||
logger.info(f" 全量数据导入完成! 总耗时: {elapsed}")
|
||
logger.info("=" * 70)
|
||
|
||
|
||
# ============================================================
|
||
# 11. 进度查询 & 断点续传辅助函数
|
||
# ============================================================
|
||
|
||
def check_daily_progress(conn=None):
|
||
"""
|
||
查看 daily 表的导入进度（按年份 + 股票数统计）
|
||
返回各年份的记录数和独立股票数
|
||
"""
|
||
own_conn = conn is None
|
||
if own_conn:
|
||
conn = get_pg_connection()
|
||
try:
|
||
cursor = conn.cursor()
|
||
cursor.execute("""
|
||
SELECT
|
||
EXTRACT(YEAR FROM trade_date)::int AS year,
|
||
COUNT(*) AS records,
|
||
COUNT(DISTINCT ts_code) AS stocks
|
||
FROM daily
|
||
GROUP BY year
|
||
ORDER BY year
|
||
""")
|
||
rows = cursor.fetchall()
|
||
if not rows:
|
||
logger.info("daily 表无数据")
|
||
return {}
|
||
|
||
logger.info(f"{'年份':<6} {'记录数':>12} {'股票数':>8}")
|
||
logger.info("-" * 30)
|
||
result = {}
|
||
for year, records, stocks in rows:
|
||
logger.info(f"{year:<6} {records:>12,} {stocks:>8,}")
|
||
result[int(year)] = {"records": int(records), "stocks": int(stocks)}
|
||
cursor.close()
|
||
return result
|
||
finally:
|
||
if own_conn:
|
||
conn.close()
|
||
|
||
|
||
def check_table_summary(conn=None):
|
||
"""
|
||
查看所有表的导入概览
|
||
"""
|
||
own_conn = conn is None
|
||
if own_conn:
|
||
conn = get_pg_connection()
|
||
try:
|
||
tables = [
|
||
("stock_basic", None),
|
||
("trade_cal", (("trade_cal", "cal_date"),)),
|
||
("daily", (("daily", "trade_date"),)),
|
||
("daily_basic", (("daily_basic", "trade_date"),)),
|
||
("adj_factor", (("adj_factor", "trade_date"),)),
|
||
("income", (("income", "end_date"),)),
|
||
("balancesheet", (("balancesheet", "end_date"),)),
|
||
("cashflow", (("cashflow", "end_date"),)),
|
||
("fina_indicator", (("fina_indicator", "end_date"),)),
|
||
("index_daily", (("index_daily", "trade_date"),)),
|
||
]
|
||
cursor = conn.cursor()
|
||
logger.info(f"{'表名':<20} {'记录数':>12} {'最早日期':>12} {'最晚日期':>12}")
|
||
logger.info("-" * 62)
|
||
for table_name, date_info in tables:
|
||
try:
|
||
cursor.execute(sql.SQL("SELECT COUNT(*) FROM {}").format(sql.Identifier(table_name)))
|
||
count = cursor.fetchone()[0]
|
||
if date_info:
|
||
tbl, col = date_info
|
||
cursor.execute(
|
||
sql.SQL("SELECT MIN({}), MAX({}) FROM {}").format(
|
||
sql.Identifier(col), sql.Identifier(col), sql.Identifier(tbl)
|
||
)
|
||
)
|
||
min_d, max_d = cursor.fetchone()
|
||
min_str = str(min_d)[:10] if min_d else "N/A"
|
||
max_str = str(max_d)[:10] if max_d else "N/A"
|
||
logger.info(f"{table_name:<20} {count:>12,} {min_str:>12} {max_str:>12}")
|
||
else:
|
||
logger.info(f"{table_name:<20} {count:>12,}")
|
||
except Exception as e:
|
||
logger.warning(f"{table_name:<20} 查询失败: {e}")
|
||
cursor.close()
|
||
finally:
|
||
if own_conn:
|
||
conn.close()
|
||
|
||
|
||
def get_missing_daily_dates(
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
conn=None,
|
||
) -> List[str]:
|
||
"""
|
||
获取 daily 表中缺失的交易日列表
|
||
对比 trade_cal 中 is_open=1 的日期和 daily 表已有的 trade_date，
|
||
返回未导入的交易日列表。
|
||
|
||
返回: 缺失交易日字符串列表 (YYYY-MM-DD 格式)
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
own_conn = conn is None
|
||
if own_conn:
|
||
conn = get_pg_connection()
|
||
|
||
try:
|
||
cursor = conn.cursor()
|
||
cursor.execute(
|
||
"""
|
||
SELECT tc.cal_date
|
||
FROM trade_cal tc
|
||
WHERE tc.is_open = 1
|
||
AND tc.cal_date >= %s
|
||
AND tc.cal_date <= %s
|
||
AND NOT EXISTS (
|
||
SELECT 1 FROM daily d
|
||
WHERE d.trade_date = tc.cal_date
|
||
)
|
||
ORDER BY tc.cal_date
|
||
""",
|
||
(start_date, end_date),
|
||
)
|
||
missing_dates = [row[0].strftime("%Y-%m-%d") if hasattr(row[0], "strftime") else str(row[0])[:10]
|
||
for row in cursor.fetchall()]
|
||
cursor.close()
|
||
return missing_dates
|
||
finally:
|
||
if own_conn:
|
||
conn.close()
|
||
|
||
|
||
def resume_daily_by_date(
|
||
start_date: Optional[str] = None,
|
||
end_date: Optional[str] = None,
|
||
sleep_interval: float = 0.3,
|
||
):
|
||
"""
|
||
按缺失日期断点续传日线行情
|
||
自动查询 daily 表已有的 trade_date 与 trade_cal 对比，
|
||
只导入缺失日期的全市场数据。
|
||
|
||
用法:
|
||
resume_daily_by_date(start_date="2010-01-01", end_date="2025-12-31")
|
||
"""
|
||
if start_date is None:
|
||
start_date = START_DATE
|
||
if end_date is None:
|
||
end_date = END_DATE
|
||
|
||
logger.info("=" * 60)
|
||
logger.info(f"[断点续传] 检测缺失日期: {start_date} ~ {end_date}")
|
||
|
||
missing_dates = get_missing_daily_dates(start_date, end_date)
|
||
|
||
if not missing_dates:
|
||
logger.info(" 所有交易日数据已完整，无需续传!")
|
||
return
|
||
|
||
total = len(missing_dates)
|
||
logger.info(f" 发现 {total} 个缺失交易日待导入")
|
||
if total <= 20:
|
||
logger.info(f" 缺失日期: {missing_dates}")
|
||
else:
|
||
logger.info(f" 缺失日期 (前20): {missing_dates[:20]}")
|
||
|
||
# 按年份分组统计
|
||
years_map: Dict[int, List[str]] = {}
|
||
for d in missing_dates:
|
||
y = int(d[:4])
|
||
years_map.setdefault(y, []).append(d)
|
||
for y in sorted(years_map):
|
||
logger.info(f" {y} 年: {len(years_map[y])} 个缺失交易日")
|
||
|
||
# 逐日期导入
|
||
conn = get_pg_connection()
|
||
pro = get_ts_pro()
|
||
success_count = 0
|
||
fail_list = []
|
||
|
||
for i, td_str in enumerate(missing_dates, 1):
|
||
try:
|
||
td_compact = td_str.replace("-", "")
|
||
|
||
def fetch():
|
||
return pro.daily_vip(trade_date=td_compact)
|
||
|
||
df = fetch_with_retry(fetch, max_retries=3)
|
||
if df is None or df.empty:
|
||
logger.warning(f" [{td_str}] 返回空数据，跳过")
|
||
continue
|
||
|
||
df = _normalize_daily_df(df)
|
||
conflict_cols = ["ts_code", "trade_date"]
|
||
batch_insert("daily", df, conn, conflict_cols)
|
||
success_count += 1
|
||
|
||
except Exception as e:
|
||
logger.error(f" [{td_str}] 导入失败: {e}")
|
||
fail_list.append(td_str)
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
|
||
if i % 50 == 0 or i == total:
|
||
logger.info(f" 续传进度: {i}/{total} 成功={success_count} 失败={len(fail_list)}")
|
||
|
||
time.sleep(sleep_interval)
|
||
|
||
conn.close()
|
||
logger.info(f" 断点续传完成: 成功 {success_count}/{total}")
|
||
if fail_list:
|
||
logger.warning(f" 失败日期({len(fail_list)}): {fail_list[:20]}...")
|
||
return fail_list
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 测试连接
|
||
try:
|
||
conn = get_pg_connection()
|
||
logger.info(f"成功连接到 PostgreSQL: {DB_CONFIG['host']}:{DB_CONFIG['port']}")
|
||
conn.close()
|
||
except Exception as e:
|
||
logger.error(f"无法连接到 PostgreSQL: {e}")
|
||
logger.error("请确认 Docker 容器已启动，且 config.py 中的连接参数正确")
|