diff --git a/airbyte_cdk/cli/source_declarative_manifest/_run.py b/airbyte_cdk/cli/source_declarative_manifest/_run.py index 1b64e51252..13b13104bf 100644 --- a/airbyte_cdk/cli/source_declarative_manifest/_run.py +++ b/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): diff --git a/airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py b/airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py index 5a20bdf1f3..2c6d5c8bd2 100644 --- a/airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py +++ b/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( diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index 055576f45f..1fafa07a62 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/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( diff --git a/airbyte_cdk/sources/declarative/yaml_declarative_source.py b/airbyte_cdk/sources/declarative/yaml_declarative_source.py index 7af07a5c8f..c397e7b1cc 100644 --- a/airbyte_cdk/sources/declarative/yaml_declarative_source.py +++ b/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 diff --git a/airbyte_cdk/utils/yaml_loader.py b/airbyte_cdk/utils/yaml_loader.py new file mode 100644 index 0000000000..6ef5e60834 --- /dev/null +++ b/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) diff --git a/unit_tests/utils/test_yaml_loader.py b/unit_tests/utils/test_yaml_loader.py new file mode 100644 index 0000000000..13ee02e173 --- /dev/null +++ b/unit_tests/utils/test_yaml_loader.py @@ -0,0 +1,244 @@ +# +# Copyright (c) 2025 Airbyte, Inc., all rights reserved. +# + +import io +import pkgutil +from pathlib import Path +from typing import Any +from unittest import mock + +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"] +""" + +# 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' +) + +ESCAPED_EMOJI_YAML = 'description: "\\ud83d\\ude00 emoji"\n' + + +class _NonSeekable(io.StringIO): + def seekable(self) -> bool: + return False + + +@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 repr(pure_python_schema) == repr(c_loader_schema) + assert c_loader_schema == _get_declarative_component_schema() + + +@requires_libyaml +def test_malformed_yaml_raises_parser_error_with_c_loader() -> None: + with pytest.raises(yaml.parser.ParserError): + safe_load_yaml(MALFORMED_YAML) + + +def test_malformed_yaml_raises_parser_error_without_libyaml( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delattr(yaml, "CSafeLoader", raising=False) + with pytest.raises(yaml.parser.ParserError): + safe_load_yaml(MALFORMED_YAML) + + +@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 + + +@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 actual_exc_info.value.__context__ is None + 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) + assert parse_exc_info.value.__context__ is None + + +@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 actual_exc_info.value.__context__ is None + assert str(path) in str(actual_exc_info.value) + + +@requires_libyaml +def test_syntax_error_does_not_chain_c_loader_error() -> None: + with pytest.raises(yaml.YAMLError) as exc_info: + safe_load_yaml(MISINDENTED_YAML) + assert exc_info.value.__context__ is None + + +@requires_libyaml +def test_non_seekable_stream_parses_and_errors_like_safe_loader() -> None: + assert safe_load_yaml(_NonSeekable(ESCAPED_EMOJI_YAML)) == yaml.load( + _NonSeekable(ESCAPED_EMOJI_YAML), Loader=yaml.SafeLoader + ) + with pytest.raises(yaml.YAMLError) as expected: + yaml.load(_NonSeekable(MISINDENTED_YAML), Loader=yaml.SafeLoader) + with pytest.raises(yaml.YAMLError) as actual: + safe_load_yaml(_NonSeekable(MISINDENTED_YAML)) + assert str(actual.value) == str(expected.value) + + +@requires_libyaml +def test_mock_open_handle_with_invalid_yaml_raises() -> None: + with pytest.raises(yaml.YAMLError): + safe_load_yaml(mock.mock_open(read_data=MISINDENTED_YAML)()) + + +def test_str_subclass_parses() -> None: + class _Str(str): + pass + + assert safe_load_yaml(_Str("a: 1")) == {"a": 1} + + +@requires_libyaml +def test_c_loader_accepts_tab_separators_that_safe_loader_rejects() -> None: + with pytest.raises(yaml.scanner.ScannerError): + yaml.load("a:\tb", Loader=yaml.SafeLoader) + assert safe_load_yaml("a:\tb") == {"a": "b"} + + +@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}