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 a44dc22b75..841511f922 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1242,7 +1242,7 @@ def create_legacy_to_per_partition_state_migration( ) def create_session_token_authenticator( - self, model: SessionTokenAuthenticatorModel, config: Config, name: str, **kwargs: Any + self, model: SessionTokenAuthenticatorModel, config: Config, *, name: str, **kwargs: Any ) -> Union[ApiKeyAuthenticator, BearerAuthenticator]: self._reject_reduce_page_size_action( model.login_requester, f"`login_requester` of the SessionTokenAuthenticator of {name}" @@ -1864,8 +1864,15 @@ def _validate_jitter_range(jitter_range_in_seconds: Optional[float]) -> None: raise ValueError("jitter_range_in_seconds must be greater than or equal to 0") def create_cursor_pagination( - self, model: CursorPaginationModel, config: Config, decoder: Decoder, **kwargs: Any + self, + model: CursorPaginationModel, + config: Config, + decoder: Optional[Decoder] = None, + **kwargs: Any, ) -> CursorPaginationStrategy: + # None when nested under a custom component, which cannot pass a decoder + if decoder is None: + decoder = JsonDecoder(parameters={}) if isinstance(decoder, PaginationDecoderDecorator): inner_decoder = decoder.decoder else: @@ -3337,12 +3344,15 @@ def create_offset_increment( self, model: OffsetIncrementModel, config: Config, - decoder: Decoder, + decoder: Optional[Decoder] = None, extractor_model: Optional[ Union[CustomRecordExtractorModel, DpathExtractorModel, CombinedExtractorModel] ] = None, **kwargs: Any, ) -> OffsetIncrement: + # None when nested under a custom component, which cannot pass a decoder + if decoder is None: + decoder = JsonDecoder(parameters={}) if isinstance(decoder, PaginationDecoderDecorator): inner_decoder = decoder.decoder else: @@ -3663,11 +3673,27 @@ def create_remove_fields( ) def create_selective_authenticator( - self, model: SelectiveAuthenticatorModel, config: Config, **kwargs: Any + self, + model: SelectiveAuthenticatorModel, + config: Config, + *, + name: Optional[str] = None, + url_base: Optional[str] = None, + **kwargs: Any, ) -> DeclarativeAuthenticator: + # Keyword-only so that _create_nested_component also fills them under a custom component, + # from the parent's kwargs or $parameters. None is not forwarded, so a missing required one + # keeps its hint. + nested_kwargs = { + key: value + for key, value in {"name": name, "url_base": url_base}.items() + if value is not None + } authenticators = { - name: self._create_component_from_model(model=auth, config=config) - for name, auth in model.authenticators.items() + key: self._create_component_from_model( + model=auth, config=config, **nested_kwargs, **kwargs + ) + for key, auth in model.authenticators.items() } # SelectiveAuthenticator will return instance of DeclarativeAuthenticator or raise ValueError error return SelectiveAuthenticator( # type: ignore[abstract] 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 2933b05754..67c15c661f 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 @@ -77,6 +77,7 @@ CompositeErrorHandler as CompositeErrorHandlerModel, ) from airbyte_cdk.sources.declarative.models import ConcurrencyLevel as ConcurrencyLevelModel +from airbyte_cdk.sources.declarative.models import CustomAuthenticator as CustomAuthenticatorModel from airbyte_cdk.sources.declarative.models import CustomErrorHandler as CustomErrorHandlerModel from airbyte_cdk.sources.declarative.models import ( CustomPartitionRouter as CustomPartitionRouterModel, @@ -3092,6 +3093,193 @@ def test_create_requester_with_selective_authenticator(input_config, expected_au assert isinstance(authenticator, expected_authenticator_class) +_SESSION_TOKEN_AUTHENTICATOR = { + "type": "SessionTokenAuthenticator", + "login_requester": { + "type": "HttpRequester", + "url_base": "https://api.sendgrid.com", + "path": "/session", + "http_method": "POST", + }, + "session_token_path": ["id"], + "request_authentication": { + "type": "ApiKey", + "inject_into": { + "type": "RequestOption", + "inject_into": "header", + "field_name": "X-Session", + }, + }, +} + +_LEGACY_SESSION_TOKEN_AUTHENTICATOR = { + "type": "LegacySessionTokenAuthenticator", + "header": "token", + "login_url": "login", + "session_token_response_key": "session", + "validate_session_url": "validate", +} + + +def _create_requester_with_authenticator(authenticator, config): + return factory.create_component( + model_type=HttpRequesterModel, + component_definition={ + "type": "HttpRequester", + "url_base": "https://api.sendgrid.com", + "path": "/v3/marketing/lists", + "authenticator": authenticator, + }, + config=config, + name="lists", + decoder=None, + ) + + +@pytest.mark.parametrize( + "auth_type, expected_authenticator_class", + [ + pytest.param("session", ApiKeyAuthenticator, id="test_session_token_selected"), + pytest.param("legacy", LegacySessionTokenAuthenticator, id="test_legacy_session_selected"), + pytest.param("token", ApiKeyAuthenticator, id="test_token_selected"), + pytest.param("oauth", DeclarativeOauth2Authenticator, id="test_oauth_selected"), + ], +) +def test_create_requester_with_selective_authenticator_nesting_session_authenticators( + auth_type, expected_authenticator_class +): + requester = _create_requester_with_authenticator( + { + "type": "SelectiveAuthenticator", + "authenticator_selection_path": ["auth", "type"], + "authenticators": { + "session": _SESSION_TOKEN_AUTHENTICATOR, + "legacy": _LEGACY_SESSION_TOKEN_AUTHENTICATOR, + "token": {"type": "ApiKeyAuthenticator", "header": "X-Key", "api_token": "key"}, + "oauth": { + "type": "OAuthAuthenticator", + "token_refresh_endpoint": "https://api.url.com", + "client_id": "some_id", + "client_secret": "some_secret", + "refresh_token": "some_token", + }, + }, + }, + {**input_config, "auth": {"type": auth_type}}, + ) + + assert isinstance(requester.authenticator, expected_authenticator_class) + if auth_type == "legacy": + assert requester.authenticator._api_url.eval(input_config) == "https://api.sendgrid.com" + + +def test_create_requester_with_selective_authenticator_passes_stream_name_to_session_token_authenticator(): + requester = _create_requester_with_authenticator( + { + "type": "SelectiveAuthenticator", + "authenticator_selection_path": ["auth", "type"], + "authenticators": {"session": _SESSION_TOKEN_AUTHENTICATOR}, + }, + {**input_config, "auth": {"type": "session"}}, + ) + + session_token_provider = requester.authenticator.token_provider.session_token_provider + # the stream name, not the selector key, like a SessionTokenAuthenticator set on the requester directly + assert session_token_provider.login_requester.name == "lists_login_requester" + + +def test_create_requester_with_custom_authenticator_nesting_session_token_authenticator(): + requester = _create_requester_with_authenticator( + { + "type": "CustomAuthenticator", + "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingCustomAuthenticator", + "inner": _SESSION_TOKEN_AUTHENTICATOR, + }, + input_config, + ) + + session_token_provider = requester.authenticator.inner.token_provider.session_token_provider + assert session_token_provider.login_requester.name == "lists_login_requester" + + +def test_create_custom_authenticator_nesting_session_token_authenticator_with_name_in_parameters(): + authenticator = factory.create_component( + model_type=CustomAuthenticatorModel, + component_definition={ + "type": "CustomAuthenticator", + "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingCustomAuthenticator", + "inner": {**_SESSION_TOKEN_AUTHENTICATOR, "$parameters": {"name": "explicit"}}, + }, + config=input_config, + ) + + session_token_provider = authenticator.inner.token_provider.session_token_provider + assert session_token_provider.login_requester.name == "explicit_login_requester" + + +def test_create_requester_with_selective_authenticator_nesting_custom_authenticator(): + requester = _create_requester_with_authenticator( + { + "type": "SelectiveAuthenticator", + "authenticator_selection_path": ["auth", "type"], + "authenticators": { + "custom": { + "type": "CustomAuthenticator", + "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingCustomAuthenticator", + "inner": _SESSION_TOKEN_AUTHENTICATOR, + }, + }, + }, + {**input_config, "auth": {"type": "custom"}}, + ) + + session_token_provider = requester.authenticator.inner.token_provider.session_token_provider + assert session_token_provider.login_requester.name == "lists_login_requester" + + +def _create_custom_requester_with_selective_authenticator(authenticator, parameters=None): + return factory.create_component( + model_type=CustomRequesterModel, + component_definition={ + "type": "CustomRequester", + "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingRequester", + "url_base": "https://api.sendgrid.com", + "path": "/v3/marketing/lists", + "authenticator": { + "type": "SelectiveAuthenticator", + "authenticator_selection_path": ["auth", "type"], + "authenticators": {"selected": authenticator}, + "$parameters": parameters or {}, + }, + }, + config={**input_config, "auth": {"type": "selected"}}, + name="lists", + ) + + +def test_create_custom_requester_with_selective_authenticator_passes_stream_name_to_session_token_authenticator(): + requester = _create_custom_requester_with_selective_authenticator(_SESSION_TOKEN_AUTHENTICATOR) + + session_token_provider = requester.authenticator.token_provider.session_token_provider + assert session_token_provider.login_requester.name == "lists_login_requester" + + +def test_create_custom_requester_with_selective_authenticator_passes_url_base_parameter_to_legacy_session_token_authenticator(): + requester = _create_custom_requester_with_selective_authenticator( + _LEGACY_SESSION_TOKEN_AUTHENTICATOR, {"url_base": "https://api.legacy.com"} + ) + + assert isinstance(requester.authenticator, LegacySessionTokenAuthenticator) + assert requester.authenticator._api_url.eval(input_config) == "https://api.legacy.com" + + +def test_create_custom_requester_with_selective_authenticator_without_url_base_for_legacy_session_token_authenticator_raises(): + with pytest.raises( + ValueError, match=r"Please provide SelectiveAuthenticator\.\$parameters\.url_base$" + ): + _create_custom_requester_with_selective_authenticator(_LEGACY_SESSION_TOKEN_AUTHENTICATOR) + + def test_create_composite_error_handler(): content = """ error_handler: @@ -3635,12 +3823,32 @@ def test_custom_components_do_not_contain_extra_fields(): assert custom_substream_partition_router.custom_pagination_strategy.page_size == 100 -def test_parse_custom_component_fields_if_subcomponent(): +@pytest.mark.parametrize( + "custom_pagination_strategy, expected_strategy_class", + [ + pytest.param( + {"type": "PageIncrement", "page_size": 100}, PageIncrement, id="test_page_increment" + ), + pytest.param( + {"type": "CursorPagination", "cursor_value": "{{ response.next }}", "page_size": 100}, + CursorPaginationStrategy, + id="test_cursor_pagination_without_decoder", + ), + pytest.param( + {"type": "OffsetIncrement", "page_size": 100}, + OffsetIncrement, + id="test_offset_increment_without_decoder", + ), + ], +) +def test_parse_custom_component_fields_if_subcomponent( + custom_pagination_strategy, expected_strategy_class +): custom_substream_partition_router_manifest = { "type": "CustomPartitionRouter", "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingCustomSubstreamPartitionRouter", "custom_field": "here", - "custom_pagination_strategy": {"type": "PageIncrement", "page_size": 100}, + "custom_pagination_strategy": custom_pagination_strategy, "parent_stream_configs": [ { "type": "ParentStreamConfig", @@ -3702,8 +3910,11 @@ def test_parse_custom_component_fields_if_subcomponent(): == "repository_id" ) - assert isinstance(custom_substream_partition_router.custom_pagination_strategy, PageIncrement) - assert custom_substream_partition_router.custom_pagination_strategy.page_size == 100 + strategy = custom_substream_partition_router.custom_pagination_strategy + assert isinstance(strategy, expected_strategy_class) + assert strategy.get_page_size() == 100 + if expected_strategy_class is not PageIncrement: + assert isinstance(strategy.decoder.decoder, JsonDecoder) class TestCreateTransformations: diff --git a/unit_tests/sources/declarative/parsers/testing_components.py b/unit_tests/sources/declarative/parsers/testing_components.py index b21b63efd6..9d5ecd9642 100644 --- a/unit_tests/sources/declarative/parsers/testing_components.py +++ b/unit_tests/sources/declarative/parsers/testing_components.py @@ -5,6 +5,7 @@ from dataclasses import dataclass, field from typing import Any, ClassVar, List, Mapping, Optional +from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator from airbyte_cdk.sources.declarative.extractors import DpathExtractor from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration from airbyte_cdk.sources.declarative.partition_routers import SubstreamPartitionRouter @@ -99,6 +100,25 @@ class TestingCustomErrorHandler(DefaultErrorHandler): __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name +@dataclass +class TestingCustomAuthenticator(DeclarativeAuthenticator): + """ + A test class wrapping a nested authenticator, used for testing custom authenticators with subcomponents. + """ + + __test__: ClassVar[bool] = False # Tell Pytest this is not a Pytest class, despite its name + + inner: DeclarativeAuthenticator + + @property + def auth_header(self) -> str: + return self.inner.auth_header + + @property + def token(self) -> str: + return self.inner.token + + @dataclass class TestingRequester(HttpRequester): request_parameters: Optional[RequestInput] = None diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index 1a65aa0b94..bbf9ac8994 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -5687,6 +5687,111 @@ def test_given_record_selector_is_filtering_when_read_then_raise_error(): list(source.read(logger=source.logger, config=input_config, catalog=catalog, state=[])) +@pytest.mark.parametrize( + "auth_type, expected_headers, expected_login_calls", + [ + pytest.param("session", {"X-Session": "a_session_token"}, 1, id="test_session_selected"), + pytest.param("token", {"X-Key": "a_key"}, 0, id="test_token_selected"), + ], +) +def test_given_selective_authenticator_nesting_session_token_authenticator_when_read_then_authenticate( + auth_type, expected_headers, expected_login_calls +): + input_config = {"auth_type": auth_type, "username": "a_user", "api_key": "a_key"} + manifest = { + "version": "0.34.2", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["Test"]}, + "streams": [ + { + "type": "DeclarativeStream", + "name": "Test", + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object"}, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://example.org", + "path": "/test", + "authenticator": { + "type": "SelectiveAuthenticator", + "authenticator_selection_path": ["auth_type"], + "authenticators": { + "session": { + "type": "SessionTokenAuthenticator", + "login_requester": { + "type": "HttpRequester", + "url_base": "https://example.org", + "path": "/login", + "http_method": "POST", + "request_body_json": { + "username": "{{ config['username'] }}" + }, + }, + "session_token_path": ["token"], + "expiration_duration": "PT1H", + "request_authentication": { + "type": "ApiKey", + "inject_into": { + "type": "RequestOption", + "inject_into": "header", + "field_name": "X-Session", + }, + }, + }, + "token": { + "type": "ApiKeyAuthenticator", + "header": "X-Key", + "api_token": "{{ config['api_key'] }}", + }, + }, + }, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + }, + } + ], + "spec": { + "type": "Spec", + "documentation_url": "https://example.org", + "connection_specification": {}, + }, + } + + catalog = create_catalog("Test") + source = ConcurrentDeclarativeSource( + source_config=manifest, + config=input_config, + catalog=catalog, + state=None, + ) + login_request = HttpRequest( + "https://example.org/login", body=json.dumps({"username": "a_user"}) + ) + + with HttpMocker() as http_mocker: + http_mocker.post(login_request, HttpResponse(json.dumps({"token": "a_session_token"}))) + http_mocker.get( + HttpRequest("https://example.org/test", headers=expected_headers), + HttpResponse(json.dumps([{"id": 1}])), + ) + messages = list( + source.read(logger=source.logger, config=input_config, catalog=catalog, state=[]) + ) + + http_mocker.assert_number_of_calls(login_request, expected_login_calls) + + assert [message.record.data for message in messages if message.type == Type.RECORD] == [ + {"id": 1} + ] + + def _make_default_stream(name: str) -> DefaultStream: """Create a minimal DefaultStream instance for testing.""" from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor