diff --git a/airbyte_cdk/sources/declarative/auth/token_provider.py b/airbyte_cdk/sources/declarative/auth/token_provider.py index f3abd85296..4a75022c11 100644 --- a/airbyte_cdk/sources/declarative/auth/token_provider.py +++ b/airbyte_cdk/sources/declarative/auth/token_provider.py @@ -4,6 +4,7 @@ import datetime +import threading from abc import abstractmethod from dataclasses import InitVar, dataclass, field from typing import Any, List, Mapping, Optional, Union @@ -40,6 +41,9 @@ class SessionTokenProvider(TokenProvider): _next_expiration_time: Optional[AirbyteDateTime] = None _token: Optional[str] = None + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + self._lock = threading.Lock() + def get_token(self) -> str: self._refresh_if_necessary() if self._token is None: @@ -47,8 +51,16 @@ def get_token(self) -> str: return self._token def _refresh_if_necessary(self) -> None: - if self._next_expiration_time is None or self._next_expiration_time < ab_datetime_now(): + if self.expiration_duration is None: + # Without an expiry nothing is cached and every call logs in, so there is nothing to lock self._refresh() + elif self._token_has_expired(): + with self._lock: + if self._token_has_expired(): + self._refresh() + + def _token_has_expired(self) -> bool: + return self._next_expiration_time is None or self._next_expiration_time < ab_datetime_now() def _refresh(self) -> None: response = self.login_requester.send_request( @@ -64,9 +76,10 @@ def _refresh(self) -> None: if response is None: raise ReadException("Failed to get session token, response got ignored by requester") session_token = dpath.get(next(self.decoder.decode(response)), self.session_token_path) + # Set the token first: `get_token` reads it without the lock as soon as the expiry is fresh + self._token = session_token # type: ignore # Returned decoded response will be Mapping and therefore session_token will be str or None if self.expiration_duration is not None: self._next_expiration_time = ab_datetime_now() + self.expiration_duration - self._token = session_token # type: ignore # Returned decoded response will be Mapping and therefore session_token will be str or None @dataclass diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index c45fb8c96f..dfabccd6b0 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -868,7 +868,7 @@ definitions: additionalProperties: true CustomPaginationStrategy: title: Custom Pagination Strategy - description: Pagination strategy component whose behavior is derived from a custom code implementation of the connector. + description: Pagination strategy component whose behavior is derived from a custom code implementation of the connector. One instance paginates every partition of the stream in several threads at once, so the next token must be derived from the method arguments, not from state kept on the instance. type: object additionalProperties: true required: @@ -953,7 +953,7 @@ definitions: additionalProperties: true CustomRetriever: title: Custom Retriever - description: Retriever component whose behavior is derived from a custom code implementation of the connector. + description: Retriever component whose behavior is derived from a custom code implementation of the connector. One instance reads every partition of the stream and its `read_records` runs in several threads at once, so per-read state must live in local variables, not on the instance. type: object additionalProperties: true required: diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py index a7a89f5400..69ed9c459c 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py @@ -15,6 +15,9 @@ class PaginationStrategy: """ Defines how to get the next page token + + One instance paginates every partition of the stream, several at once, so `next_page_token` must derive the + next token from its arguments, such as `last_page_token_value`, rather than from state kept on the instance. """ @property diff --git a/airbyte_cdk/sources/declarative/retrievers/retriever.py b/airbyte_cdk/sources/declarative/retrievers/retriever.py index 89db62c2c2..1e0b9e6c04 100644 --- a/airbyte_cdk/sources/declarative/retrievers/retriever.py +++ b/airbyte_cdk/sources/declarative/retrievers/retriever.py @@ -14,6 +14,10 @@ class Retriever: """ Responsible for fetching a stream's records from an HTTP API source. + + A stream creates one retriever and reads all its partitions through it: `read_records` is called from several + threads at once, one call per stream slice. Implementations, custom ones included, must keep per-read state such + as page tokens, counters or recursion depth in local variables rather than on the instance or its components. """ @abstractmethod diff --git a/airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py b/airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py index 593b616eec..ab5350d919 100644 --- a/airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py +++ b/airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py @@ -41,9 +41,9 @@ def __init__( max_records_limit: Optional[int] = None, ) -> None: """ - The DeclarativePartitionFactory takes a retriever_factory and not a retriever directly. The reason is that our components are not - thread safe and classes like `DefaultPaginator` may not work because multiple threads can access and modify a shared field across each other. - In order to avoid these problems, we will create one retriever per thread which should make the processing thread-safe. + Every partition this factory creates reads through the same `retriever`, and the worker threads read partitions + concurrently. Sharing it keeps one authenticator, HTTP session and async job repository per stream; see `Retriever` + for what it requires from implementations. """ self._stream_name = stream_name self._schema_loader = schema_loader diff --git a/unit_tests/sources/declarative/auth/test_token_provider.py b/unit_tests/sources/declarative/auth/test_token_provider.py index 076d39741f..de593946ca 100644 --- a/unit_tests/sources/declarative/auth/test_token_provider.py +++ b/unit_tests/sources/declarative/auth/test_token_provider.py @@ -1,8 +1,11 @@ # # Copyright (c) 2023 Airbyte, Inc., all rights reserved. # +import itertools import json -from unittest.mock import MagicMock +import threading +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock, patch import freezegun import pytest @@ -14,6 +17,7 @@ SessionTokenProvider, ) from airbyte_cdk.sources.declarative.exceptions import ReadException +from airbyte_cdk.utils.datetime_helpers import ab_datetime_now def create_session_token_provider(): @@ -82,6 +86,52 @@ def test_session_token_provider_ignored_response(): provider.get_token() +def test_session_token_provider_refreshes_once_for_concurrent_calls(): + with freezegun.freeze_time("2001-05-21T12:00:00Z"): + provider = create_session_token_provider() + provider.get_token() + provider.login_requester.send_request.reset_mock() + provider.login_requester.send_request.return_value.content = json.dumps( + {"nested": {"token": "updated_token"}} + ).encode() + + both_calls_found_token_expired = threading.Barrier(2, timeout=10) + clock_reads = itertools.count() + + def now(): + # The first two clock reads are the unlocked expiry checks of the two calls + if next(clock_reads) < 2: + both_calls_found_token_expired.wait() + return ab_datetime_now() + + with patch( + "airbyte_cdk.sources.declarative.auth.token_provider.ab_datetime_now", side_effect=now + ): + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(provider.get_token) for _ in range(2)] + + assert [future.result() for future in futures] == ["updated_token", "updated_token"] + assert provider.login_requester.send_request.call_count == 1 + + +def test_session_token_provider_sets_token_before_expiry(): + """ + `get_token` skips the lock once the expiry is fresh, so a fresh expiry must never be visible before its token. + """ + provider = create_session_token_provider() + tokens_when_expiry_was_set = [] + + def record_token_when_expiry_is_set(self, name, value): + object.__setattr__(self, name, value) + if name == "_next_expiration_time": + tokens_when_expiry_was_set.append(self._token) + + with patch.object(SessionTokenProvider, "__setattr__", record_token_when_expiry_is_set): + assert provider.get_token() == "my_token" + + assert tokens_when_expiry_was_set == ["my_token"] + + @pytest.mark.parametrize( "api_token_template,expected_token", [