Spaces:
Paused
Paused
| """ | |
| DuckDB 数据库管理模块 | |
| 实现数据库初始化、Hugging Face Dataset 同步、单例连接管理 | |
| """ | |
| import os | |
| import logging | |
| from pathlib import Path | |
| from typing import Optional | |
| from functools import lru_cache | |
| from dotenv import load_dotenv | |
| # 加载 .env 文件 | |
| load_dotenv() | |
| import duckdb | |
| from huggingface_hub import hf_hub_download, upload_file, login | |
| logger = logging.getLogger(__name__) | |
| # 环境变量配置 | |
| DUCKDB_PATH = os.getenv("DUCKDB_PATH", "/app/data/stock_data.duckdb") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| DATASET_REPO_ID = os.getenv("DATASET_REPO_ID", "") | |
| HF_HOME = os.getenv("HF_HOME", "/tmp/.huggingface") | |
| class DatabaseManager: | |
| """DuckDB 数据库管理器 - 单例模式""" | |
| _instance: Optional['DatabaseManager'] = None | |
| _connection: Optional[duckdb.DuckDBPyConnection] = None | |
| def __new__(cls) -> 'DatabaseManager': | |
| if cls._instance is None: | |
| cls._instance = super().__new__(cls) | |
| return cls._instance | |
| def conn(self) -> duckdb.DuckDBPyConnection: | |
| """获取数据库连接""" | |
| if self._connection is None: | |
| # 确保目录存在 | |
| os.makedirs(os.path.dirname(DUCKDB_PATH), exist_ok=True) | |
| self._connection = duckdb.connect(DUCKDB_PATH) | |
| # 如果在 HF 环境，配置远程访问插件 | |
| if HF_TOKEN and DATASET_REPO_ID: | |
| try: | |
| self._connection.execute("INSTALL httpfs; LOAD httpfs;") | |
| self._connection.execute(f"SET s3_region='us-east-1';") # HF 默认区域 | |
| # 这里的逻辑可以根据是否切换到 Parquet 进一步扩展 | |
| except Exception as e: | |
| logger.warning(f"Failed to load httpfs extension: {e}") | |
| logger.info(f"Database connection established: {DUCKDB_PATH}") | |
| return self._connection | |
| def close(self) -> None: | |
| """关闭数据库连接""" | |
| if self._connection is not None: | |
| self._connection.close() | |
| self._connection = None | |
| logger.info("Database connection closed") | |
| def init_db(self) -> None: | |
| """ | |
| 初始化数据库 | |
| - 在本地环境：使用本地 DuckDB | |
| - 在 HF 环境：挂载云端 Parquet 文件 | |
| """ | |
| conn = self.conn | |
| # 检查是否在 HF 环境 | |
| is_hf = os.getenv("SPACES_REPO_NAME") is not None or (HF_TOKEN and DATASET_REPO_ID) | |
| if is_hf: | |
| logger.info("HF environment detected. Downloading remote Parquet files...") | |
| try: | |
| from huggingface_hub import list_repo_files, hf_hub_download | |
| # 1. 动态获取文件列表 | |
| all_files = list_repo_files(repo_id=DATASET_REPO_ID, repo_type="dataset") | |
| # 2. 股票列表：下载并加载 | |
| list_file = "data/stock_list.parquet" | |
| if list_file in all_files: | |
| local_list_path = hf_hub_download(repo_id=DATASET_REPO_ID, filename=list_file, repo_type="dataset") | |
| conn.execute(f"CREATE OR REPLACE TABLE stock_list AS SELECT * FROM read_parquet('{local_list_path}')") | |
| logger.info("Local stock_list table created from remote parquet") | |
| # 3. 日线数据：下载并加载 (下载模式替代挂载模式) | |
| # 注意：如果文件很多，这步可能会花点时间，但比下载整个 .duckdb 快得多 | |
| parquet_files = [f for f in all_files if f.startswith("data/parquet/") and f.endswith(".parquet")] | |
| if parquet_files: | |
| local_paths = [] | |
| for f in parquet_files: | |
| # HF SDK 会自动处理缓存，已下载的文件不会重复下载 | |
| path = hf_hub_download(repo_id=DATASET_REPO_ID, filename=f, repo_type="dataset") | |
| local_paths.append(f"'{path}'") | |
| files_sql = ", ".join(local_paths) | |
| conn.execute("DROP VIEW IF EXISTS stock_daily") | |
| conn.execute("DROP TABLE IF EXISTS stock_daily") | |
| conn.execute(f"CREATE OR REPLACE VIEW stock_daily AS SELECT * FROM read_parquet([{files_sql}])") | |
| logger.info(f"Remote stock_daily view created with {len(parquet_files)} partitions") | |
| else: | |
| logger.warning("No parquet data files found in dataset") | |
| self._create_tables() | |
| # 验证 | |
| count = conn.execute("SELECT COUNT(*) FROM stock_list").fetchone()[0] | |
| logger.info(f"Verification successful: {count} stocks found") | |
| except Exception as e: | |
| logger.error(f"Failed to load remote Parquet: {e}") | |
| self._create_tables() | |
| else: | |
| self._create_tables() | |
| logger.info("Local database initialized") | |
| def _download_from_hf(self) -> None: | |
| """从 Hugging Face Dataset 下载数据库文件""" | |
| if not HF_TOKEN or not DATASET_REPO_ID: | |
| logger.warning("HF_TOKEN or DATASET_REPO_ID not set, skipping download") | |
| return | |
| try: | |
| # 登录 Hugging Face | |
| login(token=HF_TOKEN) | |
| # 下载数据库文件 | |
| downloaded_path = hf_hub_download( | |
| repo_id=DATASET_REPO_ID, | |
| filename="stock_data.duckdb", | |
| repo_type="dataset", | |
| local_dir=os.path.dirname(DUCKDB_PATH), | |
| ) | |
| # 如果下载路径不同，重命名 | |
| if downloaded_path != DUCKDB_PATH: | |
| os.rename(downloaded_path, DUCKDB_PATH) | |
| logger.info(f"Database downloaded from HF Dataset: {DATASET_REPO_ID}") | |
| except Exception as e: | |
| logger.error(f"Failed to download database from HF: {e}") | |
| def upload_db(self) -> None: | |
| """上传 Parquet 分区到 Hugging Face Dataset (不再上传全量 .duckdb)""" | |
| if not HF_TOKEN or not DATASET_REPO_ID: | |
| logger.warning("HF_TOKEN or DATASET_REPO_ID not set, skipping upload") | |
| return | |
| try: | |
| # 先关闭连接 | |
| self.close() | |
| login(token=HF_TOKEN) | |
| # 1. 上传股票列表 (这个很小，可以全量覆盖) | |
| if Path(DUCKDB_PATH).exists(): | |
| # 导出列表为 parquet 方便云端读取 | |
| conn = duckdb.connect(DUCKDB_PATH) | |
| list_path = os.path.join(os.path.dirname(DUCKDB_PATH), "stock_list.parquet") | |
| conn.execute(f"COPY stock_list TO '{list_path}' (FORMAT PARQUET)") | |
| conn.close() | |
| upload_file( | |
| path_or_fileobj=list_path, | |
| path_in_repo="data/stock_list.parquet", | |
| repo_id=DATASET_REPO_ID, | |
| repo_type="dataset", | |
| ) | |
| # 2. 上传所有 Parquet 行情文件 (增量核心) | |
| parquet_dir = Path(os.path.dirname(DUCKDB_PATH)) / "parquet" | |
| if parquet_dir.exists(): | |
| for p_file in parquet_dir.glob("*.parquet"): | |
| upload_file( | |
| path_or_fileobj=str(p_file), | |
| path_in_repo=f"data/parquet/{p_file.name}", | |
| repo_id=DATASET_REPO_ID, | |
| repo_type="dataset", | |
| ) | |
| logger.info(f"Parquet files uploaded to HF Dataset: {DATASET_REPO_ID}") | |
| except Exception as e: | |
| logger.error(f"Failed to upload to HF: {e}") | |
| finally: | |
| _ = self.conn | |
| def _create_tables(self) -> None: | |
| """创建数据库表结构""" | |
| conn = self.conn | |
| # 日线行情表 | |
| conn.execute(""" | |
| CREATE TABLE IF NOT EXISTS stock_daily ( | |
| code VARCHAR, | |
| trade_date DATE, | |
| open DOUBLE, | |
| high DOUBLE, | |
| low DOUBLE, | |
| close DOUBLE, | |
| volume BIGINT, | |
| amount DOUBLE, | |
| pct_chg DOUBLE, | |
| turnover_rate DOUBLE, | |
| PRIMARY KEY (code, trade_date) | |
| ) | |
| """) | |
| # 股票基础信息表 | |
| conn.execute(""" | |
| CREATE TABLE IF NOT EXISTS stock_list ( | |
| code VARCHAR PRIMARY KEY, | |
| name VARCHAR, | |
| market VARCHAR, | |
| list_date DATE | |
| ) | |
| """) | |
| # 创建索引 | |
| conn.execute(""" | |
| CREATE INDEX IF NOT EXISTS idx_code_date | |
| ON stock_daily (code, trade_date) | |
| """) | |
| logger.info("Database tables created/verified") | |
| # 全局单例实例 | |
| db_manager = DatabaseManager() | |
| def get_db() -> DatabaseManager: | |
| """获取数据库管理器实例""" | |
| return db_manager | |