move schema loader decorator to model_to_component_factory

Author: maxime.c

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -552,8 +552,9 @@
     InlineSchemaLoader,
     JsonFileSchemaLoader,
     SchemaTypeIdentifier,
-    TypesMap,
+    TypesMap, SchemaLoader,
 )
+from airbyte_cdk.sources.declarative.schema.caching_schema_loader_decorator import CachingSchemaLoaderDecorator
 from airbyte_cdk.sources.declarative.schema.composite_schema_loader import CompositeSchemaLoader
 from airbyte_cdk.sources.declarative.spec import ConfigMigration, Spec
 from airbyte_cdk.sources.declarative.stream_slicers import (
@@ -2095,13 +2096,7 @@ def create_default_stream(
         if isinstance(retriever, AsyncRetriever):
             stream_slicer = retriever.stream_slicer
 
-        schema_loader: Union[
-            CompositeSchemaLoader,
-            DefaultSchemaLoader,
-            DynamicSchemaLoader,
-            InlineSchemaLoader,
-            JsonFileSchemaLoader,
-        ]
+        schema_loader: SchemaLoader
         if model.schema_loader and isinstance(model.schema_loader, list):
             nested_schema_loaders = [
                 self._create_component_from_model(model=nested_schema_loader, config=config)
@@ -2120,6 +2115,7 @@ def create_default_stream(
             if "name" not in options:
                 options["name"] = model.name
             schema_loader = DefaultSchemaLoader(config=config, parameters=options)
+        schema_loader = CachingSchemaLoaderDecorator(schema_loader)
 
         stream_name = model.name or ""
         return DefaultStream(

airbyte_cdk/sources/declarative/schema/caching_schema_loader_decorator.py

@@ -0,0 +1,15 @@
+from typing import Any, Mapping, Optional
+
+from airbyte_cdk.sources.declarative.schema import SchemaLoader
+
+
+class CachingSchemaLoaderDecorator(SchemaLoader):
+    def __init__(self, schema_loader: SchemaLoader):
+        self._decorated = schema_loader
+        self._loaded_schema: Optional[Mapping[str, Any]] = None
+
+    def get_json_schema(self) -> Mapping[str, Any]:
+        if self._loaded_schema is None:
+            self._loaded_schema = self._decorated.get_json_schema()
+
+        return self._loaded_schema  # type: ignore  # at that point, we assume the schema will be populated

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -31,18 +31,6 @@ def get_total_records(self) -> int:
         return self.total_record_counter
 
 
-class SchemaLoaderCachingDecorator(SchemaLoader):
-    def __init__(self, schema_loader: SchemaLoader):
-        self._decorated = schema_loader
-        self._loaded_schema: Optional[Mapping[str, Any]] = None
-
-    def get_json_schema(self) -> Mapping[str, Any]:
-        if self._loaded_schema is None:
-            self._loaded_schema = self._decorated.get_json_schema()
-
-        return self._loaded_schema  # type: ignore  # at that point, we assume the schema will be populated
-
-
 class DeclarativePartitionFactory:
     def __init__(
         self,
@@ -58,7 +46,7 @@ def __init__(
         In order to avoid these problems, we will create one retriever per thread which should make the processing thread-safe.
         """
         self._stream_name = stream_name
-        self._schema_loader = SchemaLoaderCachingDecorator(schema_loader)
+        self._schema_loader = schema_loader
         self._retriever = retriever
         self._message_repository = message_repository
         self._max_records_limit = max_records_limit

unit_tests/sources/declarative/schema/test_caching_schema_loader_decorator.py

@@ -0,0 +1,17 @@
+from unittest import TestCase
+from unittest.mock import Mock
+
+from airbyte_cdk.sources.declarative.schema import SchemaLoader
+from airbyte_cdk.sources.declarative.schema.caching_schema_loader_decorator import CachingSchemaLoaderDecorator
+
+
+class CachingSchemaLoaderDecoratorTest(TestCase):
+    def test_given_previous_calls_when_get_json_schema_then_return_cached_schema(self):
+        decorated = Mock(spec=SchemaLoader)
+        schema_loader = CachingSchemaLoaderDecorator(decorated)
+
+        schema_loader.get_json_schema()
+        schema_loader.get_json_schema()
+        schema_loader.get_json_schema()
+
+        assert decorated.get_json_schema.call_count == 1