add configured catalog to SimpleRetriever and only fetch properties that are selected in the json_schema

Author: brianjlai

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -169,6 +169,7 @@ def __init__(
         component_factory = ModelToComponentFactory(
             emit_connector_builder_messages=emit_connector_builder_messages,
             message_repository=ConcurrentMessageRepository(queue, message_repository),
+            configured_catalog=catalog,
             connector_state_manager=self._connector_state_manager,
             max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"),
             limit_pages_fetched_per_slice=limits.max_pages_per_slice if limits else None,

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -26,6 +26,7 @@
     get_type_hints,
 )
 
+from airbyte_protocol_dataclasses.models import ConfiguredAirbyteStream
 from isodate import parse_duration
 from pydantic.v1 import BaseModel
 from requests import Response
@@ -42,6 +43,7 @@
     AirbyteStateMessage,
     AirbyteStateType,
     AirbyteStreamState,
+    ConfiguredAirbyteCatalog,
     FailureType,
     Level,
     StreamDescriptor,
@@ -668,6 +670,7 @@ def __init__(
         message_repository: Optional[MessageRepository] = None,
         connector_state_manager: Optional[ConnectorStateManager] = None,
         max_concurrent_async_job_count: Optional[int] = None,
+        configured_catalog: Optional[ConfiguredAirbyteCatalog] = None,
     ):
         self._init_mappings()
         self._limit_pages_fetched_per_slice = limit_pages_fetched_per_slice
@@ -678,6 +681,9 @@ def __init__(
         self._message_repository = message_repository or InMemoryMessageRepository(
             self._evaluate_log_level(emit_connector_builder_messages)
         )
+        self._stream_name_to_configured_stream = self._create_stream_name_to_configured_stream(
+            configured_catalog
+        )
         self._connector_state_manager = connector_state_manager or ConnectorStateManager()
         self._api_budget: Optional[Union[APIBudget, HttpAPIBudget]] = None
         self._job_tracker: JobTracker = JobTracker(max_concurrent_async_job_count or 1)
@@ -796,6 +802,14 @@ def _init_mappings(self) -> None:
         # Needed for the case where we need to perform a second parse on the fields of a custom component
         self.TYPE_NAME_TO_MODEL = {cls.__name__: cls for cls in self.PYDANTIC_MODEL_TO_CONSTRUCTOR}
 
+    @staticmethod
+    def _create_stream_name_to_configured_stream(
+        configured_catalog: Optional[ConfiguredAirbyteCatalog],
+    ) -> Mapping[str, ConfiguredAirbyteStream]:
+        if configured_catalog is None:
+            return {}
+        return {stream.stream.name: stream for stream in configured_catalog.streams}
+
     def create_component(
         self,
         model_type: Type[BaseModel],
@@ -3302,6 +3316,8 @@ def _get_url(req: Requester) -> str:
             model.ignore_stream_slicer_parameters_on_paginated_requests or False
         )
 
+        configured_stream = self._stream_name_to_configured_stream.get(name)
+
         if (
             model.partition_router
             and isinstance(model.partition_router, SubstreamPartitionRouterModel)
@@ -3337,6 +3353,7 @@ def _get_url(req: Requester) -> str:
                 request_option_provider=request_options_provider,
                 cursor=None,
                 config=config,
+                configured_airbyte_stream=configured_stream,
                 ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests,
                 parameters=model.parameters or {},
             )
@@ -3358,6 +3375,7 @@ def _get_url(req: Requester) -> str:
             request_option_provider=request_options_provider,
             cursor=None,
             config=config,
+            configured_airbyte_stream=configured_stream,
             ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests,
             additional_query_properties=query_properties,
             log_formatter=self._get_log_formatter(log_formatter, name),

airbyte_cdk/sources/declarative/requesters/query_properties/property_chunking.py

@@ -2,7 +2,7 @@
 
 from dataclasses import InitVar, dataclass
 from enum import Enum
-from typing import Any, Iterable, List, Mapping, Optional
+from typing import Any, Iterable, List, Mapping, Optional, Set
 
 from airbyte_cdk.sources.declarative.requesters.query_properties.strategies import GroupByKey
 from airbyte_cdk.sources.declarative.requesters.query_properties.strategies.merge_strategy import (
@@ -40,7 +40,10 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         )
 
     def get_request_property_chunks(
-        self, property_fields: Iterable[str], always_include_properties: Optional[List[str]]
+        self,
+        property_fields: Iterable[str],
+        always_include_properties: Optional[List[str]],
+        configured_properties: Optional[Set[str]],
     ) -> Iterable[List[str]]:
         if not self.property_limit:
             single_property_chunk = list(property_fields)
@@ -53,6 +56,8 @@ def get_request_property_chunks(
         for property_field in property_fields:
             # If property_limit_type is not defined, we default to property_count which is just an incrementing count
             # todo: Add ability to specify parameter delimiter representation and take into account in property_field_size
+            if configured_properties is not None and property_field not in configured_properties:
+                continue
             property_field_size = (
                 len(property_field)
                 + 3  # The +3 represents the extra characters for encoding the delimiter in between properties

airbyte_cdk/sources/declarative/requesters/query_properties/query_properties.py

@@ -1,8 +1,9 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
 
 from dataclasses import InitVar, dataclass
-from typing import Any, Iterable, List, Mapping, Optional, Union
+from typing import Any, Iterable, List, Mapping, Optional, Set, Union
 
+from airbyte_cdk.models import ConfiguredAirbyteStream
 from airbyte_cdk.sources.declarative.requesters.query_properties import (
     PropertiesFromEndpoint,
     PropertyChunking,
@@ -26,23 +27,44 @@ class QueryProperties:
     parameters: InitVar[Mapping[str, Any]]
 
     def get_request_property_chunks(
-        self, stream_slice: Optional[StreamSlice] = None
+        self,
+        stream_slice: Optional[StreamSlice] = None,
+        configured_stream: Optional[ConfiguredAirbyteStream] = None,
     ) -> Iterable[List[str]]:
         """
         Uses the defined property_list to fetch the total set of properties dynamically or from a static list
         and based on the resulting properties, performs property chunking if applicable.
         :param stream_slice: The StreamSlice of the current partition being processed during the sync. This is included
         because subcomponents of QueryProperties can make use of interpolation of the top-level StreamSlice object
+        :param configured_stream: The customer configured stream being synced which is needed to identify which
+        record fields to query for and emit.
         """
         fields: Union[Iterable[str], List[str]]
         if isinstance(self.property_list, PropertiesFromEndpoint):
             fields = self.property_list.get_properties_from_endpoint(stream_slice=stream_slice)
         else:
             fields = self.property_list if self.property_list else []
 
+        configured_properties = self._get_configured_properties(configured_stream)
+
         if self.property_chunking:
             yield from self.property_chunking.get_request_property_chunks(
-                property_fields=fields, always_include_properties=self.always_include_properties
+                property_fields=fields,
+                always_include_properties=self.always_include_properties,
+                configured_properties=configured_properties,
             )
         else:
-            yield list(fields)
+            # A schema might have no extra properties enabled which is valid and represented by an empty set
+            if configured_properties is not None:
+                yield from [[field for field in fields if field in configured_properties]]
+            else:
+                yield list(fields)
+
+    @staticmethod
+    def _get_configured_properties(
+        configured_stream: Optional[ConfiguredAirbyteStream] = None,
+    ) -> Optional[Set[str]]:
+        if configured_stream:
+            # todo double check that configured catalog only contains enabled fields
+            return set(configured_stream.stream.json_schema.get("properties", {}).keys())
+        return None

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -25,7 +25,10 @@
 
 from airbyte_cdk.legacy.sources.declarative.incremental import ResumableFullRefreshCursor
 from airbyte_cdk.legacy.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
-from airbyte_cdk.models import AirbyteMessage
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    ConfiguredAirbyteStream,
+)
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
@@ -97,6 +100,7 @@ class SimpleRetriever(Retriever):
     cursor: Optional[DeclarativeCursor] = None
     ignore_stream_slicer_parameters_on_paginated_requests: bool = False
     additional_query_properties: Optional[QueryProperties] = None
+    configured_airbyte_stream: Optional[ConfiguredAirbyteStream] = None
     log_formatter: Optional[Callable[[requests.Response], Any]] = None
     pagination_tracker_factory: Callable[[], PaginationTracker] = field(
         default_factory=lambda: lambda: PaginationTracker()
@@ -389,7 +393,8 @@ def _read_pages(
                     and self.additional_query_properties.property_chunking
                 ):
                     for properties in self.additional_query_properties.get_request_property_chunks(
-                        stream_slice=stream_slice
+                        stream_slice=stream_slice,
+                        configured_stream=self.configured_airbyte_stream,
                     ):
                         stream_slice = StreamSlice(
                             partition=stream_slice.partition or {},

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -12,6 +12,13 @@
 import freezegun
 import pytest
 import requests
+from airbyte_protocol_dataclasses.models.airbyte_protocol import (
+    AirbyteStream,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteStream,
+    DestinationSyncMode,
+    SyncMode,
+)
 from freezegun.api import FakeDatetime
 from pydantic.v1 import ValidationError
 
@@ -230,6 +237,28 @@ def test_create_component_type_mismatch():
         factory.create_component(CheckStreamModel, manifest["check"], {})
 
 
+def test_create_component_with_configured_catalog():
+    configured_catalog = ConfiguredAirbyteCatalog(
+        streams=[
+            ConfiguredAirbyteStream(
+                stream=AirbyteStream(
+                    name="test",
+                    json_schema={"type": "object", "properties": {"id": {"type": "string"}}},
+                    supported_sync_modes=[SyncMode.full_refresh],
+                ),
+                sync_mode=SyncMode.full_refresh,
+                destination_sync_mode=DestinationSyncMode.overwrite,
+            )
+        ]
+    )
+
+    factory_with_catalog = ModelToComponentFactory(configured_catalog=configured_catalog)
+
+    assert factory_with_catalog._stream_name_to_configured_stream == {
+        "test": configured_catalog.streams[0]
+    }
+
+
 def test_full_config_stream():
     content = """
 decoder:
@@ -1217,7 +1246,7 @@ def test_stream_with_custom_requester_and_query_properties(requests_mock):
       http_method: "GET"
       request_parameters:
         not_query: 1
-        query: 
+        query:
           type: QueryProperties
           property_list:
             - id