dex-python-sdk 0.0.2__cp311-abi3-win_amd64.whl

dex/_value_hydrator.py

@@ -0,0 +1,223 @@
+# Copyright (c) 2026 Super Durable, Inc.
+#
+# Licensed under the Super Durable Source License 1.0.
+# You may not use this file except in compliance with the License.
+# See the LICENSE file in the repository root.
+#
+# SPDX-License-Identifier: LicenseRef-Super-Durable-1.0
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+
+from dex.blob_cache import BlobCache
+from dex.dexpb import dex_pb2 as pb
+from dex.dexpb import dex_pb2_grpc
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class _PendingBlob:
+    blob_id: str
+    is_object: bool
+    request: pb.Value
+    indexes: list[int] = field(default_factory=list)
+    hydrated: pb.Value | None = None
+
+
+class ValueHydrator:
+    def __init__(
+        self,
+        service: dex_pb2_grpc.FlowServiceStub,
+        cache: BlobCache,
+    ) -> None:
+        self._service = service
+        self._cache = cache
+
+    def hydrate(self, value: pb.Value) -> pb.Value:
+        return self.hydrate_all([value])[0]
+
+    def hydrate_all(self, values: list[pb.Value]) -> list[pb.Value]:
+        hydrated = list(values)
+        pending: dict[tuple[str, bool], _PendingBlob] = {}
+        for index, value in enumerate(values):
+            key = self._blob_key(value)
+            if key is None:
+                self._validate_concrete(value)
+                continue
+            blob = pending.get(key)
+            if blob is None:
+                blob = _PendingBlob(key[0], key[1], value)
+                pending[key] = blob
+            blob.indexes.append(index)
+
+        misses: list[_PendingBlob] = []
+        for blob in pending.values():
+            concrete = self._read_cache(blob)
+            if concrete is None:
+                misses.append(blob)
+            else:
+                blob.hydrated = concrete
+        self._load_misses(misses)
+
+        for blob in pending.values():
+            if blob.hydrated is None:
+                raise RuntimeError(f"blob was not hydrated: {blob.blob_id}")
+            for index in blob.indexes:
+                hydrated[index] = blob.hydrated
+        return hydrated
+
+    def wait_for_request(
+        self,
+        request: pb.InvokeWaitForMethodRequest,
+    ) -> pb.InvokeWaitForMethodRequest:
+        result = pb.InvokeWaitForMethodRequest()
+        result.CopyFrom(request)
+        values = [request.step_input, *(entry.value for entry in request.attributes)]
+        hydrated = self.hydrate_all(values)
+        result.step_input.CopyFrom(hydrated[0])
+        for entry, value in zip(result.attributes, hydrated[1:]):
+            entry.value.CopyFrom(value)
+        return result
+
+    def execute_request(
+        self,
+        request: pb.InvokeExecuteMethodRequest,
+    ) -> pb.InvokeExecuteMethodRequest:
+        result = pb.InvokeExecuteMethodRequest()
+        result.CopyFrom(request)
+        values = [request.step_input]
+        values.extend(entry.value for entry in request.attributes)
+        values.extend(entry.value for entry in request.step_exe_locals)
+        for channel_result in request.condition_results.channel_results:
+            values.extend(channel_result.values)
+        hydrated = iter(self.hydrate_all(values))
+        result.step_input.CopyFrom(next(hydrated))
+        for entry in result.attributes:
+            entry.value.CopyFrom(next(hydrated))
+        for entry in result.step_exe_locals:
+            entry.value.CopyFrom(next(hydrated))
+        for channel_result in result.condition_results.channel_results:
+            for value in channel_result.values:
+                value.CopyFrom(next(hydrated))
+        return result
+
+    def rpc_request(
+        self,
+        request: pb.InvokeWorkerRPCRequest,
+    ) -> pb.InvokeWorkerRPCRequest:
+        result = pb.InvokeWorkerRPCRequest()
+        result.CopyFrom(request)
+        values = [request.input, *(entry.value for entry in request.attributes)]
+        hydrated = self.hydrate_all(values)
+        result.input.CopyFrom(hydrated[0])
+        for entry, value in zip(result.attributes, hydrated[1:]):
+            entry.value.CopyFrom(value)
+        return result
+
+    def step_outputs(
+        self,
+        outputs: list[pb.StepCompletionOutput],
+    ) -> list[pb.StepCompletionOutput]:
+        values = [output.completed_step_output for output in outputs]
+        hydrated = self.hydrate_all(values)
+        results: list[pb.StepCompletionOutput] = []
+        for output, value in zip(outputs, hydrated):
+            result = pb.StepCompletionOutput()
+            result.CopyFrom(output)
+            result.completed_step_output.CopyFrom(value)
+            results.append(result)
+        return results
+
+    def _load_misses(self, misses: list[_PendingBlob]) -> None:
+        if not misses:
+            return
+        response = self._service.LoadBlobs(
+            pb.LoadBlobsRequest(values=[miss.request for miss in misses])
+        )
+        for miss in misses:
+            concrete = response.values.get(miss.blob_id)
+            if concrete is None:
+                raise RuntimeError(f"LoadBlobs omitted blob {miss.blob_id}")
+            self._validate_hydrated(miss, concrete)
+            miss.hydrated = concrete
+            self._write_cache(miss, concrete)
+
+    def _read_cache(self, blob: _PendingBlob) -> pb.Value | None:
+        try:
+            payload = self._cache.get(blob.blob_id)
+            if payload is None:
+                return None
+            if blob.is_object:
+                concrete = pb.Value(obj_value=pb.EncodedObject.FromString(payload))
+            else:
+                concrete = pb.Value(string_value=payload.decode("utf-8"))
+            self._validate_hydrated(blob, concrete)
+            return concrete
+        except Exception:
+            _LOGGER.warning("cannot read cached blob %s", blob.blob_id, exc_info=True)
+            try:
+                self._cache.delete(blob.blob_id)
+            except Exception:
+                _LOGGER.warning(
+                    "cannot delete cached blob %s",
+                    blob.blob_id,
+                    exc_info=True,
+                )
+            return None
+
+    def _write_cache(self, blob: _PendingBlob, concrete: pb.Value) -> None:
+        try:
+            payload = (
+                concrete.obj_value.SerializeToString()
+                if blob.is_object
+                else concrete.string_value.encode("utf-8")
+            )
+            self._cache.put(blob.blob_id, payload)
+        except Exception:
+            _LOGGER.warning("cannot cache blob %s", blob.blob_id, exc_info=True)
+
+    @staticmethod
+    def _blob_key(value: pb.Value) -> tuple[str, bool] | None:
+        kind = value.WhichOneof("kind")
+        if kind == "internal_blob_id_for_string_value":
+            blob_id = value.internal_blob_id_for_string_value
+            if not blob_id:
+                raise ValueError("blob ID is required")
+            return blob_id, False
+        if kind == "internal_blob_id_for_obj_value":
+            blob_id = value.internal_blob_id_for_obj_value
+            if not blob_id:
+                raise ValueError("blob ID is required")
+            return blob_id, True
+        return None
+
+    @staticmethod
+    def _validate_hydrated(blob: _PendingBlob, value: pb.Value) -> None:
+        expected = "obj_value" if blob.is_object else "string_value"
+        if value.WhichOneof("kind") != expected:
+            raise RuntimeError(
+                f"blob {blob.blob_id} hydrated to {value.WhichOneof('kind')}"
+            )
+        ValueHydrator._validate_concrete(value)
+
+    @staticmethod
+    def _validate_concrete(value: pb.Value) -> None:
+        kind = value.WhichOneof("kind")
+        if kind in ("string_value", "int_value", "bool_value", "null_value"):
+            return
+        if kind == "double_value":
+            import math
+
+            if not math.isfinite(value.double_value):
+                raise ValueError("non-finite numbers are unsupported")
+            return
+        if kind == "obj_value":
+            if value.obj_value.encoding not in ("json", "rawbytes"):
+                raise ValueError(
+                    f"unsupported object encoding {value.obj_value.encoding}"
+                )
+            return
+        raise ValueError("Value has no concrete kind")

dex/_value_mapper.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2026 Super Durable, Inc.
+#
+# Licensed under the Super Durable Source License 1.0.
+# You may not use this file except in compliance with the License.
+# See the LICENSE file in the repository root.
+#
+# SPDX-License-Identifier: LicenseRef-Super-Durable-1.0
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, is_dataclass
+from datetime import datetime
+from enum import Enum
+from typing import Any, cast
+
+from google.protobuf import struct_pb2
+
+from dex.attribute import AttributeIndex, IndexType
+from dex.codec import Codec, CodecRegistry, Value, WireKind
+from dex.dexpb import dex_pb2 as pb
+
+_JSON_ENCODING = "json"
+_RAW_BYTES_ENCODING = "rawbytes"
+
+
+class ValueMapper:
+    def __init__(self, codecs: CodecRegistry) -> None:
+        self._codecs = codecs
+
+    def codec(self, value_type: object) -> Codec[Any]:
+        return self._codecs.resolve(value_type)
+
+    def encode(self, value: Any, codec: Codec[Any]) -> pb.Value:
+        if value is None:
+            return pb.Value(
+                obj_value=pb.EncodedObject(
+                    encoding=_JSON_ENCODING,
+                    payload=b"null",
+                )
+            )
+        logical = codec.encode(value)
+        if logical.kind is WireKind.STRING:
+            return pb.Value(string_value=self._require(logical.data, str))
+        if logical.kind is WireKind.BOOL:
+            return pb.Value(bool_value=self._require_exact(logical.data, bool))
+        if logical.kind is WireKind.INT64:
+            return pb.Value(int_value=self._require_exact(logical.data, int))
+        if logical.kind is WireKind.DOUBLE:
+            return pb.Value(double_value=self._require_exact(logical.data, float))
+        if logical.kind is WireKind.BYTES:
+            return pb.Value(
+                obj_value=pb.EncodedObject(
+                    encoding=_RAW_BYTES_ENCODING,
+                    payload=self._require(logical.data, bytes),
+                )
+            )
+        if logical.kind is WireKind.JSON:
+            payload = self._require(logical.data, str).encode("utf-8")
+            return pb.Value(
+                obj_value=pb.EncodedObject(
+                    encoding=_JSON_ENCODING,
+                    payload=payload,
+                )
+            )
+        raise TypeError(f"unsupported wire kind {logical.kind}")
+
+    def encode_dynamic(self, value: Any) -> pb.Value:
+        if value is None:
+            return self.encode(value, self._codecs.resolve(type(None)))
+        value_type = type(value)
+        try:
+            codec = self._codecs.resolve(value_type)
+        except TypeError:
+            if isinstance(value, (list, tuple, dict)):
+                return pb.Value(
+                    obj_value=pb.EncodedObject(
+                        encoding=_JSON_ENCODING,
+                        payload=self._encode_json_collection(value),
+                    )
+                )
+            raise
+        return self.encode(value, codec)
+
+    def decode(self, value: pb.Value, codec: Codec[Any]) -> Any:
+        kind = value.WhichOneof("kind")
+        expected = codec.wire_kind
+        if kind == "string_value" and expected is WireKind.STRING:
+            return codec.decode(Value(WireKind.STRING, value.string_value))
+        if kind == "bool_value" and expected is WireKind.BOOL:
+            return codec.decode(Value(WireKind.BOOL, value.bool_value))
+        if kind == "int_value" and expected is WireKind.INT64:
+            return codec.decode(Value(WireKind.INT64, value.int_value))
+        if kind == "double_value" and expected is WireKind.DOUBLE:
+            return codec.decode(Value(WireKind.DOUBLE, value.double_value))
+        if kind == "obj_value" and expected is WireKind.BYTES:
+            self._require_encoding(value.obj_value, _RAW_BYTES_ENCODING)
+            return codec.decode(Value(WireKind.BYTES, value.obj_value.payload))
+        if kind == "obj_value" and expected is WireKind.JSON:
+            self._require_encoding(value.obj_value, _JSON_ENCODING)
+            payload = value.obj_value.payload.decode("utf-8")
+            return codec.decode(Value(WireKind.JSON, payload))
+        if kind == "obj_value" and value.obj_value.encoding == _JSON_ENCODING:
+            if value.obj_value.payload == b"null":
+                return None
+        if kind in (
+            "internal_blob_id_for_string_value",
+            "internal_blob_id_for_obj_value",
+        ):
+            raise ValueError("blob-backed Value was not hydrated")
+        if kind == "null_value":
+            raise ValueError("attribute deletion marker cannot be decoded")
+        raise TypeError(f"cannot decode {kind or 'empty Value'} as {codec.type_name}")
+
+    @staticmethod
+    def deletion() -> pb.Value:
+        return pb.Value(null_value=cast(struct_pb2.NullValue, struct_pb2.NULL_VALUE))
+
+    @staticmethod
+    def index_config(
+        index: AttributeIndex | None,
+        dynamic: bool,
+    ) -> pb.IndexConfig | None:
+        if index is None:
+            return None
+        mapped_types = {
+            IndexType.KEYWORD: pb.INDEX_TYPE_KEYWORD,
+            IndexType.FULL_TEXT: pb.INDEX_TYPE_TEXT,
+            IndexType.KEYWORD_ARRAY: pb.INDEX_TYPE_KEYWORD_ARRAY,
+            IndexType.INT: pb.INDEX_TYPE_INT,
+            IndexType.DOUBLE: pb.INDEX_TYPE_DOUBLE,
+            IndexType.BOOL: pb.INDEX_TYPE_BOOL,
+            IndexType.DATETIME: pb.INDEX_TYPE_DATETIME,
+        }
+        config = pb.IndexConfig(enable=True, type=mapped_types[index.type])
+        if index.index_key or dynamic:
+            config.index_key = index.index_key
+        return config
+
+    @staticmethod
+    def _require(value: object, expected: type[Any]) -> Any:
+        if not isinstance(value, expected):
+            raise TypeError(f"expected {expected.__name__}")
+        return value
+
+    @staticmethod
+    def _require_exact(value: object, expected: type[Any]) -> Any:
+        if type(value) is not expected:
+            raise TypeError(f"expected {expected.__name__}")
+        return value
+
+    @staticmethod
+    def _require_encoding(value: pb.EncodedObject, expected: str) -> None:
+        if value.encoding != expected:
+            raise TypeError(f"expected {expected} encoding, got {value.encoding}")
+
+    @staticmethod
+    def _encode_json_collection(value: object) -> bytes:
+        def default(item: object) -> object:
+            if is_dataclass(item) and not isinstance(item, type):
+                return asdict(item)
+            if isinstance(item, datetime):
+                return item.isoformat()
+            if isinstance(item, Enum):
+                return item.value
+            raise TypeError(f"cannot encode {type(item).__name__} as JSON")
+
+        return json.dumps(
+            value,
+            default=default,
+            allow_nan=False,
+            separators=(",", ":"),
+            sort_keys=True,
+        ).encode("utf-8")