244 lines
12 KiB
Python
244 lines
12 KiB
Python
"""Local, resumable acquisition jobs; defaults to the original v1 protocol.
|
|
|
|
An explicit batch ID identifies one acquisition intent. A completed batch reuses
|
|
its verified capture; a new batch ID requests fresh data. No Finance submission.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from contextlib import contextmanager
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import stat
|
|
import sys
|
|
import tempfile
|
|
|
|
if __package__:
|
|
from . import collect_arr_source as source
|
|
from .audit_arr_capture import VerifiedArchive, protected_read, replay
|
|
else:
|
|
import collect_arr_source as source
|
|
from audit_arr_capture import VerifiedArchive, protected_read, replay
|
|
|
|
|
|
require = source.require
|
|
VERSION = "arr-api-capture-job/v1"
|
|
MAX_ATTEMPTS = 10
|
|
|
|
|
|
def private_directory(path: Path) -> None:
|
|
try:
|
|
path.mkdir(mode=0o700)
|
|
except FileExistsError:
|
|
pass
|
|
info = path.lstat()
|
|
require(stat.S_ISDIR(info.st_mode) and info.st_uid == os.getuid()
|
|
and stat.S_IMODE(info.st_mode) == 0o700, "unsafe_job_directory")
|
|
|
|
|
|
def sync_directory(path: Path) -> None:
|
|
fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
|
try:
|
|
os.fsync(fd)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def atomic_json(path: Path, document: dict, *, replace: bool) -> None:
|
|
"""Write private metadata atomically, then flush the directory entry."""
|
|
raw = source.json_bytes(document)
|
|
descriptor, temporary = tempfile.mkstemp(prefix=".metadata-", dir=path.parent)
|
|
temporary = Path(temporary)
|
|
try:
|
|
with os.fdopen(descriptor, "wb") as handle:
|
|
handle.write(raw)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
if replace:
|
|
os.replace(temporary, path)
|
|
else:
|
|
# Atomic exclusive publication: never replace an existing job identity.
|
|
os.link(temporary, path)
|
|
sync_directory(path.parent)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
|
|
|
|
@contextmanager
|
|
def job_lock(directory: Path):
|
|
fd = os.open(directory / ".lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK, 0o600)
|
|
try:
|
|
info = os.fstat(fd)
|
|
require(stat.S_ISREG(info.st_mode) and info.st_uid == os.getuid()
|
|
and stat.S_IMODE(info.st_mode) == 0o600, "unsafe_job_lock")
|
|
try:
|
|
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except BlockingIOError:
|
|
raise source.CollectionError("batch_busy") from None
|
|
yield
|
|
finally:
|
|
os.close(fd) # The OS releases the lock even when this process dies.
|
|
|
|
|
|
def identity(batch_id: str, options: source.Options) -> dict:
|
|
options.validate()
|
|
require(isinstance(batch_id, str) and bool(re.fullmatch(r"[a-z0-9][a-z0-9_-]{0,63}", batch_id)), "invalid_batch_id")
|
|
return {"version": VERSION, "batch_id": batch_id, "service_url": source.SERVICE,
|
|
"application_id": source.APPLICATION, "options": vars(options),
|
|
"source_contract": "arr-api-date-candidates/v1", "operations": [source.SEARCH, source.DETAIL],
|
|
"fetch_instructions": list(source.FETCH), "order_by": ["ConfirmationNo"], "sort_order": ["Asc"]}
|
|
|
|
|
|
def fingerprint(document: dict) -> str:
|
|
return hashlib.sha256(json.dumps(document, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()).hexdigest()
|
|
|
|
|
|
def read_state(directory: Path, expected_identity: dict) -> dict:
|
|
job_path, state_path = directory / "job.json", directory / "state.json"
|
|
digest = fingerprint(expected_identity)
|
|
if not os.path.lexists(job_path):
|
|
require(not os.path.lexists(state_path) and not list(directory.glob("attempt-*")), "orphan_job_state")
|
|
atomic_json(job_path, expected_identity, replace=False)
|
|
stored = source.strict_json(protected_read(job_path, 65536))
|
|
# Python equality aliases 2 with 2.0 (and True with 1). The persisted
|
|
# request must retain the exact JSON types used by its original digest.
|
|
require(fingerprint(stored) == digest, "batch_request_conflict")
|
|
if not os.path.lexists(state_path):
|
|
require(not list(directory.glob("attempt-*")), "missing_job_state")
|
|
state = {"version": expected_identity["version"], "request_sha256": digest, "attempts": []}
|
|
atomic_json(state_path, state, replace=False)
|
|
state = source.strict_json(protected_read(state_path, 65536))
|
|
require(state.get("version") == expected_identity["version"] and state.get("request_sha256") == digest,
|
|
"job_state_identity_mismatch")
|
|
attempts = state.get("attempts")
|
|
require(isinstance(attempts, list) and len(attempts) <= MAX_ATTEMPTS, "invalid_job_attempts")
|
|
for number, attempt in enumerate(attempts, 1):
|
|
require(isinstance(attempt, dict) and type(attempt.get("attempt_no")) is int and attempt["attempt_no"] == number,
|
|
"invalid_job_attempt_number")
|
|
require(attempt.get("status") in {"running", "failed", "interrupted", "complete"}, "invalid_job_attempt_status")
|
|
require(number == len(attempts) or attempt["status"] in {"failed", "interrupted"}, "invalid_job_attempt_order")
|
|
if attempt["status"] == "complete":
|
|
require(isinstance(attempt.get("manifest_sha256"), str)
|
|
and bool(re.fullmatch(r"[0-9a-f]{64}", attempt["manifest_sha256"])), "missing_completed_capture_pin")
|
|
known_directories = {f"attempt-{number:04d}" for number in range(1, len(attempts) + 1)}
|
|
require({path.name for path in directory.glob("attempt-*")} <= known_directories, "orphan_capture_attempt")
|
|
return state
|
|
|
|
|
|
def verify_completed(directory: Path, attempt: dict, options: source.Options) -> dict:
|
|
capture = directory / f"attempt-{attempt['attempt_no']:04d}"
|
|
verified = VerifiedArchive(capture, attempt["manifest_sha256"])
|
|
require(verified.options == options, "completed_capture_context_mismatch")
|
|
replay(verified)
|
|
return verified.result
|
|
|
|
|
|
def receipt(directory: Path, batch_id: str, attempt: dict, result: dict, *, reused: bool, version=VERSION) -> dict:
|
|
return {"version": version, "batch_id": batch_id, "status": "complete_candidate_capture",
|
|
"attempt_no": attempt["attempt_no"], "reused": reused,
|
|
"capture_dir": str(directory / f"attempt-{attempt['attempt_no']:04d}"),
|
|
"manifest_sha256": attempt["manifest_sha256"], "source_records": result["search_records"],
|
|
"verified_details": result["verified_details"], "candidate_capture_complete": True,
|
|
"finance_ready": False, "report_equivalence_verified": False,
|
|
**{key: result[key] for key in ("from_date", "to_date", "rate_date", "rate_responses",
|
|
"valid_rate_candidates", "all_rates_valid", "rate_issues", "profile_records", "unique_profiles",
|
|
"valid_name_candidates", "all_names_valid", "profile_http_attempts", "profile_issues") if key in result}}
|
|
|
|
|
|
def _run_locked(directory: Path, batch_id: str, options: source.Options, reader_factory, protocol=None) -> dict:
|
|
request = identity(batch_id, options) if protocol is None else protocol.identity(batch_id, options)
|
|
verify = verify_completed if protocol is None else protocol.verify_completed
|
|
state = read_state(directory, request)
|
|
attempts = state["attempts"]
|
|
if attempts and attempts[-1]["status"] == "complete":
|
|
result = verify(directory, attempts[-1], options)
|
|
return receipt(directory, batch_id, attempts[-1], result, reused=True, version=request["version"])
|
|
if attempts and attempts[-1]["status"] == "running":
|
|
# Acquiring the lock proves no cooperating local process still owns this batch.
|
|
# Even an uncommitted result.json is not promoted: its original pin is unknown.
|
|
attempts[-1].update(status="interrupted", ended_at=source.utc_now())
|
|
atomic_json(directory / "state.json", state, replace=True)
|
|
require(len(attempts) < MAX_ATTEMPTS, "batch_attempt_limit_reached")
|
|
attempt = {"attempt_no": len(attempts) + 1, "status": "running", "started_at": source.utc_now()}
|
|
attempts.append(attempt)
|
|
# Persist intent before creating a capture or calling the network.
|
|
atomic_json(directory / "state.json", state, replace=True)
|
|
capture = directory / f"attempt-{attempt['attempt_no']:04d}"
|
|
try:
|
|
archive = source.Archive(capture)
|
|
if protocol is None:
|
|
reader = reader_factory(archive, options.hotel_id)
|
|
result = source.collect(options, archive, reader)
|
|
else:
|
|
result = protocol.collect(options, archive, reader_factory)
|
|
if result.get("candidate_capture_complete") is not True:
|
|
code = result.get("error", "capture_failed")
|
|
require(isinstance(code, str) and bool(re.fullmatch(r"[a-z0-9_]{1,80}", code)), "invalid_capture_error")
|
|
raise source.CollectionError(code)
|
|
attempt["manifest_sha256"] = result["manifest_sha256"]
|
|
result = verify(directory, attempt, options)
|
|
sync_directory(capture)
|
|
except Exception as error:
|
|
code = str(error) if isinstance(error, source.CollectionError) else "capture_setup_or_storage_failure"
|
|
attempt.update(status="failed", ended_at=source.utc_now(), error=code)
|
|
atomic_json(directory / "state.json", state, replace=True)
|
|
return {"version": request["version"], "batch_id": batch_id, "status": "capture_failed", "attempt_no": attempt["attempt_no"],
|
|
"error": code, "candidate_capture_complete": False, "finance_ready": False}
|
|
attempt.update(status="complete", ended_at=source.utc_now())
|
|
atomic_json(directory / "state.json", state, replace=True)
|
|
return receipt(directory, batch_id, attempt, result, reused=False, version=request["version"])
|
|
|
|
|
|
def run_batch(root: Path, batch_id: str, options: source.Options, reader_factory, *, protocol=None) -> dict:
|
|
"""One explicit invocation starts at most one complete acquisition attempt.
|
|
|
|
reader_factory is invoked only when new reads are needed. It must return the
|
|
existing bounded Reader for v1; CLI credentials are loaded lazily in this
|
|
factory. The isolated v2 entrypoint supplies its capture/replay protocol.
|
|
"""
|
|
try:
|
|
if protocol is None:
|
|
identity(batch_id, options)
|
|
else:
|
|
protocol.identity(batch_id, options)
|
|
root = Path(root)
|
|
repository = Path(__file__).resolve().parents[2]
|
|
require(not root.resolve().is_relative_to(repository), "job_store_must_be_outside_repository")
|
|
private_directory(root)
|
|
directory = root / batch_id
|
|
private_directory(directory)
|
|
sync_directory(root)
|
|
with job_lock(directory):
|
|
return _run_locked(directory, batch_id, options, reader_factory, protocol)
|
|
except source.CollectionError as error:
|
|
return {"status": "capture_job_refused", "error": str(error), "candidate_capture_complete": False, "finance_ready": False}
|
|
except Exception:
|
|
return {"status": "capture_job_refused", "error": "job_state_or_storage_invalid",
|
|
"candidate_capture_complete": False, "finance_ready": False}
|
|
|
|
|
|
def main(argv=None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--job-store", required=True, type=Path, help="Private local directory outside the repository; parent must exist")
|
|
parser.add_argument("--batch-id", required=True, help="Stable lower-case ID for this acquisition intent; a new ID requests fresh data")
|
|
parser.add_argument("--arrival-date", required=True)
|
|
parser.add_argument("--hotel-id", required=True)
|
|
parser.add_argument("--credential-file", required=True, type=Path)
|
|
args = parser.parse_args(argv)
|
|
|
|
def reader_factory(archive, hotel_id):
|
|
key = source.load_key(args.credential_file)
|
|
return source.Reader(archive, hotel_id, source.HTTPTransport(key), key=key)
|
|
|
|
result = run_batch(args.job_store, args.batch_id, source.Options(args.arrival_date, args.hotel_id), reader_factory)
|
|
print(json.dumps(result, ensure_ascii=False, allow_nan=False))
|
|
return 0 if result.get("candidate_capture_complete") else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|