fix(low-code): parse manifest and component schema with libyaml CSafeLoader when available

airbyte_cdk/cli/source_declarative_manifest/_run.py

@@ -26,7 +26,6 @@
 from typing import Any, cast
 
 import orjson
-import yaml
 
 from airbyte_cdk.entrypoint import AirbyteEntrypoint, launch
 from airbyte_cdk.models import (
@@ -50,6 +49,7 @@
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
 from airbyte_cdk.sources.source import TState
 from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
+from airbyte_cdk.utils.yaml_loader import safe_load_yaml
 
 
 class SourceLocalYaml(YamlDeclarativeSource):
@@ -272,7 +272,7 @@ def _parse_manifest_from_file(filepath: str) -> dict[str, Any] | None:
     """Extract and parse a manifest file specified in the args."""
     try:
         with open(filepath, "r", encoding="utf-8") as manifest_file:
-            manifest_content = yaml.safe_load(manifest_file)
+            manifest_content = safe_load_yaml(manifest_file)
             if manifest_content is None:
                 raise ValueError(f"Manifest file at {filepath} is empty")
             if not isinstance(manifest_content, dict):

airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py

@@ -11,7 +11,6 @@
 from typing import Any, Dict, Iterator, List, Mapping, Optional, Set, Union
 
 import orjson
-import yaml
 from jsonschema.exceptions import ValidationError
 from jsonschema.validators import validate
 from packaging.version import InvalidVersion, Version
@@ -75,6 +74,7 @@
     SliceLogger,
 )
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+from airbyte_cdk.utils.yaml_loader import safe_load_yaml
 
 
 def _get_declarative_component_schema() -> Dict[str, Any]:
@@ -83,7 +83,7 @@ def _get_declarative_component_schema() -> Dict[str, Any]:
             "airbyte_cdk", "sources/declarative/declarative_component_schema.yaml"
         )
         if raw_component_schema is not None:
-            declarative_component_schema = yaml.load(raw_component_schema, Loader=yaml.SafeLoader)
+            declarative_component_schema = safe_load_yaml(raw_component_schema)
             return declarative_component_schema  # type: ignore
         else:
             raise RuntimeError(

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -20,7 +20,6 @@
 )
 
 import orjson
-import yaml
 from airbyte_protocol_dataclasses.models import AirbyteStreamStatus, Level, StreamDescriptor
 from jsonschema.exceptions import ValidationError
 from jsonschema.validators import validate
@@ -105,6 +104,7 @@
 from airbyte_cdk.utils.airbyte_secrets_utils import add_to_secrets, get_secrets
 from airbyte_cdk.utils.stream_status_utils import as_airbyte_message
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+from airbyte_cdk.utils.yaml_loader import safe_load_yaml
 
 
 @dataclass
@@ -128,7 +128,7 @@ def _get_declarative_component_schema() -> Dict[str, Any]:
             "airbyte_cdk", "sources/declarative/declarative_component_schema.yaml"
         )
         if raw_component_schema is not None:
-            declarative_component_schema = yaml.load(raw_component_schema, Loader=yaml.SafeLoader)
+            declarative_component_schema = safe_load_yaml(raw_component_schema)
             return declarative_component_schema  # type: ignore
         else:
             raise RuntimeError(

airbyte_cdk/sources/declarative/yaml_declarative_source.py

@@ -5,13 +5,12 @@
 import pkgutil
 from typing import Any, List, Mapping, Optional
 
-import yaml
-
 from airbyte_cdk.models import AirbyteStateMessage, ConfiguredAirbyteCatalog
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
     ConcurrentDeclarativeSource,
 )
 from airbyte_cdk.sources.types import ConnectionDefinition
+from airbyte_cdk.utils.yaml_loader import safe_load_yaml
 
 
 class YamlDeclarativeSource(ConcurrentDeclarativeSource):
@@ -44,7 +43,7 @@ def _read_and_parse_yaml_file(self, path_to_yaml_file: str) -> ConnectionDefinit
         try:
             # For testing purposes, we want to allow to just pass a file
             with open(path_to_yaml_file, "r") as f:
-                return yaml.safe_load(f)  # type: ignore  # we assume the yaml represents a ConnectionDefinition
+                return safe_load_yaml(f)  # type: ignore  # we assume the yaml represents a ConnectionDefinition
         except FileNotFoundError:
             # Running inside the container, the working directory during an operation is not structured the same as the static files
             package = self.__class__.__module__.split(".")[0]
@@ -66,4 +65,4 @@ def _parse(connection_definition_str: str) -> ConnectionDefinition:
         :param connection_definition_str: yaml string to parse
         :return: The ConnectionDefinition parsed from connection_definition_str
         """
-        return yaml.safe_load(connection_definition_str)  # type: ignore # yaml.safe_load doesn't return a type but know it is a Mapping
+        return safe_load_yaml(connection_definition_str)  # type: ignore # safe_load_yaml doesn't return a type but know it is a Mapping

airbyte_cdk/utils/yaml_loader.py

@@ -0,0 +1,55 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+from __future__ import annotations
+
+from typing import IO, Any, Optional, Union, cast
+
+import yaml
+
+
+def _get_safe_loader() -> Union[type[yaml.SafeLoader], type[yaml.CSafeLoader]]:
+    """Return libyaml's CSafeLoader when PyYAML was built with it, else the pure-Python SafeLoader."""
+    # The annotations stay unevaluated: yaml.CSafeLoader does not exist without libyaml.
+    return cast(
+        "Union[type[yaml.SafeLoader], type[yaml.CSafeLoader]]",
+        getattr(yaml, "CSafeLoader", yaml.SafeLoader),
+    )
+
+
+def _tell_if_seekable(stream: Union[IO[str], IO[bytes]]) -> Optional[int]:
+    """Return the stream's current position if it can be rewound, else None."""
+    try:
+        if stream.seekable() is not True:
+            return None
+        position = stream.tell()
+    except (AttributeError, OSError, ValueError):
+        return None
+    return position if isinstance(position, int) else None
+
+
+def safe_load_yaml(stream: Union[str, bytes, IO[str], IO[bytes]]) -> Any:
+    """Parse trusted YAML like yaml.safe_load, using libyaml's CSafeLoader when available.
+
+    Results match yaml.safe_load for documents both loaders accept. libyaml accepts
+    some input SafeLoader rejects (for example a tab after `key:`), ignores constructors
+    and resolvers registered on yaml.SafeLoader, and can crash the process on very deeply
+    nested input, so do not use this for untrusted input. When libyaml raises a
+    YAMLError, the input is parsed again with SafeLoader so callers see SafeLoader's error.
+    """
+    loader = _get_safe_loader()
+    exact_text = type(stream) in (str, bytes)
+    rewind_to = None
+    if not exact_text and not isinstance(stream, (str, bytes)):
+        rewind_to = _tell_if_seekable(stream)
+    # str/bytes subclasses, non-seekable streams and mocks keep today's behaviour.
+    if loader is yaml.SafeLoader or (not exact_text and rewind_to is None):
+        return yaml.load(stream, Loader=yaml.SafeLoader)
+    try:
+        return yaml.load(stream, Loader=loader)
+    except yaml.YAMLError:
+        pass  # Re-parse outside the handler so libyaml's error is not chained.
+    if rewind_to is not None:
+        cast(Union[IO[str], IO[bytes]], stream).seek(rewind_to)
+    return yaml.load(stream, Loader=yaml.SafeLoader)