fix:添加断点续传功能。

This commit is contained in:
2026-08-01 18:12:04 +08:00
parent cdf5582c90
commit 063f790650
2 changed files with 217 additions and 1 deletions
+144 -1
View File
@@ -982,6 +982,149 @@ def full_import(
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 resume_daily_by_year(stock_list, start_year=2010, end_year=2025):
"""
从中断点恢复按年导入日线行情
自动跳过数据库已有的年份,只导入缺失年份的数据
"""
conn = get_pg_connection()
try:
cursor = conn.cursor()
cursor.execute("""
SELECT DISTINCT EXTRACT(YEAR FROM trade_date)::int AS year
FROM daily
ORDER BY year
""")
completed_years = set(row[0] for row in cursor.fetchall())
cursor.close()
finally:
conn.close()
logger.info(f"已完成年份: {sorted(completed_years)}")
logger.info(f"待导入年份: {[y for y in range(start_year, end_year+1) if y not in completed_years]}")
for year in range(start_year, end_year + 1):
if year in completed_years:
# 检查该年的股票覆盖是否完整
conn = get_pg_connection()
try:
cursor = conn.cursor()
cursor.execute(
sql.SQL("""
SELECT COUNT(DISTINCT ts_code)
FROM {}
WHERE EXTRACT(YEAR FROM trade_date) = %s
""").format(sql.Identifier("daily")),
(year,),
)
stock_count = cursor.fetchone()[0]
cursor.close()
finally:
conn.close()
logger.info(f" {year} 年: 已有 {stock_count} 只股票, 跳过")
continue
year_start = f"{year}-01-01"
year_end = f"{year}-12-31"
logger.info(f"--- 导入 {year} 年日线行情 ---")
import_daily_batch(
stock_list,
start_date=year_start,
end_date=year_end,
sleep_interval=0.2,
)
logger.info(f" {year} 年完成\n")
if __name__ == "__main__":
# 测试连接
try:
@@ -990,4 +1133,4 @@ if __name__ == "__main__":
conn.close()
except Exception as e:
logger.error(f"无法连接到 PostgreSQL: {e}")
logger.error("请确认 Docker 容器已启动,且 config.py 中的连接参数正确")
logger.error("请确认 Docker 容器已启动,且 config.py 中的连接参数正确")