fix(declarative): speed up manifest validation at source startup

airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py

@@ -12,7 +12,6 @@
 
 import orjson
 from jsonschema.exceptions import ValidationError
-from jsonschema.validators import validate
 from packaging.version import InvalidVersion, Version
 
 from airbyte_cdk.config_observation import create_connector_config_control_message
@@ -34,6 +33,7 @@
 from airbyte_cdk.models.airbyte_protocol_serializers import AirbyteMessageSerializer
 from airbyte_cdk.sources.declarative.checks import COMPONENTS_CHECKER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import validate_manifest
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedBoolean
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConditionalStreams as ConditionalStreamsModel,
@@ -447,12 +447,7 @@ def _validate_source(self) -> None:
         Validates the connector manifest against the declarative component schema
         """
 
-        try:
-            validate(self._source_config, self._declarative_component_schema)
-        except ValidationError as e:
-            raise ValidationError(
-                "Validation against json schema defined in declarative_component_schema.yaml schema failed"
-            ) from e
+        validate_manifest(self._source_config, self._declarative_component_schema)
 
         cdk_version_str = metadata.version("airbyte_cdk")
         cdk_version = self._parse_version(cdk_version_str, "airbyte-cdk")

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -21,8 +21,8 @@
 
 import orjson
 from airbyte_protocol_dataclasses.models import AirbyteStreamStatus, Level, StreamDescriptor
-from jsonschema.exceptions import ValidationError
-from jsonschema.validators import validate
+from jsonschema.exceptions import ValidationError, best_match
+from jsonschema.validators import validator_for
 
 from airbyte_cdk.config_observation import create_connector_config_control_message
 from airbyte_cdk.connector_builder.models import (
@@ -122,7 +122,7 @@ class TestLimits:
     max_streams: int = field(default=DEFAULT_MAX_STREAMS)
 
 
-def _get_declarative_component_schema() -> Dict[str, Any]:
+def get_declarative_component_schema() -> Dict[str, Any]:
     try:
         raw_component_schema = pkgutil.get_data(
             "airbyte_cdk", "sources/declarative/declarative_component_schema.yaml"
@@ -140,6 +140,28 @@ def _get_declarative_component_schema() -> Dict[str, Any]:
         )
 
 
+def validate_manifest(manifest: Mapping[str, Any], schema: Mapping[str, Any]) -> None:
+    """Validate a manifest against the shipped component schema.
+
+    Unit tests check the schema's metaschema, so valid manifests skip `check_schema`. On failure,
+    check the schema first to surface `SchemaError`, while preserving `jsonschema.validate`'s
+    best-match error as the cause with a traceback.
+    """
+    validator_class = validator_for(schema)
+    try:
+        error = best_match(validator_class(schema).iter_errors(manifest))
+        if error is not None:
+            raise error
+    except ValidationError as e:
+        validator_class.check_schema(schema)
+        raise ValidationError(
+            "Validation against json schema defined in declarative_component_schema.yaml schema failed"
+        ) from e
+    except Exception:
+        validator_class.check_schema(schema)
+        raise
+
+
 class ConcurrentDeclarativeSource(Source):
     # By default, we defer to a value of 2. A value lower than could cause a PartitionEnqueuer to be stuck in a state of deadlock
     # because it has hit the limit of futures but not partition reader is consuming them.
@@ -204,7 +226,7 @@ def __init__(
 
         self._should_normalize = normalize_manifest
         self._should_migrate = migrate_manifest
-        self._declarative_component_schema = _get_declarative_component_schema()
+        self._declarative_component_schema = get_declarative_component_schema()
         # If custom components are needed, locate and/or register them.
         self.components_module: ModuleType | None = get_registered_components_module(config=config)
         # set additional attributes
@@ -345,12 +367,7 @@ def _validate_source(self) -> None:
         Validates the connector manifest against the declarative component schema
         """
 
-        try:
-            validate(self._source_config, self._declarative_component_schema)
-        except ValidationError as e:
-            raise ValidationError(
-                "Validation against json schema defined in declarative_component_schema.yaml schema failed"
-            ) from e
+        validate_manifest(self._source_config, self._declarative_component_schema)
 
     def _migrate_and_transform_config(
         self,

unit_tests/legacy/sources/declarative/test_manifest_declarative_source.py

@@ -15,6 +15,7 @@
 import requests
 import yaml
 from jsonschema.exceptions import ValidationError
+from jsonschema.validators import validator_for
 
 import unit_tests.sources.declarative.external_component  # Needed for dynamic imports to work
 from airbyte_cdk.legacy.sources.declarative.manifest_declarative_source import (
@@ -33,6 +34,7 @@
 )
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
     ConcurrentDeclarativeSource,
+    get_declarative_component_schema,
 )
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
@@ -43,6 +45,9 @@
 from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever
 from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from unit_tests.sources.declarative.parsers.test_model_to_component_factory import get_retriever
+from unit_tests.sources.declarative.test_concurrent_declarative_source import (
+    two_stream_schema_validation_manifest,
+)
 
 logger = logging.getLogger("airbyte")
 
@@ -2550,3 +2555,28 @@ def capturing_resolve(resolver_self, *args, **kwargs):
         "HttpComponentsResolver's requester should have api_budget set during dynamic stream "
         "discovery, but it was None. This means discovery HTTP requests are not rate-limited."
     )
+
+
+@pytest.mark.parametrize("migrate_manifest", [False, True])
+def test_building_a_source_does_not_check_the_component_schema(migrate_manifest):
+    with patch.object(
+        validator_for(get_declarative_component_schema()),
+        "check_schema",
+        side_effect=AssertionError("check_schema must not run at source build"),
+    ) as check_schema:
+        ManifestDeclarativeSource(
+            source_config=two_stream_schema_validation_manifest(),
+            migrate_manifest=migrate_manifest,
+        )
+        check_schema.assert_not_called()
+
+
+def test_legacy_source_delegates_to_public_manifest_validator():
+    with patch(
+        "airbyte_cdk.legacy.sources.declarative.manifest_declarative_source.validate_manifest"
+    ) as validate_manifest:
+        source = ManifestDeclarativeSource(source_config=two_stream_schema_validation_manifest())
+
+    validate_manifest.assert_called_once_with(
+        source._source_config, source._declarative_component_schema
+    )

unit_tests/sources/declarative/parsers/test_manifest_normalizer.py

@@ -4,7 +4,7 @@
 
 
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
-    _get_declarative_component_schema,
+    get_declarative_component_schema,
 )
 from airbyte_cdk.sources.declarative.parsers.manifest_normalizer import (
     ManifestNormalizer,
@@ -25,7 +25,7 @@ def test_when_multiple_url_base_are_resolved_and_most_frequent_is_shared(
     and the most frequent one is shared.
     """
 
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     resolved_manifest = resolver.preprocess_manifest(manifest_with_multiple_url_base)
     normalized_manifest = ManifestNormalizer(resolved_manifest, schema).normalize()
 
@@ -41,7 +41,7 @@ def test_with_shared_definitions_url_base_are_present(
     between the definitions and the `url_base` is present in the manifest.
     """
 
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     resolved_manifest = resolver.preprocess_manifest(manifest_with_url_base_linked_definition)
     normalized_manifest = ManifestNormalizer(resolved_manifest, schema).normalize()
 
@@ -75,7 +75,7 @@ def test_clean_dangling_fields_removes_unreferenced_definitions() -> None:
             }
         ],
     }
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     normalizer = ManifestNormalizer(manifest, schema)
     normalizer._clean_dangling_fields()
     assert normalizer._normalized_manifest == expected
@@ -100,7 +100,7 @@ def test_clean_dangling_fields_removes_unreferenced_schemas() -> None:
             {"name": "stream1", "schema_loader": {"schema": {"$ref": "#/schemas/referenced"}}}
         ],
     }
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     normalizer = ManifestNormalizer(manifest, schema)
     normalizer._clean_dangling_fields()
     assert normalizer._normalized_manifest == expected
@@ -134,7 +134,7 @@ def test_clean_dangling_fields_keeps_parent_paths() -> None:
             }
         ],
     }
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     normalizer = ManifestNormalizer(manifest, schema)
     normalizer._clean_dangling_fields()
     assert normalizer._normalized_manifest == expected
@@ -150,7 +150,7 @@ def test_clean_dangling_fields_removes_empty_sections() -> None:
         "streams": [{"name": "stream1", "type": "object", "properties": {}}],
     }
     expected = {"streams": [{"name": "stream1", "type": "object", "properties": {}}]}
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     normalizer = ManifestNormalizer(manifest, schema)
     normalizer._clean_dangling_fields()
     assert normalizer._normalized_manifest == expected
@@ -220,7 +220,7 @@ def test_replace_parent_streams_with_refs_replaces_with_ref():
             },
         ]
     }
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     normalizer = ManifestNormalizer(manifest, schema)
     normalizer._replace_parent_streams_with_refs()
     assert normalizer._normalized_manifest == expected
@@ -276,7 +276,7 @@ def test_replace_parent_streams_with_refs_no_match():
             },
         ]
     }
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     normalizer = ManifestNormalizer(manifest, schema)
     normalizer._replace_parent_streams_with_refs()
     assert normalizer._normalized_manifest == expected
@@ -360,7 +360,7 @@ def test_replace_parent_streams_with_refs_handles_multiple_partition_routers():
             },
         ]
     }
-    schema = _get_declarative_component_schema()
+    schema = get_declarative_component_schema()
     normalizer = ManifestNormalizer(manifest, schema)
     normalizer._replace_parent_streams_with_refs()
     assert normalizer._normalized_manifest == expected