airbyte-cdk 0.72.1__py3-none-any.whl → 6.13.1.dev4106__py3-none-any.whl

airbyte_cdk/entrypoint.py

@@ -5,6 +5,7 @@
 import argparse
 import importlib
 import ipaddress
+import json
 import logging
 import os.path
 import socket
@@ -12,35 +13,48 @@ import sys
 import tempfile
 from collections import defaultdict
 from functools import wraps
-from typing import Any, DefaultDict, Iterable, List, Mapping, MutableMapping, Optional, Union
+from typing import Any, DefaultDict, Iterable, List, Mapping, Optional
 from urllib.parse import urlparse
 
+import orjson
 import requests
+from requests import PreparedRequest, Response, Session
+
 from airbyte_cdk.connector import TConfig
 from airbyte_cdk.exception_handler import init_uncaught_exception_handler
 from airbyte_cdk.logger import init_logger
-from airbyte_cdk.models import AirbyteMessage, FailureType, Status, Type
-from airbyte_cdk.models.airbyte_protocol import AirbyteStateStats, ConnectorSpecification  # type: ignore [attr-defined]
+from airbyte_cdk.models import (
+    AirbyteConnectionStatus,
+    AirbyteMessage,
+    AirbyteMessageSerializer,
+    AirbyteStateStats,
+    ConnectorSpecification,
+    FailureType,
+    Status,
+    Type,
+)
 from airbyte_cdk.sources import Source
 from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
 from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit, split_config
+
+# from airbyte_cdk.utils import PrintBuffer, is_cloud_environment, message_utils  # add PrintBuffer back once fixed
 from airbyte_cdk.utils import is_cloud_environment, message_utils
 from airbyte_cdk.utils.airbyte_secrets_utils import get_secrets, update_secrets
 from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
-from requests import PreparedRequest, Response, Session
 
 logger = init_logger("airbyte")
 
 VALID_URL_SCHEMES = ["https"]
 CLOUD_DEPLOYMENT_MODE = "cloud"
+_HAS_LOGGED_FOR_SERIALIZATION_ERROR = False
 
 
 class AirbyteEntrypoint(object):
     def __init__(self, source: Source):
         init_uncaught_exception_handler(logger)
 
-        # deployment mode is read when instantiating the entrypoint because it is the common path shared by syncs and connector builder test requests
+        # Deployment mode is read when instantiating the entrypoint because it is the common path shared by syncs and connector builder test requests
         if is_cloud_environment():
             _init_internal_request_filter()
 
@@ -51,33 +65,54 @@ class AirbyteEntrypoint(object):
     def parse_args(args: List[str]) -> argparse.Namespace:
         # set up parent parsers
         parent_parser = argparse.ArgumentParser(add_help=False)
-        parent_parser.add_argument("--debug", action="store_true", help="enables detailed debug logs related to the sync")
+        parent_parser.add_argument(
+            "--debug", action="store_true", help="enables detailed debug logs related to the sync"
+        )
         main_parser = argparse.ArgumentParser()
         subparsers = main_parser.add_subparsers(title="commands", dest="command")
 
         # spec
-        subparsers.add_parser("spec", help="outputs the json configuration specification", parents=[parent_parser])
+        subparsers.add_parser(
+            "spec", help="outputs the json configuration specification", parents=[parent_parser]
+        )
 
         # check
-        check_parser = subparsers.add_parser("check", help="checks the config can be used to connect", parents=[parent_parser])
+        check_parser = subparsers.add_parser(
+            "check", help="checks the config can be used to connect", parents=[parent_parser]
+        )
         required_check_parser = check_parser.add_argument_group("required named arguments")
-        required_check_parser.add_argument("--config", type=str, required=True, help="path to the json configuration file")
+        required_check_parser.add_argument(
+            "--config", type=str, required=True, help="path to the json configuration file"
+        )
 
         # discover
         discover_parser = subparsers.add_parser(
-            "discover", help="outputs a catalog describing the source's schema", parents=[parent_parser]
+            "discover",
+            help="outputs a catalog describing the source's schema",
+            parents=[parent_parser],
         )
         required_discover_parser = discover_parser.add_argument_group("required named arguments")
-        required_discover_parser.add_argument("--config", type=str, required=True, help="path to the json configuration file")
+        required_discover_parser.add_argument(
+            "--config", type=str, required=True, help="path to the json configuration file"
+        )
 
         # read
-        read_parser = subparsers.add_parser("read", help="reads the source and outputs messages to STDOUT", parents=[parent_parser])
+        read_parser = subparsers.add_parser(
+            "read", help="reads the source and outputs messages to STDOUT", parents=[parent_parser]
+        )
 
-        read_parser.add_argument("--state", type=str, required=False, help="path to the json-encoded state file")
+        read_parser.add_argument(
+            "--state", type=str, required=False, help="path to the json-encoded state file"
+        )
         required_read_parser = read_parser.add_argument_group("required named arguments")
-        required_read_parser.add_argument("--config", type=str, required=True, help="path to the json configuration file")
         required_read_parser.add_argument(
-            "--catalog", type=str, required=True, help="path to the catalog used to determine which data to read"
+            "--config", type=str, required=True, help="path to the json configuration file"
+        )
+        required_read_parser.add_argument(
+            "--catalog",
+            type=str,
+            required=True,
+            help="path to the catalog used to determine which data to read",
         )
 
         return main_parser.parse_args(args)
@@ -96,12 +131,19 @@ class AirbyteEntrypoint(object):
 
         source_spec: ConnectorSpecification = self.source.spec(self.logger)
         try:
-            with tempfile.TemporaryDirectory() as temp_dir:
-                os.environ[ENV_REQUEST_CACHE_PATH] = temp_dir  # set this as default directory for request_cache to store *.sqlite files
+            with tempfile.TemporaryDirectory(
+                # Cleanup can fail on Windows due to file locks. Ignore if so,
+                # rather than failing the whole process.
+                ignore_cleanup_errors=True,
+            ) as temp_dir:
+                os.environ[ENV_REQUEST_CACHE_PATH] = (
+                    temp_dir  # set this as default directory for request_cache to store *.sqlite files
+                )
                 if cmd == "spec":
                     message = AirbyteMessage(type=Type.SPEC, spec=source_spec)
                     yield from [
-                        self.airbyte_message_to_string(queued_message) for queued_message in self._emit_queued_messages(self.source)
+                        self.airbyte_message_to_string(queued_message)
+                        for queued_message in self._emit_queued_messages(self.source)
                     ]
                     yield self.airbyte_message_to_string(message)
                 else:
@@ -109,34 +151,68 @@ class AirbyteEntrypoint(object):
                     config = self.source.configure(raw_config, temp_dir)
 
                     yield from [
-                        self.airbyte_message_to_string(queued_message) for queued_message in self._emit_queued_messages(self.source)
+                        self.airbyte_message_to_string(queued_message)
+                        for queued_message in self._emit_queued_messages(self.source)
                     ]
                     if cmd == "check":
-                        yield from map(AirbyteEntrypoint.airbyte_message_to_string, self.check(source_spec, config))
+                        yield from map(
+                            AirbyteEntrypoint.airbyte_message_to_string,
+                            self.check(source_spec, config),
+                        )
                     elif cmd == "discover":
-                        yield from map(AirbyteEntrypoint.airbyte_message_to_string, self.discover(source_spec, config))
+                        yield from map(
+                            AirbyteEntrypoint.airbyte_message_to_string,
+                            self.discover(source_spec, config),
+                        )
                     elif cmd == "read":
                         config_catalog = self.source.read_catalog(parsed_args.catalog)
                         state = self.source.read_state(parsed_args.state)
 
-                        yield from map(AirbyteEntrypoint.airbyte_message_to_string, self.read(source_spec, config, config_catalog, state))
+                        yield from map(
+                            AirbyteEntrypoint.airbyte_message_to_string,
+                            self.read(source_spec, config, config_catalog, state),
+                        )
                     else:
                         raise Exception("Unexpected command " + cmd)
         finally:
-            yield from [self.airbyte_message_to_string(queued_message) for queued_message in self._emit_queued_messages(self.source)]
+            yield from [
+                self.airbyte_message_to_string(queued_message)
+                for queued_message in self._emit_queued_messages(self.source)
+            ]
 
-    def check(self, source_spec: ConnectorSpecification, config: TConfig) -> Iterable[AirbyteMessage]:
+    def check(
+        self, source_spec: ConnectorSpecification, config: TConfig
+    ) -> Iterable[AirbyteMessage]:
         self.set_up_secret_filter(config, source_spec.connectionSpecification)
         try:
             self.validate_connection(source_spec, config)
         except AirbyteTracedException as traced_exc:
             connection_status = traced_exc.as_connection_status_message()
+            # The platform uses the exit code to surface unexpected failures so we raise the exception if the failure type not a config error
+            # If the failure is not exceptional, we'll emit a failed connection status message and return
+            if traced_exc.failure_type != FailureType.config_error:
+                raise traced_exc
             if connection_status:
                 yield from self._emit_queued_messages(self.source)
                 yield connection_status
                 return
 
-        check_result = self.source.check(self.logger, config)
+        try:
+            check_result = self.source.check(self.logger, config)
+        except AirbyteTracedException as traced_exc:
+            yield traced_exc.as_airbyte_message()
+            # The platform uses the exit code to surface unexpected failures so we raise the exception if the failure type not a config error
+            # If the failure is not exceptional, we'll emit a failed connection status message and return
+            if traced_exc.failure_type != FailureType.config_error:
+                raise traced_exc
+            else:
+                yield AirbyteMessage(
+                    type=Type.CONNECTION_STATUS,
+                    connectionStatus=AirbyteConnectionStatus(
+                        status=Status.FAILED, message=traced_exc.message
+                    ),
+                )
+                return
         if check_result.status == Status.SUCCEEDED:
             self.logger.info("Check succeeded")
         else:
@@ -145,7 +221,9 @@ class AirbyteEntrypoint(object):
         yield from self._emit_queued_messages(self.source)
         yield AirbyteMessage(type=Type.CONNECTION_STATUS, connectionStatus=check_result)
 
-    def discover(self, source_spec: ConnectorSpecification, config: TConfig) -> Iterable[AirbyteMessage]:
+    def discover(
+        self, source_spec: ConnectorSpecification, config: TConfig
+    ) -> Iterable[AirbyteMessage]:
         self.set_up_secret_filter(config, source_spec.connectionSpecification)
         if self.source.check_config_against_spec:
             self.validate_connection(source_spec, config)
@@ -155,32 +233,48 @@ class AirbyteEntrypoint(object):
         yield AirbyteMessage(type=Type.CATALOG, catalog=catalog)
 
     def read(
-        self, source_spec: ConnectorSpecification, config: TConfig, catalog: Any, state: Union[list[Any], MutableMapping[str, Any]]
+        self, source_spec: ConnectorSpecification, config: TConfig, catalog: Any, state: list[Any]
     ) -> Iterable[AirbyteMessage]:
         self.set_up_secret_filter(config, source_spec.connectionSpecification)
         if self.source.check_config_against_spec:
             self.validate_connection(source_spec, config)
 
-        stream_message_counter: DefaultDict[HashableStreamDescriptor, int] = defaultdict(int)
+        # The Airbyte protocol dictates that counts be expressed as float/double to better protect against integer overflows
+        stream_message_counter: DefaultDict[HashableStreamDescriptor, float] = defaultdict(float)
         for message in self.source.read(self.logger, config, catalog, state):
             yield self.handle_record_counts(message, stream_message_counter)
         for message in self._emit_queued_messages(self.source):
             yield self.handle_record_counts(message, stream_message_counter)
 
     @staticmethod
-    def handle_record_counts(message: AirbyteMessage, stream_message_count: DefaultDict[HashableStreamDescriptor, int]) -> AirbyteMessage:
-        if message.type == Type.RECORD:
-            stream_message_count[message_utils.get_stream_descriptor(message)] += 1
-
-        elif message.type == Type.STATE:
-            stream_descriptor = message_utils.get_stream_descriptor(message)
-
-            # Set record count from the counter onto the state message
-            message.state.sourceStats = message.state.sourceStats or AirbyteStateStats()
-            message.state.sourceStats.recordCount = stream_message_count.get(stream_descriptor, 0)
+    def handle_record_counts(
+        message: AirbyteMessage, stream_message_count: DefaultDict[HashableStreamDescriptor, float]
+    ) -> AirbyteMessage:
+        match message.type:
+            case Type.RECORD:
+                if message.record is None:
+                    raise ValueError("Record message must have a record attribute")
+
+                stream_message_count[
+                    HashableStreamDescriptor(
+                        name=message.record.stream,  # type: ignore[union-attr] # record has `stream`
+                        namespace=message.record.namespace,  # type: ignore[union-attr] # record has `namespace`
+                    )
+                ] += 1.0
+            case Type.STATE:
+                if message.state is None:
+                    raise ValueError("State message must have a state attribute")
+
+                stream_descriptor = message_utils.get_stream_descriptor(message)
+
+                # Set record count from the counter onto the state message
+                message.state.sourceStats = message.state.sourceStats or AirbyteStateStats()  # type: ignore[union-attr] # state has `sourceStats`
+                message.state.sourceStats.recordCount = stream_message_count.get(  # type: ignore[union-attr] # state has `sourceStats`
+                    stream_descriptor, 0.0
+                )
 
-            # Reset the counter
-            stream_message_count[stream_descriptor] = 0
+                # Reset the counter
+                stream_message_count[stream_descriptor] = 0.0
         return message
 
     @staticmethod
@@ -198,8 +292,18 @@ class AirbyteEntrypoint(object):
         update_secrets(config_secrets)
 
     @staticmethod
-    def airbyte_message_to_string(airbyte_message: AirbyteMessage) -> Any:
-        return airbyte_message.json(exclude_unset=True)
+    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]:
@@ -231,10 +335,13 @@ class AirbyteEntrypoint(object):
 def launch(source: Source, args: List[str]) -> None:
     source_entrypoint = AirbyteEntrypoint(source)
     parsed_args = source_entrypoint.parse_args(args)
+    # temporarily removes the PrintBuffer because we're seeing weird print behavior for concurrent syncs
+    # Refer to: https://github.com/airbytehq/oncall/issues/6235
+    # with PrintBuffer():
     for message in source_entrypoint.run(parsed_args):
         # simply printing is creating issues for concurrent CDK as Python uses different two instructions to print: one for the message and
         # the other for the break line. Adding `\n` to the message ensure that both are printed at the same time
-        print(f"{message}\n", end="")
+        print(f"{message}\n", end="", flush=True)
 
 
 def _init_internal_request_filter() -> None:
@@ -254,7 +361,9 @@ def _init_internal_request_filter() -> None:
             )
 
         if not parsed_url.hostname:
-            raise requests.exceptions.InvalidURL("Invalid URL specified: The endpoint that data is being requested from is not a valid URL")
+            raise requests.exceptions.InvalidURL(
+                "Invalid URL specified: The endpoint that data is being requested from is not a valid URL"
+            )
 
         try:
             is_private = _is_private_url(parsed_url.hostname, parsed_url.port)  # type: ignore [arg-type]

airbyte_cdk/exception_handler.py

@@ -5,12 +5,15 @@
 import logging
 import sys
 from types import TracebackType
-from typing import Any, Optional
+from typing import Any, List, Mapping, Optional
 
+from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
-def assemble_uncaught_exception(exception_type: type[BaseException], exception_value: BaseException) -> AirbyteTracedException:
+def assemble_uncaught_exception(
+    exception_type: type[BaseException], exception_value: BaseException
+) -> AirbyteTracedException:
     if issubclass(exception_type, AirbyteTracedException):
         return exception_value  # type: ignore  # validated as part of the previous line
     return AirbyteTracedException.from_exception(exception_value)
@@ -22,7 +25,11 @@ def init_uncaught_exception_handler(logger: logging.Logger) -> None:
     printed to the console without having secrets removed.
     """
 
-    def hook_fn(exception_type: type[BaseException], exception_value: BaseException, traceback_: Optional[TracebackType]) -> Any:
+    def hook_fn(
+        exception_type: type[BaseException],
+        exception_value: BaseException,
+        traceback_: Optional[TracebackType],
+    ) -> Any:
         # For developer ergonomics, we want to see the stack trace in the logs when we do a ctrl-c
         if issubclass(exception_type, KeyboardInterrupt):
             sys.__excepthook__(exception_type, exception_value, traceback_)
@@ -36,3 +43,14 @@ def init_uncaught_exception_handler(logger: logging.Logger) -> None:
         traced_exc.emit_message()
 
     sys.excepthook = hook_fn
+
+
+def generate_failed_streams_error_message(stream_failures: Mapping[str, List[Exception]]) -> str:
+    failures = "\n".join(
+        [
+            f"{stream}: {filter_secrets(exception.__repr__())}"
+            for stream, exceptions in stream_failures.items()
+            for exception in exceptions
+        ]
+    )
+    return f"During the sync, the following streams did not sync successfully: {failures}"

airbyte_cdk/logger.py

@@ -5,12 +5,18 @@
 import json
 import logging
 import logging.config
-import traceback
-from typing import Tuple
+from typing import Any, Callable, Mapping, Optional, Tuple
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage
+import orjson
+
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteMessageSerializer,
+    Level,
+    Type,
+)
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
-from deprecated import deprecated
 
 LOGGING_CONFIG = {
     "version": 1,
@@ -31,7 +37,7 @@ LOGGING_CONFIG = {
 }
 
 
-def init_logger(name: str = None):
+def init_logger(name: Optional[str] = None) -> logging.Logger:
     """Initial set up of logger"""
     logger = logging.getLogger(name)
     logger.setLevel(logging.INFO)
@@ -39,33 +45,43 @@ def init_logger(name: str = None):
     return logger
 
 
+def lazy_log(logger: logging.Logger, level: int, lazy_log_provider: Callable[[], str]) -> None:
+    """
+    This method ensure that the processing of the log message is only done if the logger is enabled for the log level.
+    """
+    if logger.isEnabledFor(level):
+        logger.log(level, lazy_log_provider())
+
+
 class AirbyteLogFormatter(logging.Formatter):
     """Output log records using AirbyteMessage"""
 
     # Transforming Python log levels to Airbyte protocol log levels
     level_mapping = {
-        logging.FATAL: "FATAL",
-        logging.ERROR: "ERROR",
-        logging.WARNING: "WARN",
-        logging.INFO: "INFO",
-        logging.DEBUG: "DEBUG",
+        logging.FATAL: Level.FATAL,
+        logging.ERROR: Level.ERROR,
+        logging.WARNING: Level.WARN,
+        logging.INFO: Level.INFO,
+        logging.DEBUG: Level.DEBUG,
     }
 
     def format(self, record: logging.LogRecord) -> str:
         """Return a JSON representation of the log message"""
         airbyte_level = self.level_mapping.get(record.levelno, "INFO")
-        if airbyte_level == "DEBUG":
+        if airbyte_level == Level.DEBUG:
             extras = self.extract_extra_args_from_record(record)
             debug_dict = {"type": "DEBUG", "message": record.getMessage(), "data": extras}
             return filter_secrets(json.dumps(debug_dict))
         else:
             message = super().format(record)
             message = filter_secrets(message)
-            log_message = AirbyteMessage(type="LOG", log=AirbyteLogMessage(level=airbyte_level, message=message))
-            return log_message.json(exclude_unset=True)
+            log_message = AirbyteMessage(
+                type=Type.LOG, log=AirbyteLogMessage(level=airbyte_level, message=message)
+            )
+            return orjson.dumps(AirbyteMessageSerializer.dump(log_message)).decode()
 
     @staticmethod
-    def extract_extra_args_from_record(record: logging.LogRecord):
+    def extract_extra_args_from_record(record: logging.LogRecord) -> Mapping[str, Any]:
         """
         The python logger conflates default args with extra args. We use an empty log record and set operations
         to isolate fields passed to the log record via extra by the developer.
@@ -88,33 +104,3 @@ def log_by_prefix(msg: str, default_level: str) -> Tuple[int, str]:
         rendered_message = msg
 
     return log_level, rendered_message
-
-
-@deprecated(version="0.1.47", reason="Use logging.getLogger('airbyte') instead")
-class AirbyteLogger:
-    def log(self, level, message):
-        log_record = AirbyteLogMessage(level=level, message=message)
-        log_message = AirbyteMessage(type="LOG", log=log_record)
-        print(log_message.json(exclude_unset=True))
-
-    def fatal(self, message):
-        self.log("FATAL", message)
-
-    def exception(self, message):
-        message = f"{message}\n{traceback.format_exc()}"
-        self.error(message)
-
-    def error(self, message):
-        self.log("ERROR", message)
-
-    def warn(self, message):
-        self.log("WARN", message)
-
-    def info(self, message):
-        self.log("INFO", message)
-
-    def debug(self, message):
-        self.log("DEBUG", message)
-
-    def trace(self, message):
-        self.log("TRACE", message)

airbyte_cdk/models/__init__.py

@@ -21,10 +21,13 @@ from .airbyte_protocol import (
     AirbyteRecordMessage,
     AirbyteStateBlob,
     AirbyteStateMessage,
+    AirbyteStateStats,
     AirbyteStateType,
     AirbyteStream,
     AirbyteStreamState,
     AirbyteStreamStatus,
+    AirbyteStreamStatusReason,
+    AirbyteStreamStatusReasonType,
     AirbyteStreamStatusTraceMessage,
     AirbyteTraceMessage,
     AuthFlowType,
@@ -36,22 +39,30 @@ from .airbyte_protocol import (
     FailureType,
     Level,
     OAuthConfigSpecification,
+    OauthConnectorInputSpecification,
     OrchestratorType,
+    State,
     Status,
     StreamDescriptor,
     SyncMode,
     TraceType,
     Type,
 )
+from .airbyte_protocol_serializers import (
+    AirbyteMessageSerializer,
+    AirbyteStateMessageSerializer,
+    AirbyteStreamStateSerializer,
+    ConfiguredAirbyteCatalogSerializer,
+    ConfiguredAirbyteStreamSerializer,
+    ConnectorSpecificationSerializer,
+)
 from .well_known_types import (
     BinaryData,
     Boolean,
     Date,
     Integer,
-    IntegerEnum,
     Model,
     Number,
-    NumberEnum,
     String,
     TimestampWithoutTimezone,
     TimestampWithTimezone,

airbyte_cdk/models/airbyte_protocol.py

@@ -2,4 +2,89 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from airbyte_protocol.models import *
+from dataclasses import InitVar, dataclass
+from typing import Annotated, Any, Dict, List, Mapping, Optional, Union
+
+from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
+from serpyco_rs.metadata import Alias
+
+from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
+
+# 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[Union[AirbyteFileTransferRecordMessage, 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]