feat(pyamber): add workflow boundary codec

amber/requirements.txt

@@ -25,6 +25,7 @@ overrides==7.7.0
 typing_extensions==4.14.1
 bidict==0.22.0
 cached_property==2.0.1
+cloudpickle==3.1.2
 psutil==7.2.2
 tzlocal==2.1
 # Not imported directly: s3fs (with aiobotocore) backs pyiceberg's

amber/src/main/python/pytexera/workflow/codec.py

@@ -0,0 +1,158 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Cloudpickle transport for explicit workflow boundary payloads."""
+
+from __future__ import annotations
+
+import pickle
+from dataclasses import dataclass
+
+import cloudpickle
+
+
+@dataclass(frozen=True)
+class BoundaryPayload:
+    """One boundary contract, its path-present fields, and encoded values."""
+
+    boundary_id: str
+    fields: tuple[str, ...]
+    present: tuple[str, ...]
+    payload: bytes
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.boundary_id, str):
+            raise TypeError("boundary ID must be a string")
+        if not self.boundary_id:
+            raise ValueError("boundary ID must be nonempty")
+        if not isinstance(self.fields, tuple) or any(
+            not isinstance(field, str) for field in self.fields
+        ):
+            raise TypeError("boundary fields must be a tuple of strings")
+        if self.fields != tuple(sorted(set(self.fields))) or any(
+            not field.isidentifier() for field in self.fields
+        ):
+            raise ValueError("boundary fields must be canonical Python names")
+        if not isinstance(self.present, tuple) or any(
+            not isinstance(field, str) for field in self.present
+        ):
+            raise TypeError("present fields must be a tuple of strings")
+        present_fields = frozenset(self.present)
+        if self.present != tuple(
+            field for field in self.fields if field in present_fields
+        ):
+            raise ValueError("present fields must be a canonical contract subset")
+        if not isinstance(self.payload, bytes):
+            raise TypeError("boundary payload must be bytes")
+
+
+@dataclass(frozen=True)
+class WorkflowEnvelope:
+    """All selected boundary payloads for one independent execution key."""
+
+    execution_key: str
+    boundaries: tuple[BoundaryPayload, ...]
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.execution_key, str):
+            raise TypeError("execution key must be a string")
+        if not self.execution_key:
+            raise ValueError("execution key must be nonempty")
+        if not isinstance(self.boundaries, tuple) or any(
+            not isinstance(row, BoundaryPayload) for row in self.boundaries
+        ):
+            raise TypeError("workflow boundaries must be a typed tuple")
+        for row in self.boundaries:
+            row.__post_init__()
+        ids = tuple(row.boundary_id for row in self.boundaries)
+        if ids != tuple(sorted(set(ids))):
+            raise ValueError("workflow boundaries must be canonical and unique")
+
+
+def encode_boundary(
+    boundary_id: str,
+    fields: tuple[str, ...],
+    values: tuple[object, ...],
+    *,
+    present: tuple[str, ...] | None = None,
+) -> BoundaryPayload:
+    """Encode values present on this path under one selected field contract."""
+
+    if not isinstance(values, tuple):
+        raise TypeError("boundary values must be a tuple")
+    present = fields if present is None else present
+    if len(present) != len(values):
+        raise ValueError("present boundary fields and values must have equal length")
+    payload = cloudpickle.dumps(values, protocol=pickle.HIGHEST_PROTOCOL)
+    return BoundaryPayload(boundary_id, fields, present, payload)
+
+
+def decode_boundary(
+    boundary: BoundaryPayload,
+    fields: tuple[str, ...],
+) -> tuple[object, ...]:
+    """Decode under the exact contract; pickle/loading exceptions propagate unchanged.
+
+    Empty payloads raise EOFError; invalid opcodes raise pickle.UnpicklingError.
+    Other loading failures also propagate. Only load trusted workflow payloads.
+    """
+
+    boundary.__post_init__()
+    if boundary.fields != fields:
+        raise ValueError("boundary fields do not match the requested contract")
+    values = cloudpickle.loads(boundary.payload)
+    if not isinstance(values, tuple) or len(values) != len(boundary.present):
+        raise ValueError("decoded boundary payload has an invalid shape")
+    return values
+
+
+def dumps_envelope(envelope: WorkflowEnvelope) -> bytes:
+    """Encode the inert envelope after validating its complete shape."""
+
+    if not isinstance(envelope, WorkflowEnvelope):
+        raise TypeError("envelope codec requires WorkflowEnvelope")
+    envelope.__post_init__()
+    return cloudpickle.dumps(envelope, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+def loads_envelope(payload: bytes) -> WorkflowEnvelope:
+    """Decode trusted bytes and revalidate envelope and nested boundary metadata.
+
+    Pickle/loading exceptions propagate unchanged, as in decode_boundary.
+    Post-load validation is not a security boundary for untrusted pickle data.
+    """
+
+    envelope = cloudpickle.loads(payload)
+    if not isinstance(envelope, WorkflowEnvelope):
+        raise TypeError("decoded payload is not WorkflowEnvelope")
+    envelope.__post_init__()
+    return envelope
+
+
+def merge_envelopes(
+    left: WorkflowEnvelope,
+    right: WorkflowEnvelope,
+) -> WorkflowEnvelope:
+    """Merge independent fan-in payloads for the same execution key."""
+
+    if left.execution_key != right.execution_key:
+        raise ValueError("cannot merge envelopes with different execution keys")
+    boundaries = (*left.boundaries, *right.boundaries)
+    return WorkflowEnvelope(
+        left.execution_key,
+        tuple(sorted(boundaries, key=lambda row: row.boundary_id)),
+    )

amber/src/test/python/pytexera/workflow/test_codec.py

@@ -0,0 +1,196 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pickle
+
+import cloudpickle
+import pytest
+from pytexera.workflow.codec import (
+    BoundaryPayload,
+    WorkflowEnvelope,
+    decode_boundary,
+    dumps_envelope,
+    encode_boundary,
+    loads_envelope,
+    merge_envelopes,
+)
+
+
+def test_boundary_cloudpickle_preserves_aliases_and_cycles() -> None:
+    shared = []
+    shared.append(shared)
+
+    boundary = encode_boundary("edge", ("left", "right"), (shared, shared))
+    values = decode_boundary(boundary, ("left", "right"))
+
+    assert values[0] is values[1]
+    assert values[0][0] is values[0]
+
+
+def test_envelope_round_trip_contains_only_explicit_boundaries() -> None:
+    boundary = encode_boundary("edge", ("value",), ({"large": [1, 2, 3]},))
+    envelope = WorkflowEnvelope("run-1", (boundary,))
+
+    decoded = loads_envelope(dumps_envelope(envelope))
+
+    assert decoded == envelope
+    assert decoded.boundaries[0].fields == ("value",)
+
+
+def test_decode_rejects_field_contract_mismatch() -> None:
+    boundary = encode_boundary("edge", ("value",), (1,))
+
+    with pytest.raises(ValueError, match="fields"):
+        decode_boundary(boundary, ("other",))
+
+
+def test_boundary_cloudpickle_preserves_an_absent_selected_field() -> None:
+    """The wire contract and the values present on one path remain distinct."""
+
+    boundary = encode_boundary(
+        "edge",
+        ("left", "right"),
+        (41,),
+        present=("left",),
+    )
+
+    assert boundary.fields == ("left", "right")
+    assert boundary.present == ("left",)
+    assert decode_boundary(boundary, ("left", "right")) == (41,)
+
+
+@pytest.mark.parametrize("size", [0, 1, 128])
+@pytest.mark.parametrize("step", [1, 2])
+def test_present_validation_scans_present_fields_linearly(size, step):
+    visits = 0
+
+    class CountedPresent(tuple):
+        def __iter__(self):
+            nonlocal visits
+            for field in super().__iter__():
+                visits += 1
+                yield field
+
+    fields = tuple(f"field{index:03d}" for index in range(size))
+    present = CountedPresent(fields[::step])
+
+    boundary = BoundaryPayload("edge", fields, present, b"payload")
+
+    assert boundary.present == present
+    assert visits <= 2 * len(present)
+
+
+@pytest.mark.parametrize("present", [("雪", "a"), ("a", "a"), ("missing",)])
+def test_present_validation_rejects_noncanonical_subsets(present):
+    with pytest.raises(ValueError, match="canonical contract subset"):
+        BoundaryPayload("edge", ("a", "雪"), present, b"payload")
+
+
+def test_envelope_rejects_duplicate_boundaries_and_cross_key_merge() -> None:
+    boundary = BoundaryPayload("edge", ("value",), ("value",), b"payload")
+    with pytest.raises(ValueError, match="canonical"):
+        WorkflowEnvelope("run", (boundary, boundary))
+
+    with pytest.raises(ValueError, match="execution key"):
+        merge_envelopes(
+            WorkflowEnvelope("left", ()),
+            WorkflowEnvelope("right", ()),
+        )
+
+
+@pytest.mark.parametrize("identity", [None, 1, True, b"edge", []])
+def test_identities_require_strings(identity):
+    with pytest.raises(TypeError, match="string"):
+        BoundaryPayload(identity, (), (), b"")
+    with pytest.raises(TypeError, match="string"):
+        WorkflowEnvelope(identity, ())
+
+
+@pytest.mark.parametrize("identity", ["", "\u96ea"])
+def test_empty_and_unicode_identities(identity):
+    if not identity:
+        with pytest.raises(ValueError, match="nonempty"):
+            WorkflowEnvelope(identity, ())
+        with pytest.raises(ValueError, match="nonempty"):
+            BoundaryPayload(identity, (), (), b"")
+    else:
+        envelope = WorkflowEnvelope(identity, (encode_boundary(identity, (), ()),))
+        assert loads_envelope(dumps_envelope(envelope)) == envelope
+
+
+@pytest.mark.parametrize(
+    ("target", "field", "value", "error"),
+    [
+        ("envelope", "execution_key", 42, TypeError),
+        ("envelope", "execution_key", "", ValueError),
+        ("envelope", "boundaries", [], TypeError),
+        ("boundary", "boundary_id", 42, TypeError),
+        ("boundary", "fields", ["value"], TypeError),
+        ("boundary", "fields", (1,), TypeError),
+        ("boundary", "present", ["value"], TypeError),
+        ("boundary", "present", ("other",), ValueError),
+        ("boundary", "payload", "bytes", TypeError),
+    ],
+)
+def test_envelope_codec_revalidates_unpickled_metadata(target, field, value, error):
+    boundary = encode_boundary("edge", ("value",), (42,))
+    envelope = WorkflowEnvelope("run", (boundary,))
+    object.__setattr__(envelope if target == "envelope" else boundary, field, value)
+    # Simulate an existing serialized instance: pickle bypasses constructors.
+    payload = cloudpickle.dumps(envelope)
+    with pytest.raises(error):
+        loads_envelope(payload)
+    with pytest.raises(error):
+        dumps_envelope(envelope)
+
+
+def test_merge_rejects_overlapping_ids_and_orders_disjoint_ids():
+    first = encode_boundary("z", ("value",), (1,))
+    second = encode_boundary("z", ("value",), (2,))
+    with pytest.raises(ValueError, match="unique"):
+        merge_envelopes(
+            WorkflowEnvelope("run", (first,)), WorkflowEnvelope("run", (second,))
+        )
+    other = encode_boundary("a", ("value",), (3,))
+    merged = merge_envelopes(
+        WorkflowEnvelope("run", (first,)), WorkflowEnvelope("run", (other,))
+    )
+    assert merged.boundaries == (other, first)
+    assert decode_boundary(merged.boundaries[0], ("value",)) == (3,)
+    assert decode_boundary(merged.boundaries[1], ("value",)) == (1,)
+    with pytest.raises(TypeError):
+        _ = first < other
+
+
+@pytest.mark.parametrize("values", [[42], {"value": 42}, "x", None])
+def test_encoder_rejects_non_tuple_values_before_serialization(values, monkeypatch):
+    def unexpected_serialization(*args, **kwargs):
+        pytest.fail("invalid values reached serialization")
+
+    monkeypatch.setattr(cloudpickle, "dumps", unexpected_serialization)
+    with pytest.raises(TypeError, match="tuple"):
+        encode_boundary("edge", ("value",), values)
+
+
+@pytest.mark.parametrize(
+    ("payload", "error"), [(b"", EOFError), (b"!", pickle.UnpicklingError)]
+)
+def test_decoders_preserve_pickle_exceptions(payload, error):
+    with pytest.raises(error):
+        loads_envelope(payload)
+    with pytest.raises(error):
+        decode_boundary(BoundaryPayload("edge", (), (), payload), ())