perf(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,48 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+from typing import IO, Any, Optional, Union, cast
+
+import yaml
+
+
+def get_safe_loader() -> Any:
+    """Return libyaml's CSafeLoader when PyYAML was built with it, else the pure-Python SafeLoader."""
+    return getattr(yaml, "CSafeLoader", yaml.SafeLoader)
+
+
+def _tell_if_seekable(stream: Any) -> Optional[int]:
+    """Return the stream's current position if it can be rewound, else None."""
+    if isinstance(stream, (str, bytes)):
+        return None
+    seekable = getattr(stream, "seekable", None)
+    if not callable(seekable) or not seekable():
+        return None
+    try:
+        return cast(int, stream.tell())
+    except (OSError, ValueError):
+        return None
+
+
+def safe_load_yaml(stream: Union[str, bytes, IO[str], IO[bytes]]) -> Any:
+    """Equivalent to yaml.safe_load, but uses the C loader when available.
+
+    If the C loader raises a yaml.YAMLError, the input is re-parsed with the
+    pure-Python SafeLoader and that error is raised instead. The C loader's
+    messages lack the source-line snippet and differ in wording, so the re-parse
+    preserves the error messages users saw before the C loader was introduced.
+    """
+    loader = get_safe_loader()
+    if loader is yaml.SafeLoader:
+        return yaml.load(stream, Loader=yaml.SafeLoader)
+    rewind_to = _tell_if_seekable(stream)
+    try:
+        return yaml.load(stream, Loader=loader)
+    except yaml.YAMLError:
+        if isinstance(stream, (str, bytes)):
+            return yaml.load(stream, Loader=yaml.SafeLoader)
+        if rewind_to is None:
+            raise
+        stream.seek(rewind_to)
+        return yaml.load(stream, Loader=yaml.SafeLoader)

unit_tests/utils/test_yaml_loader.py

@@ -0,0 +1,212 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+import io
+import pkgutil
+from pathlib import Path
+from typing import Any
+
+import pytest
+import yaml
+
+from airbyte_cdk.cli.source_declarative_manifest._run import _parse_manifest_from_file
+from airbyte_cdk.legacy.sources.declarative.manifest_declarative_source import (
+    _get_declarative_component_schema as _legacy_get_declarative_component_schema,
+)
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    _get_declarative_component_schema,
+)
+from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
+from airbyte_cdk.utils.yaml_loader import get_safe_loader, safe_load_yaml
+
+requires_libyaml = pytest.mark.skipif(
+    not yaml.__with_libyaml__, reason="PyYAML was built without libyaml"
+)
+
+MALFORMED_YAML = """
+version: "version"
+definitions:
+  this is not parsable yaml: " at all
+streams:
+  - type: DeclarativeStream
+    $parameters:
+      name: "lists"
+      primary_key: id
+      url_base: "https://api.sendgrid.com"
+check:
+  type: CheckStream
+  stream_names: ["lists"]
+"""
+
+
+@requires_libyaml
+def test_get_safe_loader_uses_c_loader_when_available() -> None:
+    assert get_safe_loader() is yaml.CSafeLoader
+
+
+def test_falls_back_to_pure_python_loader_without_libyaml(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delattr(yaml, "CSafeLoader", raising=False)
+
+    assert get_safe_loader() is yaml.SafeLoader
+    assert safe_load_yaml("a: 1\nb: [x, y]") == {"a": 1, "b": ["x", "y"]}
+    assert YamlDeclarativeSource._parse("version: 1.0.0\nstreams: []") == {
+        "version": "1.0.0",
+        "streams": [],
+    }
+
+
+@requires_libyaml
+def test_component_schema_parses_identically_with_both_loaders() -> None:
+    raw = pkgutil.get_data("airbyte_cdk", "sources/declarative/declarative_component_schema.yaml")
+    assert raw is not None
+    pure_python_schema = yaml.load(raw, Loader=yaml.SafeLoader)
+    c_loader_schema = yaml.load(raw, Loader=yaml.CSafeLoader)
+    assert pure_python_schema == c_loader_schema
+    assert c_loader_schema == _get_declarative_component_schema()
+
+
+@pytest.mark.parametrize(
+    "loader",
+    [
+        yaml.SafeLoader,
+        pytest.param(
+            yaml.CSafeLoader,
+            marks=pytest.mark.skipif(
+                not yaml.__with_libyaml__, reason="PyYAML was built without libyaml"
+            ),
+        ),
+    ],
+)
+def test_malformed_yaml_raises_parser_error_with_both_loaders(loader: Any) -> None:
+    with pytest.raises(yaml.parser.ParserError):
+        yaml.load(MALFORMED_YAML, Loader=loader)
+
+
+@requires_libyaml
+def test_declarative_call_sites_use_c_loader(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    instantiated: list[Any] = []
+
+    class _Spy(yaml.CSafeLoader):
+        def __init__(self, stream: Any) -> None:
+            instantiated.append(stream)
+            super().__init__(stream)
+
+    monkeypatch.setattr(yaml, "CSafeLoader", _Spy)
+
+    _get_declarative_component_schema()
+    assert len(instantiated) == 1
+
+    _legacy_get_declarative_component_schema()
+    assert len(instantiated) == 2
+
+    YamlDeclarativeSource._parse("a: 1")
+    assert len(instantiated) == 3
+
+    yaml_file = tmp_path / "source.yaml"
+    yaml_file.write_text("a: 1")
+    YamlDeclarativeSource._read_and_parse_yaml_file(
+        YamlDeclarativeSource.__new__(YamlDeclarativeSource), str(yaml_file)
+    )
+    assert len(instantiated) == 4
+
+    manifest_file = tmp_path / "m.yaml"
+    manifest_file.write_text("version: 1.0.0\nstreams: []")
+    _parse_manifest_from_file(str(manifest_file))
+    assert len(instantiated) == 5
+
+
+# A list item whose nested mapping is indented one space short.
+MISINDENTED_YAML = (
+    'version: "1"\nstreams:\n  - type: DeclarativeStream\n    name: a\n   primary_key: id\n'
+)
+
+
+@requires_libyaml
+def test_syntax_error_in_string_matches_pure_python_message() -> None:
+    with pytest.raises(yaml.YAMLError) as expected_exc_info:
+        yaml.load(MISINDENTED_YAML, Loader=yaml.SafeLoader)
+
+    with pytest.raises(yaml.YAMLError) as actual_exc_info:
+        safe_load_yaml(MISINDENTED_YAML)
+
+    assert str(actual_exc_info.value) == str(expected_exc_info.value)
+    assert type(actual_exc_info.value) is type(expected_exc_info.value)
+    assert "^" in str(actual_exc_info.value)
+
+    with pytest.raises(yaml.YAMLError) as parse_exc_info:
+        YamlDeclarativeSource._parse(MISINDENTED_YAML)
+
+    assert str(parse_exc_info.value) == str(expected_exc_info.value)
+    assert type(parse_exc_info.value) is type(expected_exc_info.value)
+
+
+@requires_libyaml
+def test_syntax_error_in_file_matches_pure_python_message(tmp_path: Path) -> None:
+    path = tmp_path / "manifest.yaml"
+    path.write_text(MISINDENTED_YAML)
+
+    with open(path) as f:
+        with pytest.raises(yaml.YAMLError) as expected_exc_info:
+            yaml.load(f, Loader=yaml.SafeLoader)
+
+    with open(path) as f:
+        with pytest.raises(yaml.YAMLError) as actual_exc_info:
+            safe_load_yaml(f)
+
+    assert str(actual_exc_info.value) == str(expected_exc_info.value)
+    assert str(path) in str(actual_exc_info.value)
+
+
+@requires_libyaml
+def test_syntax_error_in_non_seekable_stream_keeps_c_loader_message() -> None:
+    class _NonSeekable(io.StringIO):
+        def seekable(self) -> bool:
+            return False
+
+    with pytest.raises(yaml.YAMLError) as expected_exc_info:
+        yaml.load(_NonSeekable(MISINDENTED_YAML), Loader=yaml.CSafeLoader)
+
+    with pytest.raises(yaml.YAMLError) as actual_exc_info:
+        safe_load_yaml(_NonSeekable(MISINDENTED_YAML))
+
+    assert str(actual_exc_info.value) == str(expected_exc_info.value)
+
+
+@requires_libyaml
+def test_successful_parse_never_uses_pure_python_loader(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    instantiated: list[Any] = []
+
+    class _Spy(yaml.SafeLoader):
+        def __init__(self, stream: Any) -> None:
+            instantiated.append(stream)
+            super().__init__(stream)
+
+    monkeypatch.setattr(yaml, "SafeLoader", _Spy)
+
+    assert safe_load_yaml("a: 1") == {"a": 1}
+    assert safe_load_yaml(b"a: 1") == {"a": 1}
+
+    yaml_file = tmp_path / "source.yaml"
+    yaml_file.write_text("a: 1")
+    with open(yaml_file) as f:
+        assert safe_load_yaml(f) == {"a": 1}
+
+    assert instantiated == []
+
+
+@requires_libyaml
+def test_pure_python_result_returned_if_only_c_loader_fails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class _FailingLoader(yaml.CSafeLoader):
+        def get_single_data(self) -> Any:
+            raise yaml.YAMLError("boom")
+
+    monkeypatch.setattr(yaml, "CSafeLoader", _FailingLoader)
+
+    assert safe_load_yaml("a: 1") == {"a": 1}