From cd602bf3ac7d64dc0c1e52622003002399d259d5 Mon Sep 17 00:00:00 2001 From: jackyu66git Date: Fri, 3 Jul 2026 12:12:48 +0800 Subject: [PATCH] feat: EOD pull, EMA52 screening, DuckDB transient connections, config refactor --- src/ashare_dp/api/app.py | 324 +++++++++++++++++++++++---- src/ashare_dp/api/deps.py | 1 - src/ashare_dp/api/routers/stocks.py | 44 ++-- src/ashare_dp/cli/backfill_cmd.py | 7 - src/ashare_dp/cli/eod_cmd.py | 58 +++++ src/ashare_dp/cli/main.py | 4 + src/ashare_dp/cli/query_cmd.py | 35 ++- src/ashare_dp/cli/screening_cmd.py | 168 ++++++++++++++ src/ashare_dp/config.py | 5 + src/ashare_dp/data/backfill.py | 306 +++++++++++++++++-------- src/ashare_dp/scheduler/jobs.py | 144 ++++++++++++ src/ashare_dp/scheduler/scheduler.py | 18 +- src/ashare_dp/storage/database.py | 55 ++--- src/ashare_dp/storage/repository.py | 168 +++++++------- src/ashare_dp/storage/schema.py | 25 +++ 15 files changed, 1068 insertions(+), 294 deletions(-) create mode 100644 src/ashare_dp/cli/eod_cmd.py create mode 100644 src/ashare_dp/cli/screening_cmd.py diff --git a/src/ashare_dp/api/app.py b/src/ashare_dp/api/app.py index aec3452..9352692 100644 --- a/src/ashare_dp/api/app.py +++ b/src/ashare_dp/api/app.py @@ -4,7 +4,9 @@ from __future__ import annotations from contextlib import asynccontextmanager -from fastapi import FastAPI +from typing import Optional + +from fastapi import FastAPI, HTTPException, Query from fastapi.responses import HTMLResponse from loguru import logger @@ -82,6 +84,10 @@ DOCS_HTML = r"""
全天
自动盘后拉取
+
+ 📊 EMA52 筛选 — 扫描日线/周线价格在 EMA52 ±3% 附近的股票，按成交额排序 +
+
K线周期 股票查询 @@ -279,17 +285,142 @@ ashare-dp query stats """ +EMA52_HTML = r""" + + + + + +EMA52 筛选 — A-Share Data Platform + + + +
+

EMA52 筛选

+

每日收盘后扫描 1d / 1w 价格在 EMA52 附近的股票，按成交额降序排列

+
+ + + + 加载中… +
+
+ 价格 > EMA52 + 价格 < EMA52 + 阈值: ±3% | 按成交额↓ +
+
加载中…
+
+ ← 返回首页 +
+
+ + + + +""" + + @asynccontextmanager async def lifespan(app: FastAPI): """Startup and shutdown lifecycle.""" logger.info("Starting A-Share Data Platform...") - db = get_db() - db.connect() - for ddl in DDL_STATEMENTS: - try: - db.execute(ddl) - except Exception as e: - logger.warning(f"DDL warning: {e}") + # Run DDL with write access (short-lived) + with get_db(read_only=False) as db: + for ddl in DDL_STATEMENTS: + try: + db.execute(ddl) + except Exception as e: + logger.warning(f"DDL warning: {e}") logger.info("Database initialized") # Start realtime poller @@ -302,18 +433,6 @@ async def lifespan(app: FastAPI): logger.warning(f"Realtime poller not started: {e}") app.state.poller = None - # Start scheduler - try: - from ashare_dp.scheduler.scheduler import Scheduler - from ashare_dp.data.akshare_client import AKShareClient - scheduler = Scheduler(client=AKShareClient()) - scheduler.start() - app.state.scheduler = scheduler - logger.info("Scheduler started") - except Exception as e: - logger.warning(f"Scheduler not started: {e}") - app.state.scheduler = None - yield # Shutdown @@ -322,12 +441,6 @@ async def lifespan(app: FastAPI): await app.state.poller.stop() except Exception: pass - if app.state.scheduler: - try: - app.state.scheduler.shutdown() - except Exception: - pass - db.close() logger.info("A-Share Data Platform stopped") @@ -347,37 +460,150 @@ def create_app() -> FastAPI: @app.get("/health") async def health(): - """Basic health check.""" - db = get_db() + """Health check in data_provider format (compatible with Chan system).""" repo = KLineRepository() + from ashare_dp.core.models import BACKFILLABLE_FREQS, DERIVED_FREQS, INTRADAY_FREQS, Freq + + all_freqs = BACKFILLABLE_FREQS + INTRADAY_FREQS + DERIVED_FREQS + base_freqs = [f.value for f in BACKFILLABLE_FREQS + INTRADAY_FREQS] + derived_freqs = [f.value for f in DERIVED_FREQS] + + try: + with get_db(read_only=True) as db: + stocks = db.query("SELECT ts_code FROM stock_info") + symbols = [s[0] for s in stocks] + except Exception: + symbols = [] + return { "status": "ok", - "db_path": db.db_path, - "stocks": db.query("SELECT count(*) FROM stock_info")[0][0], - "trading_days": db.query("SELECT count(*) FROM trading_calendar")[0][0], - "daily_records": repo.count_records(Freq.d1), + "exchange": "ashare", + "symbols": symbols, + "base_timeframes": base_freqs, + "derived_timeframes": derived_freqs, + "timeframes": [f.value for f in all_freqs], + "ready": True, } - @app.get("/stats") - async def stats(): - """Full DB statistics.""" - db = get_db() - repo = KLineRepository() - from ashare_dp.core.models import Freq - freq_stats = {} - for f in [Freq.d1, Freq.w1, Freq.M1, Freq.m1, Freq.m5, Freq.m15, Freq.m30, Freq.h1]: - dr = repo.get_date_range(f) - freq_stats[f.value] = { - "records": repo.count_records(f), - "start_date": dr[0].isoformat() if dr[0] else None, - "end_date": dr[1].isoformat() if dr[1] else None, - } + @app.get("/timeframes") + async def timeframes(): + """List available timeframes in data_provider format.""" + from ashare_dp.core.models import BACKFILLABLE_FREQS, DERIVED_FREQS, INTRADAY_FREQS + + all_freqs = BACKFILLABLE_FREQS + INTRADAY_FREQS + DERIVED_FREQS return { - "stocks": db.query("SELECT count(*) FROM stock_info")[0][0], - "trading_days": db.query("SELECT count(*) FROM trading_calendar")[0][0], - "frequencies": freq_stats, + "base_timeframes": [f.value for f in BACKFILLABLE_FREQS + INTRADAY_FREQS], + "derived_timeframes": [f.value for f in DERIVED_FREQS], + "timeframes": [f.value for f in all_freqs], } + @app.get("/api/candles") + async def api_candles( + symbol: str = Query(..., description="Stock code, e.g. 000001.SZ"), + tf: str = Query("1d", description="Timeframe: 1d, 1w, 1M, 1h, 5m, 1m, etc."), + start: Optional[int] = Query(None, description="Start timestamp in ms"), + end: Optional[int] = Query(None, description="End timestamp in ms"), + limit: Optional[int] = Query(None, description="Max candles to return"), + ): + """Fetch K-line candles in data_provider-compatible format. + + Returns [{"timestamp": ..., "open": ..., "high": ..., "low": ..., "close": ..., "volume": ...}] + """ + from datetime import datetime, timezone + from ashare_dp.core.models import BACKFILLABLE_FREQS, DERIVED_FREQS, INTRADAY_FREQS, Freq + + all_freqs = BACKFILLABLE_FREQS + INTRADAY_FREQS + DERIVED_FREQS + freq_values = [f.value for f in all_freqs] + + if tf not in freq_values: + raise HTTPException(status_code=400, detail=f"Invalid timeframe: {tf}") + + freq = Freq(tf) + repo = KLineRepository() + + start_date = None + end_date = None + if start is not None: + start_date = datetime.fromtimestamp(start / 1000, tz=timezone.utc).date() + if end is not None: + end_date = datetime.fromtimestamp(end / 1000, tz=timezone.utc).date() + + df = repo.read_klines( + freq=freq, ts_code=symbol, + start_date=start_date, end_date=end_date, + limit=limit or 10000, offset=0, + ) + + if df is None or df.empty: + return [] + + candles = [] + for _, row in df.iterrows(): + ts = row.get("trade_time") + if hasattr(ts, "timestamp"): + ts_ms = int(ts.timestamp() * 1000) + else: + ts_ms = 0 + candles.append({ + "timestamp": ts_ms, + "open": float(row.get("open", 0)), + "high": float(row.get("high", 0)), + "low": float(row.get("low", 0)), + "close": float(row.get("close", 0)), + "volume": float(row.get("volume", 0)), + }) + + if limit and len(candles) > limit: + candles = candles[-limit:] + + return candles + + # ── EMA52 Screening ── + + @app.get("/screening/ema52", response_class=HTMLResponse) + async def ema52_screening_page(): + """EMA52 screening results page.""" + return EMA52_HTML + + @app.get("/api/v1/screening/ema52") + async def api_ema52_screening( + freq: str = Query("1d", description="Frequency: 1d or 1w"), + limit: int = Query(200, ge=1, le=1000), + offset: int = Query(0, ge=0), + ): + """Get EMA52 screening results from the database.""" + if freq not in ("1d", "1w"): + raise HTTPException(status_code=400, detail="freq must be 1d or 1w") + + with get_db(read_only=True) as db: + row = db.conn.execute( + "SELECT COUNT(*) FROM ema52_screening WHERE freq = ?", + [freq], + ).fetchone() + total = row[0] if row else 0 + + results = db.query( + """SELECT ts_code, name, trade_date, close_price, ema52, distance_pct, amount + FROM ema52_screening + WHERE freq = ? + ORDER BY amount DESC NULLS LAST + LIMIT ? OFFSET ?""", + [freq, limit, offset], + ) + items = [ + { + "ts_code": r[0], + "name": r[1], + "trade_date": str(r[2]) if r[2] else None, + "close_price": r[3], + "ema52": r[4], + "distance_pct": r[5], + "amount": r[6], + } + for r in results + ] + return {"total": total, "limit": limit, "offset": offset, "items": items} + # Register routers app.include_router(stocks.router, prefix="/api/v1") app.include_router(kline.router, prefix="/api/v1") diff --git a/src/ashare_dp/api/deps.py b/src/ashare_dp/api/deps.py index 57f0479..fcd87be 100644 --- a/src/ashare_dp/api/deps.py +++ b/src/ashare_dp/api/deps.py @@ -1,6 +1,5 @@ """FastAPI dependency injection.""" -from ashare_dp.storage.database import get_db from ashare_dp.storage.repository import KLineRepository diff --git a/src/ashare_dp/api/routers/stocks.py b/src/ashare_dp/api/routers/stocks.py index 4b26166..7b565dc 100644 --- a/src/ashare_dp/api/routers/stocks.py +++ b/src/ashare_dp/api/routers/stocks.py @@ -19,7 +19,6 @@ async def list_stocks( offset: int = Query(0, ge=0), ): """List stocks with optional filters.""" - db = get_db() conditions = ["1=1"] params = [] @@ -31,14 +30,15 @@ async def list_stocks( params.append(market) where = " AND ".join(conditions) - rows = db.query( - f"SELECT * FROM stock_info WHERE {where} ORDER BY ts_code LIMIT ? OFFSET ?", - tuple(params) + (limit, offset), - ) - total = db.query( - f"SELECT count(*) FROM stock_info WHERE {where}", - tuple(params), - )[0][0] + with get_db(read_only=True) as db: + rows = db.query( + f"SELECT * FROM stock_info WHERE {where} ORDER BY ts_code LIMIT ? OFFSET ?", + tuple(params) + (limit, offset), + ) + total = db.query( + f"SELECT count(*) FROM stock_info WHERE {where}", + tuple(params), + )[0][0] cols = ["ts_code", "symbol", "name", "exchange", "area", "industry", "list_date", "delist_date", "market", "updated_at"] items = [] @@ -57,11 +57,11 @@ async def list_stocks( @router.get("/search") async def search_stocks(q: str = Query(..., min_length=1, description="Search query")): """Fuzzy search stocks by name or code.""" - db = get_db() - rows = db.query( - "SELECT * FROM stock_info WHERE name LIKE ? OR symbol LIKE ? OR ts_code LIKE ? LIMIT 50", - (f"%{q}%", f"%{q}%", f"%{q}%"), - ) + with get_db(read_only=True) as db: + rows = db.query( + "SELECT * FROM stock_info WHERE name LIKE ? OR symbol LIKE ? OR ts_code LIKE ? LIMIT 50", + (f"%{q}%", f"%{q}%", f"%{q}%"), + ) cols = ["ts_code", "symbol", "name", "exchange", "area", "industry", "list_date", "delist_date", "market", "updated_at"] items = [] for row in rows: @@ -78,14 +78,14 @@ async def search_stocks(q: str = Query(..., min_length=1, description="Search qu @router.get("/{ts_code}") async def get_stock(ts_code: str): """Get single stock info by ts_code (e.g. '000001.SZ').""" - db = get_db() - row = db.query( - "SELECT * FROM stock_info WHERE ts_code = ?", - (ts_code,), - ) - if not row: - raise HTTPException(status_code=404, detail=f"Stock not found: {ts_code}") - row = row[0] + with get_db(read_only=True) as db: + row = db.query( + "SELECT * FROM stock_info WHERE ts_code = ?", + (ts_code,), + ) + if not row: + raise HTTPException(status_code=404, detail=f"Stock not found: {ts_code}") + row = row[0] cols = ["ts_code", "symbol", "name", "exchange", "area", "industry", "list_date", "delist_date", "market", "updated_at"] item = {} for i, col in enumerate(cols): diff --git a/src/ashare_dp/cli/backfill_cmd.py b/src/ashare_dp/cli/backfill_cmd.py index 713a594..1ffd01a 100644 --- a/src/ashare_dp/cli/backfill_cmd.py +++ b/src/ashare_dp/cli/backfill_cmd.py @@ -11,7 +11,6 @@ from loguru import logger from ashare_dp.core.models import BACKFILLABLE_FREQS, INTRADAY_FREQS from ashare_dp.data.akshare_client import AKShareClient from ashare_dp.data.backfill import BackfillPipeline -from ashare_dp.storage.database import get_db backfill_app = typer.Typer() @@ -19,8 +18,6 @@ backfill_app = typer.Typer() @backfill_app.command("init") def init_db(): """Initialize database schema.""" - db = get_db() - db.connect() pipeline = BackfillPipeline() pipeline.init_db() pipeline.load_stock_list() @@ -36,8 +33,6 @@ def backfill_daily( symbols: str = typer.Option(None, help="Comma-separated stock symbols (default: all)"), ): """Backfill daily/weekly/monthly K-line data.""" - db = get_db() - db.connect() start_date = datetime.strptime(start, "%Y%m%d").date() end_date = datetime.strptime(end, "%Y%m%d").date() if end else date.today() @@ -70,8 +65,6 @@ def backfill_minute( symbols: str = typer.Option(None, help="Comma-separated stock symbols (default: all)"), ): """Backfill recent minute K-line data (limited API history).""" - db = get_db() - db.connect() sym_list = [s.strip() for s in symbols.split(",")] if symbols else None diff --git a/src/ashare_dp/cli/eod_cmd.py b/src/ashare_dp/cli/eod_cmd.py new file mode 100644 index 0000000..a417e11 --- /dev/null +++ b/src/ashare_dp/cli/eod_cmd.py @@ -0,0 +1,58 @@ +"""EOD CLI subcommand: pull end-of-day data.""" + +from __future__ import annotations + +from datetime import date, datetime + +import typer +from loguru import logger + +from ashare_dp.core.calendar import BEIJING_TZ +from ashare_dp.data.eod import EODPipeline + +eod_app = typer.Typer() + + +@eod_app.command("pull") +def eod_pull( + date_str: str = typer.Option( + None, "--date", "-d", + help="Trade date in YYYY-MM-DD format (default: today in Beijing time)", + ), + daily_only: bool = typer.Option( + False, "--daily-only", + help="Only pull daily data, skip minute/weekly/monthly", + ), +): + """Pull end-of-day K-line data for all active stocks.""" + now = datetime.now(BEIJING_TZ) + today = now.date() + + if date_str: + trade_date = date.fromisoformat(date_str) + else: + trade_date = today + + # Weekend check + if now.weekday() >= 5 and date_str is None: + logger.info("Weekend, skipping EOD pull (use --date to force)") + return + + logger.info(f"EOD pull: {trade_date.isoformat()}") + + # EOD writes to Parquet files directly (via pipeline), no DuckDB lock needed. + pipeline = EODPipeline() + + symbols = pipeline.get_active_symbols() + logger.info(f"EOD: {len(symbols)} active stocks") + + # Pull daily + pipeline.pull_daily(symbols, trade_date) + + if not daily_only: + # Pull minute + pipeline.pull_minute(symbols, trade_date) + # Derive weekly/monthly + pipeline.pull_weekly_monthly(symbols, trade_date) + + logger.info(f"EOD pull complete for {trade_date.isoformat()}") diff --git a/src/ashare_dp/cli/main.py b/src/ashare_dp/cli/main.py index de4cb5b..46b6e71 100644 --- a/src/ashare_dp/cli/main.py +++ b/src/ashare_dp/cli/main.py @@ -7,6 +7,8 @@ import typer from ashare_dp.cli.backfill_cmd import backfill_app from ashare_dp.cli.serve_cmd import serve_app from ashare_dp.cli.query_cmd import query_app +from ashare_dp.cli.eod_cmd import eod_app +from ashare_dp.cli.screening_cmd import screening_app app = typer.Typer( name="ashare-dp", @@ -16,6 +18,8 @@ app = typer.Typer( app.add_typer(backfill_app, name="backfill", help="Historical data backfill") app.add_typer(serve_app, name="serve", help="Start API server") app.add_typer(query_app, name="query", help="Ad-hoc data queries") +app.add_typer(eod_app, name="eod", help="End-of-day data pull") +app.add_typer(screening_app, name="screening", help="Stock screening (EMA52, etc.)") @app.command() diff --git a/src/ashare_dp/cli/query_cmd.py b/src/ashare_dp/cli/query_cmd.py index 7dcac1e..85d3463 100644 --- a/src/ashare_dp/cli/query_cmd.py +++ b/src/ashare_dp/cli/query_cmd.py @@ -22,9 +22,6 @@ def query_kline( limit: int = typer.Option(100, help="Max records"), ): """Query K-line data from the command line.""" - db = get_db() - db.connect() - freq_enum = Freq(freq) repo = KLineRepository() @@ -48,9 +45,6 @@ def query_latest( ts_code: str = typer.Option(None, help="Stock code (optional)"), ): """Show latest K-line data.""" - db = get_db() - db.connect() - freq_enum = Freq(freq) repo = KLineRepository() df = repo.get_latest(freq=freq_enum, ts_code=ts_code) @@ -70,19 +64,17 @@ def query_stocks( limit: int = typer.Option(50, help="Max records"), ): """List stocks.""" - db = get_db() - db.connect() - - if exchange: - rows = db.query( - "SELECT ts_code, symbol, name, exchange, market, list_date FROM stock_info WHERE exchange = ? LIMIT ?", - (exchange.upper(), limit), - ) - else: - rows = db.query( - "SELECT ts_code, symbol, name, exchange, market, list_date FROM stock_info LIMIT ?", - (limit,), - ) + with get_db(read_only=True) as db: + if exchange: + rows = db.query( + "SELECT ts_code, symbol, name, exchange, market, list_date FROM stock_info WHERE exchange = ? LIMIT ?", + (exchange.upper(), limit), + ) + else: + rows = db.query( + "SELECT ts_code, symbol, name, exchange, market, list_date FROM stock_info LIMIT ?", + (limit,), + ) typer.echo(f"\n{'ts_code':<12} {'symbol':<8} {'name':<12} {'exchange':<8} {'market':<10} {'list_date'}") typer.echo("-" * 60) @@ -95,15 +87,14 @@ def query_stocks( @query_app.command("stats") def query_stats(): """Show database statistics.""" - db = get_db() - db.connect() repo = KLineRepository() typer.echo("\nDatabase Statistics:") typer.echo("-" * 40) # Stock count - n = db.query("SELECT count(*) FROM stock_info")[0][0] + with get_db(read_only=True) as db: + n = db.query("SELECT count(*) FROM stock_info")[0][0] typer.echo(f" Stocks: {n}") # Record count and date range per frequency diff --git a/src/ashare_dp/cli/screening_cmd.py b/src/ashare_dp/cli/screening_cmd.py new file mode 100644 index 0000000..34d8084 --- /dev/null +++ b/src/ashare_dp/cli/screening_cmd.py @@ -0,0 +1,168 @@ +"""Screening CLI subcommand: EMA52 scanning and other screens.""" + +from __future__ import annotations + +import typer +from loguru import logger + +from ashare_dp.storage.database import get_db + +screening_app = typer.Typer() + + +@screening_app.command("ema52") +def ema52( + threshold: float = typer.Option( + None, "--threshold", "-t", + help="EMA52 proximity threshold (e.g. 0.03 = ±3%%). Default from .env", + ), +): + """Run EMA52 screening for 1d and 1w. + + Scans all stocks, computes EMA52 from close prices, + and saves results near EMA52 to the ema52_screening table. + """ + import numpy as np + import pandas as pd + from ashare_dp.config import Settings + from ashare_dp.core.models import Freq + + settings = Settings() + if threshold is None: + threshold = settings.ema52_threshold + + with get_db(read_only=False) as db: + # Get all ts_codes and names + stocks = db.query( + "SELECT ts_code, name FROM stock_info ORDER BY ts_code" + ) + symbols = [row[0] for row in stocks] + symbol_names = {row[0]: row[1] for row in stocks} + + logger.info( + f"EMA52 screening: {len(symbols)} stocks, " + f"threshold=±{threshold * 100:.1f}%" + ) + + for freq, freq_label, bars_needed in [ + (Freq.d1, "1d", 100), + (Freq.w1, "1w", 60), + ]: + try: + from ashare_dp.storage.partitioning import partition_glob + glob = partition_glob(freq) + sql = f""" + WITH ranked AS ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY ts_code ORDER BY trade_time DESC + ) as rn + FROM read_parquet('{glob}', hive_partitioning=true, union_by_name=true) + ) + SELECT * FROM ranked WHERE rn <= {bars_needed} + ORDER BY ts_code, trade_time ASC + """ + df = _query_parquet(sql) + if df.empty: + logger.warning(f"EMA52: no data for {freq_label}") + continue + except Exception as e: + logger.error(f"EMA52: failed to read {freq_label}: {e}") + continue + + df["trade_time"] = pd.to_datetime(df["trade_time"]) + df = df.sort_values(["ts_code", "trade_time"]) + + results = [] + for ts_code, group in df.groupby("ts_code"): + if ts_code not in symbol_names: + continue + group = group.tail(bars_needed) + if len(group) < 26: + continue + + closes = group["close"].astype(float).values + ema_values = _compute_ema(closes, 52) + if len(ema_values) == 0: + continue + + latest_close = closes[-1] + latest_ema = ema_values[-1] + if latest_ema <= 0: + continue + distance = (latest_close - latest_ema) / latest_ema + + # Compute daily amount in 亿 CNY + latest_amount = float(group["amount"].iloc[-1]) + amount_yi = latest_amount / 100_000_000 # 元 → 亿 + + if abs(distance) <= threshold: + results.append({ + "trade_date": group["trade_date"].iloc[-1], + "ts_code": ts_code, + "name": symbol_names[ts_code], + "freq": freq_label, + "close_price": round(latest_close, 2), + "ema52": round(latest_ema, 2), + "distance_pct": round(distance * 100, 2), + "amount": round(amount_yi, 2), + }) + + if results: + # Sort by amount DESC (largest turnover first) + results.sort(key=lambda r: r.get("amount", 0) or 0, reverse=True) + try: + with get_db(read_only=False) as wdb: + wdb.conn.execute( + "DELETE FROM ema52_screening WHERE freq = ?", + [freq_label], + ) + wdb.conn.executemany( + """INSERT INTO ema52_screening + (trade_date, ts_code, name, freq, close_price, ema52, distance_pct, amount) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (trade_date, ts_code, freq) DO UPDATE SET + close_price=excluded.close_price, + ema52=excluded.ema52, + distance_pct=excluded.distance_pct, + amount=excluded.amount, + name=excluded.name, + updated_at=now()""", + [(r["trade_date"], r["ts_code"], r["name"], + r["freq"], r["close_price"], r["ema52"], r["distance_pct"], r["amount"]) + for r in results], + ) + except Exception as e: + logger.error(f"EMA52: failed to save {freq_label}: {e}") + continue + + logger.info( + f"EMA52 {freq_label}: {len(results)} near " + f"(±{threshold * 100:.1f}%)" + ) + + logger.info("EMA52 screening complete") + + + +def _query_parquet(sql: str) -> "pd.DataFrame": + """Execute a SQL query against Parquet files using a transient in-memory DuckDB.""" + import duckdb + import pandas as pd + conn = duckdb.connect() + try: + return conn.execute(sql).fetchdf() + finally: + conn.close() + + +def _compute_ema(values: "np.ndarray", period: int) -> "np.ndarray": + """Compute EMA for a 1D array of values.""" + import numpy as np + if len(values) < period: + return np.array([]) + alpha = 2.0 / (period + 1) + ema = np.zeros(len(values)) + ema[period - 1] = np.mean(values[:period]) + for i in range(period, len(values)): + ema[i] = alpha * values[i] + (1 - alpha) * ema[i - 1] + return ema diff --git a/src/ashare_dp/config.py b/src/ashare_dp/config.py index 2a35cdc..151aac6 100644 --- a/src/ashare_dp/config.py +++ b/src/ashare_dp/config.py @@ -28,6 +28,11 @@ class Settings(BaseSettings): # Realtime realtime_poll_interval: int = 5 + # EMA52 screening + ema52_threshold: float = 0.03 + ema52_amount_min: float = 25.0 # 最低日成交额（亿元） + ema52_amount_max: float = 35.0 # 最高日成交额（亿元） + # Logging log_level: str = "INFO" diff --git a/src/ashare_dp/data/backfill.py b/src/ashare_dp/data/backfill.py index 774b567..e130ab1 100644 --- a/src/ashare_dp/data/backfill.py +++ b/src/ashare_dp/data/backfill.py @@ -9,9 +9,12 @@ from __future__ import annotations import time from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import date, datetime, timedelta +from pathlib import Path from typing import Optional import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq from loguru import logger from ashare_dp.config import Settings @@ -23,7 +26,8 @@ from ashare_dp.core.models import ( ) from ashare_dp.data.akshare_client import AKShareClient from ashare_dp.storage.database import get_db -from ashare_dp.storage.repository import KLineRepository +from ashare_dp.storage.repository import KLineRepository, _standardize_df +from ashare_dp.storage.partitioning import partition_glob settings = Settings() @@ -74,6 +78,16 @@ LEGACY_MIN_COLUMN_MAPPING = { REQUIRED_COLS = ["ts_code", "trade_time", "open", "high", "low", "close", "volume", "amount"] +def _to_ts_code(symbol: str) -> str: + """Convert 6-digit symbol to ts_code with exchange suffix.""" + code = str(symbol).zfill(6) + if code.startswith(("4", "8")) or code.startswith("92"): + return f"{code}.BJ" + elif code.startswith("6") or code.startswith("9"): + return f"{code}.SH" + return f"{code}.SZ" + + def _normalize_hist_df(df: pd.DataFrame, symbol: str, freq: Freq) -> pd.DataFrame: """Normalize daily K-line output to standard K-line schema.""" df = df.copy() @@ -89,7 +103,7 @@ def _normalize_hist_df(df: pd.DataFrame, symbol: str, freq: Freq) -> pd.DataFram break if "ts_code" not in df.columns: - df["ts_code"] = symbol + df["ts_code"] = _to_ts_code(symbol) if "trade_time" in df.columns: df["trade_time"] = pd.to_datetime(df["trade_time"]) @@ -121,7 +135,7 @@ def _normalize_min_df(df: pd.DataFrame, symbol: str) -> pd.DataFrame: break if "ts_code" not in df.columns: - df["ts_code"] = symbol + df["ts_code"] = _to_ts_code(symbol) if "trade_time" in df.columns: df["trade_time"] = pd.to_datetime(df["trade_time"]) @@ -175,6 +189,19 @@ def _resample_daily_to_period(df: pd.DataFrame, freq: Freq) -> pd.DataFrame: return result +def _code_to_ts_code(code: str) -> str: + """Convert bare 6-digit code to ts_code format: '000001' -> '000001.SZ'.""" + code = str(code).zfill(6) + if code.startswith(("4", "8")) or code.startswith("92"): + return f"{code}.BJ" + elif code.startswith("6") or code.startswith("9"): + return f"{code}.SH" + return f"{code}.SZ" + + +WRITE_BATCH = 500 # Write to disk every N stocks to limit memory + + class BackfillPipeline: """Orchestrates historical data backfill.""" @@ -187,18 +214,18 @@ class BackfillPipeline: self.client = client or AKShareClient() self.repo = repo or KLineRepository() self.max_workers = max_workers or settings.backfill_workers - self._db = get_db() def init_db(self): """Initialize database schema (tables and views).""" from ashare_dp.storage.schema import DDL_STATEMENTS logger.info("Initializing database schema...") - for ddl in DDL_STATEMENTS: - try: - self._db.execute(ddl) - except Exception as e: - logger.warning(f"DDL warning: {e}") + with get_db(read_only=False) as db: + for ddl in DDL_STATEMENTS: + try: + db.execute(ddl) + except Exception as e: + logger.warning(f"DDL warning: {e}") logger.info("Database schema initialized") def load_stock_list(self) -> pd.DataFrame: @@ -230,21 +257,40 @@ class BackfillPipeline: df["ts_code"] = df["symbol"].apply(_make_ts_code) df["exchange"] = df["ts_code"].str[-2:] - # Upsert into DuckDB + # Upsert into DuckDB (batch insert, not row-by-row) if "ts_code" in df.columns and "symbol" in df.columns: - for _, row in df.iterrows(): - try: - self._db.execute(""" - INSERT OR REPLACE INTO stock_info (ts_code, symbol, name, exchange, updated_at) - VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP) - """, ( - row.get("ts_code"), - str(row.get("symbol", "")), - str(row.get("name", "")), - row.get("exchange", ""), - )) - except Exception: - pass + import duckdb + try: + with get_db(read_only=False) as db: + # Build column arrays for bulk insert + ts_codes = df["ts_code"].tolist() + symbols = df["symbol"].astype(str).tolist() + names = df["name"].astype(str).tolist() + exchanges = df["exchange"].tolist() + db.conn.execute(""" + INSERT OR REPLACE INTO stock_info (ts_code, symbol, name, exchange, updated_at) + SELECT ts_code, symbol, name, exchange, now() + FROM (SELECT UNNEST($1::VARCHAR[]) AS ts_code, + UNNEST($2::VARCHAR[]) AS symbol, + UNNEST($3::VARCHAR[]) AS name, + UNNEST($4::VARCHAR[]) AS exchange) + """, [ts_codes, symbols, names, exchanges]) + except Exception as e: + logger.warning(f"Batch stock_info insert failed ({e}), falling back to row-by-row") + with get_db(read_only=False) as db2: + for _, row in df.iterrows(): + try: + db2.execute(""" + INSERT OR REPLACE INTO stock_info (ts_code, symbol, name, exchange, updated_at) + VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP) + """, ( + row.get("ts_code"), + str(row.get("symbol", "")), + str(row.get("name", "")), + row.get("exchange", ""), + )) + except Exception: + pass logger.info(f"Loaded {len(df)} stocks into stock_info") return df @@ -260,20 +306,33 @@ class BackfillPipeline: df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date - for _, row in df.iterrows(): - d = row["trade_date"] - try: - self._db.execute(""" + # Batch insert into DuckDB + dates = [d for d in df["trade_date"]] + weekdays = [d.weekday() for d in dates] + years = [d.year for d in dates] + months = [d.month for d in dates] + try: + with get_db(read_only=False) as db: + db.conn.execute(""" INSERT OR REPLACE INTO trading_calendar (trade_date, is_trading_day, week_day, year, month) - VALUES (?, 1, ?, ?, ?) - """, ( - d, - d.weekday(), - d.year, - d.month, - )) - except Exception: - pass + SELECT trade_date, 1, week_day, year, month + FROM (SELECT UNNEST($1::DATE[]) AS trade_date, + UNNEST($2::INTEGER[]) AS week_day, + UNNEST($3::INTEGER[]) AS year, + UNNEST($4::INTEGER[]) AS month) + """, [dates, weekdays, years, months]) + except Exception as e: + logger.warning(f"Batch trading_calendar insert failed ({e}), falling back to row-by-row") + with get_db(read_only=False) as db2: + for _, row in df.iterrows(): + d = row["trade_date"] + try: + db2.execute(""" + INSERT OR REPLACE INTO trading_calendar (trade_date, is_trading_day, week_day, year, month) + VALUES (?, 1, ?, ?, ?) + """, (d, d.weekday(), d.year, d.month)) + except Exception: + pass logger.info(f"Loaded {len(df)} trading days into trading_calendar") @@ -288,6 +347,9 @@ class BackfillPipeline: Daily data is fetched from Sina API. Weekly and monthly are derived by resampling the daily data. + Data is written in batches (every WRITE_BATCH stocks) to avoid + holding all 5500+ stocks in memory at once. + Args: symbols: List of stock symbols (e.g. ['000001', '600000']). If None, backfills all stocks from stock_info. @@ -320,7 +382,27 @@ class BackfillPipeline: logger.info("Backfilling 1d (daily) from Sina...") completed = 0 failed = 0 - all_daily_frames = [] + total_records = 0 + batch_frames = [] + temp_dir = Path(settings.parquet_dir) / ".tmp_daily" + temp_dir.mkdir(parents=True, exist_ok=True) + batch_idx = 0 + + def _flush_batch(): + nonlocal total_records, batch_idx + if batch_frames: + chunk = pd.concat(batch_frames, ignore_index=True) + # Standardize: add freq, keep only standard columns + chunk["freq"] = Freq.d1.value + # Convert bare codes (000001) to ts_code format (000001.SZ) + chunk["ts_code"] = chunk["ts_code"].astype(str).apply(_code_to_ts_code) + chunk = _standardize_df(chunk) + total_records += len(chunk) + table = pa.Table.from_pandas(chunk, preserve_index=False) + pq.write_table(table, str(temp_dir / f"batch_{batch_idx:04d}.parquet"), compression="zstd", compression_level=3) + logger.debug(f" Flushed batch {batch_idx}: {len(chunk)} rows") + batch_idx += 1 + batch_frames.clear() def _backfill_daily(symbol: str): try: @@ -348,22 +430,47 @@ class BackfillPipeline: try: df = future.result() if df is not None and not df.empty: - all_daily_frames.append(df) + batch_frames.append(df) completed += 1 except Exception: failed += 1 + if len(batch_frames) >= WRITE_BATCH: + _flush_batch() + if (completed + failed) % 100 == 0: logger.info( f" 1d: {completed + failed}/{len(symbols)} " f"({completed} ok, {failed} fail)" ) - # Batch write all daily data at once - total_records = 0 - if all_daily_frames: - combined = pd.concat(all_daily_frames, ignore_index=True) - total_records = self.repo.write_klines(combined, Freq.d1) + _flush_batch() + + # Merge all temp batches into final Hive-partitioned output + if batch_idx > 0: + logger.info(f"Merging {batch_idx} temp batches into final parquet...") + t0 = time.monotonic() + temp_glob = str(temp_dir / "batch_*.parquet") + copy_conn = duckdb.connect() + copy_conn.execute(f""" + COPY ( + SELECT *, + YEAR(trade_date) AS year, + MONTH(trade_date) AS month, + DAY(trade_date) AS day + FROM read_parquet('{temp_glob}', union_by_name=true) + ORDER BY trade_date, ts_code + ) TO '{settings.parquet_dir}/kline_{Freq.d1.storage_dir}' + (FORMAT PARQUET, COMPRESSION ZSTD, COMPRESSION_LEVEL 3, + PARTITION_BY (year, month, day), + OVERWRITE_OR_IGNORE true, FILENAME_PATTERN 'data_{{i}}') + """) + copy_conn.close() + # Clean up temp files + import shutil + shutil.rmtree(temp_dir, ignore_errors=True) + elapsed = time.monotonic() - t0 + logger.info(f" Merged in {elapsed:.1f}s") results[Freq.d1.value] = { "records": total_records, @@ -393,56 +500,64 @@ class BackfillPipeline: end_str: str, results: dict, ): - """Derive weekly/monthly K-lines from stored daily data. + """Derive weekly/monthly K-lines from daily data using DuckDB SQL. - Reads daily data for all symbols, resamples, and writes in batch - to avoid file overwrite issues with the grouped-per-day write strategy. + Reads ALL daily parquet files in a single DuckDB scan, + groups by ts_code + truncated date, and aggregates. + Much faster than per-stock queries — avoids 5500+ individual reads. """ - total_records = 0 - completed = 0 - failed = 0 - all_frames = [] + if freq == Freq.w1: + trunc = "week" + elif freq == Freq.M1: + trunc = "month" + else: + raise ValueError(f"Unsupported derive frequency: {freq}") - for symbol in symbols: - try: - df = self.repo.read_klines( - freq=Freq.d1, - ts_code=symbol, - start_date=date.fromisoformat( - pd.Timestamp(start_str).strftime("%Y-%m-%d") - ) if len(start_str) == 8 else date.fromisoformat(start_str), - end_date=date.fromisoformat( - pd.Timestamp(end_str).strftime("%Y-%m-%d") - ) if len(end_str) == 8 else date.fromisoformat(end_str), - limit=1_000_000, - ) - if df is not None and not df.empty: - df = _resample_daily_to_period(df, freq) - all_frames.append(df) - completed += 1 - except Exception as e: - logger.error(f"Failed derive {freq.value} for {symbol}: {e}") - failed += 1 + parquet_glob = partition_glob(Freq.d1) - if (completed + failed) % 500 == 0: - logger.info( - f" {freq.value}: {completed + failed}/{len(symbols)} " - f"({completed} ok, {failed} fail)" - ) + sql = f""" + SELECT + ts_code, + date_trunc('{trunc}', trade_time) AS trade_time, + CAST(date_trunc('{trunc}', trade_time) AS DATE) AS trade_date, + FIRST(open ORDER BY trade_time) AS open, + MAX(high) AS high, + MIN(low) AS low, + LAST(close ORDER BY trade_time) AS close, + SUM(volume) AS volume, + SUM(amount) AS amount + FROM read_parquet('{parquet_glob}', + hive_partitioning=true, union_by_name=true) + GROUP BY ts_code, date_trunc('{trunc}', trade_time) + ORDER BY ts_code, trade_time + """ - # Batch write all derived data at once - if all_frames: - combined = pd.concat(all_frames, ignore_index=True) - total_records = self.repo.write_klines(combined, freq) + logger.info(f"Deriving {freq.value} with single DuckDB scan...") + t0 = time.monotonic() + try: + derive_conn = duckdb.connect() + df = derive_conn.execute(sql).fetchdf() + derive_conn.close() + except Exception as e: + logger.error(f"DuckDB derive {freq.value} failed: {e}") + results[freq.value] = {"records": 0, "completed": 0, "failed": len(symbols)} + return + elapsed = time.monotonic() - t0 + logger.info(f" Aggregated {len(df)} rows in {elapsed:.1f}s") + + if df.empty: + results[freq.value] = {"records": 0, "completed": 0, "failed": 0} + return + + total_records = self.repo.write_klines(df, freq) results[freq.value] = { "records": total_records, - "completed": completed, - "failed": failed, + "completed": len(symbols), + "failed": 0, } logger.info( - f" {freq.value} done: {total_records} records, " - f"{completed} stocks ok, {failed} failed" + f" {freq.value} done: {total_records} records" ) def backfill_minute( @@ -485,7 +600,15 @@ class BackfillPipeline: completed = 0 failed = 0 - all_min_frames = [] + total_records = 0 + batch_frames = [] + + def _write_min_batch(): + nonlocal total_records + if batch_frames: + chunk = pd.concat(batch_frames, ignore_index=True) + total_records += self.repo.write_klines(chunk, freq) + batch_frames.clear() def _backfill_min(symbol: str): try: @@ -512,22 +635,25 @@ class BackfillPipeline: try: df = future.result() if df is not None and not df.empty: - all_min_frames.append(df) + batch_frames.append(df) completed += 1 except Exception: failed += 1 + # Write batch every WRITE_BATCH stocks + if len(batch_frames) >= WRITE_BATCH: + _write_min_batch() + if (completed + failed) % 100 == 0: logger.info( f" {freq.value}: {completed + failed}/{len(symbols)} " f"({completed} ok, {failed} fail)" ) - # Batch write all data for this frequency at once - total_records = 0 - if all_min_frames: - combined = pd.concat(all_min_frames, ignore_index=True) - total_records = self.repo.write_klines(combined, freq) + # Write remaining for this frequency + if batch_frames: + chunk = pd.concat(batch_frames, ignore_index=True) + total_records += self.repo.write_klines(chunk, freq) results[freq.value] = { "records": total_records, diff --git a/src/ashare_dp/scheduler/jobs.py b/src/ashare_dp/scheduler/jobs.py index 5fbabd7..a74bbb8 100644 --- a/src/ashare_dp/scheduler/jobs.py +++ b/src/ashare_dp/scheduler/jobs.py @@ -75,3 +75,147 @@ async def health_check_job(): # Import at bottom to avoid circular from ashare_dp.core.models import Freq + + +async def ema52_screening_job(): + """EOD EMA52 screening: find stocks near EMA52 for 1d and 1w. + + Triggered after EOD data pull. Scans all stocks, computes EMA52 + from close prices, and saves results to ema52_screening table. + """ + from ashare_dp.config import Settings + + import numpy as np + import pandas as pd + from ashare_dp.storage.database import get_db + + settings = Settings() + threshold = settings.ema52_threshold + db = get_db() + + # Get all ts_codes and names + stocks = db.query( + "SELECT ts_code, name FROM stock_info ORDER BY ts_code" + ) + symbols = [row[0] for row in stocks] + symbol_names = {row[0]: row[1] for row in stocks} + + logger.info( + f"EMA52 screening: {len(symbols)} stocks, " + f"threshold=±{threshold * 100:.1f}%" + ) + + for freq, freq_label, bars_needed in [ + (Freq.d1, "1d", 100), + (Freq.w1, "1w", 60), + ]: + try: + from ashare_dp.storage.partitioning import partition_glob + glob = partition_glob(freq) + # Get last N bars per stock using ROW_NUMBER() + sql = f""" + WITH ranked AS ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY ts_code ORDER BY trade_time DESC + ) as rn + FROM read_parquet('{glob}', hive_partitioning=true, union_by_name=true) + ) + SELECT * FROM ranked WHERE rn <= {bars_needed} + ORDER BY ts_code, trade_time ASC + """ + df = db.conn.execute(sql).fetchdf() + if df.empty: + logger.warning(f"EMA52: no data for {freq_label}") + continue + except Exception as e: + logger.error(f"EMA52: failed to read {freq_label}: {e}") + continue + + df["trade_time"] = pd.to_datetime(df["trade_time"]) + df = df.sort_values(["ts_code", "trade_time"]) + + # Get last N bars per stock, compute EMA52 + results = [] + for ts_code, group in df.groupby("ts_code"): + if ts_code not in symbol_names: + continue + group = group.tail(bars_needed) + if len(group) < 26: + continue + + closes = group["close"].astype(float).values + # Compute EMA52 + ema_values = _compute_ema(closes, 52) + if len(ema_values) == 0: + continue + + latest_close = closes[-1] + latest_ema = ema_values[-1] + if latest_ema <= 0: + continue + distance = (latest_close - latest_ema) / latest_ema + + # Compute daily amount in 亿 CNY + latest_amount = float(group["amount"].iloc[-1]) + amount_yi = latest_amount / 100_000_000 # 元 → 亿 + + if abs(distance) <= threshold: + results.append({ + "trade_date": group["trade_date"].iloc[-1], + "ts_code": ts_code, + "name": symbol_names[ts_code], + "freq": freq_label, + "close_price": round(latest_close, 2), + "ema52": round(latest_ema, 2), + "distance_pct": round(distance * 100, 2), + "amount": round(amount_yi, 2), + }) + + if results: + # Sort by amount DESC (largest turnover first) + results.sort(key=lambda r: r.get("amount", 0) or 0, reverse=True) + # Upsert: delete old then insert + try: + db.conn.execute( + "DELETE FROM ema52_screening WHERE freq = ?", + [freq_label], + ) + db.conn.executemany( + """INSERT INTO ema52_screening + (trade_date, ts_code, name, freq, close_price, ema52, distance_pct, amount) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (trade_date, ts_code, freq) DO UPDATE SET + close_price=excluded.close_price, + ema52=excluded.ema52, + distance_pct=excluded.distance_pct, + amount=excluded.amount, + name=excluded.name, + updated_at=now()""", + [(r["trade_date"], r["ts_code"], r["name"], + r["freq"], r["close_price"], r["ema52"], r["distance_pct"], r["amount"]) + for r in results], + ) + except Exception as e: + logger.error(f"EMA52: failed to save {freq_label}: {e}") + continue + + near = [r for r in results if abs(r["distance_pct"]) <= threshold * 100] + logger.info( + f"EMA52 {freq_label}: {len(results)} near " + f"(±{threshold * 100:.1f}%), {len(near)} within threshold" + ) + + logger.info("EMA52 screening complete") + + +def _compute_ema(values: "np.ndarray", period: int) -> "np.ndarray": + """Compute EMA for a 1D array of values.""" + import numpy as np + if len(values) < period: + return np.array([]) + alpha = 2.0 / (period + 1) + ema = np.zeros(len(values)) + ema[period - 1] = np.mean(values[:period]) + for i in range(period, len(values)): + ema[i] = alpha * values[i] + (1 - alpha) * ema[i - 1] + return ema diff --git a/src/ashare_dp/scheduler/scheduler.py b/src/ashare_dp/scheduler/scheduler.py index 4333570..834b4e4 100644 --- a/src/ashare_dp/scheduler/scheduler.py +++ b/src/ashare_dp/scheduler/scheduler.py @@ -19,7 +19,7 @@ class Scheduler: def start(self): """Start the scheduler and register jobs.""" - from ashare_dp.scheduler.jobs import eod_pull_job, health_check_job + from ashare_dp.scheduler.jobs import eod_pull_job, ema52_screening_job, health_check_job # EOD job: 15:05 Beijing time, Mon-Fri self._scheduler.add_job( @@ -48,8 +48,22 @@ class Scheduler: replace_existing=True, ) + # EMA52 screening: 15:10 Beijing time, Mon-Fri (after EOD pull) + self._scheduler.add_job( + ema52_screening_job, + trigger=CronTrigger( + day_of_week="mon-fri", + hour=15, + minute=10, + timezone=BEIJING_TZ, + ), + id="ema52_screening", + name="EMA52 screening", + replace_existing=True, + ) + self._scheduler.start() - logger.info("Scheduler started with EOD (15:05 Mon-Fri) + health check (08:00 daily)") + logger.info("Scheduler started with EOD (15:05) + EMA52 (15:10) + health check (08:00)") def shutdown(self): """Shut down the scheduler.""" diff --git a/src/ashare_dp/storage/database.py b/src/ashare_dp/storage/database.py index 93d0e9f..8392478 100644 --- a/src/ashare_dp/storage/database.py +++ b/src/ashare_dp/storage/database.py @@ -1,14 +1,15 @@ """DuckDB connection management. -Uses a module-level connection. DuckDB connections are not thread-safe, -so all operations are serialized through a threading lock for writes. +Connections are short-lived: open, query, close. No persistent connection +so that the API process (read-only) and EOD process (write) can coexist +without lock contention. """ from __future__ import annotations -import threading +from contextlib import contextmanager from pathlib import Path -from typing import Optional +from typing import Iterator, Optional import duckdb from loguru import logger @@ -17,16 +18,14 @@ from ashare_dp.config import Settings settings = Settings() -# Module-level lock for write serialization -_write_lock = threading.Lock() - class Database: - """Manages a persistent DuckDB connection.""" + """Short-lived DuckDB connection. Use as context manager.""" - def __init__(self, db_path: str | None = None): + def __init__(self, db_path: str | None = None, read_only: bool = False): self.db_path = str(db_path or settings.duckdb_path) self._conn: Optional[duckdb.DuckDBPyConnection] = None + self._read_only = read_only @property def conn(self) -> duckdb.DuckDBPyConnection: @@ -35,11 +34,14 @@ class Database: self.connect() return self._conn + @property + def read_only(self) -> bool: + return self._read_only + def connect(self) -> duckdb.DuckDBPyConnection: - """Open a persistent connection to the DuckDB database.""" + """Open a connection to the DuckDB database.""" Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) - self._conn = duckdb.connect(self.db_path) - logger.info(f"Connected to DuckDB: {self.db_path}") + self._conn = duckdb.connect(self.db_path, read_only=self._read_only) return self._conn def close(self): @@ -47,7 +49,6 @@ class Database: if self._conn is not None: self._conn.close() self._conn = None - logger.info("DuckDB connection closed") def execute(self, sql: str, params: tuple | None = None): """Execute a SQL statement.""" @@ -74,25 +75,27 @@ class Database: ).fetchone() return result[0] > 0 - def write_lock(self): - """Acquire the write lock (use as context manager).""" - return _write_lock - def __enter__(self): + self.connect() return self def __exit__(self, *args): self.close() -# Global database instance -_db: Optional[Database] = None +@contextmanager +def get_db(read_only: bool = True) -> Iterator[Database]: + """Get a short-lived DuckDB connection. + Opens, yields, and closes. No persistent singleton. + Use as context manager: `with get_db() as db: ...` -def get_db() -> Database: - """Get or create the global database instance.""" - global _db - if _db is None: - _db = Database() - _db.connect() - return _db + Default read_only=True for API queries. + Pass read_only=False for EOD/screening writes. + """ + db = Database(read_only=read_only) + try: + db.connect() + yield db + finally: + db.close() diff --git a/src/ashare_dp/storage/repository.py b/src/ashare_dp/storage/repository.py index 53e481a..86e9ace 100644 --- a/src/ashare_dp/storage/repository.py +++ b/src/ashare_dp/storage/repository.py @@ -6,6 +6,7 @@ from datetime import date, datetime from pathlib import Path from typing import Optional +import duckdb import pandas as pd import pyarrow as pa import pyarrow.parquet as pq @@ -46,7 +47,17 @@ class KLineRepository: """Read/write K-line data from/to Parquet files.""" def __init__(self): - self._db = get_db() + pass + + def _query_parquet(self, sql: str, params: tuple | None = None) -> pd.DataFrame: + """Execute a SQL query against Parquet files using a transient DuckDB.""" + conn = duckdb.connect() + try: + if params: + return conn.execute(sql, params).fetchdf() + return conn.execute(sql).fetchdf() + finally: + conn.close() # ---- Write ---- @@ -81,31 +92,21 @@ class KLineRepository: df = _standardize_df(df) # Determine partition date - if "trade_date" not in df.columns and partition_date is None: - raise StorageError("DataFrame must have 'trade_date' column or partition_date must be provided") + if partition_date is None: + if "trade_date" not in df.columns: + raise StorageError("No trade_date column and no partition_date provided") + partition_date = pd.to_datetime(df["trade_date"].iloc[0]).date() - records_written = 0 + # Ensure partition directory + pdir = ensure_partition_dir(freq, partition_date) + fpath = pdir / "data.parquet" - if partition_date is not None: - # Single date: all records go to one file - partition_dir = ensure_partition_dir(freq, partition_date) - file_path = partition_dir / "data.parquet" - table = pa.Table.from_pandas(df, preserve_index=False) - pq.write_table(table, str(file_path), **PARQUET_WRITE_OPTIONS) - records_written = len(df) - else: - # Group by trade_date, one file per day - df["_pd"] = pd.to_datetime(df["trade_date"]).dt.date - for d, group in df.groupby("_pd"): - group = group.drop(columns=["_pd"]) - partition_dir = ensure_partition_dir(freq, d) - file_path = partition_dir / "data.parquet" - table = pa.Table.from_pandas(group, preserve_index=False) - pq.write_table(table, str(file_path), **PARQUET_WRITE_OPTIONS) - records_written += len(group) + # Write Parquet + table = pa.Table.from_pandas(df) + pq.write_table(table, str(fpath), **PARQUET_WRITE_OPTIONS) - logger.debug(f"Wrote {records_written} records to kline_{freq.storage_dir}") - return records_written + logger.debug(f"Wrote {len(df)} records to {fpath}") + return len(df) # ---- Read ---- @@ -118,25 +119,27 @@ class KLineRepository: limit: int = 10000, offset: int = 0, ) -> pd.DataFrame: - """Read K-line data from Parquet files.""" + """Read K-line data from Parquet files. + + Uses DuckDB read_parquet() with Hive partitioning. + For derived 2h freq, computes from 1h data. + """ if freq == Freq.h2: return self._read_2h(ts_code, start_date, end_date, limit, offset) glob = partition_glob(freq) - parquet_path = Path(settings.parquet_dir) / f"kline_{freq.storage_dir}" - if not parquet_path.exists(): - return pd.DataFrame(columns=STANDARD_COLS) conditions = [] params = [] + if ts_code: - conditions.append(f"ts_code = ${len(params) + 1}") + conditions.append("ts_code = ?") params.append(ts_code) if start_date: - conditions.append(f"trade_date >= ${len(params) + 1}") + conditions.append("trade_date >= ?") params.append(start_date.isoformat()) if end_date: - conditions.append(f"trade_date <= ${len(params) + 1}") + conditions.append("trade_date <= ?") params.append(end_date.isoformat()) where_clause = "" @@ -150,7 +153,7 @@ class KLineRepository: LIMIT {limit} OFFSET {offset} """ try: - return self._db.conn.execute(sql, params).fetchdf() + return self._query_parquet(sql, tuple(params) if params else None) except Exception as e: logger.warning(f"Query failed for freq={freq.value}: {e}") return pd.DataFrame(columns=STANDARD_COLS) @@ -203,19 +206,16 @@ class KLineRepository: freq: Freq, ts_code: str | None = None, ) -> pd.DataFrame: - """Get the latest K-line data for the most recent trade date.""" - if freq == Freq.h2: - return self._read_2h(ts_code, limit=100) + """Get the latest K-line for each stock (or one stock). + Uses DuckDB ROW_NUMBER() window function to get the last N bars per stock. + """ glob = partition_glob(freq) - parquet_path = Path(settings.parquet_dir) / f"kline_{freq.storage_dir}" - if not parquet_path.exists(): - return pd.DataFrame(columns=STANDARD_COLS) - conditions = [] params = [] + if ts_code: - conditions.append(f"ts_code = ${len(params) + 1}") + conditions.append("ts_code = ?") params.append(ts_code) where_clause = "" @@ -223,49 +223,67 @@ class KLineRepository: where_clause = "WHERE " + " AND ".join(conditions) sql = f""" - WITH latest AS ( - SELECT MAX(trade_date) AS max_date + WITH ranked AS ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY ts_code ORDER BY trade_time DESC + ) as rn FROM read_parquet('{glob}', hive_partitioning=true, union_by_name=true) + {where_clause} ) - SELECT k.* FROM read_parquet('{glob}', hive_partitioning=true, union_by_name=true) k - JOIN latest ON k.trade_date = latest.max_date - {where_clause} - ORDER BY k.ts_code + SELECT * FROM ranked WHERE rn = 1 + ORDER BY ts_code """ - try: - return self._db.conn.execute(sql, params).fetchdf() - except Exception as e: - logger.warning(f"get_latest failed for freq={freq.value}: {e}") - return pd.DataFrame(columns=STANDARD_COLS) + return self._query_parquet(sql, tuple(params) if params else None) - # ---- Stats ---- - - def get_date_range(self, freq: Freq) -> tuple[date | None, date | None]: - """Get the min and max trade_date for a given frequency.""" - parquet_path = Path(settings.parquet_dir) / f"kline_{freq.storage_dir}" - if not parquet_path.exists(): - return None, None + def get_latest_bars( + self, + freq: Freq, + bars: int = 100, + ts_code: str | None = None, + ) -> pd.DataFrame: + """Get the last N bars per stock.""" glob = partition_glob(freq) - try: - row = self._db.conn.execute(f""" - SELECT MIN(trade_date), MAX(trade_date) + conditions = [] + params = [] + + if ts_code: + conditions.append("ts_code = ?") + params.append(ts_code) + + where_clause = "" + if conditions: + where_clause = "WHERE " + " AND ".join(conditions) + + sql = f""" + WITH ranked AS ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY ts_code ORDER BY trade_time DESC + ) as rn FROM read_parquet('{glob}', hive_partitioning=true, union_by_name=true) - """).fetchone() - return row[0], row[1] - except Exception: - return None, None + {where_clause} + ) + SELECT * FROM ranked WHERE rn <= {bars} + ORDER BY ts_code, trade_time ASC + """ + return self._query_parquet(sql, tuple(params) if params else None) def count_records(self, freq: Freq) -> int: - """Count total records for a given frequency.""" - parquet_path = Path(settings.parquet_dir) / f"kline_{freq.storage_dir}" - if not parquet_path.exists(): - return 0 + """Count total records for a frequency.""" glob = partition_glob(freq) - try: - row = self._db.conn.execute(f""" - SELECT count(*) + result = self._query_parquet( + f"SELECT COUNT(*) FROM read_parquet('{glob}', hive_partitioning=true, union_by_name=true)" + ) + return int(result.iloc[0, 0]) if not result.empty else 0 + + def get_date_range(self, freq: Freq) -> tuple[str, str] | None: + """Get min/max trade_date for a frequency.""" + glob = partition_glob(freq) + result = self._query_parquet( + f""" + SELECT MIN(trade_date) as min_date, MAX(trade_date) as max_date FROM read_parquet('{glob}', hive_partitioning=true, union_by_name=true) - """).fetchone() - return row[0] - except Exception: - return 0 + """ + ) + if result.empty: + return None + return str(result.iloc[0, 0]), str(result.iloc[0, 1]) diff --git a/src/ashare_dp/storage/schema.py b/src/ashare_dp/storage/schema.py index 38ac3b1..7eddf4a 100644 --- a/src/ashare_dp/storage/schema.py +++ b/src/ashare_dp/storage/schema.py @@ -33,4 +33,29 @@ DDL_STATEMENTS = [ month TINYINT NOT NULL ) """, + """ + CREATE SEQUENCE IF NOT EXISTS seq_ema52_id + """, + """ + CREATE TABLE IF NOT EXISTS ema52_screening ( + id BIGINT PRIMARY KEY DEFAULT nextval('seq_ema52_id'), + trade_date DATE NOT NULL, + ts_code VARCHAR(9) NOT NULL, + name VARCHAR(40), + freq VARCHAR(3) NOT NULL, + close_price DOUBLE NOT NULL, + ema52 DOUBLE NOT NULL, + distance_pct DOUBLE NOT NULL, + amount DOUBLE, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + UNIQUE(trade_date, ts_code, freq) + ) + """, + """ + CREATE INDEX IF NOT EXISTS idx_ema52_date ON ema52_screening(trade_date) + """, + # Migration: add amount column if upgrading from older schema + """ + ALTER TABLE ema52_screening ADD COLUMN IF NOT EXISTS amount DOUBLE + """, ]