feat: operationalize local Tushare Qlib research
This commit is contained in:
@@ -0,0 +1,977 @@
|
||||
"""Create an auditable, scoped manifest from a live local Tushare mirror.
|
||||
|
||||
The command is intentionally read-only with respect to the mirror. It records
|
||||
the latest completed job for every selected ``(api_name, partition_key)`` in one
|
||||
SQLite read transaction, then optionally verifies the immutable Parquet files.
|
||||
It never reads Tushare credentials and never copies source data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from collections import Counter, defaultdict
|
||||
from datetime import date, datetime, timezone
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
from typing import Any, Iterable, Mapping, Optional, Sequence
|
||||
|
||||
|
||||
TOOL_VERSION = "2"
|
||||
DEFAULT_APIS = (
|
||||
"daily",
|
||||
"adj_factor",
|
||||
"trade_cal",
|
||||
"stock_basic",
|
||||
"index_daily",
|
||||
"index_weight",
|
||||
)
|
||||
DEFAULT_INDEX_CODES = ("000905.SH",)
|
||||
BLOCKING_LIVE_STATUSES = frozenset(("running", "deferred"))
|
||||
_MONTH_RE = re.compile(r"(?:^|[/_])month=(\d{4}-\d{2})(?:$|[/_])")
|
||||
_YEAR_RE = re.compile(r"(?:^|[/_])year=(\d{4})(?:$|[/_])")
|
||||
_INDEX_CODE_RE = re.compile(
|
||||
r"(?:^|/)(?:ts_code|index_code)=([^/]+)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_API_RE = re.compile(r"^[A-Za-z0-9_]+$")
|
||||
_INDEX_INPUT_RE = re.compile(r"^([A-Za-z0-9]+)[._]([A-Za-z0-9]+)$")
|
||||
|
||||
|
||||
class TushareSnapshotError(RuntimeError):
|
||||
"""Raised when a requested snapshot cannot be frozen safely."""
|
||||
|
||||
|
||||
def _utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat(timespec="microseconds").replace(
|
||||
"+00:00",
|
||||
"Z",
|
||||
)
|
||||
|
||||
|
||||
def _canonical_json(payload: Mapping[str, Any]) -> bytes:
|
||||
return json.dumps(
|
||||
payload,
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
allow_nan=False,
|
||||
).encode("utf-8")
|
||||
|
||||
|
||||
def _sha256_file(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _normalize_index_code(raw: str) -> str:
|
||||
value = raw.strip().upper()
|
||||
match = _INDEX_INPUT_RE.fullmatch(value)
|
||||
if match is None:
|
||||
raise ValueError(
|
||||
f"invalid index code {raw!r}; expected a value such as 000905.SH"
|
||||
)
|
||||
return f"{match.group(1)}.{match.group(2)}"
|
||||
|
||||
|
||||
def _partition_index_code(partition_key: str) -> Optional[str]:
|
||||
match = _INDEX_CODE_RE.search(partition_key)
|
||||
if match is None:
|
||||
return None
|
||||
raw = match.group(1)
|
||||
try:
|
||||
return _normalize_index_code(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _partition_month(partition_key: str) -> Optional[str]:
|
||||
match = _MONTH_RE.search(partition_key)
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def _partition_year(partition_key: str) -> Optional[int]:
|
||||
match = _YEAR_RE.search(partition_key)
|
||||
return int(match.group(1)) if match else None
|
||||
|
||||
|
||||
def _month_keys(start: date, end: date) -> list[str]:
|
||||
year = start.year
|
||||
month = start.month
|
||||
result: list[str] = []
|
||||
while (year, month) <= (end.year, end.month):
|
||||
result.append(f"{year:04d}-{month:02d}")
|
||||
month += 1
|
||||
if month == 13:
|
||||
year += 1
|
||||
month = 1
|
||||
return result
|
||||
|
||||
|
||||
def _parse_values(values: Sequence[str]) -> list[str]:
|
||||
result: list[str] = []
|
||||
for value in values:
|
||||
result.extend(part.strip() for part in value.split(",") if part.strip())
|
||||
return result
|
||||
|
||||
|
||||
def _validate_apis(apis: Iterable[str]) -> tuple[str, ...]:
|
||||
result: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for raw in apis:
|
||||
api = raw.strip()
|
||||
if not _API_RE.fullmatch(api):
|
||||
raise ValueError(f"invalid API name: {raw!r}")
|
||||
if api not in seen:
|
||||
seen.add(api)
|
||||
result.append(api)
|
||||
if not result:
|
||||
raise ValueError("at least one API must be requested")
|
||||
return tuple(result)
|
||||
|
||||
|
||||
def _scope_reason(
|
||||
api_name: str,
|
||||
partition_key: str,
|
||||
*,
|
||||
apis: set[str],
|
||||
months: set[str],
|
||||
years: set[int],
|
||||
index_weight_years: set[int],
|
||||
index_codes: set[str],
|
||||
) -> tuple[bool, str]:
|
||||
if api_name not in apis:
|
||||
return False, "api_not_requested"
|
||||
|
||||
if api_name in {"daily", "adj_factor"}:
|
||||
month = _partition_month(partition_key)
|
||||
if month is None:
|
||||
return False, "unrecognized_partition"
|
||||
return (True, "in_scope") if month in months else (
|
||||
False,
|
||||
"outside_date_range",
|
||||
)
|
||||
|
||||
if api_name == "trade_cal":
|
||||
year = _partition_year(partition_key)
|
||||
if year is None:
|
||||
return False, "unrecognized_partition"
|
||||
return (True, "in_scope") if year in years else (
|
||||
False,
|
||||
"outside_date_range",
|
||||
)
|
||||
|
||||
if api_name in {"index_daily", "index_weight"}:
|
||||
code = _partition_index_code(partition_key)
|
||||
if code not in index_codes:
|
||||
return False, "different_index_code"
|
||||
year = _partition_year(partition_key)
|
||||
if year is None:
|
||||
return False, "unrecognized_partition"
|
||||
accepted_years = (
|
||||
index_weight_years if api_name == "index_weight" else years
|
||||
)
|
||||
return (True, "in_scope") if year in accepted_years else (
|
||||
False,
|
||||
"outside_date_range",
|
||||
)
|
||||
|
||||
if api_name == "stock_basic":
|
||||
return True, "in_scope"
|
||||
|
||||
month = _partition_month(partition_key)
|
||||
if month is not None:
|
||||
return (True, "in_scope") if month in months else (
|
||||
False,
|
||||
"outside_date_range",
|
||||
)
|
||||
year = _partition_year(partition_key)
|
||||
if year is not None:
|
||||
return (True, "in_scope") if year in years else (
|
||||
False,
|
||||
"outside_date_range",
|
||||
)
|
||||
return True, "in_scope"
|
||||
|
||||
|
||||
def _state_file_identity(path: Path) -> dict[str, Any]:
|
||||
stat = path.stat()
|
||||
return {
|
||||
"device": stat.st_dev,
|
||||
"inode": stat.st_ino,
|
||||
"mtime_ns": stat.st_mtime_ns,
|
||||
"path": "data/state.sqlite3",
|
||||
"size_bytes": stat.st_size,
|
||||
}
|
||||
|
||||
|
||||
def _wal_identity(state_path: Path) -> Optional[dict[str, int]]:
|
||||
wal_path = state_path.with_name(state_path.name + "-wal")
|
||||
if not wal_path.exists():
|
||||
return None
|
||||
stat = wal_path.stat()
|
||||
return {
|
||||
"inode": stat.st_ino,
|
||||
"mtime_ns": stat.st_mtime_ns,
|
||||
"size_bytes": stat.st_size,
|
||||
}
|
||||
|
||||
|
||||
def _candidate_query(
|
||||
connection: sqlite3.Connection,
|
||||
api_name: str,
|
||||
*,
|
||||
months: Sequence[str],
|
||||
years: Sequence[int],
|
||||
index_codes: Sequence[str],
|
||||
) -> list[sqlite3.Row]:
|
||||
columns = (
|
||||
"id, api_name, partition_key, status, row_count, byte_count, "
|
||||
"sha256, file_path, completed_at, params_json, fields"
|
||||
)
|
||||
where = ["api_name = ?", "status = 'completed'"]
|
||||
parameters: list[Any] = [api_name]
|
||||
|
||||
if api_name in {"daily", "adj_factor"}:
|
||||
placeholders = ",".join("?" for _ in months)
|
||||
where.append(f"partition_key IN ({placeholders})")
|
||||
parameters.extend(f"month={month}" for month in months)
|
||||
elif api_name == "trade_cal":
|
||||
clauses = []
|
||||
for year in years:
|
||||
clauses.append("partition_key LIKE ?")
|
||||
parameters.append(f"%year={year}%")
|
||||
where.append("(" + " OR ".join(clauses) + ")")
|
||||
elif api_name in {"index_daily", "index_weight"}:
|
||||
code_clauses = []
|
||||
for code in index_codes:
|
||||
code_clauses.append("partition_key LIKE ?")
|
||||
parameters.append(f"%={code.replace('.', '_')}%")
|
||||
year_clauses = []
|
||||
for year in years:
|
||||
year_clauses.append("partition_key LIKE ?")
|
||||
parameters.append(f"%year={year}%")
|
||||
where.append("(" + " OR ".join(code_clauses) + ")")
|
||||
where.append("(" + " OR ".join(year_clauses) + ")")
|
||||
|
||||
query = f"SELECT {columns} FROM jobs WHERE " + " AND ".join(where)
|
||||
return list(connection.execute(query, parameters))
|
||||
|
||||
|
||||
def _assert_scope_complete(
|
||||
selected: Sequence[sqlite3.Row],
|
||||
*,
|
||||
apis: Sequence[str],
|
||||
months: Sequence[str],
|
||||
years: Sequence[int],
|
||||
index_weight_years: Sequence[int],
|
||||
index_codes: Sequence[str],
|
||||
) -> None:
|
||||
by_api: dict[str, list[sqlite3.Row]] = defaultdict(list)
|
||||
for row in selected:
|
||||
by_api[str(row["api_name"])].append(row)
|
||||
|
||||
missing: list[str] = []
|
||||
for api in apis:
|
||||
rows = by_api.get(api, [])
|
||||
if api in {"daily", "adj_factor"}:
|
||||
found = {
|
||||
value
|
||||
for value in (
|
||||
_partition_month(str(row["partition_key"])) for row in rows
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
missing.extend(
|
||||
f"{api}:month={month}" for month in months if month not in found
|
||||
)
|
||||
elif api == "trade_cal":
|
||||
found = {
|
||||
value
|
||||
for value in (
|
||||
_partition_year(str(row["partition_key"])) for row in rows
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
missing.extend(
|
||||
f"{api}:year={year}" for year in years if year not in found
|
||||
)
|
||||
elif api in {"index_daily", "index_weight"}:
|
||||
found = {
|
||||
(
|
||||
_partition_index_code(str(row["partition_key"])),
|
||||
_partition_year(str(row["partition_key"])),
|
||||
)
|
||||
for row in rows
|
||||
}
|
||||
required_years = (
|
||||
index_weight_years if api == "index_weight" else years
|
||||
)
|
||||
missing.extend(
|
||||
f"{api}:{code}/year={year}"
|
||||
for code in index_codes
|
||||
for year in required_years
|
||||
if (code, year) not in found
|
||||
)
|
||||
elif not rows:
|
||||
missing.append(f"{api}:no_completed_partition")
|
||||
|
||||
if missing:
|
||||
preview = ", ".join(missing[:12])
|
||||
suffix = "" if len(missing) <= 12 else f" (+{len(missing) - 12} more)"
|
||||
raise TushareSnapshotError(
|
||||
f"requested scope is incomplete; missing {preview}{suffix}"
|
||||
)
|
||||
|
||||
|
||||
def _resolve_source_file(mirror_root: Path, raw_path: str) -> tuple[Path, str]:
|
||||
path = Path(raw_path)
|
||||
if not path.is_absolute():
|
||||
path = mirror_root / path
|
||||
resolved = path.expanduser().resolve(strict=True)
|
||||
root = mirror_root.resolve(strict=True)
|
||||
try:
|
||||
relative = resolved.relative_to(root)
|
||||
except ValueError as exc:
|
||||
raise TushareSnapshotError(
|
||||
"source file escapes the declared mirror root"
|
||||
) from exc
|
||||
return resolved, relative.as_posix()
|
||||
|
||||
|
||||
def _parquet_row_count(path: Path) -> int:
|
||||
try:
|
||||
import pyarrow.parquet as parquet
|
||||
except ImportError as exc:
|
||||
raise TushareSnapshotError(
|
||||
"Parquet row verification requires pyarrow; install the Qlib "
|
||||
"research requirements or use --no-verify explicitly"
|
||||
) from exc
|
||||
try:
|
||||
metadata = parquet.ParquetFile(path).metadata
|
||||
except Exception as exc:
|
||||
raise TushareSnapshotError(
|
||||
f"cannot read Parquet metadata for {path.name}: {exc}"
|
||||
) from exc
|
||||
return int(metadata.num_rows)
|
||||
|
||||
|
||||
def _verify_job_file(
|
||||
mirror_root: Path,
|
||||
row: sqlite3.Row,
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
raw_path = row["file_path"]
|
||||
if not isinstance(raw_path, str) or not raw_path:
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has no source file path"
|
||||
)
|
||||
path, relative = _resolve_source_file(mirror_root, raw_path)
|
||||
stat = path.stat()
|
||||
observed_size = int(stat.st_size)
|
||||
expected_size = row["byte_count"]
|
||||
if expected_size is None or observed_size != int(expected_size):
|
||||
raise TushareSnapshotError(
|
||||
f"size mismatch for job {row['id']} ({relative}): "
|
||||
f"expected {expected_size}, observed {observed_size}"
|
||||
)
|
||||
|
||||
expected_sha = str(row["sha256"] or "").lower()
|
||||
if not re.fullmatch(r"[0-9a-f]{64}", expected_sha):
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has no valid SHA-256"
|
||||
)
|
||||
observed_sha = _sha256_file(path)
|
||||
if observed_sha != expected_sha:
|
||||
raise TushareSnapshotError(
|
||||
f"SHA-256 mismatch for job {row['id']} ({relative})"
|
||||
)
|
||||
|
||||
expected_rows = row["row_count"]
|
||||
observed_rows = _parquet_row_count(path)
|
||||
if expected_rows is None or observed_rows != int(expected_rows):
|
||||
raise TushareSnapshotError(
|
||||
f"row-count mismatch for job {row['id']} ({relative}): "
|
||||
f"expected {expected_rows}, observed {observed_rows}"
|
||||
)
|
||||
return relative, {
|
||||
"row_count": observed_rows,
|
||||
"sha256": observed_sha,
|
||||
"size_bytes": observed_size,
|
||||
}
|
||||
|
||||
|
||||
def _relative_job_path(mirror_root: Path, row: sqlite3.Row) -> str:
|
||||
raw_path = row["file_path"]
|
||||
if not isinstance(raw_path, str) or not raw_path:
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has no source file path"
|
||||
)
|
||||
_, relative = _resolve_source_file(mirror_root, raw_path)
|
||||
return relative
|
||||
|
||||
|
||||
def _expected_job_metadata(row: sqlite3.Row) -> tuple[int, int, str]:
|
||||
try:
|
||||
row_count = int(row["row_count"])
|
||||
byte_count = int(row["byte_count"])
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has incomplete count metadata"
|
||||
) from exc
|
||||
if row_count < 0 or byte_count < 0:
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has negative count metadata"
|
||||
)
|
||||
source_sha = str(row["sha256"] or "").lower()
|
||||
if not re.fullmatch(r"[0-9a-f]{64}", source_sha):
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has no valid SHA-256"
|
||||
)
|
||||
return row_count, byte_count, source_sha
|
||||
|
||||
|
||||
def _build_table_summaries(
|
||||
jobs: Sequence[Mapping[str, Any]],
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
grouped: dict[str, list[Mapping[str, Any]]] = defaultdict(list)
|
||||
for job in jobs:
|
||||
grouped[str(job["api_name"])].append(job)
|
||||
summaries: dict[str, dict[str, Any]] = {}
|
||||
for api_name in sorted(grouped):
|
||||
rows = grouped[api_name]
|
||||
partitions = sorted(str(row["partition_key"]) for row in rows)
|
||||
summaries[api_name] = {
|
||||
"bytes": sum(int(row["byte_count"]) for row in rows),
|
||||
"first_partition": partitions[0],
|
||||
"last_partition": partitions[-1],
|
||||
"partitions": len(rows),
|
||||
"rows": sum(int(row["row_count"]) for row in rows),
|
||||
"shadowed_completed_jobs": sum(
|
||||
int(row["shadowed_completed_count"]) for row in rows
|
||||
),
|
||||
}
|
||||
return summaries
|
||||
|
||||
|
||||
def _request_sha256(row: sqlite3.Row) -> str:
|
||||
raw_params = row["params_json"]
|
||||
if not isinstance(raw_params, str):
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has no params_json provenance"
|
||||
)
|
||||
try:
|
||||
params = json.loads(raw_params)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise TushareSnapshotError(
|
||||
f"completed job {row['id']} has invalid params_json provenance"
|
||||
) from exc
|
||||
request = {
|
||||
"api_name": str(row["api_name"]),
|
||||
"partition_key": str(row["partition_key"]),
|
||||
"fields": str(row["fields"] or ""),
|
||||
"params": params,
|
||||
}
|
||||
return hashlib.sha256(_canonical_json(request)).hexdigest()
|
||||
|
||||
|
||||
def _latest_completed_ids(
|
||||
state_path: Path,
|
||||
) -> dict[tuple[str, str], int]:
|
||||
connection = sqlite3.connect(
|
||||
state_path.as_uri() + "?mode=ro",
|
||||
uri=True,
|
||||
timeout=30,
|
||||
)
|
||||
try:
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT api_name, partition_key, MAX(id)
|
||||
FROM jobs
|
||||
WHERE status = 'completed'
|
||||
GROUP BY api_name, partition_key
|
||||
"""
|
||||
)
|
||||
return {
|
||||
(str(api_name), str(partition_key)): int(job_id)
|
||||
for api_name, partition_key, job_id in rows
|
||||
}
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def create_snapshot(
|
||||
mirror_root: str | Path,
|
||||
*,
|
||||
start: str | date,
|
||||
end: str | date,
|
||||
apis: Sequence[str] = DEFAULT_APIS,
|
||||
index_codes: Sequence[str] = DEFAULT_INDEX_CODES,
|
||||
verify: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""Read and verify one deterministic Tushare mirror selection."""
|
||||
|
||||
root = Path(mirror_root).expanduser().resolve(strict=True)
|
||||
state_path = root / "data" / "state.sqlite3"
|
||||
if not state_path.is_file():
|
||||
raise TushareSnapshotError(
|
||||
"mirror is missing its data/state.sqlite3 catalog"
|
||||
)
|
||||
|
||||
start_date = (
|
||||
start if isinstance(start, date) else date.fromisoformat(str(start))
|
||||
)
|
||||
end_date = end if isinstance(end, date) else date.fromisoformat(str(end))
|
||||
if start_date > end_date:
|
||||
raise ValueError("start must be on or before end")
|
||||
|
||||
requested_apis = _validate_apis(apis)
|
||||
requested_codes = tuple(
|
||||
dict.fromkeys(_normalize_index_code(code) for code in index_codes)
|
||||
)
|
||||
if (
|
||||
{"index_daily", "index_weight"} & set(requested_apis)
|
||||
and not requested_codes
|
||||
):
|
||||
raise ValueError("index APIs require at least one --index-codes value")
|
||||
|
||||
months = _month_keys(start_date, end_date)
|
||||
years = list(range(start_date.year, end_date.year + 1))
|
||||
index_weight_years = (
|
||||
[start_date.year - 1, *years]
|
||||
if "index_weight" in requested_apis
|
||||
else years
|
||||
)
|
||||
month_set = set(months)
|
||||
year_set = set(years)
|
||||
index_weight_year_set = set(index_weight_years)
|
||||
code_set = set(requested_codes)
|
||||
api_set = set(requested_apis)
|
||||
|
||||
identity_before = _state_file_identity(state_path)
|
||||
wal_before = _wal_identity(state_path)
|
||||
read_started_at = _utc_now()
|
||||
connection = sqlite3.connect(
|
||||
state_path.as_uri() + "?mode=ro",
|
||||
uri=True,
|
||||
timeout=30,
|
||||
)
|
||||
connection.row_factory = sqlite3.Row
|
||||
try:
|
||||
connection.execute("BEGIN")
|
||||
required_columns = {
|
||||
"id",
|
||||
"api_name",
|
||||
"partition_key",
|
||||
"status",
|
||||
"row_count",
|
||||
"byte_count",
|
||||
"sha256",
|
||||
"file_path",
|
||||
"completed_at",
|
||||
"fields",
|
||||
"params_json",
|
||||
}
|
||||
available_columns = {
|
||||
str(row["name"])
|
||||
for row in connection.execute("PRAGMA table_info(jobs)")
|
||||
}
|
||||
missing_columns = required_columns - available_columns
|
||||
if missing_columns:
|
||||
raise TushareSnapshotError(
|
||||
"jobs table is missing required columns: "
|
||||
+ ", ".join(sorted(missing_columns))
|
||||
)
|
||||
|
||||
sqlite_identity = {
|
||||
"data_version": int(
|
||||
connection.execute("PRAGMA data_version").fetchone()[0]
|
||||
),
|
||||
"journal_mode": str(
|
||||
connection.execute("PRAGMA journal_mode").fetchone()[0]
|
||||
),
|
||||
"page_count": int(
|
||||
connection.execute("PRAGMA page_count").fetchone()[0]
|
||||
),
|
||||
"page_size": int(
|
||||
connection.execute("PRAGMA page_size").fetchone()[0]
|
||||
),
|
||||
"schema_version": int(
|
||||
connection.execute("PRAGMA schema_version").fetchone()[0]
|
||||
),
|
||||
"user_version": int(
|
||||
connection.execute("PRAGMA user_version").fetchone()[0]
|
||||
),
|
||||
}
|
||||
|
||||
scoped_completed: list[sqlite3.Row] = []
|
||||
for api_name in requested_apis:
|
||||
api_years = (
|
||||
index_weight_years
|
||||
if api_name == "index_weight"
|
||||
else years
|
||||
)
|
||||
rows = _candidate_query(
|
||||
connection,
|
||||
api_name,
|
||||
months=months,
|
||||
years=api_years,
|
||||
index_codes=requested_codes,
|
||||
)
|
||||
scoped_completed.extend(
|
||||
row
|
||||
for row in rows
|
||||
if _scope_reason(
|
||||
str(row["api_name"]),
|
||||
str(row["partition_key"]),
|
||||
apis=api_set,
|
||||
months=month_set,
|
||||
years=year_set,
|
||||
index_weight_years=index_weight_year_set,
|
||||
index_codes=code_set,
|
||||
)[0]
|
||||
)
|
||||
|
||||
live_rows = list(
|
||||
connection.execute(
|
||||
"""
|
||||
SELECT id, api_name, partition_key, status
|
||||
FROM jobs
|
||||
WHERE status IN ('running', 'deferred')
|
||||
ORDER BY id
|
||||
"""
|
||||
)
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
read_finished_at = _utc_now()
|
||||
identity_after = _state_file_identity(state_path)
|
||||
wal_after = _wal_identity(state_path)
|
||||
|
||||
excluded_live_jobs: list[dict[str, Any]] = []
|
||||
blocking_live_jobs: list[sqlite3.Row] = []
|
||||
for row in live_rows:
|
||||
in_scope, reason = _scope_reason(
|
||||
str(row["api_name"]),
|
||||
str(row["partition_key"]),
|
||||
apis=api_set,
|
||||
months=month_set,
|
||||
years=year_set,
|
||||
index_weight_years=index_weight_year_set,
|
||||
index_codes=code_set,
|
||||
)
|
||||
if in_scope and str(row["status"]) in BLOCKING_LIVE_STATUSES:
|
||||
blocking_live_jobs.append(row)
|
||||
else:
|
||||
excluded_live_jobs.append(
|
||||
{
|
||||
"api_name": str(row["api_name"]),
|
||||
"id": int(row["id"]),
|
||||
"partition_key": str(row["partition_key"]),
|
||||
"reason": reason,
|
||||
"status": str(row["status"]),
|
||||
}
|
||||
)
|
||||
if blocking_live_jobs:
|
||||
details = ", ".join(
|
||||
f"{row['api_name']}:{row['partition_key']}#{row['id']}"
|
||||
f"[{row['status']}]"
|
||||
for row in blocking_live_jobs[:12]
|
||||
)
|
||||
raise TushareSnapshotError(
|
||||
"requested scope contains running/deferred jobs; retry after they "
|
||||
f"reach a terminal state: {details}"
|
||||
)
|
||||
|
||||
grouped: dict[tuple[str, str], list[sqlite3.Row]] = defaultdict(list)
|
||||
for row in scoped_completed:
|
||||
grouped[(str(row["api_name"]), str(row["partition_key"]))].append(row)
|
||||
latest_rows: list[sqlite3.Row] = []
|
||||
shadow_counts: dict[tuple[str, str], int] = {}
|
||||
for key, rows in grouped.items():
|
||||
ordered = sorted(rows, key=lambda row: int(row["id"]))
|
||||
latest_rows.append(ordered[-1])
|
||||
shadow_counts[key] = len(ordered) - 1
|
||||
latest_rows.sort(
|
||||
key=lambda row: (
|
||||
str(row["api_name"]),
|
||||
str(row["partition_key"]),
|
||||
int(row["id"]),
|
||||
)
|
||||
)
|
||||
|
||||
_assert_scope_complete(
|
||||
latest_rows,
|
||||
apis=requested_apis,
|
||||
months=months,
|
||||
years=years,
|
||||
index_weight_years=index_weight_years,
|
||||
index_codes=requested_codes,
|
||||
)
|
||||
|
||||
selected_jobs: list[dict[str, Any]] = []
|
||||
for row in latest_rows:
|
||||
key = (str(row["api_name"]), str(row["partition_key"]))
|
||||
row_count, byte_count, expected_sha = _expected_job_metadata(row)
|
||||
if verify:
|
||||
relative_path, observed = _verify_job_file(root, row)
|
||||
else:
|
||||
relative_path = _relative_job_path(root, row)
|
||||
observed = None
|
||||
job: dict[str, Any] = {
|
||||
"api_name": key[0],
|
||||
"byte_count": byte_count,
|
||||
"completed_at": row["completed_at"],
|
||||
"id": int(row["id"]),
|
||||
"partition_key": key[1],
|
||||
"path": relative_path,
|
||||
"row_count": row_count,
|
||||
"request_sha256": _request_sha256(row),
|
||||
"sha256": expected_sha,
|
||||
"shadowed_completed_count": shadow_counts[key],
|
||||
"verified": verify,
|
||||
}
|
||||
if observed is not None:
|
||||
job["observed"] = observed
|
||||
selected_jobs.append(job)
|
||||
|
||||
latest_ids_after_verification = _latest_completed_ids(state_path)
|
||||
stale_jobs = [
|
||||
{
|
||||
"api_name": job["api_name"],
|
||||
"partition_key": job["partition_key"],
|
||||
"selected_job_id": job["id"],
|
||||
"latest_job_id": latest_ids_after_verification.get(
|
||||
(job["api_name"], job["partition_key"])
|
||||
),
|
||||
}
|
||||
for job in selected_jobs
|
||||
if latest_ids_after_verification.get(
|
||||
(job["api_name"], job["partition_key"])
|
||||
)
|
||||
!= job["id"]
|
||||
]
|
||||
if stale_jobs:
|
||||
raise TushareSnapshotError(
|
||||
"selected jobs changed during file verification: "
|
||||
f"{stale_jobs[:12]}"
|
||||
)
|
||||
if verify:
|
||||
for row, job in zip(latest_rows, selected_jobs):
|
||||
relative_path, observed = _verify_job_file(root, row)
|
||||
if (
|
||||
relative_path != job["path"]
|
||||
or observed != job.get("observed")
|
||||
):
|
||||
raise TushareSnapshotError(
|
||||
"selected source bytes changed during double verification"
|
||||
)
|
||||
identity_after_verification = _state_file_identity(state_path)
|
||||
wal_after_verification = _wal_identity(state_path)
|
||||
|
||||
selection_identity = [
|
||||
{
|
||||
key: value
|
||||
for key, value in job.items()
|
||||
if key not in {"observed", "verified"}
|
||||
}
|
||||
for job in selected_jobs
|
||||
]
|
||||
selection_sha = hashlib.sha256(
|
||||
_canonical_json({"selected_jobs": selection_identity})
|
||||
).hexdigest()
|
||||
changed_while_reading = (
|
||||
identity_before != identity_after or wal_before != wal_after
|
||||
)
|
||||
status_counts = Counter(str(row["status"]) for row in live_rows)
|
||||
manifest: dict[str, Any] = {
|
||||
"as_of": read_started_at,
|
||||
"excluded_live_jobs": excluded_live_jobs,
|
||||
"format": "quant-os.tushare-scoped-snapshot/v2",
|
||||
"hash_algorithms": {
|
||||
"canonical_json": (
|
||||
"UTF-8 JSON; keys sorted; separators ',' and ':'; "
|
||||
"ensure_ascii=false; allow_nan=false"
|
||||
),
|
||||
"manifest_sha256": (
|
||||
"SHA-256 of canonical_json(the complete manifest excluding "
|
||||
"the manifest_sha256 member)"
|
||||
),
|
||||
"request_sha256": (
|
||||
"SHA-256 of canonical_json({api_name, partition_key, fields, "
|
||||
"params}), with params_json parsed as JSON and without "
|
||||
"embedding it in this manifest"
|
||||
),
|
||||
"selection_sha256": (
|
||||
"SHA-256 of canonical_json({selected_jobs: [...]}), excluding "
|
||||
"each job's observed and verified members so verify/no-verify "
|
||||
"runs over the same SQLite selection have the same identity"
|
||||
),
|
||||
},
|
||||
"mirror": {
|
||||
"label": root.name,
|
||||
"state_db": {
|
||||
"changed_while_reading": changed_while_reading,
|
||||
"changed_through_file_verification": (
|
||||
identity_before != identity_after_verification
|
||||
or wal_before != wal_after_verification
|
||||
),
|
||||
"file_after": identity_after,
|
||||
"file_after_verification": identity_after_verification,
|
||||
"file_before": identity_before,
|
||||
"read_finished_at": read_finished_at,
|
||||
"read_started_at": read_started_at,
|
||||
"sqlite_snapshot": sqlite_identity,
|
||||
"wal_after": wal_after,
|
||||
"wal_after_verification": wal_after_verification,
|
||||
"wal_before": wal_before,
|
||||
},
|
||||
},
|
||||
"scope": {
|
||||
"apis": list(requested_apis),
|
||||
"end": end_date.isoformat(),
|
||||
"index_codes": list(requested_codes),
|
||||
"partition_policy": (
|
||||
"calendar partitions overlapping [start,end]; timeless "
|
||||
"stock_basic partitions; exact requested index codes; "
|
||||
"index_weight also includes start.year-1 as the pre-start "
|
||||
"point-in-time membership anchor"
|
||||
),
|
||||
"index_weight_anchor_year": (
|
||||
start_date.year - 1
|
||||
if "index_weight" in requested_apis
|
||||
else None
|
||||
),
|
||||
"start": start_date.isoformat(),
|
||||
},
|
||||
"selected_jobs": selected_jobs,
|
||||
"selection_sha256": selection_sha,
|
||||
"summary": {
|
||||
"excluded_live_job_status_counts": dict(sorted(status_counts.items())),
|
||||
"excluded_live_jobs": len(excluded_live_jobs),
|
||||
"selected_bytes": sum(job["byte_count"] for job in selected_jobs),
|
||||
"selected_jobs": len(selected_jobs),
|
||||
"selected_rows": sum(job["row_count"] for job in selected_jobs),
|
||||
"shadowed_completed_jobs": sum(
|
||||
job["shadowed_completed_count"] for job in selected_jobs
|
||||
),
|
||||
"tables": _build_table_summaries(selected_jobs),
|
||||
},
|
||||
"tool": {
|
||||
"name": "tools/tushare_snapshot.py",
|
||||
"version": TOOL_VERSION,
|
||||
},
|
||||
"verification": {
|
||||
"mode": "size+sha256+parquet_row_count" if verify else "not_performed",
|
||||
"verified_jobs": len(selected_jobs) if verify else 0,
|
||||
"double_verified_jobs": len(selected_jobs) if verify else 0,
|
||||
"selected_jobs_still_latest_after_verification": True,
|
||||
},
|
||||
}
|
||||
manifest["manifest_sha256"] = hashlib.sha256(
|
||||
_canonical_json(manifest)
|
||||
).hexdigest()
|
||||
return manifest
|
||||
|
||||
|
||||
def write_manifest(path: str | Path, manifest: Mapping[str, Any]) -> Path:
|
||||
destination = Path(path).expanduser().resolve()
|
||||
if destination.exists() and destination.is_dir():
|
||||
raise ValueError(f"output points to a directory: {destination}")
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
payload = (
|
||||
json.dumps(
|
||||
manifest,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
allow_nan=False,
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
descriptor, temporary_name = tempfile.mkstemp(
|
||||
prefix=f".{destination.name}.",
|
||||
suffix=".tmp",
|
||||
dir=destination.parent,
|
||||
)
|
||||
temporary = Path(temporary_name)
|
||||
try:
|
||||
with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
|
||||
handle.write(payload)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(temporary, destination)
|
||||
finally:
|
||||
if temporary.exists():
|
||||
temporary.unlink()
|
||||
return destination
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--mirror-root", required=True)
|
||||
parser.add_argument("--start", required=True)
|
||||
parser.add_argument("--end", required=True)
|
||||
parser.add_argument(
|
||||
"--apis",
|
||||
nargs="+",
|
||||
default=list(DEFAULT_APIS),
|
||||
help="space- or comma-separated API names",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--index-codes",
|
||||
nargs="+",
|
||||
default=list(DEFAULT_INDEX_CODES),
|
||||
help="space- or comma-separated Tushare index codes",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verify",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="verify file size, SHA-256 and Parquet row count (default: true)",
|
||||
)
|
||||
parser.add_argument("--output", "--output-json", dest="output", required=True)
|
||||
return parser
|
||||
|
||||
|
||||
def _main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
try:
|
||||
manifest = create_snapshot(
|
||||
args.mirror_root,
|
||||
start=args.start,
|
||||
end=args.end,
|
||||
apis=_parse_values(args.apis),
|
||||
index_codes=_parse_values(args.index_codes),
|
||||
verify=args.verify,
|
||||
)
|
||||
output = write_manifest(args.output, manifest)
|
||||
except (OSError, ValueError, sqlite3.Error, TushareSnapshotError) as exc:
|
||||
parser.error(str(exc))
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"manifest_sha256": manifest["manifest_sha256"],
|
||||
"output": str(output),
|
||||
"selected_jobs": manifest["summary"]["selected_jobs"],
|
||||
"verified_jobs": manifest["verification"]["verified_jobs"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(_main())
|
||||
Reference in New Issue
Block a user