xiaobai-review/xiaobai-datahub/datahub/datasets_ext.py

"""Extended EOD datasets beyond the first-batch A/B release groups.

These publish independently (soft): a failure here must not block daily/valuation
release. Scheduler runs them after the core EOD window.
"""

from __future__ import annotations

from typing import Any

# Independent soft datasets (not part of A/B atomic groups).
EXTENDED_SOFT_DATASETS = {
    "limit_events",
    "popularity",
    "dragon_tiger",
    "sector_daily",
}

EXTENDED_SCHEMA = """
CREATE TABLE IF NOT EXISTS eod_limit_events (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, limit_type TEXT NOT NULL,
  name TEXT, industry TEXT, close REAL, pct_chg REAL, amount REAL,
  limit_amount REAL, float_mv REAL, total_mv REAL, turnover_ratio REAL,
  fd_amount REAL, first_time TEXT, last_time TEXT,
  open_times INTEGER, up_stat TEXT, limit_times INTEGER,
  batch_id TEXT NOT NULL,
  PRIMARY KEY (ts_code, trade_date, limit_type, batch_id)
) WITHOUT ROWID;

CREATE TABLE IF NOT EXISTS staging_limit_events (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, limit_type TEXT NOT NULL, batch_id TEXT NOT NULL,
  name TEXT, industry TEXT, close REAL, pct_chg REAL, amount REAL,
  limit_amount REAL, float_mv REAL, total_mv REAL, turnover_ratio REAL,
  fd_amount REAL, first_time TEXT, last_time TEXT,
  open_times INTEGER, up_stat TEXT, limit_times INTEGER,
  PRIMARY KEY (batch_id, ts_code, trade_date, limit_type)
);

CREATE TABLE IF NOT EXISTS eod_popularity (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, source TEXT NOT NULL,
  ts_name TEXT, rank INTEGER, pct_change REAL, current_price REAL,
  hot REAL, concept TEXT, data_type TEXT,
  batch_id TEXT NOT NULL,
  PRIMARY KEY (ts_code, trade_date, source, batch_id)
) WITHOUT ROWID;

CREATE TABLE IF NOT EXISTS staging_popularity (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, source TEXT NOT NULL, batch_id TEXT NOT NULL,
  ts_name TEXT, rank INTEGER, pct_change REAL, current_price REAL,
  hot REAL, concept TEXT, data_type TEXT,
  PRIMARY KEY (batch_id, ts_code, trade_date, source)
);

CREATE TABLE IF NOT EXISTS eod_dragon_tiger (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, hm_name TEXT NOT NULL,
  ts_name TEXT, buy_amount REAL, sell_amount REAL, net_amount REAL,
  hm_orgs TEXT, tag TEXT, pct_change REAL, reason TEXT,
  batch_id TEXT NOT NULL,
  PRIMARY KEY (ts_code, trade_date, hm_name, batch_id)
) WITHOUT ROWID;

CREATE TABLE IF NOT EXISTS staging_dragon_tiger (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, hm_name TEXT NOT NULL, batch_id TEXT NOT NULL,
  ts_name TEXT, buy_amount REAL, sell_amount REAL, net_amount REAL,
  hm_orgs TEXT, tag TEXT, pct_change REAL, reason TEXT,
  PRIMARY KEY (batch_id, ts_code, trade_date, hm_name)
);

CREATE TABLE IF NOT EXISTS eod_sector_daily (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, family TEXT NOT NULL,
  name TEXT, open REAL, high REAL, low REAL, close REAL, pre_close REAL,
  pct_change REAL, vol REAL, turnover_rate REAL, amount REAL,
  batch_id TEXT NOT NULL,
  PRIMARY KEY (ts_code, trade_date, family, batch_id)
) WITHOUT ROWID;

CREATE TABLE IF NOT EXISTS staging_sector_daily (
  ts_code TEXT NOT NULL, trade_date TEXT NOT NULL, family TEXT NOT NULL, batch_id TEXT NOT NULL,
  name TEXT, open REAL, high REAL, low REAL, close REAL, pre_close REAL,
  pct_change REAL, vol REAL, turnover_rate REAL, amount REAL,
  PRIMARY KEY (batch_id, ts_code, trade_date, family)
);

CREATE TABLE IF NOT EXISTS sector_master (
  ts_code TEXT PRIMARY KEY,
  name TEXT,
  family TEXT NOT NULL,
  exchange TEXT,
  list_date TEXT,
  member_count INTEGER,
  type TEXT,
  updated_at TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_eod_limit_date ON eod_limit_events(trade_date, batch_id);
CREATE INDEX IF NOT EXISTS idx_eod_pop_date ON eod_popularity(trade_date, batch_id);
CREATE INDEX IF NOT EXISTS idx_eod_lhb_date ON eod_dragon_tiger(trade_date, batch_id);
CREATE INDEX IF NOT EXISTS idx_eod_sector_date ON eod_sector_daily(trade_date, family, batch_id);
"""

EXTENDED_DATASET_TABLES = {
    "limit_events": ("eod_limit_events", "staging_limit_events"),
    "popularity": ("eod_popularity", "staging_popularity"),
    "dragon_tiger": ("eod_dragon_tiger", "staging_dragon_tiger"),
    "sector_daily": ("eod_sector_daily", "staging_sector_daily"),
}

EXTENDED_STAGING_INSERT: dict[str, tuple[str, Any]] = {
    "limit_events": (
        "INSERT INTO staging_limit_events("
        "ts_code,trade_date,limit_type,batch_id,name,industry,close,pct_chg,amount,"
        "limit_amount,float_mv,total_mv,turnover_ratio,fd_amount,first_time,last_time,"
        "open_times,up_stat,limit_times) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
        lambda r, b: (
            r["ts_code"], r["trade_date"], r["limit_type"], b,
            r.get("name"), r.get("industry"), r.get("close"), r.get("pct_chg"), r.get("amount"),
            r.get("limit_amount"), r.get("float_mv"), r.get("total_mv"), r.get("turnover_ratio"),
            r.get("fd_amount"), r.get("first_time"), r.get("last_time"),
            r.get("open_times"), r.get("up_stat"), r.get("limit_times"),
        ),
    ),
    "popularity": (
        "INSERT INTO staging_popularity("
        "ts_code,trade_date,source,batch_id,ts_name,rank,pct_change,current_price,hot,concept,data_type) "
        "VALUES (?,?,?,?,?,?,?,?,?,?,?)",
        lambda r, b: (
            r["ts_code"], r["trade_date"], r["source"], b,
            r.get("ts_name"), r.get("rank"), r.get("pct_change"), r.get("current_price"),
            r.get("hot"), r.get("concept"), r.get("data_type"),
        ),
    ),
    "dragon_tiger": (
        "INSERT INTO staging_dragon_tiger("
        "ts_code,trade_date,hm_name,batch_id,ts_name,buy_amount,sell_amount,net_amount,"
        "hm_orgs,tag,pct_change,reason) VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
        lambda r, b: (
            r["ts_code"], r["trade_date"], r["hm_name"], b,
            r.get("ts_name"), r.get("buy_amount"), r.get("sell_amount"), r.get("net_amount"),
            r.get("hm_orgs"), r.get("tag"), r.get("pct_change"), r.get("reason"),
        ),
    ),
    "sector_daily": (
        "INSERT INTO staging_sector_daily("
        "ts_code,trade_date,family,batch_id,name,open,high,low,close,pre_close,"
        "pct_change,vol,turnover_rate,amount) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
        lambda r, b: (
            r["ts_code"], r["trade_date"], r["family"], b,
            r.get("name"), r.get("open"), r.get("high"), r.get("low"), r.get("close"),
            r.get("pre_close"), r.get("pct_change"), r.get("vol"), r.get("turnover_rate"),
            r.get("amount"),
        ),
    ),
}

EXTENDED_EOD_COPY = {
    "limit_events": (
        "INSERT OR REPLACE INTO eod_limit_events "
        "SELECT ts_code,trade_date,limit_type,name,industry,close,pct_chg,amount,"
        "limit_amount,float_mv,total_mv,turnover_ratio,fd_amount,first_time,last_time,"
        "open_times,up_stat,limit_times,batch_id "
        "FROM staging_limit_events WHERE batch_id = ?"
    ),
    "popularity": (
        "INSERT OR REPLACE INTO eod_popularity "
        "SELECT ts_code,trade_date,source,ts_name,rank,pct_change,current_price,hot,concept,data_type,batch_id "
        "FROM staging_popularity WHERE batch_id = ?"
    ),
    "dragon_tiger": (
        "INSERT OR REPLACE INTO eod_dragon_tiger "
        "SELECT ts_code,trade_date,hm_name,ts_name,buy_amount,sell_amount,net_amount,"
        "hm_orgs,tag,pct_change,reason,batch_id "
        "FROM staging_dragon_tiger WHERE batch_id = ?"
    ),
    "sector_daily": (
        "INSERT OR REPLACE INTO eod_sector_daily "
        "SELECT ts_code,trade_date,family,name,open,high,low,close,pre_close,"
        "pct_change,vol,turnover_rate,amount,batch_id "
        "FROM staging_sector_daily WHERE batch_id = ?"
    ),
}