282 lines
9.6 KiB
Python
282 lines
9.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Verify exact lineage from a Tushare scoped manifest to a Qlib provider."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
from typing import Any, Mapping, Optional, Sequence
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from adapters.tushare_local import MANIFEST_NAME, verify_qlib_provider
|
|
|
|
|
|
MATCHED_FIELDS = (
|
|
("id", "job_id"),
|
|
("path", "path"),
|
|
("row_count", "row_count"),
|
|
("byte_count", "byte_count"),
|
|
("sha256", "sha256"),
|
|
("request_sha256", "request_sha256"),
|
|
)
|
|
|
|
|
|
class TushareLineageError(RuntimeError):
|
|
"""Raised when a scoped manifest and provider do not match exactly."""
|
|
|
|
|
|
def _canonical_json(value: Mapping[str, Any]) -> bytes:
|
|
return json.dumps(
|
|
dict(value),
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
allow_nan=False,
|
|
).encode("utf-8")
|
|
|
|
|
|
def _load_object(path: Path, *, label: str) -> dict[str, Any]:
|
|
try:
|
|
value = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise TushareLineageError(f"cannot read {label}: {path}") from exc
|
|
if not isinstance(value, dict):
|
|
raise TushareLineageError(f"{label} must be a JSON object")
|
|
return value
|
|
|
|
|
|
def _validate_snapshot(
|
|
snapshot: Mapping[str, Any],
|
|
) -> Sequence[Mapping[str, Any]]:
|
|
if snapshot.get("format") != "quant-os.tushare-scoped-snapshot/v2":
|
|
raise TushareLineageError("source snapshot must use scoped format v2")
|
|
manifest_sha = str(snapshot.get("manifest_sha256") or "")
|
|
without_manifest_sha = dict(snapshot)
|
|
without_manifest_sha.pop("manifest_sha256", None)
|
|
computed_manifest_sha = hashlib.sha256(
|
|
_canonical_json(without_manifest_sha)
|
|
).hexdigest()
|
|
if manifest_sha != computed_manifest_sha:
|
|
raise TushareLineageError("source snapshot manifest SHA-256 mismatch")
|
|
|
|
jobs = snapshot.get("selected_jobs")
|
|
if not isinstance(jobs, list) or not jobs:
|
|
raise TushareLineageError("source snapshot has no selected_jobs")
|
|
selection_identity = []
|
|
seen: set[tuple[str, str]] = set()
|
|
for raw_job in jobs:
|
|
if not isinstance(raw_job, Mapping):
|
|
raise TushareLineageError("source snapshot job must be an object")
|
|
key = (
|
|
str(raw_job.get("api_name") or ""),
|
|
str(raw_job.get("partition_key") or ""),
|
|
)
|
|
if not all(key) or key in seen:
|
|
raise TushareLineageError(
|
|
"source snapshot has empty or duplicate job keys"
|
|
)
|
|
seen.add(key)
|
|
path = Path(str(raw_job.get("path") or ""))
|
|
if path.is_absolute() or ".." in path.parts:
|
|
raise TushareLineageError(
|
|
"source snapshot job paths must stay relative to the mirror"
|
|
)
|
|
observed = raw_job.get("observed")
|
|
if raw_job.get("verified") is not True or not isinstance(
|
|
observed,
|
|
Mapping,
|
|
):
|
|
raise TushareLineageError(
|
|
"every source snapshot job must carry verified observations"
|
|
)
|
|
if (
|
|
observed.get("row_count") != raw_job.get("row_count")
|
|
or observed.get("size_bytes") != raw_job.get("byte_count")
|
|
or observed.get("sha256") != raw_job.get("sha256")
|
|
):
|
|
raise TushareLineageError(
|
|
"source snapshot observed file evidence is inconsistent"
|
|
)
|
|
selection_identity.append(
|
|
{
|
|
field: value
|
|
for field, value in raw_job.items()
|
|
if field not in {"observed", "verified"}
|
|
}
|
|
)
|
|
computed_selection_sha = hashlib.sha256(
|
|
_canonical_json({"selected_jobs": selection_identity})
|
|
).hexdigest()
|
|
if computed_selection_sha != snapshot.get("selection_sha256"):
|
|
raise TushareLineageError("source snapshot selection SHA-256 mismatch")
|
|
|
|
verification = snapshot.get("verification")
|
|
if (
|
|
not isinstance(verification, Mapping)
|
|
or verification.get("verified_jobs") != len(jobs)
|
|
or verification.get("double_verified_jobs") != len(jobs)
|
|
or verification.get(
|
|
"selected_jobs_still_latest_after_verification"
|
|
)
|
|
is not True
|
|
):
|
|
raise TushareLineageError(
|
|
"source snapshot lacks closed double-verification evidence"
|
|
)
|
|
return jobs
|
|
|
|
|
|
def _compare_source_jobs(
|
|
snapshot_jobs: Sequence[Mapping[str, Any]],
|
|
provider_jobs: Sequence[Mapping[str, Any]],
|
|
) -> None:
|
|
def key(job: Mapping[str, Any]) -> tuple[str, str]:
|
|
return (
|
|
str(job.get("api_name") or ""),
|
|
str(job.get("partition_key") or ""),
|
|
)
|
|
|
|
snapshot_by_key = {key(job): job for job in snapshot_jobs}
|
|
provider_by_key = {key(job): job for job in provider_jobs}
|
|
if len(snapshot_by_key) != len(snapshot_jobs):
|
|
raise TushareLineageError("source snapshot has duplicate job keys")
|
|
if len(provider_by_key) != len(provider_jobs):
|
|
raise TushareLineageError("provider manifest has duplicate job keys")
|
|
if set(snapshot_by_key) != set(provider_by_key):
|
|
missing = sorted(set(snapshot_by_key).difference(provider_by_key))
|
|
extra = sorted(set(provider_by_key).difference(snapshot_by_key))
|
|
raise TushareLineageError(
|
|
"source job sets differ: "
|
|
f"missing_from_provider={missing[:5]}, "
|
|
f"extra_in_provider={extra[:5]}"
|
|
)
|
|
mismatches: list[str] = []
|
|
for job_key in sorted(snapshot_by_key):
|
|
source = snapshot_by_key[job_key]
|
|
provider = provider_by_key[job_key]
|
|
for source_field, provider_field in MATCHED_FIELDS:
|
|
if source.get(source_field) != provider.get(provider_field):
|
|
mismatches.append(
|
|
f"{job_key[0]}:{job_key[1]}:{source_field}"
|
|
)
|
|
if mismatches:
|
|
raise TushareLineageError(
|
|
"source job fields differ: " + ", ".join(mismatches[:12])
|
|
)
|
|
|
|
|
|
def verify_lineage(
|
|
source_manifest: str | Path,
|
|
provider_dir: str | Path,
|
|
) -> dict[str, Any]:
|
|
source_path = Path(source_manifest).expanduser().resolve()
|
|
provider = Path(provider_dir).expanduser().resolve()
|
|
snapshot = _load_object(source_path, label="source snapshot")
|
|
snapshot_jobs = _validate_snapshot(snapshot)
|
|
provider_verification = verify_qlib_provider(provider)
|
|
provider_manifest = _load_object(
|
|
provider / MANIFEST_NAME,
|
|
label="provider manifest",
|
|
)
|
|
provider_jobs = provider_manifest.get("source_jobs")
|
|
if not isinstance(provider_jobs, list):
|
|
raise TushareLineageError("provider manifest has no source_jobs")
|
|
_compare_source_jobs(snapshot_jobs, provider_jobs)
|
|
return {
|
|
"artifact_type": "quant_os_tushare_source_provider_lineage",
|
|
"ok": True,
|
|
"snapshot_format": snapshot["format"],
|
|
"snapshot_manifest_sha256": snapshot["manifest_sha256"],
|
|
"snapshot_selection_sha256": snapshot["selection_sha256"],
|
|
"source_job_count": len(snapshot_jobs),
|
|
"matched_fields": [
|
|
source_field for source_field, _ in MATCHED_FIELDS
|
|
],
|
|
"mismatch_count": 0,
|
|
"provider": provider_verification,
|
|
"tool_source_sha256": hashlib.sha256(
|
|
Path(__file__).resolve().read_bytes()
|
|
).hexdigest(),
|
|
"gate_credit": [],
|
|
"investment_value_claim": False,
|
|
}
|
|
|
|
|
|
def _write_json(path: Path, payload: Mapping[str, Any]) -> None:
|
|
destination = path.expanduser().resolve()
|
|
if destination.exists() and destination.is_dir():
|
|
raise ValueError("output points to a directory")
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
provider_bytes = (
|
|
json.dumps(
|
|
dict(payload),
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
sort_keys=True,
|
|
allow_nan=False,
|
|
)
|
|
+ "\n"
|
|
).encode("utf-8")
|
|
descriptor, temporary_name = tempfile.mkstemp(
|
|
prefix=f".{destination.name}.",
|
|
suffix=".tmp",
|
|
dir=destination.parent,
|
|
)
|
|
temporary = Path(temporary_name)
|
|
try:
|
|
with os.fdopen(descriptor, "wb") as handle:
|
|
handle.write(provider_bytes)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary, destination)
|
|
finally:
|
|
if temporary.exists():
|
|
temporary.unlink()
|
|
|
|
|
|
def _main(argv: Optional[list[str]] = None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source-manifest", required=True)
|
|
parser.add_argument("--provider-dir", required=True)
|
|
parser.add_argument("--output-json")
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
result = verify_lineage(
|
|
args.source_manifest,
|
|
args.provider_dir,
|
|
)
|
|
if args.output_json:
|
|
output = Path(args.output_json).expanduser().resolve()
|
|
provider = Path(args.provider_dir).expanduser().resolve()
|
|
if output == provider or provider in output.parents:
|
|
raise ValueError(
|
|
"output-json must be outside the immutable provider "
|
|
"directory"
|
|
)
|
|
_write_json(output, result)
|
|
except (OSError, ValueError, TushareLineageError) as exc:
|
|
parser.error(str(exc))
|
|
print(
|
|
json.dumps(
|
|
result,
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
sort_keys=True,
|
|
allow_nan=False,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(_main())
|