From ef67b081a49a35e439a1ebb8d2053cfe5cfd957f Mon Sep 17 00:00:00 2001 From: hai-ld <1579471+hai-ld@users.noreply.github.com> Date: Mon, 21 Apr 2025 13:01:47 +0700 Subject: [PATCH 1/2] Allow SessionTokenAuthenticator to be created as a subcomponent to a custom component Change `name` argument in the function `create_session_token_authenticator` to a keyword-only argument so that it can be injected via `$parameters`. Other `name` arguments in the same module are already keyword-only. --- .../sources/declarative/parsers/model_to_component_factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 5d2415525f..8b8026c230 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1084,7 +1084,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]: decoder = ( self._create_component_from_model(model=model.decoder, config=config) From c8ed141ac0207cda3c2b7ee07468e6fa7f506734 Mon Sep 17 00:00:00 2001 From: Baz <22987674+bazarnov@users.noreply.github.com> Date: Mon, 28 Sep 2026 15:45:28 +0300 Subject: [PATCH 2/2] fix(low-code): pass the stream name to authenticators nested in SelectiveAuthenticator --- .../parsers/model_to_component_factory.py | 4 +- .../test_model_to_component_factory.py | 143 ++++++++++++++++++ .../declarative/parsers/testing_components.py | 20 +++ .../test_concurrent_declarative_source.py | 105 +++++++++++++ 4 files changed, 270 insertions(+), 2 deletions(-) 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 94dfd8635c..bb4255e2da 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -3460,8 +3460,8 @@ def create_selective_authenticator( self, model: SelectiveAuthenticatorModel, config: Config, **kwargs: Any ) -> DeclarativeAuthenticator: 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, **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 ea3b0cf151..4dfc95954e 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 @@ -67,6 +67,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, @@ -2779,6 +2780,148 @@ 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", + }, + }, +} + + +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": { + "type": "LegacySessionTokenAuthenticator", + "header": "token", + "login_url": "login", + "session_token_response_key": "session", + "validate_session_url": "validate", + }, + "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 test_create_composite_error_handler(): content = """ error_handler: 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 c8aaca355b..0c5089bb7e 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -5522,6 +5522,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