diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e2e8dd..7297f3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ ## 0.1.1 +* **Plugins can read job identity and scratch paths per request.** `current_invocation_identity()` + returns an `InvocationIdentity` exclusively from the bound invocation context; absent fields + are `None`. It falls back to the environment only when no context is bound. `invocation_scratch_dir(root)` returns + `root///` while a context is bound and `root` unchanged otherwise. + * **A provider rate limit no longer reaches the wire as terminal.** `plugin_error_of` declared `retryable = false` for the whole legacy `UserError` family, `RateLimitError` included, so a 429 a consumer reads off the declared field fails the record on the first throttle instead of backing diff --git a/test/api/test_invocation_scope.py b/test/api/test_invocation_scope.py new file mode 100644 index 0000000..b3e857d --- /dev/null +++ b/test/api/test_invocation_scope.py @@ -0,0 +1,133 @@ +import asyncio +from dataclasses import FrozenInstanceError, dataclass, fields +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from unstructured_platform_plugins.etl_uvicorn.api_generator import generate_fast_api +from unstructured_platform_plugins.invocation_context import InvocationContext +from unstructured_platform_plugins.invocation_scope import ( + InvocationIdentity, + current_invocation_identity, + invocation_scratch_dir, +) +from unstructured_platform_plugins.invocation_settings import invocation_envelope + + +@pytest.mark.parametrize("bound", [False, True]) +def test_identity_context_is_authoritative(monkeypatch, bound): + for field in fields(InvocationIdentity): + monkeypatch.setenv(field.name.upper(), f"env-{field.name}") + context = InvocationContext(job_id="request", tenant_id="") if bound else None + with invocation_envelope(None, context): + identity = current_invocation_identity() + assert identity.job_id == ("request" if bound else "env-job_id") + assert identity.tenant_id == ("" if bound else "env-tenant_id") + for field in fields(InvocationIdentity): + if field.name not in {"job_id", "tenant_id"}: + assert getattr(identity, field.name) == (None if bound else f"env-{field.name}") + with pytest.raises(FrozenInstanceError): + identity.job_id = "other" + + +def test_absent_identity(monkeypatch): + for field in fields(InvocationIdentity): + monkeypatch.delenv(field.name.upper(), raising=False) + assert all( + getattr(current_invocation_identity(), f.name) is None for f in fields(InvocationIdentity) + ) + + +def test_boot_path_and_creation(tmp_path): + root = tmp_path / "scratch" + assert invocation_scratch_dir(root, create=False) == root + assert not root.exists() + assert invocation_scratch_dir(str(root)) == root + assert root.is_dir() + + +@pytest.mark.parametrize("segment", ["../a", "..", ".", "a/b", "a\\b", "x" * 300, "☃"]) +def test_segments_are_safe_and_bounded(tmp_path, segment): + context = InvocationContext(job_id=segment, dag_node_id=segment, invocation_id=segment) + with invocation_envelope(None, context): + path = invocation_scratch_dir(tmp_path) + assert path.is_relative_to(tmp_path) + assert len(path.relative_to(tmp_path).parts) == 3 + assert all(p not in {".", ".."} and len(p) <= 96 for p in path.relative_to(tmp_path).parts) + assert path.is_dir() + + +def test_sanitization_does_not_alias(tmp_path): + paths = [] + for job_id in [None, "unknown", "~missing", "a/b", "a_b"]: + with invocation_envelope(None, InvocationContext(job_id=job_id, invocation_id="invoke")): + paths.append(invocation_scratch_dir(tmp_path, create=False)) + encoded_job_id = paths[3].relative_to(tmp_path).parts[0] + with invocation_envelope( + None, InvocationContext(job_id=encoded_job_id, invocation_id="invoke") + ): + paths.append(invocation_scratch_dir(tmp_path, create=False)) + assert len(set(paths)) == len(paths) + + +def test_single_record_and_batch_correlation(tmp_path): + with invocation_envelope( + None, + InvocationContext( + job_id="job", dag_node_id="node", invocation_id="invoke", record_id="record" + ), + ): + assert invocation_scratch_dir(tmp_path, create=False) == tmp_path / "job/node/invoke" + with invocation_envelope(None, InvocationContext(record_id="record")): + assert invocation_scratch_dir(tmp_path, create=False).name == "record" + with invocation_envelope(None, InvocationContext(invocation_id="", record_id="record")): + assert invocation_scratch_dir(tmp_path, create=False).name == "record" + with invocation_envelope(None, InvocationContext(record_id="record", record_ids=["record"])): + first = invocation_scratch_dir(tmp_path, create=False) + second = invocation_scratch_dir(tmp_path, create=False) + assert first != second + assert first.name != "record" + + +def test_concurrent_invocations_are_isolated(tmp_path): + async def invoke(job): + with invocation_envelope( + None, InvocationContext(job_id=job, dag_node_id="node", invocation_id="invoke") + ): + await asyncio.sleep(0) + return current_invocation_identity().job_id, invocation_scratch_dir(tmp_path) + + async def run(): + return await asyncio.gather(invoke("one"), invoke("two")) + + results = asyncio.run(run()) + assert results == [("one", tmp_path / "one/node/invoke"), ("two", tmp_path / "two/node/invoke")] + + +@dataclass +class ScratchOutput: + path: str + + +def test_streaming_retains_invocation_scope(tmp_path): + async def stream() -> ScratchOutput: + await asyncio.sleep(0) + yield ScratchOutput( + str(invocation_scratch_dir(tmp_path) / current_invocation_identity().job_id) + ) + + app = generate_fast_api(stream) + with TestClient(app) as client: + response = client.post( + "/invoke", + json={ + "invocation_context": { + "job_id": "job", + "dag_node_id": "node", + "invocation_id": "invoke", + } + }, + ) + assert response.status_code == 200 + assert str(Path(tmp_path) / "job/node/invoke/job") in response.text diff --git a/unstructured_platform_plugins/invocation_scope.py b/unstructured_platform_plugins/invocation_scope.py new file mode 100644 index 0000000..0ee7bce --- /dev/null +++ b/unstructured_platform_plugins/invocation_scope.py @@ -0,0 +1,63 @@ +"""Identity and filesystem scope for the current invocation.""" + +from __future__ import annotations + +import hashlib +import os +import re +from dataclasses import dataclass, fields +from pathlib import Path +from uuid import uuid4 + +from unstructured_platform_plugins.invocation_settings import current_invocation_context + + +@dataclass(frozen=True) +class InvocationIdentity: + job_id: str | None + tenant_id: str | None + org_id: str | None + workflow_id: str | None + attribution_id: str | None + dag_node_id: str | None + dag_node_type: str | None + dag_node_subtype: str | None + + +def current_invocation_identity() -> InvocationIdentity: + """Read identity from the bound context, or the environment when no context is bound.""" + context = current_invocation_context() + values = {} + for field in fields(InvocationIdentity): + values[field.name] = ( + getattr(context, field.name, None) + if context is not None + else os.getenv(field.name.upper()) + ) + return InvocationIdentity(**values) + + +def _segment(value: str | None) -> str: + if not value: + return "~missing" + safe = re.sub(r"[^a-zA-Z0-9_.-]", "_", value) + if safe == value and safe not in {".", ".."} and len(safe) <= 96: + return safe + digest = hashlib.sha256(value.encode()).hexdigest() + return f"~{safe[:30].strip('.') or 'unknown'}-{digest}" + + +def invocation_scratch_dir(root: str | os.PathLike[str], *, create: bool = True) -> Path: + """Return an invocation directory, or the unchanged root when no context is bound.""" + path = Path(root) + context = current_invocation_context() + if context is not None: + identity = current_invocation_identity() + invocation_id = context.invocation_id + if not invocation_id and context.record_ids is None and context.invocation_ids is None: + invocation_id = context.record_id + path = path / _segment(identity.job_id) / _segment(identity.dag_node_id) + path = path / _segment(invocation_id or uuid4().hex) + if create: + path.mkdir(parents=True, exist_ok=True) + return path