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

airbyte_ops_mcp/mcp/{github.py → github_actions.py}

@@ -15,12 +15,22 @@ from fastmcp import FastMCP
 from pydantic import BaseModel, Field
 
 from airbyte_ops_mcp.github_actions import (
-    GITHUB_API_BASE,
     get_workflow_jobs,
+    trigger_workflow_dispatch,
+)
+from airbyte_ops_mcp.github_api import (
+    GITHUB_API_BASE,
+    get_pr_head_ref,
     resolve_github_token,
 )
 from airbyte_ops_mcp.mcp._mcp_utils import mcp_tool, register_mcp_tools
 
+# Token env vars for workflow triggering (in order of preference)
+WORKFLOW_TRIGGER_TOKEN_ENV_VARS = [
+    "GITHUB_CI_WORKFLOW_TRIGGER_PAT",
+    "GITHUB_TOKEN",
+]
+
 DOCKERHUB_API_BASE = "https://hub.docker.com/v2"
 
 
@@ -36,7 +46,7 @@ class JobInfo(BaseModel):
 
 
 class WorkflowRunStatus(BaseModel):
-    """Response model for check_workflow_status MCP tool."""
+    """Response model for check_ci_workflow_status MCP tool."""
 
     run_id: int
     status: str
@@ -116,7 +126,7 @@ def _get_workflow_run(
     idempotent=True,
     open_world=True,
 )
-def check_workflow_status(
+def check_ci_workflow_status(
     workflow_url: Annotated[
         str | None,
         Field(
@@ -196,6 +206,108 @@ def check_workflow_status(
     )
 
 
+class TriggerCIWorkflowResult(BaseModel):
+    """Response model for trigger_ci_workflow MCP tool."""
+
+    success: bool
+    message: str
+    workflow_url: str
+    run_id: int | None = None
+    run_url: str | None = None
+
+
+@mcp_tool(
+    read_only=False,
+    idempotent=False,
+    open_world=True,
+)
+def trigger_ci_workflow(
+    owner: Annotated[
+        str,
+        Field(description="Repository owner (e.g., 'airbytehq')"),
+    ],
+    repo: Annotated[
+        str,
+        Field(description="Repository name (e.g., 'airbyte')"),
+    ],
+    workflow_file: Annotated[
+        str,
+        Field(description="Workflow file name (e.g., 'connector-regression-test.yml')"),
+    ],
+    workflow_definition_ref: Annotated[
+        str | None,
+        Field(
+            description="Branch name or PR number for the workflow definition to use. "
+            "If a PR number (integer string) is provided, it resolves to the PR's head branch name. "
+            "If a branch name is provided, it is used directly. "
+            "Defaults to the repository's default branch if not specified."
+        ),
+    ] = None,
+    inputs: Annotated[
+        dict[str, str] | None,
+        Field(
+            description="Workflow inputs as a dictionary of string key-value pairs. "
+            "These are passed to the workflow_dispatch event."
+        ),
+    ] = None,
+) -> TriggerCIWorkflowResult:
+    """Trigger a GitHub Actions CI workflow via workflow_dispatch.
+
+    This tool triggers a workflow in any GitHub repository that has workflow_dispatch
+    enabled. It resolves PR numbers to branch names automatically since GitHub's
+    workflow_dispatch API only accepts branch names, not refs/pull/{pr}/head format.
+
+    Requires GITHUB_CI_WORKFLOW_TRIGGER_PAT or GITHUB_TOKEN environment variable
+    with 'actions:write' permission.
+    """
+    # Guard: Check for required token
+    token = resolve_github_token(WORKFLOW_TRIGGER_TOKEN_ENV_VARS)
+
+    # Resolve workflow definition ref
+    # If a PR number is provided (integer string), resolve to the PR's head branch name
+    # Otherwise use the provided branch name or default to repo's default branch
+    if workflow_definition_ref:
+        if workflow_definition_ref.isdigit():
+            # Resolve PR number to branch name via GitHub API
+            pr_head_info = get_pr_head_ref(
+                owner,
+                repo,
+                int(workflow_definition_ref),
+                token,
+            )
+            resolved_ref = pr_head_info.ref
+        else:
+            resolved_ref = workflow_definition_ref
+    else:
+        # Default to main (most common default branch)
+        resolved_ref = "main"
+
+    # Trigger the workflow
+    result = trigger_workflow_dispatch(
+        owner=owner,
+        repo=repo,
+        workflow_file=workflow_file,
+        ref=resolved_ref,
+        inputs=inputs or {},
+        token=token,
+        find_run=True,
+    )
+
+    # Build response message
+    if result.run_id:
+        message = f"Successfully triggered workflow {workflow_file} on {owner}/{repo} (ref: {resolved_ref}). Run ID: {result.run_id}"
+    else:
+        message = f"Successfully triggered workflow {workflow_file} on {owner}/{repo} (ref: {resolved_ref}). Run ID not yet available."
+
+    return TriggerCIWorkflowResult(
+        success=True,
+        message=message,
+        workflow_url=result.workflow_url,
+        run_id=result.run_id,
+        run_url=result.run_url,
+    )
+
+
 class DockerImageInfo(BaseModel):
     """Response model for get_docker_image_info MCP tool."""
 
@@ -282,8 +394,8 @@ def get_docker_image_info(
     )
 
 
-def register_github_tools(app: FastMCP) -> None:
-    """Register GitHub tools with the FastMCP app.
+def register_github_actions_tools(app: FastMCP) -> None:
+    """Register GitHub Actions tools with the FastMCP app.
 
     Args:
         app: FastMCP application instance

airbyte_ops_mcp/mcp/prerelease.py

@@ -2,12 +2,15 @@
 """MCP tools for triggering connector pre-release workflows.
 
 This module provides MCP tools for triggering the publish-connectors-prerelease
-workflow in the airbytehq/airbyte repository via GitHub's workflow dispatch API.
+workflow in the airbytehq/airbyte repository (for OSS connectors) or the
+publish_enterprise_connectors workflow in airbytehq/airbyte-enterprise
+(for enterprise connectors) via GitHub's workflow dispatch API.
 """
 
 from __future__ import annotations
 
 import base64
+from enum import StrEnum
 from typing import Annotated, Literal
 
 import requests
@@ -15,15 +18,32 @@ import yaml
 from fastmcp import FastMCP
 from pydantic import BaseModel, Field
 
-from airbyte_ops_mcp.github_actions import GITHUB_API_BASE, resolve_github_token
+from airbyte_ops_mcp.github_api import (
+    GITHUB_API_BASE,
+    get_pr_head_ref,
+    resolve_github_token,
+)
 from airbyte_ops_mcp.mcp._mcp_utils import mcp_tool, register_mcp_tools
 
+
+class ConnectorRepo(StrEnum):
+    """Repository where connector code is located."""
+
+    AIRBYTE = "airbyte"
+    AIRBYTE_ENTERPRISE = "airbyte-enterprise"
+
+
 DEFAULT_REPO_OWNER = "airbytehq"
-DEFAULT_REPO_NAME = "airbyte"
+DEFAULT_REPO_NAME = ConnectorRepo.AIRBYTE
 DEFAULT_BRANCH = "master"
 PRERELEASE_WORKFLOW_FILE = "publish-connectors-prerelease-command.yml"
 CONNECTOR_PATH_PREFIX = "airbyte-integrations/connectors"
 
+# Enterprise repository constants
+ENTERPRISE_REPO_NAME = ConnectorRepo.AIRBYTE_ENTERPRISE
+ENTERPRISE_DEFAULT_BRANCH = "main"
+ENTERPRISE_PRERELEASE_WORKFLOW_FILE = "publish_enterprise_connectors.yml"
+
 # Token env vars for prerelease publishing (in order of preference)
 PRERELEASE_TOKEN_ENV_VARS = [
     "GITHUB_CONNECTOR_PUBLISHING_PAT",
@@ -71,14 +91,6 @@ def compute_prerelease_docker_image_tag(base_version: str, sha: str) -> str:
     return f"{base_version}-{PRERELEASE_TAG_PREFIX}.{short_sha}"
 
 
-class PRHeadInfo(BaseModel):
-    """Information about a PR's head commit."""
-
-    ref: str
-    sha: str
-    short_sha: str
-
-
 class PrereleaseWorkflowResult(BaseModel):
     """Response model for publish_connector_to_airbyte_registry MCP tool."""
 
@@ -91,47 +103,6 @@ class PrereleaseWorkflowResult(BaseModel):
     docker_image_tag: str | None = None
 
 
-def _get_pr_head_info(
-    owner: str,
-    repo: str,
-    pr_number: int,
-    token: str,
-) -> PRHeadInfo:
-    """Get the head ref and SHA for a PR.
-
-    Args:
-        owner: Repository owner (e.g., "airbytehq")
-        repo: Repository name (e.g., "airbyte")
-        pr_number: Pull request number
-        token: GitHub API token
-
-    Returns:
-        PRHeadInfo with ref, sha, and short_sha.
-
-    Raises:
-        ValueError: If PR not found or API error.
-    """
-    url = f"{GITHUB_API_BASE}/repos/{owner}/{repo}/pulls/{pr_number}"
-    headers = {
-        "Authorization": f"Bearer {token}",
-        "Accept": "application/vnd.github+json",
-        "X-GitHub-Api-Version": "2022-11-28",
-    }
-
-    response = requests.get(url, headers=headers, timeout=30)
-    if response.status_code == 404:
-        raise ValueError(f"PR {owner}/{repo}#{pr_number} not found")
-    response.raise_for_status()
-
-    pr_data = response.json()
-    sha = pr_data["head"]["sha"]
-    return PRHeadInfo(
-        ref=pr_data["head"]["ref"],
-        sha=sha,
-        short_sha=sha[:7],
-    )
-
-
 def _get_connector_metadata(
     owner: str,
     repo: str,
@@ -238,6 +209,14 @@ def publish_connector_to_airbyte_registry(
         int,
         Field(description="The pull request number containing the connector changes"),
     ],
+    repo: Annotated[
+        ConnectorRepo,
+        Field(
+            default=ConnectorRepo.AIRBYTE,
+            description="Repository where the connector PR is located. "
+            "Use 'airbyte' for OSS connectors (default) or 'airbyte-enterprise' for enterprise connectors.",
+        ),
+    ],
     prerelease: Annotated[
         Literal[True],
         Field(
@@ -249,8 +228,10 @@ def publish_connector_to_airbyte_registry(
     """Publish a connector to the Airbyte registry.
 
     Currently only supports pre-release publishing. This tool triggers the
-    publish-connectors-prerelease workflow in the airbytehq/airbyte repository,
-    which publishes a pre-release version of the specified connector from the PR branch.
+    publish-connectors-prerelease workflow in the airbytehq/airbyte repository
+    (for OSS connectors) or the publish_enterprise_connectors workflow in
+    airbytehq/airbyte-enterprise (for enterprise connectors), which publishes
+    a pre-release version of the specified connector from the PR branch.
 
     Pre-release versions are tagged with the format: {version}-preview.{7-char-git-sha}
     These versions are available for version pinning via the scoped_configuration API.
@@ -267,17 +248,25 @@ def publish_connector_to_airbyte_registry(
     # Guard: Check for required token
     token = resolve_github_token(PRERELEASE_TOKEN_ENV_VARS)
 
+    # Determine repo-specific settings
+    is_enterprise = repo == ConnectorRepo.AIRBYTE_ENTERPRISE
+    target_repo_name = ENTERPRISE_REPO_NAME if is_enterprise else DEFAULT_REPO_NAME
+    target_branch = ENTERPRISE_DEFAULT_BRANCH if is_enterprise else DEFAULT_BRANCH
+    target_workflow = (
+        ENTERPRISE_PRERELEASE_WORKFLOW_FILE
+        if is_enterprise
+        else PRERELEASE_WORKFLOW_FILE
+    )
+
     # Get the PR's head SHA for computing the docker image tag
     # Note: We no longer pass gitref to the workflow - it derives the ref from PR number
-    head_info = _get_pr_head_info(
-        DEFAULT_REPO_OWNER, DEFAULT_REPO_NAME, pr_number, token
-    )
+    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}/{DEFAULT_REPO_NAME}",
+        "repo": f"{DEFAULT_REPO_OWNER}/{target_repo_name}",
         "pr": str(pr_number),
     }
 
@@ -285,9 +274,9 @@ def publish_connector_to_airbyte_registry(
     # The workflow will checkout the PR branch via inputs.gitref
     workflow_url = _trigger_workflow_dispatch(
         owner=DEFAULT_REPO_OWNER,
-        repo=DEFAULT_REPO_NAME,
-        workflow_file=PRERELEASE_WORKFLOW_FILE,
-        ref=DEFAULT_BRANCH,
+        repo=target_repo_name,
+        workflow_file=target_workflow,
+        ref=target_branch,
         inputs=workflow_inputs,
         token=token,
     )
@@ -297,7 +286,7 @@ def publish_connector_to_airbyte_registry(
     docker_image_tag: str | None = None
     metadata = _get_connector_metadata(
         DEFAULT_REPO_OWNER,
-        DEFAULT_REPO_NAME,
+        target_repo_name,
         connector_name,
         head_info.sha,
         token,
@@ -311,9 +300,10 @@ def publish_connector_to_airbyte_registry(
                 base_version, head_info.sha
             )
 
+    repo_info = f" from {repo}" if is_enterprise else ""
     return PrereleaseWorkflowResult(
         success=True,
-        message=f"Successfully triggered pre-release workflow for {connector_name} from PR #{pr_number}",
+        message=f"Successfully triggered pre-release workflow for {connector_name}{repo_info} from PR #{pr_number}",
         workflow_url=workflow_url,
         connector_name=connector_name,
         pr_number=pr_number,

airbyte_ops_mcp/mcp/prod_db_queries.py

@@ -7,7 +7,7 @@ airbyte_ops_mcp.prod_db_access.queries for use by AI agents.
 
 from __future__ import annotations
 
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import StrEnum
 from typing import Annotated, Any
 
@@ -24,9 +24,11 @@ from airbyte_ops_mcp.prod_db_access.queries import (
     query_connections_by_destination_connector,
     query_connector_versions,
     query_dataplanes_list,
+    query_destination_connection_stats,
     query_failed_sync_attempts_for_connector,
     query_new_connector_releases,
     query_recent_syncs_for_connector,
+    query_source_connection_stats,
     query_syncs_for_version_pinned_connector,
     query_workspace_info,
     query_workspaces_by_email_domain,
@@ -90,6 +92,94 @@ class WorkspacesByEmailDomainResult(BaseModel):
     )
 
 
+class LatestAttemptBreakdown(BaseModel):
+    """Breakdown of connections by latest attempt status."""
+
+    succeeded: int = Field(
+        default=0, description="Connections where latest attempt succeeded"
+    )
+    failed: int = Field(
+        default=0, description="Connections where latest attempt failed"
+    )
+    cancelled: int = Field(
+        default=0, description="Connections where latest attempt was cancelled"
+    )
+    running: int = Field(
+        default=0, description="Connections where latest attempt is still running"
+    )
+    unknown: int = Field(
+        default=0,
+        description="Connections with no recent attempts in the lookback window",
+    )
+
+
+class VersionPinStats(BaseModel):
+    """Stats for connections pinned to a specific version."""
+
+    pinned_version_id: str | None = Field(
+        description="The connector version UUID (None for unpinned connections)"
+    )
+    docker_image_tag: str | None = Field(
+        default=None, description="The docker image tag for this version"
+    )
+    total_connections: int = Field(description="Total number of connections")
+    enabled_connections: int = Field(
+        description="Number of enabled (active status) connections"
+    )
+    active_connections: int = Field(
+        description="Number of connections with recent sync activity"
+    )
+    latest_attempt: LatestAttemptBreakdown = Field(
+        description="Breakdown by latest attempt status"
+    )
+
+
+class ConnectorConnectionStats(BaseModel):
+    """Aggregate connection stats for a connector."""
+
+    connector_definition_id: str = Field(description="The connector definition UUID")
+    connector_type: str = Field(description="'source' or 'destination'")
+    canonical_name: str | None = Field(
+        default=None, description="The canonical connector name if resolved"
+    )
+    total_connections: int = Field(
+        description="Total number of non-deprecated connections"
+    )
+    enabled_connections: int = Field(
+        description="Number of enabled (active status) connections"
+    )
+    active_connections: int = Field(
+        description="Number of connections with recent sync activity"
+    )
+    pinned_connections: int = Field(
+        description="Number of connections with explicit version pins"
+    )
+    unpinned_connections: int = Field(
+        description="Number of connections on default version"
+    )
+    latest_attempt: LatestAttemptBreakdown = Field(
+        description="Overall breakdown by latest attempt status"
+    )
+    by_version: list[VersionPinStats] = Field(
+        description="Stats broken down by pinned version"
+    )
+
+
+class ConnectorConnectionStatsResponse(BaseModel):
+    """Response containing connection stats for multiple connectors."""
+
+    sources: list[ConnectorConnectionStats] = Field(
+        default_factory=list, description="Stats for source connectors"
+    )
+    destinations: list[ConnectorConnectionStats] = Field(
+        default_factory=list, description="Stats for destination connectors"
+    )
+    active_within_days: int = Field(
+        description="Lookback window used for 'active' connections"
+    )
+    generated_at: datetime = Field(description="When this response was generated")
+
+
 # Cloud registry URL for resolving canonical names
 CLOUD_REGISTRY_URL = (
     "https://connectors.airbyte.com/files/registries/v0/cloud_registry.json"
@@ -851,6 +941,213 @@ def query_prod_workspaces_by_email_domain(
     )
 
 
+def _build_connector_stats(
+    connector_definition_id: str,
+    connector_type: str,
+    canonical_name: str | None,
+    rows: list[dict[str, Any]],
+    version_tags: dict[str, str | None],
+) -> ConnectorConnectionStats:
+    """Build ConnectorConnectionStats from query result rows."""
+    # Aggregate totals across all version groups
+    total_connections = 0
+    enabled_connections = 0
+    active_connections = 0
+    pinned_connections = 0
+    unpinned_connections = 0
+    total_succeeded = 0
+    total_failed = 0
+    total_cancelled = 0
+    total_running = 0
+    total_unknown = 0
+
+    by_version: list[VersionPinStats] = []
+
+    for row in rows:
+        version_id = row.get("pinned_version_id")
+        row_total = int(row.get("total_connections", 0))
+        row_enabled = int(row.get("enabled_connections", 0))
+        row_active = int(row.get("active_connections", 0))
+        row_pinned = int(row.get("pinned_connections", 0))
+        row_unpinned = int(row.get("unpinned_connections", 0))
+        row_succeeded = int(row.get("succeeded_connections", 0))
+        row_failed = int(row.get("failed_connections", 0))
+        row_cancelled = int(row.get("cancelled_connections", 0))
+        row_running = int(row.get("running_connections", 0))
+        row_unknown = int(row.get("unknown_connections", 0))
+
+        total_connections += row_total
+        enabled_connections += row_enabled
+        active_connections += row_active
+        pinned_connections += row_pinned
+        unpinned_connections += row_unpinned
+        total_succeeded += row_succeeded
+        total_failed += row_failed
+        total_cancelled += row_cancelled
+        total_running += row_running
+        total_unknown += row_unknown
+
+        by_version.append(
+            VersionPinStats(
+                pinned_version_id=str(version_id) if version_id else None,
+                docker_image_tag=version_tags.get(str(version_id))
+                if version_id
+                else None,
+                total_connections=row_total,
+                enabled_connections=row_enabled,
+                active_connections=row_active,
+                latest_attempt=LatestAttemptBreakdown(
+                    succeeded=row_succeeded,
+                    failed=row_failed,
+                    cancelled=row_cancelled,
+                    running=row_running,
+                    unknown=row_unknown,
+                ),
+            )
+        )
+
+    return ConnectorConnectionStats(
+        connector_definition_id=connector_definition_id,
+        connector_type=connector_type,
+        canonical_name=canonical_name,
+        total_connections=total_connections,
+        enabled_connections=enabled_connections,
+        active_connections=active_connections,
+        pinned_connections=pinned_connections,
+        unpinned_connections=unpinned_connections,
+        latest_attempt=LatestAttemptBreakdown(
+            succeeded=total_succeeded,
+            failed=total_failed,
+            cancelled=total_cancelled,
+            running=total_running,
+            unknown=total_unknown,
+        ),
+        by_version=by_version,
+    )
+
+
+@mcp_tool(
+    read_only=True,
+    idempotent=True,
+    open_world=True,
+)
+def query_prod_connector_connection_stats(
+    source_definition_ids: Annotated[
+        list[str] | None,
+        Field(
+            description=(
+                "List of source connector definition IDs (UUIDs) to get stats for. "
+                "Example: ['afa734e4-3571-11ec-991a-1e0031268139']"
+            ),
+            default=None,
+        ),
+    ] = None,
+    destination_definition_ids: Annotated[
+        list[str] | None,
+        Field(
+            description=(
+                "List of destination connector definition IDs (UUIDs) to get stats for. "
+                "Example: ['94bd199c-2ff0-4aa2-b98e-17f0acb72610']"
+            ),
+            default=None,
+        ),
+    ] = None,
+    active_within_days: Annotated[
+        int,
+        Field(
+            description=(
+                "Number of days to look back for 'active' connections (default: 7). "
+                "Connections with sync activity within this window are counted as active."
+            ),
+            default=7,
+        ),
+    ] = 7,
+) -> ConnectorConnectionStatsResponse:
+    """Get aggregate connection stats for multiple connectors.
+
+    Returns counts of connections grouped by pinned version for each connector,
+    including:
+    - Total, enabled, and active connection counts
+    - Pinned vs unpinned breakdown
+    - Latest attempt status breakdown (succeeded, failed, cancelled, running, unknown)
+
+    This tool is designed for release monitoring workflows. It allows you to:
+    1. Query recently released connectors to identify which ones to monitor
+    2. Get aggregate stats showing how many connections are using each version
+    3. See health metrics (pass/fail) broken down by version
+
+    The 'active_within_days' parameter controls the lookback window for:
+    - Counting 'active' connections (those with recent sync activity)
+    - Determining 'latest attempt status' (most recent attempt within the window)
+
+    Connections with no sync activity in the lookback window will have
+    'unknown' status in the latest_attempt breakdown.
+    """
+    # Initialize empty lists if None
+    source_ids = source_definition_ids or []
+    destination_ids = destination_definition_ids or []
+
+    if not source_ids and not destination_ids:
+        raise PyAirbyteInputError(
+            message=(
+                "At least one of source_definition_ids or destination_definition_ids "
+                "must be provided."
+            ),
+        )
+
+    sources: list[ConnectorConnectionStats] = []
+    destinations: list[ConnectorConnectionStats] = []
+
+    # Process source connectors
+    for source_def_id in source_ids:
+        # Get version info for tag lookup
+        versions = query_connector_versions(source_def_id)
+        version_tags = {
+            str(v["version_id"]): v.get("docker_image_tag") for v in versions
+        }
+
+        # Get aggregate stats
+        rows = query_source_connection_stats(source_def_id, days=active_within_days)
+
+        sources.append(
+            _build_connector_stats(
+                connector_definition_id=source_def_id,
+                connector_type="source",
+                canonical_name=None,
+                rows=rows,
+                version_tags=version_tags,
+            )
+        )
+
+    # Process destination connectors
+    for dest_def_id in destination_ids:
+        # Get version info for tag lookup
+        versions = query_connector_versions(dest_def_id)
+        version_tags = {
+            str(v["version_id"]): v.get("docker_image_tag") for v in versions
+        }
+
+        # Get aggregate stats
+        rows = query_destination_connection_stats(dest_def_id, days=active_within_days)
+
+        destinations.append(
+            _build_connector_stats(
+                connector_definition_id=dest_def_id,
+                connector_type="destination",
+                canonical_name=None,
+                rows=rows,
+                version_tags=version_tags,
+            )
+        )
+
+    return ConnectorConnectionStatsResponse(
+        sources=sources,
+        destinations=destinations,
+        active_within_days=active_within_days,
+        generated_at=datetime.now(timezone.utc),
+    )
+
+
 def register_prod_db_query_tools(app: FastMCP) -> None:
     """Register prod DB query tools with the FastMCP app."""
     register_mcp_tools(app, domain=__name__)