diff --git a/airbyte_cdk/config_observation.py b/airbyte_cdk/config_observation.py index ae85e8277..7abb3c3d4 100644 --- a/airbyte_cdk/config_observation.py +++ b/airbyte_cdk/config_observation.py @@ -87,12 +87,19 @@ def observe_connector_config( def emit_configuration_as_airbyte_control_message(config: MutableMapping[str, Any]) -> None: - """ - WARNING: deprecated - emit_configuration_as_airbyte_control_message is being deprecated in favor of the MessageRepository mechanism. - See the airbyte_cdk.sources.message package + """Print the connector config as a CONNECTOR_CONFIG control message directly to stdout. + + Single-use OAuth refresh token rotation relies on this direct emission: it reaches the + platform for every command (including `check` and `discover`) and does not depend on a + message repository being drained. The payload and its newline are written in one `write()` + call so a concurrent writer to the shared stdout buffer cannot split the JSON line. """ airbyte_message = create_connector_config_control_message(config) - print(orjson.dumps(AirbyteMessageSerializer.dump(airbyte_message)).decode()) + print( + f"{orjson.dumps(AirbyteMessageSerializer.dump(airbyte_message)).decode()}\n", + end="", + flush=True, + ) def create_connector_config_control_message(config: MutableMapping[str, Any]) -> AirbyteMessage: 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 aa546c73a..364782c39 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -2984,6 +2984,7 @@ def create_oauth_authenticator( token_expiry_date_format=model.token_expiry_date_format, token_expiry_is_time_of_expiration=bool(model.token_expiry_date_format), message_repository=self._message_repository, + emit_control_message_to_message_repository=self._emit_connector_builder_messages, refresh_token_error_status_codes=refresh_token_error_status_codes, refresh_token_error_key=refresh_token_error_key, refresh_token_error_values=refresh_token_error_values, diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py b/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py index e4243b3a2..1bacced2c 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +++ b/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py @@ -181,6 +181,7 @@ def __init__( refresh_token_error_status_codes: Tuple[int, ...] = (), refresh_token_error_key: str = "", refresh_token_error_values: Tuple[str, ...] = (), + emit_control_message_to_message_repository: bool = False, ) -> None: """ Args: @@ -201,7 +202,8 @@ def __init__( token_expiry_date_config_path (Sequence[str]): Dpath to the token_expiry_date field in the connector configuration. Defaults to ("credentials", "token_expiry_date"). token_expiry_date_format (Optional[str]): Date format of the token expiry date field (set by expires_in_name). If not specified the token expiry date is interpreted as number of seconds until expiration. token_expiry_is_time_of_expiration bool: set True it if expires_in is returned as time of expiration instead of the number seconds until expiration - message_repository (MessageRepository): the message repository used to emit logs on HTTP requests and control message on config update + message_repository (MessageRepository): The message repository used to emit logs on HTTP requests and control message on config update. + emit_control_message_to_message_repository (bool): Whether the message repository additionally receives the CONNECTOR_CONFIG message. This is only for in-process consumers such as the Connector Builder, because messages in the repository are otherwise printed to stdout by the entrypoint and would duplicate the direct stdout emission. Defaults to False. """ self._connector_config = connector_config self._client_id: str = self._get_config_value_by_path( @@ -220,6 +222,9 @@ def __init__( self._grant_type_name = grant_type_name self._connector_config = connector_config self.__message_repository = message_repository + self._emit_control_message_to_message_repository = ( + emit_control_message_to_message_repository + ) super().__init__( token_refresh_endpoint=token_refresh_endpoint, client_id_name=self._client_id_name, @@ -407,24 +412,20 @@ def _get_config_value_by_path( ) def _emit_control_message(self) -> None: + """Emit the updated connector config as a CONNECTOR_CONFIG control message. + + The message is always printed to stdout, which is the delivery the platform relies on: it + is immediate and independent of whether the source exposes a message repository and of + when that repository is drained (for `ConcurrentDeclarativeSource` the internal repository + is only drained during `read`, after partitions are processed). The message repository is + only used in addition when explicitly requested, for in-process consumers such as the + Connector Builder; repository messages are otherwise printed by the entrypoint too and + would duplicate the stdout emission. """ - Emits a control message based on the connector configuration. - - Control messages for config updates (like refreshed tokens) must be printed directly - to stdout so the platform can process them immediately. The message repository is - also used to queue the message for any additional processing. - - Note: - The function `emit_configuration_as_airbyte_control_message` prints directly to - stdout, which is required for the platform to detect and persist config changes. - """ - # Always emit to stdout so the platform can process the config update immediately. - # This is critical for single-use refresh tokens where the new token must be persisted - # before subsequent operations try to use the old (now invalid) token. emit_configuration_as_airbyte_control_message(self._connector_config) # type: ignore[arg-type] - - # Also emit to the message repository for any additional processing (e.g., logging) - if not isinstance(self._message_repository, NoopMessageRepository): + if self._emit_control_message_to_message_repository and not isinstance( + self._message_repository, NoopMessageRepository + ): self._message_repository.emit_message( create_connector_config_control_message(self._connector_config) # type: ignore[arg-type] ) diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index e2fef671c..06046e2d6 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -44,6 +44,7 @@ DestinationSyncMode, FailureType, Level, + OrchestratorType, Status, StreamDescriptor, SyncMode, @@ -52,6 +53,7 @@ from airbyte_cdk.sources.declarative.async_job.job_tracker import ConcurrentJobLimitReached from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( ConcurrentDeclarativeSource, + TestLimits, ) from airbyte_cdk.sources.declarative.extractors.record_filter import ( ClientSideIncrementalRecordFilterDecorator, @@ -6113,3 +6115,173 @@ def capturing_resolve(resolver_self, *args, **kwargs): "HttpComponentsResolver's requester should have api_budget set during dynamic stream " "discovery, but it was None. This means discovery HTTP requests are not rate-limited." ) + + +def _single_use_refresh_token_manifest(include_check: bool = True) -> Dict[str, Any]: + manifest: Dict[str, Any] = { + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": [ + { + "type": "DeclarativeStream", + "name": "items", + "primary_key": "id", + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": { + "type": "object", + "properties": {"id": {"type": "string"}}, + }, + }, + "retriever": { + "type": "SimpleRetriever", + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "paginator": {"type": "NoPagination"}, + "requester": { + "type": "HttpRequester", + "url_base": "https://api.example.com", + "path": "/items", + "http_method": "GET", + "authenticator": { + "type": "OAuthAuthenticator", + "token_refresh_endpoint": "https://api.example.com/oauth/token", + "client_id": "{{ config['credentials']['client_id'] }}", + "client_secret": "{{ config['credentials']['client_secret'] }}", + "refresh_token": "{{ config['credentials']['refresh_token'] }}", + "refresh_token_updater": {}, + }, + }, + }, + } + ], + } + if include_check: + manifest["check"] = {"type": "CheckStream", "stream_names": ["items"]} + return manifest + + +def _single_use_refresh_token_config() -> Dict[str, Any]: + return { + "credentials": { + "client_id": "client_id", + "client_secret": "client_secret", + "refresh_token": "old_refresh", + } + } + + +def _single_use_refresh_token_catalog() -> ConfiguredAirbyteCatalog: + return ConfiguredAirbyteCatalog( + streams=[ + ConfiguredAirbyteStream( + stream=AirbyteStream( + name="items", + json_schema={"type": "object", "properties": {"id": {"type": "string"}}}, + supported_sync_modes=[SyncMode.full_refresh], + ), + sync_mode=SyncMode.full_refresh, + destination_sync_mode=DestinationSyncMode.append, + ) + ] + ) + + +def _mock_single_use_refresh_token_requests(http_mocker: HttpMocker) -> None: + http_mocker.post( + HttpRequest( + "https://api.example.com/oauth/token", + body="grant_type=refresh_token&client_id=client_id&client_secret=client_secret&refresh_token=old_refresh", + ), + HttpResponse( + json.dumps( + { + "access_token": "new_access", + "refresh_token": "new_refresh", + "expires_in": 3600, + } + ) + ), + ) + http_mocker.get( + HttpRequest("https://api.example.com/items"), + HttpResponse(json.dumps([{"id": "item-1"}])), + ) + + +def _connector_config_lines(output: str) -> List[Dict[str, Any]]: + parsed = [json.loads(line) for line in output.splitlines() if line.strip()] + return [message for message in parsed if message.get("type") == "CONTROL"] + + +def test_given_refresh_token_updater_when_read_then_exactly_one_connector_config_message_on_stdout( + capsys, +): + config = _single_use_refresh_token_config() + catalog = _single_use_refresh_token_catalog() + source = ConcurrentDeclarativeSource( + source_config=_single_use_refresh_token_manifest(), + config=config, + catalog=catalog, + state=None, + ) + + with HttpMocker() as http_mocker: + _mock_single_use_refresh_token_requests(http_mocker) + messages = list(source.read(logger, config, catalog, [])) + + assert not any(message.type == Type.CONTROL for message in messages) + assert any(message.type == Type.RECORD for message in messages) + connector_config_messages = _connector_config_lines(capsys.readouterr().out) + assert len(connector_config_messages) == 1 + assert ( + connector_config_messages[0]["control"]["connectorConfig"]["config"]["credentials"][ + "refresh_token" + ] + == "new_refresh" + ) + + +def test_given_refresh_token_updater_and_connector_builder_mode_when_read_then_control_message_yielded_once( + capsys, +): + config = _single_use_refresh_token_config() + catalog = _single_use_refresh_token_catalog() + source = ConcurrentDeclarativeSource( + source_config=_single_use_refresh_token_manifest(), + config=config, + catalog=catalog, + state=None, + emit_connector_builder_messages=True, + limits=TestLimits(max_records=10, max_pages_per_slice=10, max_slices=10, max_streams=10), + ) + + with HttpMocker() as http_mocker: + _mock_single_use_refresh_token_requests(http_mocker) + messages = list(source.read(logger, config, catalog, [])) + + yielded_control_messages = [message for message in messages if message.type == Type.CONTROL] + assert len(yielded_control_messages) == 1 + assert yielded_control_messages[0].control.type == OrchestratorType.CONNECTOR_CONFIG + connector_config_messages = _connector_config_lines(capsys.readouterr().out) + assert len(connector_config_messages) == 1 + + +def test_given_refresh_token_updater_when_check_then_connector_config_message_on_stdout(capsys): + config = _single_use_refresh_token_config() + source = ConcurrentDeclarativeSource( + source_config=_single_use_refresh_token_manifest(), + config=config, + catalog=None, + state=None, + ) + + with HttpMocker() as http_mocker: + _mock_single_use_refresh_token_requests(http_mocker) + connection_status = source.check(logger, config) + + assert connection_status.status == Status.SUCCEEDED + connector_config_messages = _connector_config_lines(capsys.readouterr().out) + assert len(connector_config_messages) == 1 diff --git a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py index 3a7cb8d3d..723a916c7 100644 --- a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py +++ b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py @@ -17,6 +17,7 @@ from requests.exceptions import RequestException from airbyte_cdk.models import FailureType, OrchestratorType, Type +from airbyte_cdk.sources.message import InMemoryMessageRepository, NoopMessageRepository from airbyte_cdk.sources.streams.http.requests_native_auth import ( BasicHttpAuthenticator, MultipleTokenAuthenticator, @@ -726,7 +727,9 @@ def test_given_no_message_repository_get_access_token( authenticator.token_has_expired = mocker.Mock(return_value=True) access_token = authenticator.get_access_token() captured = capsys.readouterr() - airbyte_message = json.loads(captured.out) + messages = [json.loads(line) for line in captured.out.splitlines() if line.strip()] + assert len(messages) == 1 + airbyte_message = messages[0] expected_new_config = connector_config.copy() expected_new_config["credentials"]["access_token"] = "new_access_token" expected_new_config["credentials"]["refresh_token"] = "new_refresh_token" @@ -741,10 +744,10 @@ def test_given_no_message_repository_get_access_token( assert not captured.out assert authenticator.access_token == access_token == "new_access_token" - def test_given_message_repository_when_get_access_token_then_emit_message( - self, mocker, connector_config + def test_given_message_repository_when_get_access_token_then_emit_exactly_one_control_message_to_stdout( + self, mocker, connector_config, capsys ): - message_repository = Mock() + message_repository = InMemoryMessageRepository() authenticator = SingleUseRefreshTokenOauth2Authenticator( connector_config, token_refresh_endpoint="https://refresh_endpoint.com", @@ -771,8 +774,66 @@ def test_given_message_repository_when_get_access_token_then_emit_message( authenticator.get_access_token() - emitted_message = message_repository.emit_message.call_args_list[0].args[0] - assert emitted_message.type == Type.CONTROL + messages = [ + json.loads(line) for line in capsys.readouterr().out.splitlines() if line.strip() + ] + assert len(messages) == 1 + emitted_message = messages[0] + assert emitted_message["type"] == "CONTROL" + assert emitted_message["control"]["type"] == "CONNECTOR_CONFIG" + assert ( + emitted_message["control"]["connectorConfig"]["config"]["credentials"]["access_token"] + == "new_access_token" + ) + assert ( + emitted_message["control"]["connectorConfig"]["config"]["credentials"]["refresh_token"] + == "new_refresh_token" + ) + assert [m for m in message_repository.consume_queue() if m.type == Type.CONTROL] == [] + + def test_given_emit_control_message_to_message_repository_when_get_access_token_then_emit_to_stdout_and_repository( + self, mocker, connector_config, capsys + ): + message_repository = InMemoryMessageRepository() + authenticator = SingleUseRefreshTokenOauth2Authenticator( + connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + client_id=connector_config["credentials"]["client_id"], + client_secret=connector_config["credentials"]["client_secret"], + token_expiry_is_time_of_expiration=True, + token_expiry_date_format="YYYY-MM-DD", + message_repository=message_repository, + emit_control_message_to_message_repository=True, + ) + resp.status_code = 200 + mocker.patch.object( + resp, + "json", + return_value={ + authenticator.get_access_token_name(): "new_access_token", + authenticator.get_expires_in_name(): "2023-04-04", + authenticator.get_refresh_token_name(): "new_refresh_token", + }, + ) + mocker.patch.object(requests, "request", side_effect=mock_request, autospec=True) + + authenticator.token_has_expired = mocker.Mock(return_value=True) + + authenticator.get_access_token() + + messages = [ + json.loads(line) for line in capsys.readouterr().out.splitlines() if line.strip() + ] + assert len(messages) == 1 + assert messages[0]["type"] == "CONTROL" + assert messages[0]["control"]["type"] == "CONNECTOR_CONFIG" + emitted_messages = [ + message + for message in message_repository.consume_queue() + if message.type == Type.CONTROL + ] + assert len(emitted_messages) == 1 + emitted_message = emitted_messages[0] assert emitted_message.control.type == OrchestratorType.CONNECTOR_CONFIG assert ( emitted_message.control.connectorConfig.config["credentials"]["access_token"] @@ -795,6 +856,40 @@ def test_given_message_repository_when_get_access_token_then_emit_message( == "my_client_secret" ) + def test_given_noop_message_repository_and_emit_flag_when_get_access_token_then_only_stdout( + self, mocker, connector_config, capsys + ): + authenticator = SingleUseRefreshTokenOauth2Authenticator( + connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + client_id=connector_config["credentials"]["client_id"], + client_secret=connector_config["credentials"]["client_secret"], + message_repository=NoopMessageRepository(), + emit_control_message_to_message_repository=True, + ) + resp.status_code = 200 + mocker.patch.object( + resp, + "json", + return_value={ + authenticator.get_access_token_name(): "new_access_token", + authenticator.get_expires_in_name(): 3600, + authenticator.get_refresh_token_name(): "new_refresh_token", + }, + ) + mocker.patch.object(requests, "request", side_effect=mock_request, autospec=True) + + authenticator.token_has_expired = mocker.Mock(return_value=True) + + authenticator.get_access_token() + + messages = [ + json.loads(line) for line in capsys.readouterr().out.splitlines() if line.strip() + ] + assert len(messages) == 1 + assert messages[0]["type"] == "CONTROL" + assert messages[0]["control"]["type"] == "CONNECTOR_CONFIG" + def test_given_message_repository_when_get_access_token_then_log_request( self, mocker, connector_config ):