feat: Add MCP tool and Python API for connection enable/disable and scheduling #923
Conversation
Co-Authored-By: AJ Steers <[email protected]>
Original prompt from AJ Steers |
🤖 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. Testing 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/1766175960-connection-enable-disable-schedule' pyairbyte --help
# Install PyAirbyte from this branch for development:
pip install 'git+https://github.com/airbytehq/PyAirbyte.git@devin/1766175960-connection-enable-disable-schedule'Helpful ResourcesPR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
Community SupportQuestions? Join the #pyairbyte channel in our Slack workspace. |
📝 WalkthroughWalkthroughAdds enable/disable and CRON/manual scheduling methods to CloudConnection (refreshing state, patching via API, updating internal info, returning self) and a single MCP tool function Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant MCP_Tool as update_cloud_connection
participant WorkspaceResolver as WorkspaceResolver
participant CloudWorkspace as CloudWorkspace
participant CloudConnection as CloudConnection
participant AirbyteAPI as AirbyteAPI
User->>MCP_Tool: call update_cloud_connection(params)
MCP_Tool->>WorkspaceResolver: resolve workspace (workspace_id?)
WorkspaceResolver-->>MCP_Tool: CloudWorkspace
MCP_Tool->>CloudWorkspace: get connection by id
CloudWorkspace-->>MCP_Tool: CloudConnection
MCP_Tool->>CloudConnection: refresh state / fetch current info
CloudConnection->>AirbyteAPI: GET /connections/{id}
AirbyteAPI-->>CloudConnection: connection info
alt enabled change provided
MCP_Tool->>CloudConnection: set_enabled(enabled)
CloudConnection->>AirbyteAPI: PATCH /connections/{id} (status payload)
AirbyteAPI-->>CloudConnection: updated connection info
end
alt cron_expression provided
MCP_Tool->>CloudConnection: set_schedule(cron_expression)
CloudConnection->>AirbyteAPI: PATCH /connections/{id} (schedule payload)
AirbyteAPI-->>CloudConnection: updated connection info
end
alt manual_schedule provided
MCP_Tool->>CloudConnection: set_manual_schedule()
CloudConnection->>AirbyteAPI: PATCH /connections/{id} (manual schedule payload)
AirbyteAPI-->>CloudConnection: updated connection info
end
CloudConnection-->>MCP_Tool: returns updated instance
MCP_Tool-->>User: formatted success message with connection URL and summary
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ 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: 1
🧹 Nitpick comments (1)
airbyte/cloud/connections.py (1)
438-468: Consider validating the cron expression format, wdyt?The implementation follows the established pattern—no input validation before making the API call. This is consistent across all CloudConnection setter methods (
rename(),set_table_prefix(),set_selected_streams(), etc.), suggesting a deliberate design choice to delegate validation to the API. Invalid formats will be rejected by the API, though users may see less detailed error messages upstream.If you'd prefer to add client-side validation for better UX, that would be a nice enhancement. Otherwise, leaving it to the API is a reasonable approach given the existing architecture.
Since no unit tests were found for this method, it would be good to verify that:
- Invalid cron expressions are rejected gracefully by the API
- The schedule is correctly applied and persisted
📜 Review details
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
airbyte/cloud/connections.py(2 hunks)airbyte/mcp/cloud_ops.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
airbyte/mcp/cloud_ops.py (4)
airbyte/mcp/_tool_utils.py (1)
mcp_tool(189-235)tests/integration_tests/cloud/conftest.py (1)
workspace_id(49-50)airbyte/cloud/workspaces.py (1)
get_connection(315-327)airbyte/cloud/connections.py (4)
enable(397-414)disable(416-434)set_schedule(438-468)set_manual_schedule(470-490)
airbyte/cloud/connections.py (1)
airbyte/_util/api_util.py (1)
patch_connection(1337-1399)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Pytest (All, Python 3.10, Windows)
- GitHub Check: Pytest (All, Python 3.11, Windows)
- GitHub Check: Pytest (All, Python 3.11, Ubuntu)
- GitHub Check: Pytest (All, Python 3.10, Ubuntu)
- GitHub Check: Pytest (Fast)
- GitHub Check: Pytest (No Creds)
🔇 Additional comments (3)
airbyte/cloud/connections.py (3)
19-19: LGTM - appropriate lint suppression.The noqa comment is justified given the intentional expansion of the class's public API surface with these new lifecycle and scheduling methods.
397-434: LGTM - enable/disable implementation looks solid.Both methods follow the established pattern in this class (similar to
rename()andset_table_prefix()), and I've verified from the provided code snippets thatapi_util.patch_connectiondoes accept thestatusparameter with typemodels.ConnectionStatusEnum | None.The docstrings clearly explain the behavior, particularly the note in
disable()that manual syncs can still be triggered.
470-490: LGTM - clean implementation.The method is straightforward and correctly implements manual scheduling by setting
schedule_type=MANUAL. The docstring clearly explains that this disables automatic syncs while still allowing manual triggers.
| @mcp_tool( | ||
| domain="cloud", | ||
| open_world=True, | ||
| destructive=True, | ||
| extra_help_text=CLOUD_AUTH_TIP_TEXT, | ||
| ) | ||
| def enable_cloud_connection( | ||
| connection_id: Annotated[ | ||
| str, | ||
| Field(description="The ID of the connection to enable."), | ||
| ], | ||
| *, | ||
| workspace_id: Annotated[ | ||
| str | None, | ||
| Field( | ||
| description=WORKSPACE_ID_TIP_TEXT, | ||
| default=None, | ||
| ), | ||
| ], | ||
| ) -> str: | ||
| """Enable a connection on Airbyte Cloud. | ||
| Sets the connection status to 'active', allowing scheduled syncs to run. | ||
| """ | ||
| workspace: CloudWorkspace = _get_cloud_workspace(workspace_id) | ||
| connection = workspace.get_connection(connection_id=connection_id) | ||
| connection.enable() | ||
| return ( | ||
| f"Successfully enabled connection '{connection_id}'. " f"URL: {connection.connection_url}" | ||
| ) | ||
|
|
||
|
|
||
| @mcp_tool( | ||
| domain="cloud", | ||
| open_world=True, | ||
| destructive=True, | ||
| extra_help_text=CLOUD_AUTH_TIP_TEXT, | ||
| ) | ||
| def disable_cloud_connection( | ||
| connection_id: Annotated[ | ||
| str, | ||
| Field(description="The ID of the connection to disable."), | ||
| ], | ||
| *, | ||
| workspace_id: Annotated[ | ||
| str | None, | ||
| Field( | ||
| description=WORKSPACE_ID_TIP_TEXT, | ||
| default=None, | ||
| ), | ||
| ], | ||
| ) -> str: | ||
| """Disable a connection on Airbyte Cloud. | ||
| Sets the connection status to 'inactive', preventing scheduled syncs from running. | ||
| Manual syncs can still be triggered using the run_cloud_sync tool. | ||
| """ | ||
| workspace: CloudWorkspace = _get_cloud_workspace(workspace_id) | ||
| connection = workspace.get_connection(connection_id=connection_id) | ||
| connection.disable() | ||
| return ( | ||
| f"Successfully disabled connection '{connection_id}'. " f"URL: {connection.connection_url}" | ||
| ) | ||
|
|
||
|
|
||
| @mcp_tool( | ||
| domain="cloud", | ||
| open_world=True, | ||
| destructive=True, | ||
| extra_help_text=CLOUD_AUTH_TIP_TEXT, | ||
| ) | ||
| def set_cloud_connection_schedule( | ||
| connection_id: Annotated[ | ||
| str, | ||
| Field(description="The ID of the connection to update."), | ||
| ], | ||
| cron_expression: Annotated[ | ||
| str, | ||
| Field( | ||
| description=( | ||
| "A cron expression defining when syncs should run. " | ||
| "Examples: '0 0 * * *' (daily at midnight UTC), " | ||
| "'0 */6 * * *' (every 6 hours), " | ||
| "'0 0 * * 0' (weekly on Sunday at midnight UTC)." | ||
| ) | ||
| ), | ||
| ], | ||
| *, | ||
| workspace_id: Annotated[ | ||
| str | None, | ||
| Field( | ||
| description=WORKSPACE_ID_TIP_TEXT, | ||
| default=None, | ||
| ), | ||
| ], | ||
| ) -> str: | ||
| """Set a cron schedule for a connection on Airbyte Cloud. | ||
| Configures the connection to automatically sync on the specified schedule. | ||
| """ | ||
| workspace: CloudWorkspace = _get_cloud_workspace(workspace_id) | ||
| connection = workspace.get_connection(connection_id=connection_id) | ||
| connection.set_schedule(cron_expression=cron_expression) | ||
| return ( | ||
| f"Successfully set schedule for connection '{connection_id}' " | ||
| f"to '{cron_expression}'. URL: {connection.connection_url}" | ||
| ) | ||
|
|
||
|
|
||
| @mcp_tool( | ||
| domain="cloud", | ||
| open_world=True, | ||
| destructive=True, | ||
| extra_help_text=CLOUD_AUTH_TIP_TEXT, | ||
| ) | ||
| def set_cloud_connection_manual_schedule( | ||
| connection_id: Annotated[ | ||
| str, | ||
| Field(description="The ID of the connection to update."), | ||
| ], | ||
| *, | ||
| workspace_id: Annotated[ | ||
| str | None, | ||
| Field( | ||
| description=WORKSPACE_ID_TIP_TEXT, | ||
| default=None, | ||
| ), | ||
| ], | ||
| ) -> str: | ||
| """Set a connection to manual scheduling on Airbyte Cloud. | ||
| Disables automatic syncs. Syncs will only run when manually triggered | ||
| using the run_cloud_sync tool. | ||
| """ | ||
| workspace: CloudWorkspace = _get_cloud_workspace(workspace_id) | ||
| connection = workspace.get_connection(connection_id=connection_id) | ||
| connection.set_manual_schedule() | ||
| return ( | ||
| f"Successfully set connection '{connection_id}' to manual scheduling. " | ||
| f"URL: {connection.connection_url}" | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Should these tools include session guards for safety, wdyt?
I notice that other destructive connection operations in this file use check_guid_created_in_session(connection_id) to prevent modifying resources not created in the current session:
set_cloud_connection_table_prefix(line 2239)set_cloud_connection_selected_streams(line 2283)
However, these four new tools don't include this guard. For consistency and safety in MCP safe mode, should we add:
check_guid_created_in_session(connection_id)at the start of each tool (after the workspace/connection retrieval)?
This would align with the pattern used by similar destructive connection operations and provide the session-based safety mechanism mentioned in the register_cloud_ops_tools docstring at line 2495.
That said, I also notice rename_cloud_connection (line 2182) doesn't use this guard either, so there's some inconsistency in the file already. It might be worth standardizing the approach across all connection modification tools.
🔎 Proposed fix to add session guards
For enable_cloud_connection:
def enable_cloud_connection(
connection_id: Annotated[
str,
Field(description="The ID of the connection to enable."),
],
*,
workspace_id: Annotated[
str | None,
Field(
description=WORKSPACE_ID_TIP_TEXT,
default=None,
),
],
) -> str:
"""Enable a connection on Airbyte Cloud.
Sets the connection status to 'active', allowing scheduled syncs to run.
"""
+ check_guid_created_in_session(connection_id)
workspace: CloudWorkspace = _get_cloud_workspace(workspace_id)
connection = workspace.get_connection(connection_id=connection_id)
connection.enable()Apply the same pattern to disable_cloud_connection, set_cloud_connection_schedule, and set_cloud_connection_manual_schedule.
🤖 Prompt for AI Agents
In airbyte/mcp/cloud_ops.py around lines 2296 to 2436, the four tools
enable_cloud_connection, disable_cloud_connection,
set_cloud_connection_schedule, and set_cloud_connection_manual_schedule are
missing the session guard used by other destructive connection ops; add a call
to check_guid_created_in_session(connection_id) immediately after retrieving the
connection (i.e., after connection = workspace.get_connection(...)) in each of
the four functions so they follow the same safety pattern as
set_cloud_connection_table_prefix and set_cloud_connection_selected_streams.
airbyte/cloud/connections.py
Outdated
|
|
||
| # Enable/Disable | ||
|
|
||
| def enable(self) -> CloudConnection: |
There was a problem hiding this comment.
Let's combine these into a single method to set enabled status. Let's also add an input arg to ignore no op, meaning don't raise if the status is already what we are asking for. Default behavior is to Not raise If the request turns out to be a no-op.
Then, on top of that, add an implementation that's a simpler enabled property, calling the method with the default option to not raise if it's tried to set to something that's already existing. When retrieving this value, we will not use any cached values, because we won't know if another process or human has toggled the setting in between us checking it
There was a problem hiding this comment.
Done! I've refactored the implementation based on your feedback:
- Combined
enable()anddisable()into a singleset_enabled(*, enabled: bool, ignore_noop: bool = True)method - Added an
enabledproperty with:- Getter: Always fetches fresh data from the API (doesn't use cached values)
- Setter: Calls
set_enabled()with the defaultignore_noop=True
- The
ignore_noopflag defaults toTrue(silently returns if already in requested state) - Updated the MCP tools to use
connection.set_enabled(enabled=True/False)
The enabled parameter is keyword-only to satisfy the FBT001 lint rule (no boolean positional arguments).
…ag and enabled property Co-Authored-By: AJ Steers <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
airbyte/mcp/cloud_ops.py (1)
2296-2325: Missing session guard for consistency with other destructive tools.This concern was already raised in a previous review -
set_cloud_connection_table_prefixandset_cloud_connection_selected_streamsincludecheck_guid_created_in_session(connection_id)but these new tools don't. Should we add it for consistency, wdyt?
🧹 Nitpick comments (1)
airbyte/cloud/connections.py (1)
472-502: Solid scheduling implementation!The methods follow the established patterns nicely. One thought: the
cron_expressionis passed directly to the API without client-side validation. This is fine since the API will reject invalid formats, but would you consider adding a quick format check (e.g., ensuring 5 space-separated fields) to provide earlier feedback to users, wdyt? Not blocking - just a thought for UX.
📜 Review details
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
airbyte/cloud/connections.py(2 hunks)airbyte/mcp/cloud_ops.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
airbyte/mcp/cloud_ops.py (2)
airbyte/cloud/workspaces.py (2)
CloudWorkspace(88-886)get_connection(315-327)airbyte/cloud/connections.py (3)
set_enabled(420-468)enabled(398-408)enabled(411-418)
airbyte/cloud/connections.py (1)
airbyte/_util/api_util.py (1)
patch_connection(1337-1399)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Pytest (All, Python 3.10, Windows)
- GitHub Check: Pytest (All, Python 3.11, Windows)
- GitHub Check: Pytest (All, Python 3.10, Ubuntu)
- GitHub Check: Pytest (All, Python 3.11, Ubuntu)
- GitHub Check: Pytest (No Creds)
- GitHub Check: Pytest (Fast)
🔇 Additional comments (7)
airbyte/cloud/connections.py (4)
19-19: LGTM on the noqa suppression!The
# noqa: PLR0904is reasonable given the new methods being added. The class is growing to provide a complete API surface for connection management, so suppressing this lint warning makes sense here.
397-408: Nice implementation of theenabledproperty getter!The fresh API fetch on every access is the right call here since connection status can change externally. This aligns perfectly with the requirement to avoid stale cached values.
410-468: Well-designedset_enabledimplementation!The combined method with
ignore_noopflag is clean, and the property setter that delegates withignore_noop=Trueprovides a nice ergonomic API. The fresh fetch before patching ensures we're always working with current state - good defensive design.
504-524: Cleanset_manual_scheduleimplementation!Straightforward and consistent with the
set_schedulemethod. The docstring clearly explains the behavior.airbyte/mcp/cloud_ops.py (3)
2328-2358: Good docstring addition!The clarification that "manual syncs can still be triggered using the run_cloud_sync tool" is helpful context for MCP consumers. Implementation looks correct.
2361-2402: Nice inline examples in the field description!The cron expression examples (
'0 0 * * *','0 */6 * * *','0 0 * * 0') in the Field description are really helpful for LLM agents that will consume this tool. Good design choice.
2405-2436: Clean implementation!The docstring's mention of
run_cloud_syncfor manual triggering creates a nice cross-reference for tool discovery. LGTM.
…d_connection tool Co-Authored-By: AJ Steers <[email protected]>
devin-ai-integration
bot
changed the title
|
I've consolidated the 4 separate tools into a single
The tool validates that at least one setting is provided and that |
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
airbyte/mcp/cloud_ops.py (1)
2302-2401: Should this tool include a session guard for safety, wdyt?Similar to the concern raised in previous reviews, this consolidated tool doesn't call
check_guid_created_in_session(connection_id). Other destructive connection operations in this file use this guard for consistency and safety in MCP safe mode:
set_cloud_connection_table_prefix(line 2239)set_cloud_connection_selected_streams(line 2283)Adding the session guard would align with the pattern used by similar destructive connection operations and provide the session-based safety mechanism mentioned in the
register_cloud_ops_toolsdocstring.🔎 Proposed fix to add session guard
def update_cloud_connection( connection_id: Annotated[ str, Field(description="The ID of the connection to update."), ], *, enabled: Annotated[ bool | None, Field( description=( "Set the connection's enabled status. " "True enables the connection (status='active'), " "False disables it (status='inactive'). " "Leave unset to keep the current status." ), default=None, ), ], cron_expression: Annotated[ str | None, Field( description=( "A cron expression defining when syncs should run. " "Examples: '0 0 * * *' (daily at midnight UTC), " "'0 */6 * * *' (every 6 hours), " "'0 0 * * 0' (weekly on Sunday at midnight UTC). " "Leave unset to keep the current schedule. " "Cannot be used together with 'manual_schedule'." ), default=None, ), ], manual_schedule: Annotated[ bool | None, Field( description=( "Set to True to disable automatic syncs (manual scheduling only). " "Syncs will only run when manually triggered. " "Cannot be used together with 'cron_expression'." ), default=None, ), ], workspace_id: Annotated[ str | None, Field( description=WORKSPACE_ID_TIP_TEXT, default=None, ), ], ) -> str: """Update a connection's settings on Airbyte Cloud. This tool allows updating multiple connection settings in a single call: - Enable or disable the connection - Set a cron schedule for automatic syncs - Switch to manual scheduling (no automatic syncs) At least one setting must be provided. The 'cron_expression' and 'manual_schedule' parameters are mutually exclusive. """ + check_guid_created_in_session(connection_id) + # Validate that at least one setting is provided if enabled is None and cron_expression is None and manual_schedule is None:
🧹 Nitpick comments (1)
airbyte/mcp/cloud_ops.py (1)
2365-2368: Consider using PyAirbyteInputError for consistency, wdyt?The validation errors here raise
ValueError, but other similar validation errors in this file usePyAirbyteInputError(see lines 710-713, 1039-1042, 1270-1276, 1703-1712). Using the consistent exception type would align with the established pattern.🔎 Proposed fix
# Validate that at least one setting is provided if enabled is None and cron_expression is None and manual_schedule is None: - raise ValueError( - "At least one setting must be provided: 'enabled', 'cron_expression', " - "or 'manual_schedule'." + raise PyAirbyteInputError( + message=( + "At least one setting must be provided: 'enabled', 'cron_expression', " + "or 'manual_schedule'." + ), ) # Validate mutually exclusive schedule options if cron_expression is not None and manual_schedule is True: - raise ValueError( - "Cannot specify both 'cron_expression' and 'manual_schedule=True'. " - "Use 'cron_expression' for scheduled syncs or 'manual_schedule=True' " - "for manual-only syncs." + raise PyAirbyteInputError( + message=( + "Cannot specify both 'cron_expression' and 'manual_schedule=True'. " + "Use 'cron_expression' for scheduled syncs or 'manual_schedule=True' " + "for manual-only syncs." + ), )Also applies to: 2372-2376
📜 Review details
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
airbyte/mcp/cloud_ops.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
airbyte/mcp/cloud_ops.py (3)
airbyte/mcp/_tool_utils.py (1)
mcp_tool(189-235)airbyte/cloud/connections.py (6)
enabled(398-408)enabled(411-418)set_enabled(420-468)set_schedule(472-502)set_manual_schedule(504-524)connection_url(160-162)airbyte/cloud/workspaces.py (1)
get_connection(315-327)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Pytest (All, Python 3.10, Ubuntu)
- GitHub Check: Pytest (All, Python 3.10, Windows)
- GitHub Check: Pytest (All, Python 3.11, Windows)
- GitHub Check: Pytest (All, Python 3.11, Ubuntu)
- GitHub Check: Pytest (Fast)
- GitHub Check: Pytest (No Creds)
feat: Add MCP tool and Python API for connection enable/disable and scheduling
Summary
Adds a consolidated MCP tool and underlying Python methods for enabling/disabling connections and managing connection schedules on Airbyte Cloud. The MCP tool is marked as
destructive=Trueso it will be disabled when the MCP server runs in read-only mode.New CloudConnection methods:
enabledproperty - Gets/sets connection enabled status (getter always fetches fresh data from API)set_enabled(*, enabled: bool, ignore_noop: bool = True)- Sets connection status with no-op handlingset_schedule(cron_expression)- Sets a cron schedule for automatic syncsset_manual_schedule()- Disables automatic syncsNew MCP tool:
update_cloud_connection- Consolidated tool accepting optionalenabled,cron_expression, andmanual_scheduleparametersUpdates Since Last Revision
Per reviewer feedback:
enable()anddisable()methods into a singleset_enabled()method withignore_noopflag (defaults to True)enabledproperty with fresh-read getter (never uses cached values) and setterupdate_cloud_connectiontool that accepts multiple optional parametersReview & Testing Checklist for Human
api_util.patch_connectionactually acceptsstatusandscheduleparameters - I followed the existing pattern but didn't test against the live APIenabledproperty getter fetches fresh data on each accessRecommended Test Plan
Notes
# noqa: PLR0904to CloudConnection class to suppress "too many public methods" lint warning (approved by user)Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.