feat(cloud): add connection state and catalog management to CloudConnection#975
Conversation
Add typed connection state get/set operations with stream-level filtering: - New connection_state module with Pydantic models (ConnectionStateResponse, StreamState, GlobalState, StreamDescriptor) with camelCase API aliases - CloudConnection.get_state() - typed state retrieval with optional stream filter - CloudConnection.set_state() - full state update via safe endpoint (HTTP 423 guard) - CloudConnection.get_stream_state() - convenience single-stream state getter - CloudConnection.set_stream_state() - fetch-modify-push single stream update - API layer: create_or_update_connection_state_safe() in api_util - Unit tests for models and helper functions (15 tests) Co-Authored-By: AJ Steers <aj@airbyte.io>
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
👋 Greetings, Airbyte Team Member!Here are some helpful tips and reminders for your convenience. 💡 Show Tips and TricksTesting This PyAirbyte VersionYou can test this version of PyAirbyte using the following: # Run PyAirbyte CLI from this branch:
uvx --from 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1770419251-connection-state-management' pyairbyte --help
# Install PyAirbyte from this branch for development:
pip install 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1770419251-connection-state-management'PR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
📚 Show Repo GuidanceHelpful ResourcesCommunity SupportQuestions? Join the #pyairbyte channel in our Slack workspace. |
Co-Authored-By: AJ Steers <aj@airbyte.io>
📝 WalkthroughWalkthroughAdds structured connection-state models and helpers, API helpers to replace full connection state and catalog (mapping HTTP 423 to a new exception), CloudConnection methods for dumping/importing and per-stream edits, a new AirbyteConnectionSyncActiveError, and unit tests for state models and helpers. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant CloudConnection
participant StateModels
participant API as "api_util"
participant Backend as "Airbyte Backend"
User->>CloudConnection: dump_raw_state(connection_id)
CloudConnection->>API: GET connection state
API->>Backend: GET /state/{connectionId}
Backend-->>API: ConnectionStateResponse JSON
API-->>CloudConnection: dict
CloudConnection->>StateModels: parse -> ConnectionStateResponse
User->>CloudConnection: set_stream_state(name, new_stream_state)
CloudConnection->>CloudConnection: dump_raw_state()
CloudConnection->>StateModels: _get_stream_list() / _match_stream()
CloudConnection->>CloudConnection: build full state envelope
CloudConnection->>API: POST replace_connection_state(connection_id, full_state)
API->>Backend: POST /v1/state/create_or_update_safe
alt success
Backend-->>API: updated ConnectionStateResponse JSON
API-->>CloudConnection: dict
CloudConnection-->>User: updated state dict
else 423 LOCKED
Backend-->>API: HTTP 423
API-->>CloudConnection: raises AirbyteConnectionSyncActiveError
CloudConnection-->>User: exception
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Would you like me to add suggested unit tests for the new catalog-replace paths similar to the state tests, wdyt? 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@airbyte/cloud/connection_state.py`:
- Around line 110-120: The current _match_stream treats stream_namespace=None as
a wildcard which lets callers (e.g., set_stream_state) match any namespace for a
given stream name; update the _match_stream docstring to explicitly state that
None is treated as a wildcard and will match streams with any namespace, and
update set_stream_state to detect when stream_namespace is None and more than
one matching StreamState exists and emit a warning (or raise/log) so callers are
alerted to potential ambiguity; reference _match_stream and set_stream_state
when making these changes.
In `@airbyte/cloud/connections.py`:
- Around line 543-555: The streamDescriptor currently omits namespace when
stream_namespace is an empty string because the condition uses "if
stream_namespace"; update the conditional to explicitly check for presence using
"if stream_namespace is not None" so that empty-string namespaces are preserved;
modify the construction of new_stream_entry (the streamDescriptor dict creation
around stream_namespace) to use this explicit None check, referencing the
variables new_stream_entry, streamDescriptor, and stream_namespace.
🧹 Nitpick comments (3)
airbyte/_util/api_util.py (1)
2095-2135: Clean addition following existing patterns.The function mirrors the structure of
get_connection_statenicely. One thought: the docstring mentions HTTP 423 behavior, but_make_config_api_requestwill surface it as a genericAirbyteError. Would it be helpful to catch 423 specifically and raise a more descriptive error (e.g., "Cannot update state while a sync is running") so users get actionable feedback, wdyt?airbyte/cloud/connections.py (2)
453-478:set_statedoesn't validate thatconnection_statecontains a matchingconnectionId— could this lead to silent mismatches?The docstring says
connectionIdinside the dict "Must match this connection's ID," but there's no enforcement. If a caller passes a state dict from a different connection (e.g., copy-paste error), the API would receive a top-levelconnectionId(fromself.connection_id) that might conflict with the one insideconnection_state. Depending on which one the API respects, this could either silently do the right thing or corrupt another connection's state. A quick guard here could prevent a painful debugging session, wdyt?🛡️ Proposed validation
def set_state( self, connection_state: dict[str, Any], ) -> ConnectionStateResponse: + if ( + "connectionId" in connection_state + and connection_state["connectionId"] != self.connection_id + ): + raise PyAirbyteInputError( + message="connectionId in state payload does not match this connection.", + context={ + "expected": self.connection_id, + "actual": connection_state["connectionId"], + }, + ) updated_data = api_util.create_or_update_connection_state_safe(
504-583:set_stream_stateuses a fetch-modify-push pattern — inherent TOCTOU, but acceptable given the safe endpoint.The read-modify-write on lines 529→583 means two concurrent
set_stream_statecalls could race, with the last writer winning. The "safe" endpoint guards against sync-time corruption (423), but not against concurrent client updates. This is likely acceptable for the use case, but maybe worth a brief note in the docstring so callers are aware, wdyt?
Co-Authored-By: AJ Steers <aj@airbyte.io>
…m merge Co-Authored-By: AJ Steers <aj@airbyte.io>
…of Pydantic models Co-Authored-By: AJ Steers <aj@airbyte.io>
…state Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@airbyte/cloud/connections.py`:
- Around line 537-548: The loop that replaces matched streams uses _match_stream
with stream_namespace=None as a wildcard, causing multiple streams with the same
name but different namespaces to each be replaced by the one new_stream_entry
and corrupt state; change the logic in the block that iterates over streams
(uses streams, _match_stream, updated_streams_raw, new_stream_entry) to first
collect all matching indices, and if more than one match is found when
stream_namespace is None either raise a clear exception (e.g., ValueError) or
log a warning and abort the replace to avoid duplicating namespace-less entries;
otherwise proceed to replace exactly one matched entry (or append if none) so
you never insert multiple copies of new_stream_entry.
🧹 Nitpick comments (1)
airbyte/_util/api_util.py (1)
2095-2142: Clean implementation, consistent with existing patterns.The docstring does a great job explaining the full-replacement semantics and the HTTP 423 behavior. One minor thought: the
connection_statedict likely contains its ownconnectionIdfield, and there's no validation that it matches theconnection_idparameter passed to this function. A mismatch could lead to confusing API errors. Would it be worth adding a quick sanity check, or is that better left to the caller, wdyt?💡 Optional validation
+ # Sanity check: if connectionState includes a connectionId, it should match + inner_conn_id = connection_state.get("connectionId") + if inner_conn_id and inner_conn_id != connection_id: + raise PyAirbyteInputError( + message="connectionId in connection_state does not match the connection_id parameter.", + context={ + "connection_id": connection_id, + "connection_state_connection_id": inner_conn_id, + }, + ) return _make_config_api_request( path="/state/create_or_update_safe",
…tate Co-Authored-By: AJ Steers <aj@airbyte.io>
…update_connection_state Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
…api_util Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
…tionSyncActiveError Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@airbyte/_util/api_util.py`:
- Around line 2222-2233: The call to _make_config_api_request in the function
that updates connections currently sets "skipReset": False which will trigger a
destination reset; change this to a safer default by making skipReset True or,
better, add a caller-configurable parameter (e.g., skip_reset) to the
surrounding function and pass it through to _make_config_api_request instead of
the hardcoded False; update the call site to set "skipReset": skip_reset
(default True) and adjust any callers/tests accordingly so the reset behavior is
explicit.
🧹 Nitpick comments (1)
airbyte/_util/api_util.py (1)
38-43: DuplicateTYPE_CHECKINGblock — could you merge these two?Lines 38-40 and 42-47 both guard
if TYPE_CHECKING:imports. The first one importsCallableand the second importsCallableagain (plusDestinationConfiguration). This results in a redundant import ofCallableand two separate blocks that could be consolidated, wdyt?🔧 Proposed fix
if TYPE_CHECKING: from collections.abc import Callable - - -if TYPE_CHECKING: - from collections.abc import Callable from airbyte_api.models import ( DestinationConfiguration, )
…ethods Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
…//git-manager.devin.ai/proxy/github.com/airbytehq/PyAirbyte into devin/1770419251-connection-state-management
…ict, state_blob_dict, configured_catalog_dict) Co-Authored-By: AJ Steers <aj@airbyte.io>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@airbyte/cloud/connections.py`:
- Around line 547-563: The zip(raw_streams, streams, strict=False) can drop
trailing raw entries when parsed streams are shorter; replace the zip loop with
an index-based loop over raw_streams (e.g., for i, raw_s in
enumerate(raw_streams): parsed_s = streams[i] if i < len(streams) else None) and
then use _match_stream(parsed_s, stream_name, stream_namespace) (or check
parsed_s is not None before matching) to decide whether to append
new_stream_entry or raw_s to updated_streams_raw; this ensures no raw_streams
are silently dropped if _get_stream_list(current) returns fewer entries.
- Around line 618-640: import_raw_catalog currently calls
replace_connection_catalog directly and will surface a generic AirbyteError when
the backend returns HTTP 423; update replace_connection_catalog (or wrap its
call inside import_raw_catalog) to mirror the HTTP 423 handling used by
replace_connection_state/import_raw_state by catching the error from
replace_connection_catalog and raising AirbyteConnectionSyncActiveError when the
response indicates a sync is active (HTTP 423). Locate the call to
replace_connection_catalog in import_raw_catalog and either add the same
try/except that maps 423 → AirbyteConnectionSyncActiveError (reusing the exact
logic from replace_connection_state) or modify replace_connection_catalog to
perform that mapping so both import_raw_catalog and existing callers receive the
clearer sync-active exception.
🧹 Nitpick comments (1)
airbyte/_util/api_util.py (1)
1467-1472: No timeout on therequests.request()call — could block indefinitely.
_make_config_api_requestdoesn't pass atimeouttorequests.request(). If the config API hangs, callers (including the new state/catalog methods) would block forever. Would adding a reasonable default timeout (e.g.,timeout=60) be worth considering here, wdyt?🔧 Proposed fix
response = requests.request( method="POST", url=full_url, headers=headers, json=json, + timeout=60, )
devin-ai-integration
Bot
changed the title
Co-Authored-By: AJ Steers <aj@airbyte.io>
Co-Authored-By: AJ Steers <aj@airbyte.io>
…t destructive resets Co-Authored-By: AJ Steers <aj@airbyte.io>
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
Responses to remaining CodeRabbit nitpick suggestionsDuplicate No timeout on Multiple matches debug log in |
Co-Authored-By: AJ Steers <aj@airbyte.io>
Aaron ("AJ") Steers (aaronsteers)
deleted the
devin/1770419251-connection-state-management
branch
1 participant
Summary
Adds connection state and catalog dump/import operations to
CloudConnection, extracted from airbytehq/airbyte-ops-mcp#338 as reusable core library functionality. This enables both MCP tools and CLI commands to share the same implementation via PyAirbyte.Private module
airbyte/cloud/_connection_state.py:ConnectionStateResponse,StreamState,GlobalState,StreamDescriptor) with camelCase aliases — used internally only, not exported in the public API_match_stream()and_get_stream_list()for stream-level filteringNew
CloudConnectionmethods (all inputs/outputs are raw dicts):dump_raw_state()import_raw_state(connection_state_dict)connectionIdfor cross-connection portability. Safe endpoint rejects updates during active syncs (HTTP 423).get_stream_state(stream_name)set_stream_state(stream_name, state_blob_dict)dump_raw_catalog()syncCatalog) as a dictimport_raw_catalog(catalog)API layer (
api_util.py):replace_connection_state()— callsPOST /v1/state/create_or_update_safe, catches HTTP 423 →AirbyteConnectionSyncActiveErrorreplace_connection_catalog()— callsPOST /v1/web_backend/connections/updatewithskipReset: FalseOther changes:
get_state_artifacts()andget_catalog_artifact()deprecated via@deprecateddecoratorAirbyteConnectionSyncActiveErrorexception for HTTP 423get_connection_artifacttool updated to usedump_raw_state()/dump_raw_catalog()set_stream_statepreserves raw dicts for non-matching streams (avoids Pydantic round-trip data loss)Review & Testing Checklist for Human
set_stream_statefetch-modify-push correctness (HIGH RISK): This method useszip(raw_streams, parsed_streams)to pair raw API dicts with Pydantic-parsed models for matching. Non-matching streams are preserved as raw dicts (no Pydantic round-trip), but the zip pairing assumes the raw and parsed lists are in the same order and same length. Must be tested against a real connection with bothstreamandglobalstate types.import_raw_catalogwithskipReset: False: Catalog replacement triggers destination table resets. For backup/restore workflows where the catalog hasn't actually changed, this could unnecessarily reset destination tables. Evaluate whetherskipResetshould be a caller-configurable parameter.connectionIdoverride behavior:replace_connection_statedoes{**connection_state_dict, "connectionId": connection_id}. Verify the API accepts this and doesn't requireconnectionIdin a nested location._match_stream()treatsNoneand""as equivalent (both mean "no namespace"). There is no wildcard behavior — caller must match the actual namespace. Verify this is acceptable.Suggested test plan:
Notes
CloudConnectionmethods — they require real API credentialsLink to Devin run: https://app.devin.ai/sessions/5ce605c67c5f436d94d5d76a2c6b95fc
Requested by: Aaron ("AJ") Steers (@aaronsteers)
Summary by CodeRabbit
New Features
Bug Fixes / Reliability
Deprecated
Tests
API
Important
Auto-merge enabled.
This PR is set to merge automatically when all requirements are met.