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

airbyte_cdk/cli/source_declarative_manifest/_run.py

@@ -32,13 +32,13 @@ from airbyte_cdk.entrypoint import AirbyteEntrypoint, launch
 from airbyte_cdk.models import (
     AirbyteErrorTraceMessage,
     AirbyteMessage,
+    AirbyteMessageSerializer,
     AirbyteStateMessage,
     AirbyteTraceMessage,
     ConfiguredAirbyteCatalog,
+    ConnectorSpecificationSerializer,
     TraceType,
     Type,
-    ab_connector_spec_from_string,
-    ab_message_to_string,
 )
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
     ConcurrentDeclarativeSource,
@@ -105,19 +105,21 @@ def _get_local_yaml_source(args: list[str]) -> SourceLocalYaml:
         )
     except Exception as error:
         print(
-            ab_message_to_string(
-                AirbyteMessage(
-                    type=Type.TRACE,
-                    trace=AirbyteTraceMessage(
-                        type=TraceType.ERROR,
-                        emitted_at=ab_datetime_now().to_epoch_millis(),
-                        error=AirbyteErrorTraceMessage(
-                            message=f"Error starting the sync. This could be due to an invalid configuration or catalog. Please contact Support for assistance. Error: {error}",
-                            stack_trace=traceback.format_exc(),
+            orjson.dumps(
+                AirbyteMessageSerializer.dump(
+                    AirbyteMessage(
+                        type=Type.TRACE,
+                        trace=AirbyteTraceMessage(
+                            type=TraceType.ERROR,
+                            emitted_at=ab_datetime_now().to_epoch_millis(),
+                            error=AirbyteErrorTraceMessage(
+                                message=f"Error starting the sync. This could be due to an invalid configuration or catalog. Please contact Support for assistance. Error: {error}",
+                                stack_trace=traceback.format_exc(),
+                            ),
                         ),
-                    ),
-                ),
-            )
+                    )
+                )
+            ).decode()
         )
         raise error
 
@@ -147,10 +149,11 @@ def handle_remote_manifest_command(args: list[str]) -> None:
                 "Could not find `spec.json` file for source-declarative-manifest"
             )
 
-        spec = ab_connector_spec_from_string(json_spec.decode("utf-8"))
+        spec_obj = json.loads(json_spec)
+        spec = ConnectorSpecificationSerializer.load(spec_obj)
 
         message = AirbyteMessage(type=Type.SPEC, spec=spec)
-        print(ab_message_to_string(message))
+        print(AirbyteEntrypoint.airbyte_message_to_string(message))
     else:
         source = create_declarative_source(args)
         launch(
@@ -212,19 +215,21 @@ def create_declarative_source(
         )
     except Exception as error:
         print(
-            ab_message_to_string(
-                AirbyteMessage(
-                    type=Type.TRACE,
-                    trace=AirbyteTraceMessage(
-                        type=TraceType.ERROR,
-                        emitted_at=ab_datetime_now().to_epoch_millis(),
-                        error=AirbyteErrorTraceMessage(
-                            message=f"Error starting the sync. This could be due to an invalid configuration or catalog. Please contact Support for assistance. Error: {error}",
-                            stack_trace=traceback.format_exc(),
+            orjson.dumps(
+                AirbyteMessageSerializer.dump(
+                    AirbyteMessage(
+                        type=Type.TRACE,
+                        trace=AirbyteTraceMessage(
+                            type=TraceType.ERROR,
+                            emitted_at=ab_datetime_now().to_epoch_millis(),
+                            error=AirbyteErrorTraceMessage(
+                                message=f"Error starting the sync. This could be due to an invalid configuration or catalog. Please contact Support for assistance. Error: {error}",
+                                stack_trace=traceback.format_exc(),
+                            ),
                         ),
-                    ),
-                ),
-            ),
+                    )
+                )
+            ).decode()
         )
         raise error
 
@@ -293,10 +298,10 @@ def _register_components_from_file(filepath: str) -> None:
     spec.loader.exec_module(module)
 
 
-def run(args: list[str] | None = None) -> None:
+def run() -> None:
     """Run the `source-declarative-manifest` CLI.
 
     Args are detected from the command line, and the appropriate command is executed.
     """
-    args = args or sys.argv[1:]
+    args: list[str] = sys.argv[1:]
     handle_command(args)

airbyte_cdk/config_observation.py

@@ -16,9 +16,9 @@ from airbyte_cdk.models import (
     AirbyteControlConnectorConfigMessage,
     AirbyteControlMessage,
     AirbyteMessage,
+    AirbyteMessageSerializer,
     OrchestratorType,
     Type,
-    ab_message_to_string,
 )
 
 
@@ -92,7 +92,7 @@ def emit_configuration_as_airbyte_control_message(config: MutableMapping[str, An
     See the airbyte_cdk.sources.message package
     """
     airbyte_message = create_connector_config_control_message(config)
-    print(ab_message_to_string(airbyte_message))
+    print(orjson.dumps(AirbyteMessageSerializer.dump(airbyte_message)).decode())
 
 
 def create_connector_config_control_message(config: MutableMapping[str, Any]) -> AirbyteMessage:

airbyte_cdk/connector.py

@@ -15,6 +15,7 @@ import yaml
 from airbyte_cdk.models import (
     AirbyteConnectionStatus,
     ConnectorSpecification,
+    ConnectorSpecificationSerializer,
 )
 
 
@@ -94,7 +95,7 @@ class BaseConnector(ABC, Generic[TConfig]):
         else:
             raise FileNotFoundError("Unable to find spec.yaml or spec.json in the package.")
 
-        return ConnectorSpecification.model_validate(spec_obj)
+        return ConnectorSpecificationSerializer.load(spec_obj)
 
     @abstractmethod
     def check(self, logger: logging.Logger, config: TConfig) -> AirbyteConnectionStatus:

airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -108,7 +108,12 @@ def read_stream(
         stream_name = configured_catalog.streams[0].stream.name
 
         stream_read = test_read_handler.run_test_read(
-            source, config, configured_catalog, state, limits.max_records
+            source,
+            config,
+            configured_catalog,
+            stream_name,
+            state,
+            limits.max_records,
         )
 
         return AirbyteMessage(

airbyte_cdk/connector_builder/main.py

@@ -20,10 +20,11 @@ from airbyte_cdk.connector_builder.connector_builder_handler import (
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
 from airbyte_cdk.models import (
     AirbyteMessage,
+    AirbyteMessageSerializer,
     AirbyteStateMessage,
     ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
 )
-from airbyte_cdk.models.airbyte_protocol_serializers import ab_message_to_string
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.source import Source
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
@@ -52,7 +53,7 @@ def get_config_and_catalog_from_args(
 
     command = config["__command"]
     if command == "test_read":
-        catalog = ConfiguredAirbyteCatalog.model_validate(BaseConnector.read_config(catalog_path))
+        catalog = ConfiguredAirbyteCatalogSerializer.load(BaseConnector.read_config(catalog_path))
         state = Source.read_state(state_path)
     else:
         catalog = None
@@ -91,27 +92,19 @@ def handle_request(args: List[str]) -> str:
     command, config, catalog, state = get_config_and_catalog_from_args(args)
     limits = get_limits(config)
     source = create_source(config, limits)
-    return ab_message_to_string(
-        handle_connector_builder_request(source, command, config, catalog, state, limits)
-    )
+    return orjson.dumps(
+        AirbyteMessageSerializer.dump(
+            handle_connector_builder_request(source, command, config, catalog, state, limits)
+        )
+    ).decode()  # type: ignore[no-any-return] # Serializer.dump() always returns AirbyteMessage
 
-def run(args: list[str] | None) -> None:
-    """Run the connector builder handler."""
-    if args is None:
-        args = sys.argv[1:]
 
+if __name__ == "__main__":
     try:
-        result = handle_request(args)
-        print(result)
+        print(handle_request(sys.argv[1:]))
     except Exception as exc:
         error = AirbyteTracedException.from_exception(
             exc, message=f"Error handling request: {str(exc)}"
         )
         m = error.as_airbyte_message()
-        print(ab_message_to_string(m))
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    run(sys.argv[1:])
-    sys.exit(1)
+        print(orjson.dumps(AirbyteMessageSerializer.dump(m)).decode())

airbyte_cdk/connector_builder/test_reader/helpers.py

@@ -269,6 +269,37 @@ def should_close_page_for_slice(at_least_one_page_in_group: bool, message: Airby
     return at_least_one_page_in_group and should_process_slice_descriptor(message)
 
 
+def is_page_http_request_for_different_stream(
+    json_message: Optional[Dict[str, Any]], stream_name: str
+) -> bool:
+    """
+    Determines whether a given JSON message represents a page HTTP request for a different stream.
+
+    This function checks if the provided JSON message is a page HTTP request, and if the stream name in the log is
+    different from the provided stream name.
+
+    This is needed because dynamic streams result in extra page HTTP requests for the dynamic streams that we want to ignore
+    when they do not match the stream that is being read.
+
+    Args:
+        json_message (Optional[Dict[str, Any]]): The JSON message to evaluate.
+        stream_name (str): The name of the stream to compare against.
+
+    Returns:
+        bool: True if the JSON message is a page HTTP request for a different stream, False otherwise.
+    """
+    if not json_message or not is_page_http_request(json_message):
+        return False
+
+    message_stream_name: str | None = (
+        json_message.get("airbyte_cdk", {}).get("stream", {}).get("name", None)
+    )
+    if message_stream_name is None:
+        return False
+
+    return message_stream_name != stream_name
+
+
 def is_page_http_request(json_message: Optional[Dict[str, Any]]) -> bool:
     """
     Determines whether a given JSON message represents a page HTTP request.

airbyte_cdk/connector_builder/test_reader/message_grouper.py

@@ -28,6 +28,7 @@ from .helpers import (
     is_async_auxiliary_request,
     is_config_update_message,
     is_log_message,
+    is_page_http_request_for_different_stream,
     is_record_message,
     is_state_message,
     is_trace_with_error,
@@ -44,6 +45,7 @@ def get_message_groups(
     schema_inferrer: SchemaInferrer,
     datetime_format_inferrer: DatetimeFormatInferrer,
     limit: int,
+    stream_name: str,
 ) -> MESSAGE_GROUPS:
     """
     Processes an iterator of AirbyteMessage objects to group and yield messages based on their type and sequence.
@@ -96,6 +98,9 @@ def get_message_groups(
     while records_count < limit and (message := next(messages, None)):
         json_message = airbyte_message_to_json(message)
 
+        if is_page_http_request_for_different_stream(json_message, stream_name):
+            continue
+
         if should_close_page(at_least_one_page_in_group, message, json_message):
             current_page_request, current_page_response = handle_current_page(
                 current_page_request,

airbyte_cdk/connector_builder/test_reader/reader.py

@@ -86,6 +86,7 @@ class TestReader:
         source: DeclarativeSource,
         config: Mapping[str, Any],
         configured_catalog: ConfiguredAirbyteCatalog,
+        stream_name: str,
         state: List[AirbyteStateMessage],
         record_limit: Optional[int] = None,
     ) -> StreamRead:
@@ -112,14 +113,17 @@ class TestReader:
 
         record_limit = self._check_record_limit(record_limit)
         # The connector builder currently only supports reading from a single stream at a time
-        stream = source.streams(config)[0]
+        streams = source.streams(config)
+        stream = next((stream for stream in streams if stream.name == stream_name), None)
 
         # get any deprecation warnings during the component creation
         deprecation_warnings: List[LogMessage] = source.deprecation_warnings()
 
         schema_inferrer = SchemaInferrer(
-            self._pk_to_nested_and_composite_field(stream.primary_key),
-            self._cursor_field_to_nested_and_composite_field(stream.cursor_field),
+            self._pk_to_nested_and_composite_field(stream.primary_key) if stream else None,
+            self._cursor_field_to_nested_and_composite_field(stream.cursor_field)
+            if stream
+            else None,
         )
         datetime_format_inferrer = DatetimeFormatInferrer()
 
@@ -128,6 +132,7 @@ class TestReader:
             schema_inferrer,
             datetime_format_inferrer,
             record_limit,
+            stream_name,
         )
 
         slices, log_messages, auxiliary_requests, latest_config_update = self._categorise_groups(

airbyte_cdk/destinations/destination.py

@@ -15,12 +15,10 @@ from airbyte_cdk.connector import Connector
 from airbyte_cdk.exception_handler import init_uncaught_exception_handler
 from airbyte_cdk.models import (
     AirbyteMessage,
+    AirbyteMessageSerializer,
     ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
     Type,
-    ab_configured_catalog_from_string,
-    ab_configured_catalog_to_string,
-    ab_message_from_string,
-    ab_message_to_string,
 )
 from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
@@ -48,7 +46,7 @@ class Destination(Connector, ABC):
         """Reads from stdin, converting to Airbyte messages"""
         for line in input_stream:
             try:
-                yield ab_message_from_string(line)
+                yield AirbyteMessageSerializer.load(orjson.loads(line))
             except orjson.JSONDecodeError:
                 logger.info(
                     f"ignoring input which can't be deserialized as Airbyte Message: {line}"
@@ -60,7 +58,9 @@ class Destination(Connector, ABC):
         configured_catalog_path: str,
         input_stream: io.TextIOWrapper,
     ) -> Iterable[AirbyteMessage]:
-        catalog = ab_configured_catalog_from_string(open(configured_catalog_path).read())
+        catalog = ConfiguredAirbyteCatalogSerializer.load(
+            orjson.loads(open(configured_catalog_path).read())
+        )
         input_messages = self._parse_input_stream(input_stream)
         logger.info("Begin writing to the destination...")
         yield from self.write(
@@ -151,4 +151,4 @@ class Destination(Connector, ABC):
         parsed_args = self.parse_args(args)
         output_messages = self.run_cmd(parsed_args)
         for message in output_messages:
-            print(ab_message_to_string(message))
+            print(orjson.dumps(AirbyteMessageSerializer.dump(message)).decode())

airbyte_cdk/entrypoint.py

@@ -26,12 +26,12 @@ from airbyte_cdk.logger import PRINT_BUFFER, init_logger
 from airbyte_cdk.models import (
     AirbyteConnectionStatus,
     AirbyteMessage,
+    AirbyteMessageSerializer,
     AirbyteStateStats,
     ConnectorSpecification,
     FailureType,
     Status,
     Type,
-    ab_message_to_string,
 )
 from airbyte_cdk.sources import Source
 from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
@@ -47,6 +47,7 @@ logger = init_logger("airbyte")
 
 VALID_URL_SCHEMES = ["https"]
 CLOUD_DEPLOYMENT_MODE = "cloud"
+_HAS_LOGGED_FOR_SERIALIZATION_ERROR = False
 
 
 class AirbyteEntrypoint(object):
@@ -177,26 +178,26 @@ class AirbyteEntrypoint(object):
                 if cmd == "spec":
                     message = AirbyteMessage(type=Type.SPEC, spec=source_spec)
                     yield from [
-                        ab_message_to_string(queued_message)
+                        self.airbyte_message_to_string(queued_message)
                         for queued_message in self._emit_queued_messages(self.source)
                     ]
-                    yield ab_message_to_string(message)
+                    yield self.airbyte_message_to_string(message)
                 else:
                     raw_config = self.source.read_config(parsed_args.config)
                     config = self.source.configure(raw_config, temp_dir)
 
                     yield from [
-                        ab_message_to_string(queued_message)
+                        self.airbyte_message_to_string(queued_message)
                         for queued_message in self._emit_queued_messages(self.source)
                     ]
                     if cmd == "check":
                         yield from map(
-                            ab_message_to_string,
+                            AirbyteEntrypoint.airbyte_message_to_string,
                             self.check(source_spec, config),
                         )
                     elif cmd == "discover":
                         yield from map(
-                            ab_message_to_string,
+                            AirbyteEntrypoint.airbyte_message_to_string,
                             self.discover(source_spec, config),
                         )
                     elif cmd == "read":
@@ -204,14 +205,14 @@ class AirbyteEntrypoint(object):
                         state = self.source.read_state(parsed_args.state)
 
                         yield from map(
-                            ab_message_to_string,
+                            AirbyteEntrypoint.airbyte_message_to_string,
                             self.read(source_spec, config, config_catalog, state),
                         )
                     else:
                         raise Exception("Unexpected command " + cmd)
         finally:
             yield from [
-                ab_message_to_string(queued_message)
+                self.airbyte_message_to_string(queued_message)
                 for queued_message in self._emit_queued_messages(self.source)
             ]
 
@@ -326,6 +327,20 @@ class AirbyteEntrypoint(object):
         config_secrets = get_secrets(connection_specification, config)
         update_secrets(config_secrets)
 
+    @staticmethod
+    def airbyte_message_to_string(airbyte_message: AirbyteMessage) -> str:
+        global _HAS_LOGGED_FOR_SERIALIZATION_ERROR
+        serialized_message = AirbyteMessageSerializer.dump(airbyte_message)
+        try:
+            return orjson.dumps(serialized_message).decode()
+        except Exception as exception:
+            if not _HAS_LOGGED_FOR_SERIALIZATION_ERROR:
+                logger.warning(
+                    f"There was an error during the serialization of an AirbyteMessage: `{exception}`. This might impact the sync performances."
+                )
+                _HAS_LOGGED_FOR_SERIALIZATION_ERROR = True
+            return json.dumps(serialized_message)
+
     @classmethod
     def extract_state(cls, args: List[str]) -> Optional[Any]:
         parsed_args = cls.parse_args(args)

airbyte_cdk/logger.py

@@ -12,9 +12,9 @@ import orjson
 from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
+    AirbyteMessageSerializer,
     Level,
     Type,
-    ab_message_to_string,
 )
 from airbyte_cdk.utils import PrintBuffer
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
@@ -81,7 +81,7 @@ class AirbyteLogFormatter(logging.Formatter):
             log_message = AirbyteMessage(
                 type=Type.LOG, log=AirbyteLogMessage(level=airbyte_level, message=message)
             )
-            return ab_message_to_string(log_message)
+            return orjson.dumps(AirbyteMessageSerializer.dump(log_message)).decode()
 
     @staticmethod
     def extract_extra_args_from_record(record: logging.LogRecord) -> Mapping[str, Any]:

airbyte_cdk/models/__init__.py

@@ -50,13 +50,12 @@ from .airbyte_protocol import (
     Type,
 )
 from .airbyte_protocol_serializers import (
-    ab_configured_catalog_from_string,
-    ab_configured_catalog_to_string,
-    ab_connector_spec_from_string,
-    ab_connector_spec_to_string,
-    ab_message_from_string,
-    ab_message_to_string,
-    ab_state_message_to_string,
+    AirbyteMessageSerializer,
+    AirbyteStateMessageSerializer,
+    AirbyteStreamStateSerializer,
+    ConfiguredAirbyteCatalogSerializer,
+    ConfiguredAirbyteStreamSerializer,
+    ConnectorSpecificationSerializer,
 )
 from .well_known_types import (
     BinaryData,

airbyte_cdk/models/airbyte_protocol.py

@@ -2,8 +2,87 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import sys
 from dataclasses import InitVar, dataclass
 from typing import Annotated, Any, Dict, List, Mapping, Optional, Union
 
-from airbyte_protocol.models import *  # noqa: F403  # Allow '*'
+from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
+from serpyco_rs.metadata import Alias
+
+# ruff: noqa: F405  # ignore fuzzy import issues with 'import *'
+
+
+@dataclass
+class AirbyteStateBlob:
+    """
+    A dataclass that dynamically sets attributes based on provided keyword arguments and positional arguments.
+    Used to "mimic" pydantic Basemodel with ConfigDict(extra='allow') option.
+
+    The `AirbyteStateBlob` class allows for flexible instantiation by accepting any number of keyword arguments
+    and positional arguments. These are used to dynamically update the instance's attributes. This class is useful
+    in scenarios where the attributes of an object are not known until runtime and need to be set dynamically.
+
+    Attributes:
+        kwargs (InitVar[Mapping[str, Any]]): A dictionary of keyword arguments used to set attributes dynamically.
+
+    Methods:
+        __init__(*args: Any, **kwargs: Any) -> None:
+            Initializes the `AirbyteStateBlob` by setting attributes from the provided arguments.
+
+        __eq__(other: object) -> bool:
+            Checks equality between two `AirbyteStateBlob` instances based on their internal dictionaries.
+            Returns `False` if the other object is not an instance of `AirbyteStateBlob`.
+    """
+
+    kwargs: InitVar[Mapping[str, Any]]
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        # Set any attribute passed in through kwargs
+        for arg in args:
+            self.__dict__.update(arg)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            False
+            if not isinstance(other, AirbyteStateBlob)
+            else bool(self.__dict__ == other.__dict__)
+        )
+
+
+# The following dataclasses have been redeclared to include the new version of AirbyteStateBlob
+@dataclass
+class AirbyteStreamState:
+    stream_descriptor: StreamDescriptor  # type: ignore [name-defined]
+    stream_state: Optional[AirbyteStateBlob] = None
+
+
+@dataclass
+class AirbyteGlobalState:
+    stream_states: List[AirbyteStreamState]
+    shared_state: Optional[AirbyteStateBlob] = None
+
+
+@dataclass
+class AirbyteStateMessage:
+    type: Optional[AirbyteStateType] = None  # type: ignore [name-defined]
+    stream: Optional[AirbyteStreamState] = None
+    global_: Annotated[AirbyteGlobalState | None, Alias("global")] = (
+        None  # "global" is a reserved keyword in python ⇒ Alias is used for (de-)serialization
+    )
+    data: Optional[Dict[str, Any]] = None
+    sourceStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]
+    destinationStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]
+
+
+@dataclass
+class AirbyteMessage:
+    type: Type  # type: ignore [name-defined]
+    log: Optional[AirbyteLogMessage] = None  # type: ignore [name-defined]
+    spec: Optional[ConnectorSpecification] = None  # type: ignore [name-defined]
+    connectionStatus: Optional[AirbyteConnectionStatus] = None  # type: ignore [name-defined]
+    catalog: Optional[AirbyteCatalog] = None  # type: ignore [name-defined]
+    record: Optional[AirbyteRecordMessage] = None  # type: ignore [name-defined]
+    state: Optional[AirbyteStateMessage] = None
+    trace: Optional[AirbyteTraceMessage] = None  # type: ignore [name-defined]
+    control: Optional[AirbyteControlMessage] = None  # type: ignore [name-defined]