新增一键运行功能

daily.py

@@ -2,29 +2,29 @@ import tushare as ts
 import pymysql
 from datetime import datetime
 import time
+import argparse
 
-# ===== 配置 =====
-TUSHARE_TOKEN = "0fad3cf498757089e2630028455d5bbe1637475788bcdaa1f2175e93"
-
-DB_CONFIG = {
-    "host": "8.159.129.156",
-    "port": 10836,
-    "user": "yangfan",
-    "password": "aA%8023321088",
-    "database": "level",
-    "charset": "utf8mb4"
-}
+from stock_config import DB_CONFIG, TUSHARE_TOKEN
 
 # ===== 初始化 =====
 ts.set_token(TUSHARE_TOKEN)
 pro = ts.pro_api()
 
-conn = pymysql.connect(**DB_CONFIG)
-cursor = conn.cursor()
+conn = None
+cursor = None
+
+
+def get_connection():
+    global conn, cursor
+    if conn is None or not conn.open:
+        conn = pymysql.connect(**DB_CONFIG)
+        cursor = conn.cursor()
+    return conn, cursor
 
 
 # ===== 建表 =====
 def create_tables():
+    db_conn, db_cursor = get_connection()
     sql = """
     CREATE TABLE IF NOT EXISTS stock_daily (
         id BIGINT AUTO_INCREMENT PRIMARY KEY,
@@ -43,8 +43,8 @@ def create_tables():
         UNIQUE KEY uniq_code_date (ts_code, trade_date)
     );
     """
-    cursor.execute(sql)
-    conn.commit()
+    db_cursor.execute(sql)
+    db_conn.commit()
     print("表检查/创建完成")
 
 
@@ -55,6 +55,7 @@ def str_to_date(date_str):
 
 # ===== 获取交易日 =====
 def get_trade_dates(start_date, end_date):
+    _, db_cursor = get_connection()
     sql = """
     SELECT cal_date 
     FROM trade_calendar
@@ -62,14 +63,15 @@ def get_trade_dates(start_date, end_date):
       AND is_open = 1
     ORDER BY cal_date
     """
-    cursor.execute(sql, (str_to_date(start_date), str_to_date(end_date)))
-    results = cursor.fetchall()
+    db_cursor.execute(sql, (str_to_date(start_date), str_to_date(end_date)))
+    results = db_cursor.fetchall()
 
     return [row[0].strftime("%Y%m%d") for row in results]
 
 
 # ===== 插入数据 =====
 def insert_data(df):
+    db_conn, db_cursor = get_connection()
     if df is None or df.empty:
         return 0, 0
 
@@ -99,10 +101,10 @@ def insert_data(df):
             row['amount']
         ))
 
-    cursor.executemany(sql, data)
-    conn.commit()
+    db_cursor.executemany(sql, data)
+    db_conn.commit()
 
-    inserted = cursor.rowcount
+    inserted = db_cursor.rowcount
     skipped = len(data) - inserted
 
     return inserted, skipped
@@ -140,11 +142,14 @@ def run(start_date, end_date):
     print("\n====== 汇总 ======")
     print(f"总新增: {total_inserted}")
     print(f"总跳过: {total_skipped}")
+    return True
 
 
 # ===== 输入参数 =====
 if __name__ == "__main__":
-    START_DATE = "20260424"
-    END_DATE = "20260424"
+    parser = argparse.ArgumentParser(description="同步日行情")
+    parser.add_argument("--start-date", default="20260429", help="开始日期，格式 YYYYMMDD")
+    parser.add_argument("--end-date", default="20260429", help="结束日期，格式 YYYYMMDD")
+    args = parser.parse_args()
 
-    run(START_DATE, END_DATE)
+    run(args.start_date, args.end_date)