diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index 67240c69dd..05c993805d 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py @@ -412,9 +412,27 @@ def read( yield from self._concurrent_source.read(selected_concurrent_streams) def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog: - return AirbyteCatalog( - streams=[stream.as_airbyte_stream() for stream in self.streams(config=self._config)] + streams = ( + self._catalog_streams() + if not config and not self.check_config_during_discover + else self.streams(config=self._config) ) + return AirbyteCatalog(streams=[stream.as_airbyte_stream() for stream in streams]) + + def _catalog_streams(self) -> List[AbstractStream]: + """Create streams that can only describe their catalog entry, for config-free discover.""" + return [ + self._constructor.create_catalog_stream( + ( + StateDelegatingStreamModel + if stream_config.get("type") == StateDelegatingStreamModel.__name__ + else DeclarativeStreamModel + ), + stream_config, + self._config, + ) + for stream_config in self._stream_configs(self._source_config) + self.dynamic_streams + ] # todo: add PR comment about whether we can change the signature to List[AbstractStream] def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]: # type: ignore # we are migrating away from the AbstractSource and are expecting that this will only be called by ConcurrentDeclarativeSource or the Connector Builder diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 219e0f70ad..52f53b1991 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -17,6 +17,7 @@ Any, Callable, Dict, + Iterable, List, Mapping, MutableMapping, @@ -701,6 +702,10 @@ ) from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream from airbyte_cdk.sources.streams.concurrent.helpers import get_primary_key_from_stream +from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition +from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import ( + PartitionGenerator, +) from airbyte_cdk.sources.streams.concurrent.partitions.stream_slicer import ( StreamSlicer as ConcurrentStreamSlicer, ) @@ -1450,22 +1455,7 @@ def create_concurrent_cursor_from_datetime_based_cursor( model_parameters = datetime_based_cursor_model.parameters or {} - cursor_field = self._get_catalog_defined_cursor_field( - stream_name=stream_name, - allow_catalog_defined_cursor_field=datetime_based_cursor_model.allow_catalog_defined_cursor_field - or False, - ) - - if not cursor_field: - interpolated_cursor_field = InterpolatedString.create( - datetime_based_cursor_model.cursor_field, - parameters=model_parameters, - ) - cursor_field = CursorField( - cursor_field_key=interpolated_cursor_field.eval(config=config), - supports_catalog_defined_cursor_field=datetime_based_cursor_model.allow_catalog_defined_cursor_field - or False, - ) + cursor_field = self._resolve_cursor_field(datetime_based_cursor_model, stream_name, config) interpolated_partition_field_start = InterpolatedString.create( datetime_based_cursor_model.partition_field_start or "start_time", @@ -1691,23 +1681,10 @@ def create_concurrent_cursor_from_incrementing_count_cursor( else: evaluated_start_value = 0 - cursor_field = self._get_catalog_defined_cursor_field( - stream_name=stream_name, - allow_catalog_defined_cursor_field=incrementing_count_cursor_model.allow_catalog_defined_cursor_field - or False, + cursor_field = self._resolve_cursor_field( + incrementing_count_cursor_model, stream_name, config ) - if not cursor_field: - interpolated_cursor_field = InterpolatedString.create( - incrementing_count_cursor_model.cursor_field, - parameters=incrementing_count_cursor_model.parameters or {}, - ) - cursor_field = CursorField( - cursor_field_key=interpolated_cursor_field.eval(config=config), - supports_catalog_defined_cursor_field=incrementing_count_cursor_model.allow_catalog_defined_cursor_field - or False, - ) - connector_state_converter = IncrementingCountStreamStateConverter( is_sequential_state=True, # ConcurrentPerPartitionCursor only works with sequential state ) @@ -1776,26 +1753,7 @@ def create_concurrent_cursor_from_perpartition_cursor( f"Expected {model_type.__name__} component, but received {datetime_based_cursor_model.__class__.__name__}" ) - cursor_field = self._get_catalog_defined_cursor_field( - stream_name=stream_name, - allow_catalog_defined_cursor_field=datetime_based_cursor_model.allow_catalog_defined_cursor_field - or False, - ) - - if not cursor_field: - interpolated_cursor_field = InterpolatedString.create( - datetime_based_cursor_model.cursor_field, - # FIXME the interfaces of the concurrent cursor are kind of annoying as they take a `ComponentDefinition` instead of the actual model. This was done because the ConcurrentDeclarativeSource didn't have access to the models [here for example](https://github.com/airbytehq/airbyte-python-cdk/blob/f525803b3fec9329e4cc8478996a92bf884bfde9/airbyte_cdk/sources/declarative/concurrent_declarative_source.py#L354C54-L354C91). So now we have two cases: - # * The ComponentDefinition comes from model.__dict__ in which case we have `parameters` - # * The ComponentDefinition comes from the manifest as a dict in which case we have `$parameters` - # We should change those interfaces to use the model once we clean up the code in CDS at which point the parameter propagation should happen as part of the ModelToComponentFactory. - parameters=datetime_based_cursor_model.parameters or {}, - ) - cursor_field = CursorField( - cursor_field_key=interpolated_cursor_field.eval(config=config), - supports_catalog_defined_cursor_field=datetime_based_cursor_model.allow_catalog_defined_cursor_field - or False, - ) + cursor_field = self._resolve_cursor_field(datetime_based_cursor_model, stream_name, config) datetime_format = datetime_based_cursor_model.datetime_format @@ -2239,6 +2197,34 @@ def create_default_stream( if isinstance(retriever, AsyncRetriever): stream_slicer = retriever.stream_slicer + schema_loader = self._build_schema_loader(model, config) + + stream_name = model.name or "" + return DefaultStream( + partition_generator=StreamSlicerPartitionGenerator( + DeclarativePartitionFactory( + stream_name, + schema_loader, + retriever, + self._message_repository, + ), + stream_slicer, + slice_limit=self._limit_slices_fetched, + ), + name=stream_name, + json_schema=schema_loader.get_json_schema, + primary_key=get_primary_key_from_stream(primary_key), + cursor_field=( + concurrent_cursor.cursor_field + if hasattr(concurrent_cursor, "cursor_field") + else None + ), + logger=logging.getLogger(f"airbyte.{stream_name}"), + cursor=concurrent_cursor, + supports_file_transfer=hasattr(model, "file_uploader") and bool(model.file_uploader), + ) + + def _build_schema_loader(self, model: DeclarativeStreamModel, config: Config) -> SchemaLoader: schema_loader: SchemaLoader if model.schema_loader and isinstance(model.schema_loader, list): nested_schema_loaders = [ @@ -2258,31 +2244,59 @@ 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) + return CachingSchemaLoaderDecorator(schema_loader) + + def create_catalog_stream( + self, + model_type: Type[BaseModel], + component_definition: ComponentDefinition, + config: Config, + ) -> DefaultStream: + """Create a stream that can only describe itself in a catalog. + + The stream gets the same name, schema loader, primary key, cursor field and file transfer + flag as `create_default_stream`, so `as_airbyte_stream()` returns the same catalog entry. + Requesters, authenticators, partition routers and cursors are never created, so a config is + not needed unless the schema loader or cursor field interpolates it. Reading a config value + that is not set raises `MissingConfigValueError` rather than producing an incomplete catalog + entry. Reading from the stream raises an error. + """ + model = model_type.parse_obj(component_definition) + if isinstance(model, StateDelegatingStreamModel): + # Mirrors create_state_delegating_stream: without state, the full refresh stream is used. + model = ( + model.incremental_stream + if self._connector_state_manager.get_stream_state(model.name, None) + else model.full_refresh_stream + ) + if not isinstance(model, DeclarativeStreamModel): + raise ValueError( + f"Expected a DeclarativeStream or StateDelegatingStream component, but received {model.__class__.__name__}" + ) + config = _StrictConfig(config) stream_name = model.name or "" + cursor_field: Optional[CursorField] = None + if isinstance( + model.incremental_sync, (DatetimeBasedCursorModel, IncrementingCountCursorModel) + ): + cursor_field = self._resolve_cursor_field(model.incremental_sync, stream_name, config) + elif model.incremental_sync: + raise ValueError( + f"Incremental sync of type {type(model.incremental_sync)} is not supported" + ) + return DefaultStream( - partition_generator=StreamSlicerPartitionGenerator( - DeclarativePartitionFactory( - stream_name, - schema_loader, - retriever, - self._message_repository, - ), - stream_slicer, - slice_limit=self._limit_slices_fetched, - ), + partition_generator=_CatalogOnlyPartitionGenerator(stream_name), name=stream_name, - json_schema=schema_loader.get_json_schema, - primary_key=get_primary_key_from_stream(primary_key), - cursor_field=( - concurrent_cursor.cursor_field - if hasattr(concurrent_cursor, "cursor_field") - else None + json_schema=self._build_schema_loader(model, config).get_json_schema, + primary_key=get_primary_key_from_stream( + model.primary_key.__root__ if model.primary_key else None ), + cursor_field=cursor_field, logger=logging.getLogger(f"airbyte.{stream_name}"), - cursor=concurrent_cursor, - supports_file_transfer=hasattr(model, "file_uploader") and bool(model.file_uploader), + cursor=FinalStateCursor(stream_name, None, self._message_repository), + supports_file_transfer=bool(model.file_uploader), ) def _warn_on_ineffective_incremental_dependency(self, model: DeclarativeStreamModel) -> None: @@ -5787,6 +5801,30 @@ def _ensure_query_properties_to_model( request_parameter ) + def _resolve_cursor_field( + self, + cursor_model: Union[DatetimeBasedCursorModel, IncrementingCountCursorModel], + stream_name: str, + config: Config, + ) -> CursorField: + """Return the catalog-defined cursor field if allowed and set, else the manifest's cursor field.""" + allow_catalog_defined = cursor_model.allow_catalog_defined_cursor_field or False + cursor_field = self._get_catalog_defined_cursor_field( + stream_name=stream_name, + allow_catalog_defined_cursor_field=allow_catalog_defined, + ) + if cursor_field: + return cursor_field + + interpolated_cursor_field = InterpolatedString.create( + cursor_model.cursor_field, + parameters=cursor_model.parameters or {}, + ) + return CursorField( + cursor_field_key=interpolated_cursor_field.eval(config=config), + supports_catalog_defined_cursor_field=allow_catalog_defined, + ) + def _get_catalog_defined_cursor_field( self, stream_name: str, allow_catalog_defined_cursor_field: bool ) -> Optional[CursorField]: @@ -5814,3 +5852,31 @@ def _get_catalog_defined_cursor_field( cursor_field_key=configured_stream.cursor_field[0], supports_catalog_defined_cursor_field=allow_catalog_defined_cursor_field, ) + + +class _CatalogOnlyPartitionGenerator(PartitionGenerator): + """Partition generator for streams created by `create_catalog_stream`, which cannot be read.""" + + def __init__(self, stream_name: str) -> None: + self._stream_name = stream_name + + def generate(self) -> Iterable[Partition]: + raise RuntimeError( + f"Stream `{self._stream_name}` was created for discover only and cannot be read." + ) + + +class MissingConfigValueError(ValueError): + """Raised when a catalog-only stream reads a config value that is not set.""" + + +class _StrictConfig(Dict[str, Any]): + """Config that raises on missing keys instead of letting interpolation render them as empty. + + Explicit defaults such as `config.get('key', 'default')` still work. + """ + + def __missing__(self, key: str) -> Any: + raise MissingConfigValueError( + f"Config value `{key}` is required to build this stream's catalog entry." + ) diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source_config_free_discover.py b/unit_tests/sources/declarative/test_concurrent_declarative_source_config_free_discover.py index f08027619d..2e65a6447e 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source_config_free_discover.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source_config_free_discover.py @@ -9,6 +9,9 @@ from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( ConcurrentDeclarativeSource, ) +from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ( + MissingConfigValueError, +) _RETRIEVER = { "type": "SimpleRetriever", @@ -182,3 +185,124 @@ def test_discover_static_schema_without_config() -> None: catalog = source.discover(source.logger, {}) assert sorted(stream.name for stream in catalog.streams) == ["dynamic_items", "items"] + + +_INCREMENTAL_SYNC = { + "type": "DatetimeBasedCursor", + "cursor_field": "{{ parameters['cursor'] }}", + "datetime_format": "%Y-%m-%dT%H:%M:%SZ", + "start_datetime": "{{ config['start_date'] }}", +} + + +def _incremental_stream(name: str) -> dict[str, Any]: + return { + **_stream(name, _INLINE_SCHEMA_LOADER), + "primary_key": ["id"], + "incremental_sync": _INCREMENTAL_SYNC, + "$parameters": {"cursor": "updated_at"}, + } + + +def _substream(name: str, parent: Mapping[str, Any]) -> dict[str, Any]: + return { + **_incremental_stream(name), + "retriever": { + **_RETRIEVER, + "partition_router": { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "type": "ParentStreamConfig", + "stream": parent, + "parent_key": "id", + "partition_field": "parent_id", + } + ], + }, + }, + } + + +def _state_delegating_stream(name: str) -> dict[str, Any]: + return { + "type": "StateDelegatingStream", + "name": name, + "full_refresh_stream": {**_stream(name, _INLINE_SCHEMA_LOADER), "primary_key": ["id"]}, + "incremental_stream": _incremental_stream(name), + } + + +@pytest.mark.parametrize( + "streams", + [ + pytest.param([_incremental_stream("items")], id="config_dependent_cursor"), + pytest.param( + [_substream("children", _incremental_stream("parents"))], + id="substream_with_config_dependent_parent", + ), + pytest.param([_state_delegating_stream("items")], id="state_delegating_stream"), + ], +) +def test_config_free_discover_matches_configured_discover(streams: list[Mapping[str, Any]]) -> None: + config = {"api_key": "dummy", "start_date": "2024-01-01T00:00:00Z"} + configured_source = ConcurrentDeclarativeSource(source_config=_manifest(streams), config=config) + config_free_source = ConcurrentDeclarativeSource(source_config=_manifest(streams), config={}) + + configured_catalog = configured_source.discover(configured_source.logger, config) + config_free_catalog = config_free_source.discover(config_free_source.logger, {}) + + assert config_free_catalog == configured_catalog + + +def test_config_free_discover_streams_cannot_be_read() -> None: + source = ConcurrentDeclarativeSource( + source_config=_manifest([_incremental_stream("items")]), config={} + ) + + (stream,) = source._catalog_streams() + + with pytest.raises(RuntimeError, match="discover only"): + list(stream.generate_partitions()) + + +def test_discover_without_config_fails_for_config_dependent_cursor_without_catalog_streams() -> ( + None +): + source = ConcurrentDeclarativeSource( + source_config=_manifest([_incremental_stream("items")]), config={} + ) + + with pytest.raises(ValueError): + source.streams(config={}) + + +@pytest.mark.parametrize( + "cursor_field", + [ + pytest.param("{{ config['record_cursor'] }}", id="whole_field"), + pytest.param("updated_{{ config['suffix'] }}", id="partial_field"), + ], +) +def test_config_free_discover_fails_when_cursor_field_reads_missing_config( + cursor_field: str, +) -> None: + stream = _incremental_stream("items") + stream["incremental_sync"] = {**_INCREMENTAL_SYNC, "cursor_field": cursor_field} + source = ConcurrentDeclarativeSource(source_config=_manifest([stream]), config={}) + + with pytest.raises(MissingConfigValueError): + source.discover(source.logger, {}) + + +def test_config_free_discover_uses_cursor_field_template_default() -> None: + stream = _incremental_stream("items") + stream["incremental_sync"] = { + **_INCREMENTAL_SYNC, + "cursor_field": "{{ config.get('record_cursor', 'updated_at') }}", + } + source = ConcurrentDeclarativeSource(source_config=_manifest([stream]), config={}) + + (catalog_stream,) = source.discover(source.logger, {}).streams + + assert catalog_stream.default_cursor_field == ["updated_at"]