airbyte-internal-ops 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

airbyte_ops_mcp/mcp/cloud_connector_versions.py

@@ -25,7 +25,9 @@ from airbyte_ops_mcp.cloud_admin.auth import (
 )
 from airbyte_ops_mcp.cloud_admin.models import (
     ConnectorVersionInfo,
+    OrganizationVersionOverrideResult,
     VersionOverrideOperationResult,
+    WorkspaceVersionOverrideResult,
 )
 from airbyte_ops_mcp.github_api import (
     GitHubAPIError,
@@ -434,6 +436,464 @@ def set_cloud_connector_version_override(
         )
 
 
+@mcp_tool(
+    destructive=True,
+    idempotent=False,
+    open_world=True,
+)
+def set_workspace_connector_version_override(
+    workspace_id: Annotated[
+        str,
+        Field(description="The Airbyte Cloud workspace ID."),
+    ],
+    connector_name: Annotated[
+        str,
+        Field(
+            description="The connector name (e.g., 'source-github', 'destination-bigquery')."
+        ),
+    ],
+    connector_type: Annotated[
+        Literal["source", "destination"],
+        "The type of connector (source or destination)",
+    ],
+    approval_comment_url: Annotated[
+        str,
+        Field(
+            description="URL to a GitHub comment where the admin has explicitly "
+            "requested or authorized this deployment. Must be a valid GitHub comment URL. "
+            "Required for authorization. The admin email is automatically derived from "
+            "the comment author's GitHub profile.",
+        ),
+    ],
+    version: Annotated[
+        str | None,
+        Field(
+            description="The semver version string to pin to (e.g., '0.1.0'). "
+            "Must be None if unset is True.",
+            default=None,
+        ),
+    ],
+    unset: Annotated[
+        bool,
+        Field(
+            description="If True, removes any existing version override. "
+            "Cannot be True if version is provided.",
+            default=False,
+        ),
+    ],
+    override_reason: Annotated[
+        str | None,
+        Field(
+            description="Required when setting a version. "
+            "Explanation for the override (min 10 characters).",
+            default=None,
+        ),
+    ],
+    override_reason_reference_url: Annotated[
+        str | None,
+        Field(
+            description="Optional URL with more context (e.g., issue link).",
+            default=None,
+        ),
+    ],
+    issue_url: Annotated[
+        str | None,
+        Field(
+            description="URL to the GitHub issue providing context for this operation. "
+            "Must be a valid GitHub URL (https://github.com/...). Required for authorization.",
+            default=None,
+        ),
+    ],
+    ai_agent_session_url: Annotated[
+        str | None,
+        Field(
+            description="URL to the AI agent session driving this operation, if applicable. "
+            "Provides additional auditability for AI-driven operations.",
+            default=None,
+        ),
+    ],
+) -> WorkspaceVersionOverrideResult:
+    """Set or clear a workspace-level version override for a connector type.
+
+    This pins ALL instances of a connector type within a workspace to a specific version.
+    For example, pinning 'source-github' at workspace level means all GitHub sources
+    in that workspace will use the pinned version.
+
+    **Admin-only operation** - Requires:
+    - AIRBYTE_INTERNAL_ADMIN_FLAG=airbyte.io environment variable
+    - issue_url parameter (GitHub issue URL for context)
+    - approval_comment_url parameter (GitHub comment URL with approval from an @airbyte.io user)
+
+    You must specify EXACTLY ONE of `version` OR `unset=True`, but not both.
+    When setting a version, `override_reason` is required.
+    """
+    # Validate admin access (check env var flag)
+    try:
+        require_internal_admin_flag_only()
+    except CloudAuthError as e:
+        return WorkspaceVersionOverrideResult(
+            success=False,
+            message=f"Admin authentication failed: {e}",
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+    # Validate authorization parameters
+    validation_errors: list[str] = []
+
+    if not issue_url:
+        validation_errors.append(
+            "issue_url is required for authorization (GitHub issue URL)"
+        )
+    elif not issue_url.startswith("https://github.com/"):
+        validation_errors.append(
+            f"issue_url must be a valid GitHub URL (https://github.com/...), got: {issue_url}"
+        )
+
+    if not approval_comment_url.startswith("https://github.com/"):
+        validation_errors.append(
+            f"approval_comment_url must be a valid GitHub URL, got: {approval_comment_url}"
+        )
+    elif (
+        "#issuecomment-" not in approval_comment_url
+        and "#discussion_r" not in approval_comment_url
+    ):
+        validation_errors.append(
+            "approval_comment_url must be a GitHub comment URL "
+            "(containing #issuecomment- or #discussion_r)"
+        )
+
+    if validation_errors:
+        return WorkspaceVersionOverrideResult(
+            success=False,
+            message="Authorization validation failed: " + "; ".join(validation_errors),
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+    # Derive admin email from approval comment URL
+    try:
+        admin_user_email = get_admin_email_from_approval_comment(approval_comment_url)
+    except GitHubCommentParseError as e:
+        return WorkspaceVersionOverrideResult(
+            success=False,
+            message=f"Failed to parse approval comment URL: {e}",
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+    except GitHubAPIError as e:
+        return WorkspaceVersionOverrideResult(
+            success=False,
+            message=f"Failed to fetch approval comment from GitHub: {e}",
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+    except GitHubUserEmailNotFoundError as e:
+        return WorkspaceVersionOverrideResult(
+            success=False,
+            message=str(e),
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+    # Build enhanced override reason with audit fields (only for 'set' operations)
+    enhanced_override_reason = override_reason
+    if not unset and override_reason:
+        audit_parts = [override_reason]
+        audit_parts.append(f"Issue: {issue_url}")
+        audit_parts.append(f"Approval: {approval_comment_url}")
+        if ai_agent_session_url:
+            audit_parts.append(f"AI Session: {ai_agent_session_url}")
+        enhanced_override_reason = " | ".join(audit_parts)
+
+    # Resolve auth and call API client
+    try:
+        auth = _resolve_cloud_auth()
+
+        result = api_client.set_workspace_connector_version_override(
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+            api_root=constants.CLOUD_CONFIG_API_ROOT,
+            client_id=auth.client_id,
+            client_secret=auth.client_secret,
+            bearer_token=auth.bearer_token,
+            version=version,
+            unset=unset,
+            override_reason=enhanced_override_reason,
+            override_reason_reference_url=override_reason_reference_url,
+            user_email=admin_user_email,
+        )
+
+        if unset:
+            if result:
+                message = f"Successfully cleared workspace-level version override for {connector_name}."
+            else:
+                message = f"No workspace-level version override was active for {connector_name} (nothing to clear)"
+        else:
+            message = f"Successfully pinned {connector_name} to version {version} at workspace level."
+
+        return WorkspaceVersionOverrideResult(
+            success=True,
+            message=message,
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+            version=version if not unset else None,
+        )
+
+    except PyAirbyteInputError as e:
+        return WorkspaceVersionOverrideResult(
+            success=False,
+            message=str(e),
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+    except CloudAuthError as e:
+        return WorkspaceVersionOverrideResult(
+            success=False,
+            message=f"Authentication failed: {e}",
+            workspace_id=workspace_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+
+@mcp_tool(
+    destructive=True,
+    idempotent=False,
+    open_world=True,
+)
+def set_organization_connector_version_override(
+    organization_id: Annotated[
+        str,
+        Field(description="The Airbyte Cloud organization ID."),
+    ],
+    connector_name: Annotated[
+        str,
+        Field(
+            description="The connector name (e.g., 'source-github', 'destination-bigquery')."
+        ),
+    ],
+    connector_type: Annotated[
+        Literal["source", "destination"],
+        "The type of connector (source or destination)",
+    ],
+    approval_comment_url: Annotated[
+        str,
+        Field(
+            description="URL to a GitHub comment where the admin has explicitly "
+            "requested or authorized this deployment. Must be a valid GitHub comment URL. "
+            "Required for authorization. The admin email is automatically derived from "
+            "the comment author's GitHub profile.",
+        ),
+    ],
+    version: Annotated[
+        str | None,
+        Field(
+            description="The semver version string to pin to (e.g., '0.1.0'). "
+            "Must be None if unset is True.",
+            default=None,
+        ),
+    ],
+    unset: Annotated[
+        bool,
+        Field(
+            description="If True, removes any existing version override. "
+            "Cannot be True if version is provided.",
+            default=False,
+        ),
+    ],
+    override_reason: Annotated[
+        str | None,
+        Field(
+            description="Required when setting a version. "
+            "Explanation for the override (min 10 characters).",
+            default=None,
+        ),
+    ],
+    override_reason_reference_url: Annotated[
+        str | None,
+        Field(
+            description="Optional URL with more context (e.g., issue link).",
+            default=None,
+        ),
+    ],
+    issue_url: Annotated[
+        str | None,
+        Field(
+            description="URL to the GitHub issue providing context for this operation. "
+            "Must be a valid GitHub URL (https://github.com/...). Required for authorization.",
+            default=None,
+        ),
+    ],
+    ai_agent_session_url: Annotated[
+        str | None,
+        Field(
+            description="URL to the AI agent session driving this operation, if applicable. "
+            "Provides additional auditability for AI-driven operations.",
+            default=None,
+        ),
+    ],
+) -> OrganizationVersionOverrideResult:
+    """Set or clear an organization-level version override for a connector type.
+
+    This pins ALL instances of a connector type across an entire organization to a
+    specific version. For example, pinning 'source-github' at organization level means
+    all GitHub sources in all workspaces within that organization will use the pinned version.
+
+    **Admin-only operation** - Requires:
+    - AIRBYTE_INTERNAL_ADMIN_FLAG=airbyte.io environment variable
+    - issue_url parameter (GitHub issue URL for context)
+    - approval_comment_url parameter (GitHub comment URL with approval from an @airbyte.io user)
+
+    You must specify EXACTLY ONE of `version` OR `unset=True`, but not both.
+    When setting a version, `override_reason` is required.
+    """
+    # Validate admin access (check env var flag)
+    try:
+        require_internal_admin_flag_only()
+    except CloudAuthError as e:
+        return OrganizationVersionOverrideResult(
+            success=False,
+            message=f"Admin authentication failed: {e}",
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+    # Validate authorization parameters
+    validation_errors: list[str] = []
+
+    if not issue_url:
+        validation_errors.append(
+            "issue_url is required for authorization (GitHub issue URL)"
+        )
+    elif not issue_url.startswith("https://github.com/"):
+        validation_errors.append(
+            f"issue_url must be a valid GitHub URL (https://github.com/...), got: {issue_url}"
+        )
+
+    if not approval_comment_url.startswith("https://github.com/"):
+        validation_errors.append(
+            f"approval_comment_url must be a valid GitHub URL, got: {approval_comment_url}"
+        )
+    elif (
+        "#issuecomment-" not in approval_comment_url
+        and "#discussion_r" not in approval_comment_url
+    ):
+        validation_errors.append(
+            "approval_comment_url must be a GitHub comment URL "
+            "(containing #issuecomment- or #discussion_r)"
+        )
+
+    if validation_errors:
+        return OrganizationVersionOverrideResult(
+            success=False,
+            message="Authorization validation failed: " + "; ".join(validation_errors),
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+    # Derive admin email from approval comment URL
+    try:
+        admin_user_email = get_admin_email_from_approval_comment(approval_comment_url)
+    except GitHubCommentParseError as e:
+        return OrganizationVersionOverrideResult(
+            success=False,
+            message=f"Failed to parse approval comment URL: {e}",
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+    except GitHubAPIError as e:
+        return OrganizationVersionOverrideResult(
+            success=False,
+            message=f"Failed to fetch approval comment from GitHub: {e}",
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+    except GitHubUserEmailNotFoundError as e:
+        return OrganizationVersionOverrideResult(
+            success=False,
+            message=str(e),
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+    # Build enhanced override reason with audit fields (only for 'set' operations)
+    enhanced_override_reason = override_reason
+    if not unset and override_reason:
+        audit_parts = [override_reason]
+        audit_parts.append(f"Issue: {issue_url}")
+        audit_parts.append(f"Approval: {approval_comment_url}")
+        if ai_agent_session_url:
+            audit_parts.append(f"AI Session: {ai_agent_session_url}")
+        enhanced_override_reason = " | ".join(audit_parts)
+
+    # Resolve auth and call API client
+    try:
+        auth = _resolve_cloud_auth()
+
+        result = api_client.set_organization_connector_version_override(
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+            api_root=constants.CLOUD_CONFIG_API_ROOT,
+            client_id=auth.client_id,
+            client_secret=auth.client_secret,
+            bearer_token=auth.bearer_token,
+            version=version,
+            unset=unset,
+            override_reason=enhanced_override_reason,
+            override_reason_reference_url=override_reason_reference_url,
+            user_email=admin_user_email,
+        )
+
+        if unset:
+            if result:
+                message = f"Successfully cleared organization-level version override for {connector_name}."
+            else:
+                message = f"No organization-level version override was active for {connector_name} (nothing to clear)"
+        else:
+            message = f"Successfully pinned {connector_name} to version {version} at organization level."
+
+        return OrganizationVersionOverrideResult(
+            success=True,
+            message=message,
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+            version=version if not unset else None,
+        )
+
+    except PyAirbyteInputError as e:
+        return OrganizationVersionOverrideResult(
+            success=False,
+            message=str(e),
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+    except CloudAuthError as e:
+        return OrganizationVersionOverrideResult(
+            success=False,
+            message=f"Authentication failed: {e}",
+            organization_id=organization_id,
+            connector_name=connector_name,
+            connector_type=connector_type,
+        )
+
+
 def register_cloud_connector_version_tools(app: FastMCP) -> None:
     """Register cloud connector version management tools with the FastMCP app.
 

airbyte_ops_mcp/mcp/prerelease.py

@@ -18,6 +18,7 @@ import yaml
 from fastmcp import FastMCP
 from pydantic import BaseModel, Field
 
+from airbyte_ops_mcp.github_actions import trigger_workflow_dispatch
 from airbyte_ops_mcp.github_api import (
     GITHUB_API_BASE,
     get_pr_head_ref,
@@ -149,49 +150,6 @@ def _get_connector_metadata(
     return yaml.safe_load(content)
 
 
-def _trigger_workflow_dispatch(
-    owner: str,
-    repo: str,
-    workflow_file: str,
-    ref: str,
-    inputs: dict,
-    token: str,
-) -> str:
-    """Trigger a GitHub Actions workflow via workflow_dispatch.
-
-    Args:
-        owner: Repository owner (e.g., "airbytehq")
-        repo: Repository name (e.g., "airbyte")
-        workflow_file: Workflow file name (e.g., "publish-connectors-prerelease-command.yml")
-        ref: Git ref to run the workflow on (branch name)
-        inputs: Workflow inputs dictionary
-        token: GitHub API token
-
-    Returns:
-        URL to view workflow runs.
-
-    Raises:
-        requests.HTTPError: If API request fails.
-    """
-    url = f"{GITHUB_API_BASE}/repos/{owner}/{repo}/actions/workflows/{workflow_file}/dispatches"
-    headers = {
-        "Authorization": f"Bearer {token}",
-        "Accept": "application/vnd.github+json",
-        "X-GitHub-Api-Version": "2022-11-28",
-    }
-    payload = {
-        "ref": ref,
-        "inputs": inputs,
-    }
-
-    response = requests.post(url, headers=headers, json=payload, timeout=30)
-    response.raise_for_status()
-
-    # workflow_dispatch returns 204 No Content on success
-    # Return URL to view workflow runs
-    return f"https://github.com/{owner}/{repo}/actions/workflows/{workflow_file}"
-
-
 @mcp_tool(
     read_only=False,
     destructive=False,
@@ -263,23 +221,25 @@ def publish_connector_to_airbyte_registry(
     head_info = get_pr_head_ref(DEFAULT_REPO_OWNER, target_repo_name, pr_number, token)
 
     # Prepare workflow inputs
-    # The workflow uses refs/pull/{pr}/head directly - no gitref needed
-    # Note: The workflow auto-detects modified connectors from the PR
     workflow_inputs = {
         "repo": f"{DEFAULT_REPO_OWNER}/{target_repo_name}",
         "pr": str(pr_number),
+        "connector": connector_name,
     }
 
     # Trigger the workflow on the default branch
     # The workflow will checkout the PR branch via inputs.gitref
-    workflow_url = _trigger_workflow_dispatch(
+    dispatch_result = trigger_workflow_dispatch(
         owner=DEFAULT_REPO_OWNER,
         repo=target_repo_name,
         workflow_file=target_workflow,
         ref=target_branch,
         inputs=workflow_inputs,
         token=token,
+        find_run=True,
     )
+    # Use the specific run URL if found, otherwise fall back to the workflow URL
+    workflow_url = dispatch_result.run_url or dispatch_result.workflow_url
 
     # Try to compute docker_image and docker_image_tag from connector metadata
     docker_image: str | None = None