diff --git a/airbyte_cdk/models/airbyte_protocol_serializers.py b/airbyte_cdk/models/airbyte_protocol_serializers.py index 52cab83d21..f0e6c5bc0a 100644 --- a/airbyte_cdk/models/airbyte_protocol_serializers.py +++ b/airbyte_cdk/models/airbyte_protocol_serializers.py @@ -1,5 +1,5 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. -from typing import Any, Dict +from typing import Any, Dict, cast from serpyco_rs import CustomType, Serializer @@ -28,10 +28,61 @@ def get_json_schema(self) -> Dict[str, Any]: return {"type": "object"} -def custom_type_resolver(t: type) -> CustomType[AirbyteStateBlob, Dict[str, Any]] | None: +def _blob_custom_type_resolver( + t: type, +) -> CustomType[AirbyteStateBlob, Dict[str, Any]] | None: return AirbyteStateBlobType() if t is AirbyteStateBlob else None +class AirbyteStateMessageType(CustomType[AirbyteStateMessage, Dict[str, Any]]): + _KNOWN_PROPERTIES = { + "type", + "stream", + "global", + "data", + "sourceStats", + "destinationStats", + } + + def __init__(self) -> None: + self._inner = Serializer( + AirbyteStateMessage, + omit_none=True, + custom_type_resolver=_blob_custom_type_resolver, + ) + + def serialize(self, value: AirbyteStateMessage) -> Dict[str, Any]: + result = cast(Dict[str, Any], self._inner.dump(value)) + result.update( + { + key: property_value + for key, property_value in getattr(value, "additional_properties", {}).items() + if key not in self._KNOWN_PROPERTIES + } + ) + return result + + def deserialize(self, value: Dict[str, Any]) -> AirbyteStateMessage: + message = self._inner.load(value) + additional_properties = { + key: property_value + for key, property_value in value.items() + if key not in self._KNOWN_PROPERTIES + } + if additional_properties: + message.additional_properties = additional_properties # type: ignore[attr-defined] + return message + + def get_json_schema(self) -> Dict[str, Any]: + return self._inner.get_json_schema() + + +def custom_type_resolver(t: type) -> CustomType[Any, Any] | None: + if t is AirbyteStateMessage: + return AirbyteStateMessageType() + return _blob_custom_type_resolver(t) + + AirbyteCatalogSerializer = Serializer(AirbyteCatalog, omit_none=True) AirbyteStreamSerializer = Serializer(AirbyteStream, omit_none=True) AirbyteStreamStateSerializer = Serializer( diff --git a/unit_tests/test_airbyte_protocol_serializers.py b/unit_tests/test_airbyte_protocol_serializers.py new file mode 100644 index 0000000000..66f5bcb1eb --- /dev/null +++ b/unit_tests/test_airbyte_protocol_serializers.py @@ -0,0 +1,74 @@ +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. + +from airbyte_cdk.models import AirbyteMessageSerializer, AirbyteStateMessageSerializer + + +def test_stream_state_message_preserves_additional_properties(): + message = { + "type": "STATE", + "state": { + "type": "STREAM", + "stream": { + "stream_descriptor": {"name": "users", "namespace": "public"}, + "stream_state": {"cursor": "value"}, + }, + "id": 42, + }, + } + + loaded = AirbyteMessageSerializer.load(message) + + assert AirbyteMessageSerializer.dump(loaded) == message + + +def test_global_state_message_preserves_additional_properties(): + message = { + "type": "STATE", + "state": { + "type": "GLOBAL", + "global": { + "stream_states": [ + { + "stream_descriptor": {"name": "users", "namespace": "public"}, + "stream_state": {"cursor": "value"}, + } + ], + "shared_state": {"shared": "value"}, + }, + "id": 42, + }, + } + + loaded = AirbyteMessageSerializer.load(message) + + assert AirbyteMessageSerializer.dump(loaded) == message + + +def test_direct_state_message_serializer_preserves_additional_properties(): + message = { + "type": "STREAM", + "stream": { + "stream_descriptor": {"name": "users"}, + "stream_state": {"cursor": "value"}, + }, + "id": 42, + } + + loaded = AirbyteStateMessageSerializer.load(message) + + assert AirbyteStateMessageSerializer.dump(loaded) == message + + +def test_state_message_without_additional_properties_does_not_leak_attribute(): + message = { + "type": "STREAM", + "stream": { + "stream_descriptor": {"name": "users"}, + "stream_state": {"cursor": "value"}, + }, + } + + loaded = AirbyteStateMessageSerializer.load(message) + + assert AirbyteStateMessageSerializer.dump(loaded) == message + assert "additional_properties" not in AirbyteStateMessageSerializer.dump(loaded)