airbyte-cdk 0.36.0__py3-none-any.whl → 0.36.1__py3-none-any.whl

airbyte_cdk/sources/declarative/declarative_stream.py

@@ -5,14 +5,14 @@
 from dataclasses import InitVar, dataclass, field
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteTraceMessage, SyncMode
+from airbyte_cdk.models import AirbyteMessage, SyncMode
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.declarative.schema import DefaultSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
 from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.declarative.types import Config, StreamSlice
-from airbyte_cdk.sources.streams.core import Stream
+from airbyte_cdk.sources.streams.core import Stream, StreamData
 
 
 @dataclass
@@ -102,17 +102,27 @@ class DeclarativeStream(Stream):
 
     def _apply_transformations(
         self,
-        message_or_record_data: Union[AirbyteMessage, AirbyteLogMessage, AirbyteTraceMessage, Mapping[str, Any]],
+        message_or_record_data: StreamData,
         config: Config,
         stream_slice: StreamSlice,
     ):
-        # If the input is an AirbyteRecord, transform the record's data
-        # If the input is another type of Airbyte Message, return it as is
+        # If the input is an AirbyteMessage with a record, transform the record's data
+        # If the input is another type of AirbyteMessage, return it as is
         # If the input is a dict, transform it
-        if isinstance(message_or_record_data, AirbyteLogMessage) or isinstance(message_or_record_data, AirbyteTraceMessage):
-            return message_or_record_data
+        if isinstance(message_or_record_data, AirbyteMessage):
+            if message_or_record_data.record:
+                record = message_or_record_data.record.data
+            else:
+                return message_or_record_data
+        elif isinstance(message_or_record_data, dict):
+            record = message_or_record_data
+        else:
+            # Raise an error because this is unexpected and indicative of a typing problem in the CDK
+            raise ValueError(
+                f"Unexpected record type. Expected {StreamData}. Got {type(message_or_record_data)}. This is probably due to a bug in the CDK."
+            )
         for transformation in self.transformations:
-            transformation.transform(message_or_record_data, config=config, stream_state=self.state, stream_slice=stream_slice)
+            transformation.transform(record, config=config, stream_state=self.state, stream_slice=stream_slice)
 
         return message_or_record_data
 

airbyte_cdk/sources/streams/core.py

@@ -11,7 +11,7 @@ from functools import lru_cache
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
 
 import airbyte_cdk.sources.utils.casing as casing
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteStream, AirbyteTraceMessage, SyncMode
+from airbyte_cdk.models import AirbyteMessage, AirbyteStream, SyncMode
 
 # list of all possible HTTP methods which can be used for sending of request bodies
 from airbyte_cdk.sources.utils.schema_helpers import ResourceSchemaLoader
@@ -24,10 +24,8 @@ if typing.TYPE_CHECKING:
 
 # A stream's read method can return one of the following types:
 # Mapping[str, Any]: The content of an AirbyteRecordMessage
-# AirbyteRecordMessage: An AirbyteRecordMessage
-# AirbyteLogMessage: A log message
-# AirbyteTraceMessage: A trace message
-StreamData = Union[Mapping[str, Any], AirbyteLogMessage, AirbyteTraceMessage]
+# AirbyteMessage: An AirbyteMessage. Could be of any type
+StreamData = Union[Mapping[str, Any], AirbyteMessage]
 
 
 def package_name_from_class(cls: object) -> str:

{airbyte_cdk-0.36.0.dist-info → airbyte_cdk-0.36.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 0.36.0
+Version: 0.36.1
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://github.com/airbytehq/airbyte
 Author: Airbyte

{airbyte_cdk-0.36.0.dist-info → airbyte_cdk-0.36.1.dist-info}/RECORD

@@ -24,7 +24,7 @@ airbyte_cdk/sources/declarative/__init__.py,sha256=ZnqYNxHsKCgO38IwB34RQyRMXTs4G
 airbyte_cdk/sources/declarative/create_partial.py,sha256=sUJOwD8hBzW4pxw2XhYlSTMgl-WMc5WpP5Oq_jo3fHw,3371
 airbyte_cdk/sources/declarative/declarative_component_schema.yaml,sha256=DJIV9HgHfBgv4w8zg6DhLIqXjZe235UVTe79IT7O51Q,74761
 airbyte_cdk/sources/declarative/declarative_source.py,sha256=U2As9PDKmcWDgbsWUo-RetJ9fxQOBlwntWZ0NOgs5Ac,1453
-airbyte_cdk/sources/declarative/declarative_stream.py,sha256=dhd228ddbG2YIptSw688HSWq9Hhh8GKLRxZS0VAz6rw,6211
+airbyte_cdk/sources/declarative/declarative_stream.py,sha256=0iZSpypxt8bhO3Lmf3BpGRTO7Fp0Q2GI8m8xyJJUjeM,6580
 airbyte_cdk/sources/declarative/exceptions.py,sha256=kTPUA4I2NV4J6HDz-mKPGMrfuc592akJnOyYx38l_QM,176
 airbyte_cdk/sources/declarative/manifest_declarative_source.py,sha256=vTbRNM8D9P_ChOu1GNvtNRt-PM2L9N5Y0pNRyfVFuZg,9759
 airbyte_cdk/sources/declarative/types.py,sha256=b_RJpL9TyAgxJIRYZx5BxpC39p-WccHKxbAqxWrn9oE,482
@@ -128,7 +128,7 @@ airbyte_cdk/sources/singer/singer_helpers.py,sha256=q1LmgjFxSnN-dobMy7nikUwcK-9F
 airbyte_cdk/sources/singer/source.py,sha256=3YY8UTOXmctvMVUnYmIegmL3_IxF55iGP_bc_s2MZdY,8530
 airbyte_cdk/sources/streams/__init__.py,sha256=XGrzYjIkqItvnMshsOUzYhi4lC4M9kFHhxG0oCAoAyE,176
 airbyte_cdk/sources/streams/availability_strategy.py,sha256=7BM0qLvXS0QrlKvnVkBEw4Cw8i7PCENCBLcIAcuD3nY,1007
-airbyte_cdk/sources/streams/core.py,sha256=2yN4jCVEY1nAEpHCsRIERjMqLn6XpgEfEioLpzXlRo8,11133
+airbyte_cdk/sources/streams/core.py,sha256=G0MxhVkrF1DAPoPgmRBd_wnhrI7S6ZAdtOG0h9N_zHU,11021
 airbyte_cdk/sources/streams/http/__init__.py,sha256=6hRmA0P_RhB7X54xQbtj2RTz1RGlP52AG9V4pPWLaEQ,261
 airbyte_cdk/sources/streams/http/availability_strategy.py,sha256=uYYjanBG0f-CIrUYn1SnMyJqHoQQHnTggeJiB0m4i-Y,6344
 airbyte_cdk/sources/streams/http/exceptions.py,sha256=OokLDI7W8hZvq9e15sL3em2AdwmzmcAl72Ms-i5l0Nw,1334
@@ -177,7 +177,7 @@ unit_tests/sources/test_source.py,sha256=eVtU9Zuc9gBsg11Pb5xjDtyU0gVrbYqbZ4RmzPv
 unit_tests/sources/declarative/__init__.py,sha256=ZnqYNxHsKCgO38IwB34RQyRMXTs4GTvlRi3ImKnIioo,61
 unit_tests/sources/declarative/external_component.py,sha256=lU2gL736bLEWtmrGm1B2k83RXt_3XkROimLIahZd5dg,293
 unit_tests/sources/declarative/test_create_partial.py,sha256=s_KIywQqt8RlauOCWNJVk3HC3KBTAtSwFTN6JVQgu80,2636
-unit_tests/sources/declarative/test_declarative_stream.py,sha256=o-KNEH1NXOQc3ZkThHlasMHk9dPD1pbWXtX7O1BvN9w,4815
+unit_tests/sources/declarative/test_declarative_stream.py,sha256=3leJnZIYHiFq8XI4jb3TjPXTubGJmvNGzABt4c01EkQ,5436
 unit_tests/sources/declarative/test_manifest_declarative_source.py,sha256=GckUc3nepzZkD1UM24woHlYCVZb5DP4IAQC3IeMyZF0,58924
 unit_tests/sources/declarative/test_yaml_declarative_source.py,sha256=6HhsUFgB7ueN0yOUHWb4gpPYLng5jasxN_plvz3x37g,5097
 unit_tests/sources/declarative/auth/__init__.py,sha256=4Hw-PX1-VgESLF16cDdvuYCzGJtHntThLF4qIiULWeo,61
@@ -259,8 +259,8 @@ unit_tests/utils/test_schema_inferrer.py,sha256=ckl17GlNOZInqgxni7Z2A0bg_p6JDy0G
 unit_tests/utils/test_secret_utils.py,sha256=XKe0f1RHYii8iwE6ATmBr5JGDI1pzzrnZUGdUSMJQP4,4886
 unit_tests/utils/test_stream_status_utils.py,sha256=NpV155JMXA6CG-2Zvofa14lItobyh3Onttc59X4m5DI,3382
 unit_tests/utils/test_traced_exception.py,sha256=bDFP5zMBizFenz6V2WvEZTRCKGB5ijh3DBezjbfoYIs,4198
-airbyte_cdk-0.36.0.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
-airbyte_cdk-0.36.0.dist-info/METADATA,sha256=Ca9A-Ln_vN5vBxLQCRsoXABK7QcKZl8E2UozABuvjfs,8902
-airbyte_cdk-0.36.0.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-airbyte_cdk-0.36.0.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
-airbyte_cdk-0.36.0.dist-info/RECORD,,
+airbyte_cdk-0.36.1.dist-info/LICENSE.txt,sha256=Wfe61S4BaGPj404v8lrAbvhjYR68SHlkzeYrg3_bbuM,1051
+airbyte_cdk-0.36.1.dist-info/METADATA,sha256=julG0JMgo00RRk6h3nZf_h-2-wVJffJ3e_fEEsV3dhQ,8902
+airbyte_cdk-0.36.1.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+airbyte_cdk-0.36.1.dist-info/top_level.txt,sha256=edvsDKTnE6sD2wfCUaeTfKf5gQIL6CPVMwVL2sWZzqo,51
+airbyte_cdk-0.36.1.dist-info/RECORD,,

unit_tests/sources/declarative/test_declarative_stream.py

@@ -5,7 +5,16 @@
 from unittest import mock
 from unittest.mock import MagicMock, call
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteTraceMessage, Level, SyncMode, TraceType
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteRecordMessage,
+    AirbyteTraceMessage,
+    Level,
+    SyncMode,
+    TraceType,
+    Type,
+)
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.transformations import AddFields, RecordTransformation
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
@@ -24,8 +33,8 @@ def test_declarative_stream():
     records = [
         {"pk": 1234, "field": "value"},
         {"pk": 4567, "field": "different_value"},
-        AirbyteLogMessage(level=Level.INFO, message="This is a log  message"),
-        AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345),
+        AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is a log  message")),
+        AirbyteMessage(type=Type.TRACE, trace=AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345)),
     ]
     stream_slices = [
         {"date": "2021-01-01"},
@@ -84,15 +93,17 @@ def test_declarative_stream_with_add_fields_transform():
     retriever_records = [
         {"pk": 1234, "field": "value"},
         {"pk": 4567, "field": "different_value"},
-        AirbyteLogMessage(level=Level.INFO, message="This is a log  message"),
-        AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345),
+        AirbyteMessage(type=Type.RECORD, record=AirbyteRecordMessage(data={"pk": 1357, "field": "a_value"}, emitted_at=12344, stream="stream")),
+        AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is a log  message")),
+        AirbyteMessage(type=Type.TRACE, trace=AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345)),
     ]
 
     expected_records = [
         {"pk": 1234, "field": "value", "added_key": "added_value"},
         {"pk": 4567, "field": "different_value", "added_key": "added_value"},
-        AirbyteLogMessage(level=Level.INFO, message="This is a log  message"),
-        AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345),
+        AirbyteMessage(type=Type.RECORD, record=AirbyteRecordMessage(data={"pk": 1357, "field": "a_value", "added_key": "added_value"}, emitted_at=12344, stream="stream")),
+        AirbyteMessage(type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="This is a log  message")),
+        AirbyteMessage(type=Type.TRACE, trace=AirbyteTraceMessage(type=TraceType.ERROR, emitted_at=12345)),
     ]
     stream_slices = [
         {"date": "2021-01-01"},