Show durable reservation acquisition progress in download tasks

integrations/ohip/arr_data.py

@@ -19,10 +19,11 @@ import stat
 import time
 
 from . import collect_arr_source as base
+from . import acquisition_progress as progress
 from . import profile_summary as profiles
 from . import rate_info, room_calendar_evidence as calendar, source_fields as fields
 from .audit_arr_capture import protected_read
-from .capture_job import atomic_json, job_lock, private_directory, sync_directory
+from .capture_job import atomic_json, fingerprint, job_lock, private_directory, sync_directory
 from .data_client import DataReader, document, json_bytes, typed_id
 
 
@@ -379,7 +380,7 @@ def _record(search, detail, reader, sequence):
             "fields": values, "related": related}
 
 
-def collect(reader):
+def collect(reader, report_progress=None):
     """Collect all source rows; return no Finance/business success claims."""
     options = reader.options
     rows, fatal = [], None
@@ -387,8 +388,18 @@ def collect(reader):
                                           "service_url": base.SERVICE, "application_id": base.APPLICATION,
                                           "source_kind": reader.source_kind, "max_requests": reader.max_requests})
     complete = False
+    def report(phase, completed=0, total=None):
+        if report_progress is not None:
+            try:
+                report_progress(phase, completed, total)
+            except Exception:
+                pass
+    total = None
+    report("searching")
     try:
         searches = base.search_day(reader, options, server_sort=False)
+        total = len(searches)
+        report("fetching", 0, total)
         search_refs = list(reader.used)
         for index, search in enumerate(searches, 1):
             before = len(reader.used)
@@ -415,6 +426,8 @@ def collect(reader):
             row = _record(search, detail, reader, index)
             row["sources"] = list(dict.fromkeys(search_refs + reader.used[before:]))
             rows.append(row)
+            report("fetching", len(rows), total)
+        report("verifying", len(rows), total)
         pending = [r for r in rows if r["fields"]["DISP_ROOM_NO"]["state"] not in {"available", "empty"}]
         if pending:
             before = len(reader.used)
@@ -443,6 +456,8 @@ def collect(reader):
         f["state"] == "available" or (key in OPTIONAL_FIELDS and f["state"] == "empty")
         for row in rows for key, f in row["fields"].items())
     failed = fatal is not None or states["failed"] > 0
+    if failed:
+        report("failed", len(rows), total)
     status = "failed" if failed else "collected" if ready else "collected_with_gaps"
     payload = {"version": VERSION, "hotel_id": options.hotel_id, "report_date": options.arrival_date,
                "source_kind": reader.source_kind,
@@ -481,13 +496,74 @@ class ARRDataSource:
         self.page_size, self.max_pages, self.max_records = page_size, max_pages, max_records
         self.max_requests, self.sleep = max_requests, sleep
 
-    def fetch(self, report_date: str, request_id: str) -> dict:
+    def _identity(self, report_date, request_id):
         options = base.Options(report_date, self.hotel_id, self.page_size, self.max_pages, self.max_records)
         options.validate()
         require(type(request_id) is str and re.fullmatch(r"[0-9a-f]{32}", request_id), "invalid_data_request_id")
         identity = {"version": VERSION, "options": vars(options), "max_requests": self.max_requests,
                     "request_id": request_id, "service_url": base.SERVICE, "application_id": base.APPLICATION,
                     "source_kind": "ohip_platform" if self.transport_factory is None else "test_transport"}
+        return options, identity
+
+    def get_acquisition_progress(self, *, report_date: str, request_id: str):
+        """Read local observations/legacy completion metadata without an execution lock or HTTP."""
+        try:
+            _, identity = self._identity(report_date, request_id)
+            folder = self.root / request_id
+            self._safe_directory(self.root)
+            self._safe_directory(folder)
+            stored = document(protected_read(folder / "request.json", 65536))
+            require(fingerprint(stored) == fingerprint(identity), "data_progress_context_mismatch")
+            pointer = folder / "ready.json"
+            if pointer.exists():
+                return self._completed_progress(folder, document(protected_read(pointer, 65536)), identity)
+            stored_progress = document(protected_read(folder / "progress.json", 4096))
+            require(set(stored_progress) == {"version", "identity_sha256", "progress"}
+                    and stored_progress["version"] == progress.VERSION
+                    and stored_progress["identity_sha256"] == fingerprint(identity), "data_progress_context_mismatch")
+            observation = progress.validate(stored_progress["progress"])
+            # Only a published completion checkpoint establishes completion.
+            return observation if observation["phase"] != "completed" else None
+        except Exception:
+            return None
+
+    @staticmethod
+    def _safe_directory(path):
+        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_data_directory")
+
+    def _completed_progress(self, folder, pointer, identity):
+        require(set(pointer) == {"attempt", "manifest_sha256", "data_sha256"}
+                and type(pointer["attempt"]) is str and re.fullmatch(r"attempt-[0-9]{4}", pointer["attempt"]),
+                "invalid_data_checkpoint")
+        directory = folder / pointer["attempt"]
+        self._safe_directory(directory)
+        raw = protected_read(directory / "result.json", base.MAX_MANIFEST_BYTES)
+        require(hashlib.sha256(raw).hexdigest() == pointer["manifest_sha256"], "data_manifest_changed")
+        result = document(raw)
+        require(result.get("version") == VERSION and result.get("status") in {"collected", "collected_with_gaps"}
+                and result.get("collection_complete") is True
+                and result.get("source_kind") == identity["source_kind"]
+                and result.get("data_sha256") == pointer["data_sha256"], "invalid_completed_data")
+        capture_raw = protected_read(directory / "capture.json", 65536)
+        capture = document(capture_raw)
+        require(set(capture) == {"version", "options", "service_url", "application_id", "max_requests", "source_kind"}
+                and json_bytes(capture) == json_bytes({k: identity[k] for k in capture}), "data_capture_context_mismatch")
+        files = result.get("files")
+        require(type(files) is list and sum(item.get("name") == "capture.json" for item in files if type(item) is dict) == 1,
+                "invalid_data_inventory")
+        entry = next(item for item in files if type(item) is dict and item.get("name") == "capture.json")
+        require(entry.get("bytes") == len(capture_raw) and entry.get("sha256") == hashlib.sha256(capture_raw).hexdigest(),
+                "data_capture_changed")
+        count = result.get("records")
+        require(type(count) is int and 0 <= count <= self.max_records, "invalid_data_count")
+        stamp = directory.joinpath("result.json").stat().st_mtime
+        from datetime import datetime, timezone
+        return progress.snapshot("completed", count, count, updated_at=datetime.fromtimestamp(stamp, timezone.utc).isoformat())
+
+    def fetch(self, report_date: str, request_id: str) -> dict:
+        options, identity = self._identity(report_date, request_id)
         private_directory(self.root)
         folder = self.root / request_id
         private_directory(folder)
@@ -504,21 +580,31 @@ class ARRDataSource:
             attempts = [p for p in folder.iterdir() if re.fullmatch(r"attempt-[0-9]{4}", p.name)]
             attempt = max([int(p.name[-4:]) for p in attempts], default=0) + 1
             require(attempt <= 9999, "attempt_limit_exceeded")
-            key = ""
-            if self.transport_factory is None:
-                key = base.load_key(Path(self.credential_file))
-                transport = base.HTTPTransport(key)
-            else:
-                transport = self.transport_factory()
-            archive = base.Archive(folder / f"attempt-{attempt:04d}")
-            reader = DataReader(archive, options, transport, key=key, sleep=self.sleep, max_requests=self.max_requests,
-                                source_kind=identity["source_kind"])
-            summary = collect(reader)
-            summary.update(data_path=str(archive.path / "arr-data.json"), attempt=attempt)
-            if summary["status"] != "failed":
-                atomic_json(pointer, {"attempt": archive.path.name, "manifest_sha256": summary["manifest_sha256"],
-                                      "data_sha256": summary["data_sha256"]}, replace=False)
-            return summary
+            progress.publish(folder, identity, "searching")
+            latest_counts = [0, None]
+            def report_progress(phase, completed, total):
+                latest_counts[:] = [completed, total]
+                progress.publish(folder, identity, phase, completed, total)
+            try:
+                key = ""
+                if self.transport_factory is None:
+                    key = base.load_key(Path(self.credential_file))
+                    transport = base.HTTPTransport(key)
+                else:
+                    transport = self.transport_factory()
+                archive = base.Archive(folder / f"attempt-{attempt:04d}")
+                reader = DataReader(archive, options, transport, key=key, sleep=self.sleep, max_requests=self.max_requests,
+                                    source_kind=identity["source_kind"])
+                summary = collect(reader, report_progress)
+                summary.update(data_path=str(archive.path / "arr-data.json"), attempt=attempt)
+                if summary["status"] != "failed":
+                    atomic_json(pointer, {"attempt": archive.path.name, "manifest_sha256": summary["manifest_sha256"],
+                                          "data_sha256": summary["data_sha256"]}, replace=False)
+                    progress.publish(folder, identity, "completed", summary["records"], summary["records"])
+                return summary
+            except Exception:
+                progress.publish(folder, identity, "failed", *latest_counts)
+                raise
 
     @staticmethod
     def _replay(folder, pointer, identity):