airbyte-cdk 6.60.0.post35.dev16509779638__py3-none-any.whl → 6.60.1__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/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/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/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:

airbyte_cdk/test/entrypoint_wrapper.py

@@ -17,7 +17,6 @@ than that, there are integrations point that are annoying to integrate with usin
 import json
 import logging
 import re
-import sys
 import tempfile
 import traceback
 from collections import deque
@@ -29,6 +28,7 @@ from typing import Any, List, Literal, Optional, Union, final, overload
 
 import orjson
 from pydantic import ValidationError as V2ValidationError
+from serpyco_rs import SchemaValidationError
 
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
 from airbyte_cdk.exception_handler import assemble_uncaught_exception
@@ -36,31 +36,20 @@ from airbyte_cdk.logger import AirbyteLogFormatter
 from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
+    AirbyteMessageSerializer,
     AirbyteStateMessage,
+    AirbyteStateMessageSerializer,
     AirbyteStreamState,
     AirbyteStreamStatus,
     ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
     Level,
     TraceType,
     Type,
-    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,
 )
 from airbyte_cdk.sources import Source
 from airbyte_cdk.test.models.scenario import ExpectedOutcome
 
-JsonValidationErrors: tuple[type[Exception], ...] = (orjson.JSONDecodeError,)
-# Conditionally import and create a union type for exception handling
-if sys.platform != "emscripten":
-    from serpyco_rs import SchemaValidationError
-
-    JsonValidationErrors = (orjson.JSONDecodeError, SchemaValidationError)
-
 
 class AirbyteEntrypointException(Exception):
     """Exception raised for errors in the AirbyteEntrypoint execution.
@@ -128,8 +117,8 @@ class EntrypointOutput:
     @staticmethod
     def _parse_message(message: str) -> AirbyteMessage:
         try:
-            return ab_message_from_string(message)
-        except JsonValidationErrors:
+            return AirbyteMessageSerializer.load(orjson.loads(message))
+        except (orjson.JSONDecodeError, SchemaValidationError):
             # The platform assumes that logs that are not of AirbyteMessage format are log messages
             return AirbyteMessage(
                 type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message=message)
@@ -452,7 +441,7 @@ def read(
         config_file = make_file(tmp_directory_path / "config.json", config)
         catalog_file = make_file(
             tmp_directory_path / "catalog.json",
-            ab_configured_catalog_to_string(catalog),
+            orjson.dumps(ConfiguredAirbyteCatalogSerializer.dump(catalog)).decode(),
         )
         args = [
             "read",
@@ -464,13 +453,15 @@ def read(
         if debug:
             args.append("--debug")
         if state is not None:
-            args.extend([
-                "--state",
-                make_file(
-                    tmp_directory_path / "state.json",
-                    f"[{','.join([ab_state_message_to_string(stream_state) for stream_state in state])}]",
-                ),
-            ])
+            args.extend(
+                [
+                    "--state",
+                    make_file(
+                        tmp_directory_path / "state.json",
+                        f"[{','.join([orjson.dumps(AirbyteStateMessageSerializer.dump(stream_state)).decode() for stream_state in state])}]",
+                    ),
+                ]
+            )
 
         return _run_command(
             source,

airbyte_cdk/utils/datetime_helpers.py

@@ -86,6 +86,7 @@ from typing import Any, Optional, Union, overload
 
 from dateutil import parser
 from typing_extensions import Never
+from whenever import Instant, LocalDateTime, ZonedDateTime
 
 
 class AirbyteDateTime(datetime):
@@ -137,18 +138,6 @@ class AirbyteDateTime(datetime):
             dt.tzinfo or timezone.utc,
         )
 
-    @classmethod
-    def from_timestamp(cls, timestamp: float) -> "AirbyteDateTime":
-        """Creates an AirbyteDateTime from a Unix timestamp in seconds.
-
-        Args:
-            timestamp: A Unix timestamp in seconds (float).
-
-        Returns:
-            AirbyteDateTime: A new timezone-aware datetime instance (UTC).
-        """
-        return AirbyteDateTime.from_datetime(datetime.fromtimestamp(timestamp, tz=timezone.utc))
-
     def to_datetime(self) -> datetime:
         """Converts this AirbyteDateTime to a standard datetime object.
 
@@ -411,7 +400,8 @@ def ab_datetime_parse(dt_str: str | int) -> AirbyteDateTime:
                 raise ValueError("Timestamp cannot be negative")
             if len(str(abs(timestamp))) > 10:
                 raise ValueError("Timestamp value too large")
-            return AirbyteDateTime.from_timestamp(timestamp)
+            instant = Instant.from_timestamp(timestamp)
+            return AirbyteDateTime.from_datetime(instant.py_datetime())
 
         if not isinstance(dt_str, str):
             raise ValueError(
@@ -424,7 +414,8 @@ def ab_datetime_parse(dt_str: str | int) -> AirbyteDateTime:
                 year, month, day = map(int, dt_str.split("-"))
                 if not (1 <= month <= 12 and 1 <= day <= 31):
                     raise ValueError(f"Invalid date format: {dt_str}")
-                return AirbyteDateTime(year, month, day, 0, 0, 0)
+                instant = Instant.from_utc(year, month, day, 0, 0, 0)
+                return AirbyteDateTime.from_datetime(instant.py_datetime())
             except (ValueError, TypeError):
                 raise ValueError(f"Invalid date format: {dt_str}")
 

airbyte_cdk/utils/traced_exception.py

@@ -11,12 +11,12 @@ from airbyte_cdk.models import (
     AirbyteConnectionStatus,
     AirbyteErrorTraceMessage,
     AirbyteMessage,
+    AirbyteMessageSerializer,
     AirbyteTraceMessage,
     FailureType,
     Status,
     StreamDescriptor,
     TraceType,
-    ab_message_to_string,
 )
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
@@ -95,7 +95,7 @@ class AirbyteTracedException(Exception):
         Prints the exception as an AirbyteTraceMessage.
         Note that this will be called automatically on uncaught exceptions when using the airbyte_cdk entrypoint.
         """
-        message = ab_message_to_string(self.as_airbyte_message())
+        message = orjson.dumps(AirbyteMessageSerializer.dump(self.as_airbyte_message())).decode()
         filtered_message = filter_secrets(message)
         print(filtered_message)