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"""
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
+ """,
]