Feat: Make CDK runnable in WASM

.gitignore

@@ -23,3 +23,6 @@ dist
 
 # ASDF tool versions files
 .tool-versions
+
+# Pyodide build artifacts
+.pyodide-xbuildenv

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_connector_spec_from_string,
+    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
 
@@ -149,11 +147,10 @@ def handle_remote_manifest_command(args: list[str]) -> None:
                 "Could not find `spec.json` file for source-declarative-manifest"
             )
 
-        spec_obj = json.loads(json_spec)
-        spec = ConnectorSpecificationSerializer.load(spec_obj)
+        spec = ab_connector_spec_from_string(json_spec.decode("utf-8"))
 
         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 +212,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
 
@@ -298,10 +293,10 @@ def _register_components_from_file(filepath: str) -> None:
     spec.loader.exec_module(module)
 
 
-def run() -> None:
+def run(args: list[str] | None = None) -> None:
     """Run the `source-declarative-manifest` CLI.
 
     Args are detected from the command line, and the appropriate command is executed.
     """
-    args: list[str] = sys.argv[1:]
+    args = args or sys.argv[1:]
     handle_command(args)

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.py

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

airbyte_cdk/connector_builder/main.py

@@ -20,11 +20,10 @@
 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
@@ -53,7 +52,7 @@ def get_config_and_catalog_from_args(
 
     command = config["__command"]
     if command == "test_read":
-        catalog = ConfiguredAirbyteCatalogSerializer.load(BaseConnector.read_config(catalog_path))
+        catalog = ConfiguredAirbyteCatalog.model_validate(BaseConnector.read_config(catalog_path))
         state = Source.read_state(state_path)
     else:
         catalog = None
@@ -92,19 +91,27 @@ 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)
+    )
 
+def run(args: list[str] | None) -> None:
+    """Run the connector builder handler."""
+    if args is None:
+        args = sys.argv[1:]
 
-if __name__ == "__main__":
     try:
-        print(handle_request(sys.argv[1:]))
+        result = handle_request(args)
+        print(result)
     except Exception as exc:
         error = AirbyteTracedException.from_exception(
             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))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run(sys.argv[1:])
+    sys.exit(1)

airbyte_cdk/destinations/destination.py

@@ -15,10 +15,12 @@
 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
@@ -46,7 +48,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}"
@@ -58,9 +60,7 @@ def _run_write(
         configured_catalog_path: str,
         input_stream: io.TextIOWrapper,
     ) -> Iterable[AirbyteMessage]:
-        catalog = ConfiguredAirbyteCatalogSerializer.load(
-            orjson.loads(open(configured_catalog_path).read())
-        )
+        catalog = ab_configured_catalog_from_string(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 @@ 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,13 @@
     Type,
 )
 from .airbyte_protocol_serializers import (
-    AirbyteMessageSerializer,
-    AirbyteStateMessageSerializer,
-    AirbyteStreamStateSerializer,
-    ConfiguredAirbyteCatalogSerializer,
-    ConfiguredAirbyteStreamSerializer,
-    ConnectorSpecificationSerializer,
+    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 .well_known_types import (
     BinaryData,