modify scripts

src/config/config.py

@@ -1,21 +1,43 @@
 import os
 from pathlib import Path
 
-# MySQL 配置
-db_config = {
+db_config_dev = {
     'host': 'testdb',
     'user': 'root',
     'password': 'mysqlpw',
     'database': 'stockdb'
 }
-# 读取环境变量，来修改 db_config 中host的值 
-if 'DB_HOST' in os.environ:
-    db_config['host'] = os.environ['DB_HOST']
+
+db_config_nas = {
+    'host': 'mysql',
+    'user': 'root',
+    'password': 'root',
+    'database': 'stockdb'
+}
+
+db_config_run = {
+    'host': 'testdb',
+    'user': 'root',
+    'password': 'mysqlpw',
+    'database': 'stockdb'
+}
+
+db_config = {}
+current_env = os.environ.get('DB_ENV', 'dev')
+if current_env == 'dev':
+    db_config = db_config_dev
+elif current_env == 'nas':
+    db_config = db_config_nas
+elif current_env == 'run':
+    db_config = db_config_run
+else:
+    db_config = db_config_dev  # 默认使用开发环境配置
 
 home_dir = os.path.expanduser("~")
 global_host_data_dir = f'{home_dir}/hostdir/stock_data'
 global_share_db_dir  = f'{home_dir}/sharedata/sqlite'
 global_host_input_dir = f'{global_host_data_dir}/input'
+global_pdf_data_dir = f'{home_dir}/sharedata/stock_data'
 
 # 兼容以前的定义
 global_stock_data_dir = global_host_data_dir

src/em_reports/fetch.py

@@ -11,7 +11,7 @@ from datetime import datetime, timedelta
 from functools import partial
 import src.crawler.em.reports as em
 import src.utils.utils as utils
-from src.config.config import global_host_data_dir, global_share_db_dir, db_config
+from src.config.config import global_host_data_dir, global_share_db_dir, db_config, global_pdf_data_dir
 from src.db_utils.reports import StockReportDB, DatabaseConnectionError
 from src.db_utils.reports_mysql import StockReportMysql
 from src.logger.logger import setup_logging
@@ -22,7 +22,9 @@ setup_logging()
 
 debug = False
 force = False
-pdf_base_dir = f"{global_host_data_dir}/pdfs"  # 下载 PDF 存放目录
+pdf_base_dir = f"{global_pdf_data_dir}/pdfs"  # 下载 PDF 存放目录
+pdf_tmp_dir = f"{global_pdf_data_dir}/tmp"  # 下载 PDF 存放目录
+reports_db_name = "em_repors"
 
 # 定义下载页面的链接
 map_pdf_page = {
@@ -139,7 +141,8 @@ def parse_func_general(row, tbl_name):
     # 拼目录
     dir_year = publish_date[:4] if len(publish_date)>4 else ''
     #dir_path = f'{pdf_base_dir}/{dir_year}/{map_tbl_name[tbl_name]}'
-    dir_path = f'{pdf_base_dir}/{dir_year}'
+    #dir_path = f'{pdf_base_dir}/{dir_year}'
+    dir_path = f'{pdf_base_dir}'
     os.makedirs(dir_path, exist_ok=True)
     return url, os.path.join(dir_path, file_name)
 
@@ -147,15 +150,10 @@ def parse_func_general(row, tbl_name):
 def check_pdf_pages(file_path, row, tbl_name):
     pages = get_pdf_page_count(file_path)
     if pages is None or pages < min_down_pages:
-        # 获取文件所在目录
-        file_dir = os.path.dirname(file_path)
-        # 创建 tmp 子目录
-        tmp_dir = os.path.join(file_dir, 'tmp')
-        if not os.path.exists(tmp_dir):
-            os.makedirs(tmp_dir)
         # 移动文件到 tmp 子目录
+        os.makedirs(pdf_tmp_dir, exist_ok=True)
         file_name = os.path.basename(file_path)
-        new_path = os.path.join(tmp_dir, file_name)
+        new_path = os.path.join(pdf_tmp_dir, file_name)
         shutil.move(file_path, new_path)
         logging.debug(f"move {file_name} to {tmp_dir}")
 
@@ -339,7 +337,7 @@ def main(cmd, mode, args_debug, args_force, begin, end):
     global db_tools
     try:
         #db_tools = StockReportDB(db_path)
-        db_tools = StockReportMysql(db_host=db_config['host'], db_user=db_config['user'], db_password=db_config['password'], db_name = db_config['database'], port=3306)  # 使用配置文件中的数据库配置
+        db_tools = StockReportMysql(db_host=db_config['host'], db_user=db_config['user'], db_password=db_config['password'], db_name = reports_db_name, port=3306)  # 使用配置文件中的数据库配置
 
         # 进行数据库操作
     except DatabaseConnectionError as e:

src/sample/futu_add_stock.py

@@ -0,0 +1,23 @@
+import time
+
+from futu import *
+quote_ctx = OpenQuoteContext(host='127.0.0.1', port=11111)
+
+code_list = ["SH.1B0888","SH.1A0001","SH.1B0300",
+            "SH.399001","SH.399006",
+            "SH.510880",
+            "SH.513920","SH.159331","SH.159719","SH.513910","SH.159601","SH.510050","SH.515650","SH.588000","SH.159883","SH.588080","SH.512680","SH.513850","SH.513180","SH.513010","SH.513330","SH.512800","SH.512880","SH.159857","SH.512010","SH.512760","SH.512480","SH.588200","SH.513120","SH.159712","SH.159941","SH.159605","SH.159687","SH.513100","SH.159995","SH.515700","SH.159869","SH.562000","SH.512500","SH.510500","SH.159845","SH.512100","SH.563300"
+            ]
+# 使用code_list 的倒序添加自选股，防止因已满而添加失败
+code_list.reverse()
+for code in code_list:
+    tmp_list = []
+    tmp_list.append(code)
+    ret, data = quote_ctx.modify_user_security("ETF", ModifyUserSecurityOp.ADD, tmp_list)
+    if ret == RET_OK:
+        print(f'添加自选股成功: {code}')
+    else:
+        print(f'添加自选股失败: {code}, error: {data}')
+    time.sleep(3.5)  # 添加间隔，防止请求过快被风控
+
+quote_ctx.close() # 结束后记得关闭当条连接，防止连接条数用尽