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 8940f33fa3..ad53c31682 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2760,14 +2760,22 @@ def create_dynamic_schema_loader( self._create_component_from_model(model=transformation_model, config=config) ) name = "dynamic_properties" + partition_router = self._build_stream_slicer_from_partition_router(model.retriever, config) + retriever_kwargs: Dict[str, Any] = {} + if isinstance(model.retriever, AsyncRetrieverModel): + if model.retriever.partition_router: + raise ValueError( + "DynamicSchemaLoader does not support an AsyncRetriever with a partition_router: only the first record is read, so jobs for other partitions would be orphaned. Remove the partition_router from the schema loader's AsyncRetriever." + ) + # create_async_retriever requires stream_slicer; passing it to a CustomRetriever would override its own + retriever_kwargs["stream_slicer"] = partition_router retriever = self._create_component_from_model( model=model.retriever, config=config, name=name, primary_key=None, - partition_router=self._build_stream_slicer_from_partition_router( - model.retriever, config - ), + partition_router=partition_router, + **retriever_kwargs, transformations=[], use_cache=True, log_formatter=( diff --git a/airbyte_cdk/sources/declarative/retrievers/async_retriever.py b/airbyte_cdk/sources/declarative/retrievers/async_retriever.py index 986a11d39b..3a79f04b93 100644 --- a/airbyte_cdk/sources/declarative/retrievers/async_retriever.py +++ b/airbyte_cdk/sources/declarative/retrievers/async_retriever.py @@ -82,6 +82,12 @@ def read_records( records_schema: Mapping[str, Any], stream_slice: Optional[StreamSlice] = None, ) -> Iterable[StreamData]: + if stream_slice is None: + # Slice-less callers such as DynamicSchemaLoader read this way: create and poll the jobs here + for job_slice in self.stream_slicer.stream_slices(): + yield from self.read_records(records_schema, job_slice) + return + # emit the slice_descriptor log message, for connector builder TestRead yield self.slice_logger.create_slice_log_message(stream_slice.cursor_slice) # type: ignore diff --git a/airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py b/airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py index 7d5946a125..37a434b74e 100644 --- a/airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py +++ b/airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py @@ -10,6 +10,7 @@ import dpath from typing_extensions import deprecated +from airbyte_cdk.models import AirbyteMessage, Type from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString @@ -120,6 +121,10 @@ def _update_pointer( class DynamicSchemaLoader(SchemaLoader): """ Dynamically loads a JSON Schema by extracting data from retrieved records. + + With an AsyncRetriever, each schema load creates one async job (retried on failure) and stops at + the first record, so the job's delete endpoint is not called. During a read the schema loads on + the first partition read and its job takes a slot from max_concurrent_async_job_count. """ retriever: Retriever @@ -134,7 +139,15 @@ def get_json_schema(self) -> Mapping[str, Any]: Constructs a JSON Schema based on retrieved data. """ properties = {} - retrieved_record = next(self.retriever.read_records({}), None) # type: ignore[call-overload] # read_records return Iterable data type + retrieved_record = next( + ( + record + for record in self.retriever.read_records({}) + # AsyncRetriever emits a slice log message before its records + if not (isinstance(record, AirbyteMessage) and record.type == Type.LOG) + ), + None, + ) raw_schema = ( self._extract_data( diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index ea3b0cf151..3f6ad4b75c 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -79,6 +79,9 @@ from airbyte_cdk.sources.declarative.models import DeclarativeStream as DeclarativeStreamModel from airbyte_cdk.sources.declarative.models import DefaultPaginator as DefaultPaginatorModel from airbyte_cdk.sources.declarative.models import DpathExtractor as DpathExtractorModel +from airbyte_cdk.sources.declarative.models import ( + DynamicSchemaLoader as DynamicSchemaLoaderModel, +) from airbyte_cdk.sources.declarative.models import FileUploader as FileUploaderModel from airbyte_cdk.sources.declarative.models import ( GroupingPartitionRouter as GroupingPartitionRouterModel, @@ -205,7 +208,11 @@ PageSizeReduction, PageSizeResetPolicy, ) -from airbyte_cdk.sources.declarative.schema import InlineSchemaLoader, JsonFileSchemaLoader +from airbyte_cdk.sources.declarative.schema import ( + DynamicSchemaLoader, + InlineSchemaLoader, + JsonFileSchemaLoader, +) from airbyte_cdk.sources.declarative.schema.caching_schema_loader_decorator import ( CachingSchemaLoaderDecorator, ) @@ -5124,6 +5131,137 @@ def test_create_async_retriever(): assert download_retriever_record_selector.schema_normalization._config.name == "NoTransform" +def test_dynamic_schema_loader_with_async_retriever(): + """ + Verifies that DynamicSchemaLoader can be created with an AsyncRetriever without raising + TypeError due to missing stream_slicer kwarg. + Regression test for https://github.com/airbytehq/airbyte-python-cdk/issues/766 + """ + content = """ +stream_with_dynamic_schema: + type: DeclarativeStream + name: test_stream + schema_loader: + type: DynamicSchemaLoader + retriever: + type: AsyncRetriever + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: ["data"] + status_mapping: + failed: + - Error + running: + - Pending + completed: + - Success + timeout: [] + status_extractor: + type: DpathExtractor + field_path: + - status + download_target_extractor: + type: DpathExtractor + field_path: + - download_url + creation_requester: + type: HttpRequester + url_base: https://api.test.com + path: /reports/create + http_method: POST + polling_requester: + type: HttpRequester + url_base: https://api.test.com + path: /reports/status + http_method: GET + download_requester: + type: HttpRequester + url_base: "{{download_target}}" + http_method: GET + schema_type_identifier: + type: SchemaTypeIdentifier + key_pointer: + - name + type_pointer: + - type + retriever: + type: SimpleRetriever + requester: + type: HttpRequester + url_base: https://api.test.com + path: /data + http_method: GET + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: [] + paginator: + type: NoPagination + $parameters: + name: "test_stream" + primary_key: "id" +""" + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + stream_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["stream_with_dynamic_schema"], {} + ) + + stream = factory.create_component( + model_type=DeclarativeStreamModel, + component_definition=stream_manifest, + config=input_config, + ) + + assert isinstance(stream, DefaultStream) + schema_loader = get_schema_loader(stream) + assert isinstance(schema_loader, DynamicSchemaLoader) + assert isinstance(schema_loader.retriever, AsyncRetriever) + + +def test_dynamic_schema_loader_keeps_custom_retriever_stream_slicer(): + content = """ +schema_loader: + type: DynamicSchemaLoader + retriever: + type: CustomRetriever + class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever + stream_slicer: + type: ListPartitionRouter + cursor_field: p + values: ["a"] + requester: + type: HttpRequester + url: https://api.test.com/schema + record_selector: + type: RecordSelector + extractor: + type: DpathExtractor + field_path: [] + schema_type_identifier: + type: SchemaTypeIdentifier + key_pointer: + - name +""" + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + schema_loader_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["schema_loader"], {} + ) + + schema_loader = factory.create_component( + model_type=DynamicSchemaLoaderModel, + component_definition=schema_loader_manifest, + config=input_config, + ) + + assert isinstance(schema_loader.retriever, TestingCustomRetriever) + assert isinstance(schema_loader.retriever.stream_slicer, ListPartitionRouter) + + def test_api_budget(): manifest = { "type": "DeclarativeSource", diff --git a/unit_tests/sources/declarative/schema/test_dynamic_schema_loader.py b/unit_tests/sources/declarative/schema/test_dynamic_schema_loader.py index 20147465fe..ec20333079 100644 --- a/unit_tests/sources/declarative/schema/test_dynamic_schema_loader.py +++ b/unit_tests/sources/declarative/schema/test_dynamic_schema_loader.py @@ -4,16 +4,20 @@ import json from copy import deepcopy -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest +from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level, SyncMode, Type +from airbyte_cdk.sources.declarative.async_job.job_orchestrator import AsyncJobOrchestrator from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( ConcurrentDeclarativeSource, TestLimits, ) from airbyte_cdk.sources.declarative.schema import DynamicSchemaLoader, SchemaTypeIdentifier +from airbyte_cdk.test.catalog_builder import CatalogBuilder from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse +from airbyte_cdk.utils.traced_exception import AirbyteTracedException _CONFIG = { "start_date": "2024-07-01T00:00:00.000Z", @@ -414,3 +418,218 @@ def test_dynamic_schema_loader_with_type_conditions(): assert len(actual_catalog.streams) == 1 assert actual_catalog.streams[0].json_schema == expected_schema + + +def _async_schema_loader_manifest(download_decoder, schema_type_identifier, partition_router=None): + retriever = { + "type": "AsyncRetriever", + "status_mapping": { + "failed": ["failed"], + "running": ["pending"], + "timeout": ["timeout"], + "completed": ["ready"], + }, + "status_extractor": {"type": "DpathExtractor", "field_path": ["status"]}, + "download_target_extractor": {"type": "DpathExtractor", "field_path": ["urls"]}, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "creation_requester": { + "type": "HttpRequester", + "url": "https://api.test.com/schema_job", + "http_method": "POST", + }, + "polling_requester": { + "type": "HttpRequester", + "url": "https://api.test.com/schema_job/{{ creation_response['id'] }}", + "http_method": "GET", + }, + "download_requester": { + "type": "HttpRequester", + "url": "{{ download_target }}", + "http_method": "GET", + }, + } + if download_decoder: + retriever["download_decoder"] = download_decoder + if partition_router: + retriever["partition_router"] = partition_router + manifest = deepcopy(_MANIFEST) + manifest["definitions"]["party_members_stream"]["schema_loader"] = { + "type": "DynamicSchemaLoader", + "retriever": retriever, + "schema_type_identifier": schema_type_identifier, + } + return manifest + + +@pytest.mark.parametrize( + "download_decoder, schema_type_identifier, download_body, expected_properties", + [ + pytest.param( + None, + {"schema_pointer": ["fields"], "key_pointer": ["name"], "type_pointer": ["type"]}, + json.dumps( + {"fields": [{"name": "id", "type": "integer"}, {"name": "name", "type": "string"}]} + ), + {"id": {"type": ["null", "integer"]}, "name": {"type": ["null", "string"]}}, + id="json_download", + ), + pytest.param( + {"type": "CsvDecoder", "encoding": "utf-8", "delimiter": ","}, + {"schema_pointer": [], "key_pointer": []}, + "id,name,amount\n1,foo,3\n2,bar,4\n", + { + "id": {"type": ["null", "string"]}, + "name": {"type": ["null", "string"]}, + "amount": {"type": ["null", "string"]}, + }, + id="csv_download_header_as_schema", + ), + pytest.param( + None, + {"schema_pointer": ["fields"], "key_pointer": ["name"], "type_pointer": ["type"]}, + "", + {}, + id="empty_download", + ), + ], +) +@patch.object(AsyncJobOrchestrator, "_WAIT_TIME_BETWEEN_STATUS_UPDATE_IN_SECONDS", 0) +def test_dynamic_schema_loader_with_async_retriever( + download_decoder, schema_type_identifier, download_body, expected_properties +): + manifest = _async_schema_loader_manifest(download_decoder, schema_type_identifier) + source = ConcurrentDeclarativeSource( + source_config=manifest, config=_CONFIG, catalog=None, state=None + ) + creation = HttpRequest(url="https://api.test.com/schema_job") + polling = HttpRequest(url="https://api.test.com/schema_job/job_1") + download = HttpRequest(url="https://api.test.com/download/1") + + with HttpMocker() as http_mocker: + http_mocker.post(creation, HttpResponse(body=json.dumps({"id": "job_1"}))) + http_mocker.get( + polling, + HttpResponse( + body=json.dumps({"status": "ready", "urls": ["https://api.test.com/download/1"]}) + ), + ) + http_mocker.get(download, HttpResponse(body=download_body)) + + actual_catalog = source.discover(logger=source.logger, config=_CONFIG) + + http_mocker.assert_number_of_calls(creation, 1) + http_mocker.assert_number_of_calls(polling, 1) + http_mocker.assert_number_of_calls(download, 1) + + assert actual_catalog.streams[0].json_schema["properties"] == expected_properties + + +@patch.object(AsyncJobOrchestrator, "_WAIT_TIME_BETWEEN_STATUS_UPDATE_IN_SECONDS", 0) +def test_read_with_async_retriever_and_async_dynamic_schema_loader(): + manifest = _async_schema_loader_manifest( + None, {"schema_pointer": ["fields"], "key_pointer": ["name"], "type_pointer": ["type"]} + ) + stream = manifest["definitions"]["party_members_stream"] + data_retriever = deepcopy(stream["schema_loader"]["retriever"]) + data_retriever["creation_requester"]["url"] = "https://api.test.com/data_job" + data_retriever["polling_requester"]["url"] = ( + "https://api.test.com/data_job/{{ creation_response['id'] }}" + ) + # casts records to the loaded schema, so the output proves which schema the read used + data_retriever["record_selector"]["schema_normalization"] = "Default" + stream["retriever"] = data_retriever + catalog = CatalogBuilder().with_stream("party_members", SyncMode.full_refresh).build() + source = ConcurrentDeclarativeSource( + source_config=manifest, config=_CONFIG, catalog=catalog, state=None + ) + schema_creation = HttpRequest(url="https://api.test.com/schema_job") + + with HttpMocker() as http_mocker: + http_mocker.post(schema_creation, HttpResponse(body=json.dumps({"id": "s1"}))) + http_mocker.get( + HttpRequest(url="https://api.test.com/schema_job/s1"), + HttpResponse( + body=json.dumps( + {"status": "ready", "urls": ["https://api.test.com/download/schema"]} + ) + ), + ) + http_mocker.get( + HttpRequest(url="https://api.test.com/download/schema"), + HttpResponse(body=json.dumps({"fields": [{"name": "id", "type": "integer"}]})), + ) + http_mocker.post( + HttpRequest(url="https://api.test.com/data_job"), + HttpResponse(body=json.dumps({"id": "d1"})), + ) + http_mocker.get( + HttpRequest(url="https://api.test.com/data_job/d1"), + HttpResponse( + body=json.dumps({"status": "ready", "urls": ["https://api.test.com/download/data"]}) + ), + ) + http_mocker.get( + HttpRequest(url="https://api.test.com/download/data"), + HttpResponse(body=json.dumps({"id": "1"})), + ) + + messages = list(source.read(source.logger, _CONFIG, catalog, None)) + + http_mocker.assert_number_of_calls(schema_creation, 1) + + assert [message.record.data for message in messages if message.type == Type.RECORD] == [ + {"id": 1} + ] + + +@patch.object(AsyncJobOrchestrator, "_WAIT_TIME_BETWEEN_STATUS_UPDATE_IN_SECONDS", 0) +def test_discover_with_failed_async_schema_job_raises(): + manifest = _async_schema_loader_manifest( + None, {"schema_pointer": ["fields"], "key_pointer": ["name"]} + ) + source = ConcurrentDeclarativeSource( + source_config=manifest, config=_CONFIG, catalog=None, state=None + ) + + with HttpMocker() as http_mocker: + http_mocker.post( + HttpRequest(url="https://api.test.com/schema_job"), + HttpResponse(body=json.dumps({"id": "job_1"})), + ) + http_mocker.get( + HttpRequest(url="https://api.test.com/schema_job/job_1"), + HttpResponse(body=json.dumps({"status": "failed"})), + ) + + with pytest.raises(AirbyteTracedException, match="async jobs failed"): + source.discover(logger=source.logger, config=_CONFIG) + + +def test_dynamic_schema_loader_with_async_retriever_and_partition_router_raises(): + manifest = _async_schema_loader_manifest( + None, + {"schema_pointer": ["fields"], "key_pointer": ["name"]}, + partition_router={"type": "ListPartitionRouter", "cursor_field": "p", "values": ["a", "b"]}, + ) + source = ConcurrentDeclarativeSource( + source_config=manifest, config=_CONFIG, catalog=None, state=None + ) + + with pytest.raises(ValueError, match="DynamicSchemaLoader does not support an AsyncRetriever"): + source.streams(config=_CONFIG) + + +def test_dynamic_schema_loader_skips_non_record_items(dynamic_schema_loader): + log_message = AirbyteMessage( + type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message="slice") + ) + dynamic_schema_loader.retriever.read_records.return_value = iter( + [log_message, {"schema": [{"key": "name", "type": "string"}]}] + ) + + schema = dynamic_schema_loader.get_json_schema() + + assert schema["properties"] == {"name": {"type": ["null", "string"]}}