diff --git a/quantitative_data/importer.py b/quantitative_data/importer.py index 5a39d61..32aeba2 100644 --- a/quantitative_data/importer.py +++ b/quantitative_data/importer.py @@ -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 中的连接参数正确") \ No newline at end of file + logger.error("请确认 Docker 容器已启动,且 config.py 中的连接参数正确") diff --git a/quantitative_data/数据批量导入.ipynb b/quantitative_data/数据批量导入.ipynb index 73a6788..ba684be 100644 --- a/quantitative_data/数据批量导入.ipynb +++ b/quantitative_data/数据批量导入.ipynb @@ -134,6 +134,9 @@ " get_stock_codes_from_db,\n", " full_import,\n", " batch_insert,\n", + " check_daily_progress,\n", + " check_table_summary,\n", + " resume_daily_by_year,\n", " logger,\n", ")\n", "from config import DB_CONFIG, TUSHARE_TOKEN, START_DATE, END_DATE\n", @@ -345,6 +348,76 @@ ")" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.2.1 查看日线导入进度\n", + "\n", + "按年份统计 daily 表中已导入的记录数和独立股票数,用于确认导入到哪个年份了。" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 查看每天的日线导入进度(按年份统计)\n", + "check_daily_progress()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 或者查看所有表的整体概览\n", + "check_table_summary()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.2.2 断点续传 — 从中断处继续导入\n", + "\n", + "**方法一(推荐):** 使用 `resume_daily_by_year` 自动跳过已完成的年份,只导入缺失年份。\n", + "\n", + "**方法二:** 手动修改 `start_year` 参数重新调用 `import_daily_by_year`(因为 UPSERT 幂等,重复导入不会造成数据问题)。" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 方法一:自动检测已有年份,只导入缺失年份(推荐)\n", + "stock_list = get_stock_codes_from_db()\n", + "resume_daily_by_year(stock_list, start_year=2010, end_year=2025)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# 方法二:手动指定断点年份重新调用 import_daily_by_year\n", + "# 例如假设 2010~2020 已完成,从 2021 年开始继续\n", + "# stock_list = get_stock_codes_from_db()\n", + "# import_daily_by_year(stock_list, start_year=2021, end_year=2025)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> **提示:** 以上两个方法都可以安全使用。因为 daily 表使用 `ON CONFLICT (ts_code, trade_date) DO UPDATE`,重复导入已存在的数据不会产生重复记录。另外也可查看 `import_data.log` 文件获取最后一次成功的日志输出。" + ] + }, { "cell_type": "markdown", "metadata": {},