cherry-pick-me: replace serializer classes with helper functions

Author: aaronsteers

airbyte_cdk/cli/source_declarative_manifest/_run.py

@@ -32,13 +32,13 @@
 from airbyte_cdk.models import (
     AirbyteErrorTraceMessage,
     AirbyteMessage,
-    AirbyteMessageSerializer,
     AirbyteStateMessage,
     AirbyteTraceMessage,
     ConfiguredAirbyteCatalog,
     ConnectorSpecificationSerializer,
     TraceType,
     Type,
+    ab_message_to_string,
 )
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
     ConcurrentDeclarativeSource,
@@ -105,21 +105,19 @@ def _get_local_yaml_source(args: list[str]) -> SourceLocalYaml:
         )
     except Exception as error:
         print(
-            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(),
-                            ),
+            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(),
                         ),
-                    )
-                )
-            ).decode()
+                    ),
+                ),
+            )
         )
         raise error
 
@@ -153,7 +151,7 @@ def handle_remote_manifest_command(args: list[str]) -> None:
         spec = ConnectorSpecificationSerializer.load(spec_obj)
 
         message = AirbyteMessage(type=Type.SPEC, spec=spec)
-        print(AirbyteEntrypoint.airbyte_message_to_string(message))
+        print(ab_message_to_string(message))
     else:
         source = create_declarative_source(args)
         launch(
@@ -215,21 +213,19 @@ def create_declarative_source(
         )
     except Exception as error:
         print(
-            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(),
-                            ),
+            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(),
                         ),
-                    )
-                )
-            ).decode()
+                    ),
+                ),
+            ),
         )
         raise error
 

airbyte_cdk/config_observation.py

@@ -16,9 +16,9 @@
     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(orjson.dumps(AirbyteMessageSerializer.dump(airbyte_message)).decode())
+    print(ab_message_to_string(airbyte_message))
 
 
 def create_connector_config_control_message(config: MutableMapping[str, Any]) -> AirbyteMessage:

airbyte_cdk/connector_builder/main.py

@@ -20,10 +20,10 @@
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
 from airbyte_cdk.models import (
     AirbyteMessage,
-    AirbyteMessageSerializer,
     AirbyteStateMessage,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteCatalogSerializer,
+    ab_message_to_string,
 )
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.source import Source
@@ -92,11 +92,9 @@ 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 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
+    return ab_message_to_string(
+        handle_connector_builder_request(source, command, config, catalog, state, limits)
+    )
 
 
 if __name__ == "__main__":
@@ -107,4 +105,4 @@ def handle_request(args: List[str]) -> str:
             exc, message=f"Error handling request: {str(exc)}"
         )
         m = error.as_airbyte_message()
-        print(orjson.dumps(AirbyteMessageSerializer.dump(m)).decode())
+        print(ab_message_to_string(m))

airbyte_cdk/destinations/destination.py

@@ -15,10 +15,11 @@
 from airbyte_cdk.exception_handler import init_uncaught_exception_handler
 from airbyte_cdk.models import (
     AirbyteMessage,
-    AirbyteMessageSerializer,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteCatalogSerializer,
     Type,
+    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
@@ -46,7 +47,7 @@ def _parse_input_stream(self, input_stream: io.TextIOWrapper) -> Iterable[Airbyt
         """Reads from stdin, converting to Airbyte messages"""
         for line in input_stream:
             try:
-                yield AirbyteMessageSerializer.load(orjson.loads(line))
+                yield ab_message_from_string(line)
             except orjson.JSONDecodeError:
                 logger.info(
                     f"ignoring input which can't be deserialized as Airbyte Message: {line}"
@@ -151,4 +152,4 @@ def run(self, args: List[str]) -> None:
         parsed_args = self.parse_args(args)
         output_messages = self.run_cmd(parsed_args)
         for message in output_messages:
-            print(orjson.dumps(AirbyteMessageSerializer.dump(message)).decode())
+            print(ab_message_to_string(message))

airbyte_cdk/entrypoint.py

@@ -26,12 +26,12 @@
 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,7 +47,6 @@
 
 VALID_URL_SCHEMES = ["https"]
 CLOUD_DEPLOYMENT_MODE = "cloud"
-_HAS_LOGGED_FOR_SERIALIZATION_ERROR = False
 
 
 class AirbyteEntrypoint(object):
@@ -178,41 +177,41 @@ def run(self, parsed_args: argparse.Namespace) -> Iterable[str]:
                 if cmd == "spec":
                     message = AirbyteMessage(type=Type.SPEC, spec=source_spec)
                     yield from [
-                        self.airbyte_message_to_string(queued_message)
+                        ab_message_to_string(queued_message)
                         for queued_message in self._emit_queued_messages(self.source)
                     ]
-                    yield self.airbyte_message_to_string(message)
+                    yield ab_message_to_string(message)
                 else:
                     raw_config = self.source.read_config(parsed_args.config)
                     config = self.source.configure(raw_config, temp_dir)
 
                     yield from [
-                        self.airbyte_message_to_string(queued_message)
+                        ab_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,
+                            ab_message_to_string,
                             self.check(source_spec, config),
                         )
                     elif cmd == "discover":
                         yield from map(
-                            AirbyteEntrypoint.airbyte_message_to_string,
+                            ab_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,
+                            ab_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)
+                ab_message_to_string(queued_message)
                 for queued_message in self._emit_queued_messages(self.source)
             ]
 
@@ -327,20 +326,6 @@ def set_up_secret_filter(config: TConfig, connection_specification: Mapping[str,
         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 @@
 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 @@ def format(self, record: logging.LogRecord) -> str:
             log_message = AirbyteMessage(
                 type=Type.LOG, log=AirbyteLogMessage(level=airbyte_level, message=message)
             )
-            return orjson.dumps(AirbyteMessageSerializer.dump(log_message)).decode()
+            return ab_message_to_string(log_message)
 
     @staticmethod
     def extract_extra_args_from_record(record: logging.LogRecord) -> Mapping[str, Any]:

airbyte_cdk/models/__init__.py

@@ -50,12 +50,11 @@
     Type,
 )
 from .airbyte_protocol_serializers import (
-    AirbyteMessageSerializer,
     AirbyteStateMessageSerializer,
-    AirbyteStreamStateSerializer,
     ConfiguredAirbyteCatalogSerializer,
-    ConfiguredAirbyteStreamSerializer,
     ConnectorSpecificationSerializer,
+    ab_message_from_string,
+    ab_message_to_string,
 )
 from .well_known_types import (
     BinaryData,