airbyte-cdk 6.60.0.post35.dev16509779638__py3-none-any.whl → 6.60.2__py3-none-any.whl

airbyte_cdk/models/airbyte_protocol_serializers.py

@@ -1,13 +1,7 @@
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
-import json
-import logging
-import sys
-from enum import Enum
-from typing import Any, Callable, Dict, Type, TypeVar, cast
+from typing import Any, Dict
 
-import dacite
-import orjson
-from pydantic import ValidationError
+from serpyco_rs import CustomType, Serializer
 
 from .airbyte_protocol import (  # type: ignore[attr-defined] # all classes are imported to airbyte_protocol via *
     AirbyteCatalog,
@@ -21,154 +15,34 @@ from .airbyte_protocol import (  # type: ignore[attr-defined] # all classes are
     ConnectorSpecification,
 )
 
-USE_RUST_BACKEND = sys.platform != "emscripten"
-"""When run in WASM, use the pure Python backend for serpyco."""
 
-_HAS_LOGGED_FOR_SERIALIZATION_ERROR = False
-"""Track if we have logged an error for serialization issues."""
+class AirbyteStateBlobType(CustomType[AirbyteStateBlob, Dict[str, Any]]):
+    def serialize(self, value: AirbyteStateBlob) -> Dict[str, Any]:
+        # cant use orjson.dumps() directly because private attributes are excluded, e.g. "__ab_full_refresh_sync_complete"
+        return {k: v for k, v in value.__dict__.items()}
 
-T = TypeVar("T")
+    def deserialize(self, value: Dict[str, Any]) -> AirbyteStateBlob:
+        return AirbyteStateBlob(value)
 
-logger = logging.getLogger("airbyte")
+    def get_json_schema(self) -> Dict[str, Any]:
+        return {"type": "object"}
 
-# Making this a no-op for now:
 
+def custom_type_resolver(t: type) -> CustomType[AirbyteStateBlob, Dict[str, Any]] | None:
+    return AirbyteStateBlobType() if t is AirbyteStateBlob else None
 
-def ab_message_to_string(
-    message: AirbyteMessage,
-) -> str:
-    """
-    Convert an AirbyteMessage to a JSON string.
 
-    Args:
-        message (AirbyteMessage): The Airbyte message to convert.
-
-    Returns:
-        str: JSON string representation of the AirbyteMessage.
-    """
-    return message.model_dump_json()
-
-
-def ab_message_from_string(
-    message_json: str,
-) -> AirbyteMessage:
-    """
-    Convert a JSON string to an AirbyteMessage.
-
-    Args:
-        message_str (str): The JSON string to convert.
-
-    Returns:
-        AirbyteMessage: The deserialized AirbyteMessage.
-    """
-    try:
-        return AirbyteMessage.model_validate_json(message_json)
-    except ValidationError as e:
-        raise ValueError(f"Invalid AirbyteMessage format: {e}") from e
-    except orjson.JSONDecodeError as e:
-        raise ValueError(f"Failed to decode JSON: {e}") from e
-
-
-def ab_connector_spec_from_string(
-    spec_json: str,
-) -> ConnectorSpecification:
-    """
-    Convert a JSON string to a ConnectorSpecification.
-
-    Args:
-        spec_str (str): The JSON string to convert.
-
-    Returns:
-        ConnectorSpecification: The deserialized ConnectorSpecification.
-    """
-    try:
-        return ConnectorSpecification.model_validate_json(spec_json)
-    except ValidationError as e:
-        raise ValueError(f"Invalid ConnectorSpecification format: {e}") from e
-    except orjson.JSONDecodeError as e:
-        raise ValueError(f"Failed to decode JSON: {e}") from e
-
-
-def ab_connector_spec_to_string(
-    spec: ConnectorSpecification,
-) -> str:
-    """
-    Convert a ConnectorSpecification to a JSON string.
-
-    Args:
-        spec (ConnectorSpecification): The ConnectorSpecification to convert.
-
-    Returns:
-        str: JSON string representation of the ConnectorSpecification.
-    """
-    return spec.model_dump_json()
-
-
-def ab_configured_catalog_to_string(
-    catalog: ConfiguredAirbyteCatalog,
-) -> str:
-    """
-    Convert a ConfiguredAirbyteCatalog to a JSON string.
-
-    Args:
-        catalog (ConfiguredAirbyteCatalog): The ConfiguredAirbyteCatalog to convert.
-
-    Returns:
-        str: JSON string representation of the ConfiguredAirbyteCatalog.
-    """
-    return catalog.model_dump_json()
-
-
-def ab_configured_catalog_from_string(
-    catalog_json: str,
-) -> ConfiguredAirbyteCatalog:
-    """
-    Convert a JSON string to a ConfiguredAirbyteCatalog.
-
-    Args:
-        catalog_json (str): The JSON string to convert.
-
-    Returns:
-        ConfiguredAirbyteCatalog: The deserialized ConfiguredAirbyteCatalog.
-    """
-    try:
-        return ConfiguredAirbyteCatalog.model_validate_json(catalog_json)
-    except ValidationError as e:
-        raise ValueError(f"Invalid ConfiguredAirbyteCatalog format: {e}") from e
-    except orjson.JSONDecodeError as e:
-        raise ValueError(f"Failed to decode JSON: {e}") from e
-
-
-def ab_state_message_from_string(
-    state_json: str,
-) -> AirbyteStateMessage:
-    """
-    Convert a JSON string to an AirbyteStateMessage.
-
-    Args:
-        state_json (str): The JSON string to convert.
-
-    Returns:
-        AirbyteStateMessage: The deserialized AirbyteStateMessage.
-    """
-    try:
-        return AirbyteStateMessage.model_validate_json(state_json)
-    except ValidationError as e:
-        raise ValueError(f"Invalid AirbyteStateMessage format: {e}") from e
-    except orjson.JSONDecodeError as e:
-        raise ValueError(f"Failed to decode JSON: {e}") from e
-
-
-def ab_state_message_to_string(
-    state: AirbyteStateMessage,
-) -> str:
-    """
-    Convert an AirbyteStateMessage to a JSON string.
-
-    Args:
-        state (AirbyteStateMessage): The AirbyteStateMessage to convert.
-
-    Returns:
-        str: JSON string representation of the AirbyteStateMessage.
-    """
-    return state.model_dump_json()
+AirbyteCatalogSerializer = Serializer(AirbyteCatalog, omit_none=True)
+AirbyteStreamSerializer = Serializer(AirbyteStream, omit_none=True)
+AirbyteStreamStateSerializer = Serializer(
+    AirbyteStreamState, omit_none=True, custom_type_resolver=custom_type_resolver
+)
+AirbyteStateMessageSerializer = Serializer(
+    AirbyteStateMessage, omit_none=True, custom_type_resolver=custom_type_resolver
+)
+AirbyteMessageSerializer = Serializer(
+    AirbyteMessage, omit_none=True, custom_type_resolver=custom_type_resolver
+)
+ConfiguredAirbyteCatalogSerializer = Serializer(ConfiguredAirbyteCatalog, omit_none=True)
+ConfiguredAirbyteStreamSerializer = Serializer(ConfiguredAirbyteStream, omit_none=True)
+ConnectorSpecificationSerializer = Serializer(ConnectorSpecification, omit_none=True)

airbyte_cdk/models/well_known_types.py

@@ -2,4 +2,4 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from airbyte_protocol.models.well_known_types import *  # noqa: F403  # Allow '*'
+from airbyte_protocol_dataclasses.models.well_known_types import *  # noqa: F403  # Allow '*'

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -63,7 +63,7 @@ from airbyte_cdk.sources.streams.concurrent.helpers import get_primary_key_from_
 class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
     # By default, we defer to a value of 2. A value lower than than could cause a PartitionEnqueuer to be stuck in a state of deadlock
     # because it has hit the limit of futures but not partition reader is consuming them.
-    _LOWEST_SAFE_CONCURRENCY_LEVEL = 1  # TODO: revert-me: 2
+    _LOWEST_SAFE_CONCURRENCY_LEVEL = 2
 
     def __init__(
         self,
@@ -119,9 +119,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
             )  # Partition_generation iterates using range based on this value. If this is floored to zero we end up in a dead lock during start up
         else:
             concurrency_level = self._LOWEST_SAFE_CONCURRENCY_LEVEL
-            initial_number_of_partitions_to_generate = max(
-                self._LOWEST_SAFE_CONCURRENCY_LEVEL // 2, 1
-            )
+            initial_number_of_partitions_to_generate = self._LOWEST_SAFE_CONCURRENCY_LEVEL // 2
 
         self._concurrent_source = ConcurrentSource.create(
             num_workers=concurrency_level,

airbyte_cdk/sources/declarative/incremental/datetime_based_cursor.py

@@ -429,17 +429,6 @@ class DatetimeBasedCursor(DeclarativeCursor):
                 )
             )
 
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        cursor_field = self.cursor_field.eval(self.config)  # type: ignore  # cursor_field is converted to an InterpolatedString in __post_init__
-        first_cursor_value = first.get(cursor_field)
-        second_cursor_value = second.get(cursor_field)
-        if first_cursor_value and second_cursor_value:
-            return self.parse_date(first_cursor_value) >= self.parse_date(second_cursor_value)
-        elif first_cursor_value:
-            return True
-        else:
-            return False
-
     def set_runtime_lookback_window(self, lookback_window_in_seconds: int) -> None:
         """
         Updates the lookback window based on a given number of seconds if the new duration

airbyte_cdk/sources/declarative/incremental/global_substream_cursor.py

@@ -338,12 +338,6 @@ class GlobalSubstreamCursor(DeclarativeCursor):
     def should_be_synced(self, record: Record) -> bool:
         return self._stream_cursor.should_be_synced(self._convert_record_to_cursor_record(record))
 
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        return self._stream_cursor.is_greater_than_or_equal(
-            self._convert_record_to_cursor_record(first),
-            self._convert_record_to_cursor_record(second),
-        )
-
     @staticmethod
     def _convert_record_to_cursor_record(record: Record) -> Record:
         return Record(

airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py

@@ -315,21 +315,6 @@ class PerPartitionCursor(DeclarativeCursor):
             self._convert_record_to_cursor_record(record)
         )
 
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        if not first.associated_slice or not second.associated_slice:
-            raise ValueError(
-                f"Both records should have an associated slice but got {first.associated_slice} and {second.associated_slice}"
-            )
-        if first.associated_slice.partition != second.associated_slice.partition:
-            raise ValueError(
-                f"To compare records, partition should be the same but got {first.associated_slice.partition} and {second.associated_slice.partition}"
-            )
-
-        return self._get_cursor(first).is_greater_than_or_equal(
-            self._convert_record_to_cursor_record(first),
-            self._convert_record_to_cursor_record(second),
-        )
-
     @staticmethod
     def _convert_record_to_cursor_record(record: Record) -> Record:
         return Record(

airbyte_cdk/sources/declarative/incremental/per_partition_with_global.py

@@ -195,6 +195,3 @@ class PerPartitionWithGlobalCursor(DeclarativeCursor):
 
     def should_be_synced(self, record: Record) -> bool:
         return self._get_active_cursor().should_be_synced(record)
-
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        return self._global_cursor.is_greater_than_or_equal(first, second)

airbyte_cdk/sources/declarative/incremental/resumable_full_refresh_cursor.py

@@ -42,12 +42,6 @@ class ResumableFullRefreshCursor(DeclarativeCursor):
         """
         return True
 
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        """
-        RFR record don't have ordering to be compared between one another.
-        """
-        return False
-
     def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
         # A top-level RFR cursor only manages the state of a single partition
         return self._cursor

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -31,7 +31,7 @@ from airbyte_cdk.models import (
     ConnectorSpecification,
     FailureType,
 )
-from airbyte_cdk.models.airbyte_protocol_serializers import ab_message_to_string
+from airbyte_cdk.models.airbyte_protocol_serializers import AirbyteMessageSerializer
 from airbyte_cdk.sources.declarative.checks import COMPONENTS_CHECKER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
@@ -234,7 +234,7 @@ class ManifestDeclarativeSource(DeclarativeSource):
             )
             # We have no mechanism for consuming the queue, so we print the messages to stdout
             for message in self.message_repository.consume_queue():
-                print(ab_message_to_string(message))
+                print(orjson.dumps(AirbyteMessageSerializer.dump(message)).decode())
         self._spec_component.transform_config(mutable_config)
         return mutable_config
 
@@ -542,11 +542,19 @@ class ManifestDeclarativeSource(DeclarativeSource):
                 components_resolver_config["retriever"]["requester"]["use_cache"] = True
 
             # Create a resolver for dynamic components based on type
-            components_resolver = self._constructor.create_component(
-                COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type],
-                components_resolver_config,
-                config,
-            )
+            if resolver_type == "HttpComponentsResolver":
+                components_resolver = self._constructor.create_component(
+                    model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type],
+                    component_definition=components_resolver_config,
+                    config=config,
+                    stream_name=dynamic_definition.get("name"),
+                )
+            else:
+                components_resolver = self._constructor.create_component(
+                    model_type=COMPONENTS_RESOLVER_TYPE_MAPPING[resolver_type],
+                    component_definition=components_resolver_config,
+                    config=config,
+                )
 
             stream_template_config = dynamic_definition["stream_template"]
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -3493,10 +3493,11 @@ class ModelToComponentFactory:
                 requester=download_requester,
                 record_selector=record_selector,
                 primary_key=None,
-                name=job_download_components_name,
+                name=name,
                 paginator=paginator,
                 config=config,
                 parameters={},
+                log_formatter=self._get_log_formatter(None, name),
             )
 
         def _get_job_timeout() -> datetime.timedelta:
@@ -3805,7 +3806,7 @@ class ModelToComponentFactory:
         )
 
     def create_http_components_resolver(
-        self, model: HttpComponentsResolverModel, config: Config
+        self, model: HttpComponentsResolverModel, config: Config, stream_name: Optional[str] = None
     ) -> Any:
         stream_slicer = self._build_stream_slicer_from_partition_router(model.retriever, config)
         combined_slicers = self._build_resumable_cursor(model.retriever, stream_slicer)
@@ -3813,7 +3814,7 @@ class ModelToComponentFactory:
         retriever = self._create_component_from_model(
             model=model.retriever,
             config=config,
-            name="",
+            name=f"{stream_name if stream_name else '__http_components_resolver'}",
             primary_key=None,
             stream_slicer=stream_slicer if stream_slicer else combined_slicers,
             transformations=[],
@@ -3890,7 +3891,9 @@ class ModelToComponentFactory:
         )
 
     def create_parametrized_components_resolver(
-        self, model: ParametrizedComponentsResolverModel, config: Config
+        self,
+        model: ParametrizedComponentsResolverModel,
+        config: Config,
     ) -> ParametrizedComponentsResolver:
         stream_parameters = StreamParametersDefinition(
             list_of_parameters_for_stream=model.stream_parameters.list_of_parameters_for_stream

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -41,7 +41,6 @@ from airbyte_cdk.sources.declarative.requesters.request_options import (
 from airbyte_cdk.sources.declarative.requesters.requester import Requester
 from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
-from airbyte_cdk.sources.http_logger import format_http_message
 from airbyte_cdk.sources.source import ExperimentalClassWarning
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState
@@ -528,35 +527,13 @@ class SimpleRetriever(Retriever):
                 if self.cursor and current_record:
                     self.cursor.observe(_slice, current_record)
 
-                # Latest record read, not necessarily within slice boundaries.
-                # TODO Remove once all custom components implement `observe` method.
-                # https://github.com/airbytehq/airbyte-internal-issues/issues/6955
-                most_recent_record_from_slice = self._get_most_recent_record(
-                    most_recent_record_from_slice, current_record, _slice
-                )
                 yield stream_data
 
             if self.cursor:
-                self.cursor.close_slice(_slice, most_recent_record_from_slice)
+                self.cursor.close_slice(_slice)
         return
 
-    def _get_most_recent_record(
-        self,
-        current_most_recent: Optional[Record],
-        current_record: Optional[Record],
-        stream_slice: StreamSlice,
-    ) -> Optional[Record]:
-        if self.cursor and current_record:
-            if not current_most_recent:
-                return current_record
-            else:
-                return (
-                    current_most_recent
-                    if self.cursor.is_greater_than_or_equal(current_most_recent, current_record)
-                    else current_record
-                )
-        else:
-            return None
+    # FIXME based on the comment above in SimpleRetriever.read_records, it seems like we can tackle https://github.com/airbytehq/airbyte-internal-issues/issues/6955 and remove this
 
     def _extract_record(
         self, stream_data: StreamData, stream_slice: StreamSlice

airbyte_cdk/sources/declarative/spec/spec.py

@@ -8,8 +8,8 @@ from typing import Any, List, Mapping, MutableMapping, Optional
 from airbyte_cdk.models import (
     AdvancedAuth,
     ConnectorSpecification,
+    ConnectorSpecificationSerializer,
 )
-from airbyte_cdk.models.airbyte_protocol_serializers import ab_connector_spec_from_string
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import AuthFlow
 from airbyte_cdk.sources.declarative.transformations.config_transformations.config_transformation import (
     ConfigTransformation,
@@ -59,7 +59,7 @@ class Spec:
             obj["advanced_auth"] = self.advanced_auth.dict()
 
         # We remap these keys to camel case because that's the existing format expected by the rest of the platform
-        return ConnectorSpecification.model_validate(obj)
+        return ConnectorSpecificationSerializer.load(obj)
 
     def migrate_config(self, config: MutableMapping[str, Any]) -> None:
         """

airbyte_cdk/sources/file_based/file_based_source.py

@@ -77,9 +77,9 @@ from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
 from airbyte_cdk.utils.analytics_message import create_analytics_message
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
-DEFAULT_CONCURRENCY = 1.0  # TODO: Revert me: 100
-MAX_CONCURRENCY = 1  # TODO: Revert me: 100
-INITIAL_N_PARTITIONS = max(MAX_CONCURRENCY // 2, 1)
+DEFAULT_CONCURRENCY = 100
+MAX_CONCURRENCY = 100
+INITIAL_N_PARTITIONS = MAX_CONCURRENCY // 2
 IDENTITIES_STREAM = "identities"
 
 

airbyte_cdk/sources/source.py

@@ -12,7 +12,9 @@ from airbyte_cdk.models import (
     AirbyteCatalog,
     AirbyteMessage,
     AirbyteStateMessage,
+    AirbyteStateMessageSerializer,
     ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
 )
 
 TState = TypeVar("TState")
@@ -70,7 +72,7 @@ class Source(
             state_obj = BaseConnector._read_json_file(state_path)
             if state_obj:
                 for state in state_obj:  # type: ignore  # `isinstance(state_obj, List)` ensures that this is a list
-                    parsed_message = AirbyteStateMessage.model_validate(state)
+                    parsed_message = AirbyteStateMessageSerializer.load(state)
                     if (
                         not parsed_message.stream
                         and not parsed_message.data
@@ -85,7 +87,7 @@ class Source(
     # can be overridden to change an input catalog
     @classmethod
     def read_catalog(cls, catalog_path: str) -> ConfiguredAirbyteCatalog:
-        return ConfiguredAirbyteCatalog.model_validate(cls._read_json_file(catalog_path))
+        return ConfiguredAirbyteCatalogSerializer.load(cls._read_json_file(catalog_path))
 
     @property
     def name(self) -> str:

airbyte_cdk/sources/streams/checkpoint/cursor.py

@@ -62,12 +62,6 @@ class Cursor(ABC):
         Evaluating if a record should be synced allows for filtering and stop condition on pagination
         """
 
-    @abstractmethod
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        """
-        Evaluating which record is greater in terms of cursor. This is used to avoid having to capture all the records to close a slice
-        """
-
     @abstractmethod
     def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
         """

airbyte_cdk/sources/streams/checkpoint/resumable_full_refresh_cursor.py

@@ -40,12 +40,6 @@ class ResumableFullRefreshCursor(Cursor):
         """
         return True
 
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        """
-        RFR record don't have ordering to be compared between one another.
-        """
-        return False
-
     def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
         # A top-level RFR cursor only manages the state of a single partition
         return self._cursor

airbyte_cdk/sources/streams/checkpoint/substream_resumable_full_refresh_cursor.py

@@ -89,12 +89,6 @@ class SubstreamResumableFullRefreshCursor(Cursor):
         """
         return True
 
-    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
-        """
-        RFR record don't have ordering to be compared between one another.
-        """
-        return False
-
     def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
         if not stream_slice:
             raise ValueError("A partition needs to be provided in order to extract a state")

airbyte_cdk/sources/streams/http/http_client.py

@@ -14,12 +14,12 @@ import requests_cache
 from requests.auth import AuthBase
 
 from airbyte_cdk.models import (
+    AirbyteMessageSerializer,
     AirbyteStreamStatus,
     AirbyteStreamStatusReason,
     AirbyteStreamStatusReasonType,
     Level,
     StreamDescriptor,
-    ab_message_to_string,
 )
 from airbyte_cdk.sources.http_config import MAX_CONNECTION_POOL_SIZE
 from airbyte_cdk.sources.message import MessageRepository
@@ -396,11 +396,13 @@ class HttpClient:
         if error_resolution.response_action == ResponseAction.RATE_LIMITED:
             # TODO: Update to handle with message repository when concurrent message repository is ready
             reasons = [AirbyteStreamStatusReason(type=AirbyteStreamStatusReasonType.RATE_LIMITED)]
-            message = ab_message_to_string(
-                stream_status_as_airbyte_message(
-                    StreamDescriptor(name=self._name), AirbyteStreamStatus.RUNNING, reasons
+            message = orjson.dumps(
+                AirbyteMessageSerializer.dump(
+                    stream_status_as_airbyte_message(
+                        StreamDescriptor(name=self._name), AirbyteStreamStatus.RUNNING, reasons
+                    )
                 )
-            )
+            ).decode()
 
             # Simply printing the stream status is a temporary solution and can cause future issues. Currently, the _send method is
             # wrapped with backoff decorators, and we can only emit messages by iterating record_iterator in the abstract source at the

airbyte_cdk/sources/streams/permissions/identities_stream.py

@@ -6,7 +6,7 @@ import traceback
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Iterable, List, Mapping, MutableMapping, Optional
 
-from airbyte_protocol.models import SyncMode
+from airbyte_protocol_dataclasses.models import SyncMode
 
 from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level
 from airbyte_cdk.models import Type as MessageType

airbyte_cdk/sql/shared/sql_processor.py

@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, final
 import pandas as pd
 import sqlalchemy
 import ulid
-from airbyte_protocol.models import AirbyteStateMessage
+from airbyte_protocol_dataclasses.models import AirbyteStateMessage
 from pandas import Index
 from pydantic import BaseModel, Field
 from sqlalchemy import Column, Table, and_, create_engine, insert, null, select, text, update

airbyte_cdk/test/catalog_builder.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, List, Union, overload
 from airbyte_cdk.models import (
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
+    ConfiguredAirbyteStreamSerializer,
     SyncMode,
 )
 
@@ -41,7 +42,7 @@ class ConfiguredAirbyteStreamBuilder:
         return self
 
     def build(self) -> ConfiguredAirbyteStream:
-        return ConfiguredAirbyteStream(**self._stream)
+        return ConfiguredAirbyteStreamSerializer.load(self._stream)
 
 
 class CatalogBuilder: