from __future__ import annotations
|
|
import argparse
|
import os
|
import re
|
import sys
|
import time
|
from datetime import datetime
|
from pathlib import Path
|
|
import pymysql
|
|
|
MARKET_SUFFIX = {
|
"SH": ".SH",
|
"SZ": ".SZ",
|
"BJ": ".BJ",
|
}
|
|
|
def connect(args: argparse.Namespace):
|
password = args.password or os.environ.get("TIANXIA_MYSQL_PASSWORD")
|
if not password:
|
raise RuntimeError("missing MySQL password: pass --password or set TIANXIA_MYSQL_PASSWORD")
|
return pymysql.connect(
|
host=args.host,
|
port=args.port,
|
user=args.user,
|
password=password,
|
database=args.database,
|
charset="utf8mb4",
|
autocommit=False,
|
local_infile=True,
|
)
|
|
|
def ensure_schema(conn) -> None:
|
ddl = [
|
"""
|
CREATE TABLE IF NOT EXISTS a_share_minute_price_full_adj (
|
trade_date date NOT NULL,
|
trade_time time NOT NULL,
|
bar_time datetime NOT NULL,
|
symbol char(9) NOT NULL,
|
market varchar(8) NOT NULL,
|
open_price decimal(14,4) NULL,
|
high_price decimal(14,4) NULL,
|
low_price decimal(14,4) NULL,
|
close_price decimal(14,4) NULL,
|
volume decimal(24,4) NULL,
|
amount decimal(24,4) NULL,
|
source_file varchar(255) NOT NULL,
|
source_archive varchar(255) NOT NULL,
|
import_run_id varchar(96) NOT NULL,
|
imported_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
PRIMARY KEY (symbol, trade_date, trade_time),
|
KEY idx_full_adj_minute_date_symbol_time (trade_date, symbol, trade_time),
|
KEY idx_full_adj_minute_run (import_run_id)
|
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
""",
|
"""
|
CREATE TABLE IF NOT EXISTS a_share_minute_full_adj_import_run (
|
run_id varchar(96) NOT NULL,
|
started_at datetime NOT NULL,
|
finished_at datetime NULL,
|
status varchar(32) NOT NULL,
|
source_archive varchar(512) NOT NULL,
|
source_root varchar(512) NOT NULL,
|
files_total int NOT NULL DEFAULT 0,
|
files_loaded int NOT NULL DEFAULT 0,
|
rows_loaded bigint NOT NULL DEFAULT 0,
|
min_trade_date date NULL,
|
max_trade_date date NULL,
|
error_message text NULL,
|
PRIMARY KEY (run_id)
|
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
""",
|
"""
|
CREATE TABLE IF NOT EXISTS a_share_minute_full_adj_import_file_log (
|
id bigint NOT NULL AUTO_INCREMENT,
|
run_id varchar(96) NOT NULL,
|
symbol char(9) NOT NULL,
|
market varchar(8) NOT NULL,
|
source_file varchar(512) NOT NULL,
|
file_size bigint NOT NULL,
|
rows_loaded bigint NOT NULL,
|
started_at datetime NOT NULL,
|
finished_at datetime NOT NULL,
|
duration_ms bigint NOT NULL,
|
status varchar(32) NOT NULL,
|
error_message text NULL,
|
PRIMARY KEY (id),
|
UNIQUE KEY uq_full_adj_file_run (run_id, source_file),
|
KEY idx_full_adj_file_symbol (symbol),
|
KEY idx_full_adj_file_status (status)
|
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
""",
|
]
|
with conn.cursor() as cur:
|
for sql in ddl:
|
cur.execute(sql)
|
conn.commit()
|
|
|
def discover_files(root: Path) -> list[Path]:
|
files = sorted(root.rglob("price_*.csv"))
|
if not files:
|
raise RuntimeError(f"no price_*.csv files found under {root}")
|
return files
|
|
|
def symbol_for(path: Path) -> tuple[str, str]:
|
market = path.parent.name.upper()
|
suffix = MARKET_SUFFIX.get(market)
|
if not suffix:
|
raise RuntimeError(f"unsupported market directory {market}: {path}")
|
m = re.match(r"price_(\d{6})\.csv$", path.name, re.IGNORECASE)
|
if not m:
|
raise RuntimeError(f"unsupported file name: {path}")
|
return m.group(1) + suffix, market
|
|
|
def init_run(conn, run_id: str, source_archive: str, source_root: str, files_total: int) -> None:
|
with conn.cursor() as cur:
|
cur.execute(
|
"""
|
INSERT INTO a_share_minute_full_adj_import_run
|
(run_id, started_at, status, source_archive, source_root, files_total)
|
VALUES (%s, NOW(), 'running', %s, %s, %s)
|
ON DUPLICATE KEY UPDATE
|
status='running',
|
source_archive=VALUES(source_archive),
|
source_root=VALUES(source_root),
|
files_total=VALUES(files_total),
|
error_message=NULL
|
""",
|
(run_id, source_archive, source_root, files_total),
|
)
|
conn.commit()
|
|
|
def finish_run(conn, run_id: str, status: str, error: str | None = None) -> None:
|
with conn.cursor() as cur:
|
cur.execute(
|
"""
|
UPDATE a_share_minute_full_adj_import_run r
|
LEFT JOIN (
|
SELECT
|
import_run_id,
|
COUNT(*) AS rows_loaded,
|
MIN(trade_date) AS min_trade_date,
|
MAX(trade_date) AS max_trade_date
|
FROM a_share_minute_price_full_adj
|
WHERE import_run_id=%s
|
GROUP BY import_run_id
|
) x ON x.import_run_id = r.run_id
|
SET
|
r.finished_at=NOW(),
|
r.status=%s,
|
r.rows_loaded=COALESCE(x.rows_loaded, r.rows_loaded),
|
r.min_trade_date=COALESCE(x.min_trade_date, r.min_trade_date),
|
r.max_trade_date=COALESCE(x.max_trade_date, r.max_trade_date),
|
r.error_message=%s
|
WHERE r.run_id=%s
|
""",
|
(run_id, status, error, run_id),
|
)
|
conn.commit()
|
|
|
def load_file(conn, path: Path, root: Path, archive: str, run_id: str) -> int:
|
symbol, market = symbol_for(path)
|
rel = str(path.relative_to(root)).replace("\\", "/")
|
started = datetime.now()
|
t0 = time.perf_counter()
|
with conn.cursor() as cur:
|
cur.execute("DROP TEMPORARY TABLE IF EXISTS tmp_full_adj_minute_csv")
|
cur.execute(
|
"""
|
CREATE TEMPORARY TABLE tmp_full_adj_minute_csv (
|
timetag varchar(19) NOT NULL,
|
open_price decimal(14,4) NULL,
|
high_price decimal(14,4) NULL,
|
low_price decimal(14,4) NULL,
|
close_price decimal(14,4) NULL,
|
volume decimal(24,4) NULL,
|
amount decimal(24,4) NULL
|
) ENGINE=InnoDB
|
"""
|
)
|
cur.execute(
|
"""
|
LOAD DATA LOCAL INFILE %s
|
INTO TABLE tmp_full_adj_minute_csv
|
CHARACTER SET utf8mb4
|
FIELDS TERMINATED BY ',' OPTIONALLY ENCLOSED BY '"'
|
LINES TERMINATED BY '\n'
|
IGNORE 1 LINES
|
(timetag, open_price, high_price, low_price, close_price, volume, amount)
|
""",
|
(str(path),),
|
)
|
loaded_tmp = cur.rowcount
|
cur.execute(
|
"""
|
INSERT INTO a_share_minute_price_full_adj
|
(trade_date, trade_time, bar_time, symbol, market,
|
open_price, high_price, low_price, close_price, volume, amount,
|
source_file, source_archive, import_run_id)
|
SELECT
|
DATE(STR_TO_DATE(timetag, '%%Y%%m%%d %%H:%%i:%%s')),
|
TIME(STR_TO_DATE(timetag, '%%Y%%m%%d %%H:%%i:%%s')),
|
STR_TO_DATE(timetag, '%%Y%%m%%d %%H:%%i:%%s'),
|
%s,
|
%s,
|
open_price,
|
high_price,
|
low_price,
|
close_price,
|
volume,
|
amount,
|
%s,
|
%s,
|
%s
|
FROM tmp_full_adj_minute_csv
|
ON DUPLICATE KEY UPDATE
|
open_price=VALUES(open_price),
|
high_price=VALUES(high_price),
|
low_price=VALUES(low_price),
|
close_price=VALUES(close_price),
|
volume=VALUES(volume),
|
amount=VALUES(amount),
|
source_file=VALUES(source_file),
|
source_archive=VALUES(source_archive),
|
import_run_id=VALUES(import_run_id),
|
imported_at=CURRENT_TIMESTAMP
|
""",
|
(symbol, market, rel, archive, run_id),
|
)
|
inserted = cur.rowcount
|
duration_ms = int((time.perf_counter() - t0) * 1000)
|
cur.execute(
|
"""
|
INSERT INTO a_share_minute_full_adj_import_file_log
|
(run_id, symbol, market, source_file, file_size, rows_loaded,
|
started_at, finished_at, duration_ms, status)
|
VALUES (%s, %s, %s, %s, %s, %s, %s, NOW(), %s, 'loaded')
|
ON DUPLICATE KEY UPDATE
|
rows_loaded=VALUES(rows_loaded),
|
file_size=VALUES(file_size),
|
finished_at=VALUES(finished_at),
|
duration_ms=VALUES(duration_ms),
|
status='loaded',
|
error_message=NULL
|
""",
|
(run_id, symbol, market, rel, path.stat().st_size, loaded_tmp, started, duration_ms),
|
)
|
cur.execute(
|
"""
|
UPDATE a_share_minute_full_adj_import_run
|
SET files_loaded = (
|
SELECT COUNT(*) FROM a_share_minute_full_adj_import_file_log
|
WHERE run_id=%s AND status='loaded'
|
)
|
WHERE run_id=%s
|
""",
|
(run_id, run_id),
|
)
|
conn.commit()
|
return loaded_tmp if loaded_tmp >= 0 else inserted
|
|
|
def main() -> int:
|
parser = argparse.ArgumentParser()
|
parser.add_argument("--source-root", required=True)
|
parser.add_argument("--source-archive", required=True)
|
parser.add_argument("--run-id")
|
parser.add_argument("--host", default="127.0.0.1")
|
parser.add_argument("--port", type=int, default=3306)
|
parser.add_argument("--user", default="root")
|
parser.add_argument("--password")
|
parser.add_argument("--database", default="tianxia")
|
parser.add_argument("--limit", type=int)
|
parser.add_argument("--skip-loaded", action="store_true")
|
args = parser.parse_args()
|
|
root = Path(args.source_root).resolve()
|
files = discover_files(root)
|
if args.limit:
|
files = files[: args.limit]
|
run_id = args.run_id or "full_adj_minute_" + datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
conn = connect(args)
|
try:
|
ensure_schema(conn)
|
init_run(conn, run_id, args.source_archive, str(root), len(files))
|
loaded_total = 0
|
with conn.cursor() as cur:
|
loaded_files = set()
|
if args.skip_loaded:
|
cur.execute(
|
"SELECT source_file FROM a_share_minute_full_adj_import_file_log WHERE run_id=%s AND status='loaded'",
|
(run_id,),
|
)
|
loaded_files = {row[0] for row in cur.fetchall()}
|
for idx, path in enumerate(files, 1):
|
rel = str(path.relative_to(root)).replace("\\", "/")
|
if rel in loaded_files:
|
continue
|
rows = load_file(conn, path, root, args.source_archive, run_id)
|
loaded_total += rows
|
if idx % 10 == 0 or idx == 1 or idx == len(files):
|
print(f"{datetime.now():%Y-%m-%d %H:%M:%S} loaded {idx}/{len(files)} files; last={rel}; rows={rows}; total_rows_seen={loaded_total}", flush=True)
|
finish_run(conn, run_id, "completed")
|
print(f"completed run_id={run_id} files={len(files)} rows_seen={loaded_total}")
|
return 0
|
except Exception as exc:
|
conn.rollback()
|
try:
|
finish_run(conn, run_id, "failed", str(exc))
|
except Exception:
|
pass
|
print(f"failed run_id={run_id}: {exc}", file=sys.stderr)
|
return 1
|
finally:
|
conn.close()
|
|
|
if __name__ == "__main__":
|
raise SystemExit(main())
|
