Skip to content

fix(low-code): document shared retrievers, lock session token refresh - #1201

Merged
Baz (bazarnov) merged 1 commit into
mainfrom
bazarnov/fix-shared-retriever-contract
Oct 6, 2026
Merged

Baz (bazarnov) merged 1 commit into
mainfrom
bazarnov/fix-shared-retriever-contract

Conversation

@bazarnov

@bazarnov Baz (bazarnov) commented Oct 6, 2026 •

Copy link
Copy Markdown
Contributor

Overview

👉 TL;DR: Documents that a stream reads all its partitions in parallel through one shared set of components, so custom code must not keep per-read state on itself. It also fixes the built-in session login, which could log in several times at once, or fail with a missing token, when partitions started together.

Specifically, this states the contract on DeclarativePartitionFactory, Retriever, PaginationStrategy, CustomRetriever and CustomPaginationStrategy, and makes SessionTokenProvider refresh its cached token under a double-checked lock and set the token before its expiry.

Pairs with airbytehq/airbyte#88062, which updates the Custom Components docs page to the same contract:

Changes

  • DeclarativePartitionFactory.__init__ docstring: replaced the stale text, which said every thread gets its own retriever, with what the code does. Every partition reads through the one retriever, concurrently, which keeps one authenticator, HTTP session and async job repository per stream.
  • Retriever and PaginationStrategy docstrings: one instance serves every partition at once, so per-read state such as page tokens, counters or recursion depth belongs in local variables or in method arguments such as last_page_token_value.
  • CustomRetriever and CustomPaginationStrategy descriptions in declarative_component_schema.yaml: one sentence each with the same rule. The generated models do not include class descriptions, so declarative_component_schema.py is unchanged.
  • SessionTokenProvider: with expiration_duration set, the refresh runs under a per-instance lock and re-checks the expiry once the lock is held, and _token is assigned before _next_expiration_time. Without expiration_duration, every call still logs in, without the lock.
  • test_token_provider.py: two tests. Two calls that find the token expired together send one login request and both get the new token. The token is already stored when its expiry is assigned.

Review Spotlight

Reviewers with limited time, please review first:

Risks

  • Behaviour changes only for connectors that use SessionTokenAuthenticator with expiration_duration. Concurrent partitions of a stream now share one login instead of each sending their own. A call can no longer fail with Failed to get session token, token is None while another call is logging in. That is 12 connectors, counted with git grep over the manifest.yaml files on airbyte master: 18 use SessionTokenAuthenticator (source-metabase uses the legacy class), and 12 of them set expiration_duration.
  • The other 6 do not set expiration_duration and keep logging in on every call, without the lock, as before.
  • If a login fails, the calls waiting on the lock then try their own logins one after another instead of in parallel. A failing login therefore takes longer to fail every partition of the stream.
  • The lock is per provider instance, and each stream builds its own provider. Logins of different streams, or of different connectors in one process such as the manifest server, never wait for each other.
  • A SessionTokenProvider can no longer be deep-copied or pickled (TypeError: cannot pickle '_thread.lock' object). A factory-built one could before. Nothing copies one: the CDK only deep-copies manifests, configs, state, schemas and records. The 9 deepcopy, copy.copy and pickle.dump calls in non-test source connector code on airbyte master all copy data (git grep). RateLimitedMultipleTokenAuthenticator in the same package already holds instance locks.
  • Everything else is docs: no manifest field, type or default changes. Raw datamodel-codegen 0.26.3 output for the schema is byte-identical before and after.

Test plan

  • Red/green. With the source files reverted to origin/main (git checkout origin/main -- <source files>), both new tests fail in 100 out of 100 runs:

    • test_session_token_provider_refreshes_once_for_concurrent_calls: assert 2 == 1 (two login requests)
    • test_session_token_provider_sets_token_before_expiry: assert [None] == ['my_token']

    With the fix, both pass in 200 out of 200 runs.

  • python -m pytest unit_tests/sources/declarative/auth unit_tests/sources/declarative/stream_slicers unit_tests/sources/declarative/retrievers unit_tests/sources/declarative/requesters/paginators unit_tests/sources/declarative/parsers -q: 866 passed.

  • python -m pytest unit_tests/sources/declarative/test_declarative_component_schema_is_loadable.py unit_tests/legacy/sources/declarative/test_manifest_declarative_source.py unit_tests/sources/declarative/test_concurrent_declarative_source.py unit_tests/sources/declarative/test_concurrent_declarative_source_config_overrides.py unit_tests/sources/declarative/interpolation/test_jinja.py unit_tests/sources/declarative/spec/test_spec.py -k "not test_read_with_concurrent_and_synchronous_streams" -q: 237 passed, 3 deselected. The deselected tests can crash on macOS because of the SQLite cache race tracked in Cached parent stream with more than one partition corrupts the shared sqlite request cache under concurrent reads (no such table: responses, file is not a database, no more rows available) #1150.

  • ruff format, ruff check and mypy --config-file mypy.ini on the changed files: clean.

  • prettier 3.0.3 on declarative_component_schema.yaml: the same 5 findings as on main, none on the edited lines.

Design notes

Show/Hide Content
  • Per-instance lock instead of the class-level RLock that AbstractOauth2Authenticator uses: the guarded state belongs to the instance. A class-level lock would make logins of unrelated streams and connectors wait for each other, and a plain class-level Lock would deadlock a login requester that authenticates with its own SessionTokenAuthenticator.
  • No lock without expiration_duration: nothing is cached, so every call needs its own login anyway, and a lock would only make those logins run one at a time.
  • The tests hold both calls at their unlocked expiry check with a Barrier on the first two clock reads, and record the token at the moment the expiry is assigned. A call waiting on the lock cannot be observed, so a test that blocks the login on an Event would need a timeout to tell a correct run from a broken one.

Fixes #1129

A stream reads all its partitions through one retriever, from several
threads at once. The DeclarativePartitionFactory docstring still said
the opposite and no extension point told custom component authors, so
per-read state kept on a custom retriever or pagination strategy can
silently drop records.

- State the contract on DeclarativePartitionFactory, Retriever,
  PaginationStrategy and the CustomRetriever and CustomPaginationStrategy
  schema descriptions.
- SessionTokenProvider broke it: concurrent calls each sent a login
  request, and the expiry was set before the token, so a call in between
  got "Failed to get session token, token is None". Refresh the cached
  token under a double-checked lock and set the token before its expiry.
  Without expiration_duration every call still logs in, unlocked.

Fixes #1129
@bazarnov
Baz (bazarnov) requested a review from a team as a code owner October 6, 2026 07:58
@github-actions

github-actions Bot commented Oct 6, 2026

Copy link
Copy Markdown

👋 Greetings, Airbyte Team Member!

Here are some helpful tips and reminders for your convenience.

💡 Show Tips and Tricks

Testing This CDK Version

You can test this version of the CDK using the following:

# Run the CLI from this branch:
uvx 'git+https://github.com/airbytehq/airbyte-python-cdk.git@bazarnov/fix-shared-retriever-contract#egg=airbyte-python-cdk[dev]' --help

# Update a connector to use the CDK from this branch ref:
cd airbyte-integrations/connectors/source-example
poe use-cdk-branch bazarnov/fix-shared-retriever-contract

PR Slash Commands

Airbyte Maintainers can execute the following slash commands on your PR:

  • /autofix - Fixes most formatting and linting issues
  • /poetry-lock - Updates poetry.lock file
  • /test - Runs connector tests with the updated CDK
  • /prerelease - Triggers a prerelease publish with default arguments
  • /poe build - Regenerate git-committed build artifacts, such as the pydantic models which are generated from the manifest JSON schema in YAML.
  • /poe <command> - Runs any poe command in the CDK environment
📚 Show Repo Guidance

Helpful Resources

📝 Edit this welcome message.

@github-actions

github-actions Bot commented Oct 6, 2026

Copy link
Copy Markdown

PyTest Results (Fast)

4 984 tests  +2   4 972 ✅ +2   11m 7s ⏱️ + 1m 34s
    1 suites ±0      12 💤 ±0 
    1 files   ±0       0 ❌ ±0 

Results for commit 0919b1e. ± Comparison against base commit d214b0c.

@github-actions

github-actions Bot commented Oct 6, 2026

Copy link
Copy Markdown

PyTest Results (Full)

4 987 tests  +2   4 975 ✅ +2   15m 55s ⏱️ -9s
    1 suites ±0      12 💤 ±0 
    1 files   ±0       0 ❌ ±0 

Results for commit 0919b1e. ± Comparison against base commit d214b0c.

@bazarnov
Baz (bazarnov) merged commit bc90c8d into main Oct 6, 2026
31 of 32 checks passed
@bazarnov
Baz (bazarnov) deleted the bazarnov/fix-shared-retriever-contract branch October 6, 2026 09:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

DeclarativePartitionFactory shares one retriever instance across worker threads despite per-thread docstring

2 participants