airbyte-internal-ops 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

airbyte_ops_mcp/mcp/prod_db_queries.py

@@ -7,7 +7,8 @@ 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
 
 import requests
@@ -23,13 +24,25 @@ 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_sync_results_for_version,
+    query_recent_syncs_for_connector,
+    query_source_connection_stats,
+    query_syncs_for_version_pinned_connector,
     query_workspace_info,
     query_workspaces_by_email_domain,
 )
 
+
+class StatusFilter(StrEnum):
+    """Filter for job status in sync queries."""
+
+    ALL = "all"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+
+
 # Cloud UI base URL for building connection URLs
 CLOUD_UI_BASE_URL = "https://cloud.airbyte.com"
 
@@ -79,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"
@@ -293,7 +394,7 @@ def query_prod_actors_by_connector_version(
     read_only=True,
     idempotent=True,
 )
-def query_prod_connector_version_sync_results(
+def query_prod_recent_syncs_for_version_pinned_connector(
     connector_version_id: Annotated[
         str,
         Field(description="Connector version UUID to find sync results for"),
@@ -314,11 +415,16 @@ def query_prod_connector_version_sync_results(
         ),
     ] = False,
 ) -> list[dict[str, Any]]:
-    """List sync job results for actors pinned to a specific connector version.
+    """List sync job results for actors PINNED to a specific connector version.
+
+    IMPORTANT: This tool ONLY returns results for actors that have been explicitly
+    pinned to the specified version via scoped_configuration. Most connections run
+    unpinned and will NOT appear in these results.
 
-    Returns sync job results for connections using actors that are pinned
-    to the specified version. Useful for monitoring rollout health and
-    identifying issues with specific connector versions.
+    Use this tool when you want to monitor rollout health for actors that have been
+    explicitly pinned to a pre-release or specific version. For finding healthy
+    connections across ALL actors using a connector type (regardless of pinning),
+    use query_prod_recent_syncs_for_connector instead.
 
     The actor_id field is the actor ID (superset of source_id/destination_id).
 
@@ -327,7 +433,7 @@ def query_prod_connector_version_sync_results(
     pin_origin_type, pin_origin, workspace_id, workspace_name, organization_id,
     dataplane_group_id, dataplane_name
     """
-    return query_sync_results_for_version(
+    return query_syncs_for_version_pinned_connector(
         connector_version_id,
         days=days,
         limit=limit,
@@ -335,6 +441,163 @@ def query_prod_connector_version_sync_results(
     )
 
 
+@mcp_tool(
+    read_only=True,
+    idempotent=True,
+    open_world=True,
+)
+def query_prod_recent_syncs_for_connector(
+    source_definition_id: Annotated[
+        str | None,
+        Field(
+            description=(
+                "Source connector definition ID (UUID) to search for. "
+                "Provide this OR source_canonical_name OR destination_definition_id "
+                "OR destination_canonical_name (exactly one required). "
+                "Example: 'afa734e4-3571-11ec-991a-1e0031268139' for YouTube Analytics."
+            ),
+            default=None,
+        ),
+    ],
+    source_canonical_name: Annotated[
+        str | None,
+        Field(
+            description=(
+                "Canonical source connector name to search for. "
+                "Provide this OR source_definition_id OR destination_definition_id "
+                "OR destination_canonical_name (exactly one required). "
+                "Examples: 'source-youtube-analytics', 'YouTube Analytics'."
+            ),
+            default=None,
+        ),
+    ],
+    destination_definition_id: Annotated[
+        str | None,
+        Field(
+            description=(
+                "Destination connector definition ID (UUID) to search for. "
+                "Provide this OR destination_canonical_name OR source_definition_id "
+                "OR source_canonical_name (exactly one required). "
+                "Example: '94bd199c-2ff0-4aa2-b98e-17f0acb72610' for DuckDB."
+            ),
+            default=None,
+        ),
+    ],
+    destination_canonical_name: Annotated[
+        str | None,
+        Field(
+            description=(
+                "Canonical destination connector name to search for. "
+                "Provide this OR destination_definition_id OR source_definition_id "
+                "OR source_canonical_name (exactly one required). "
+                "Examples: 'destination-duckdb', 'DuckDB'."
+            ),
+            default=None,
+        ),
+    ],
+    status_filter: Annotated[
+        StatusFilter,
+        Field(
+            description=(
+                "Filter by job status: 'all' (default), 'succeeded', or 'failed'. "
+                "Use 'succeeded' to find healthy connections with recent successful syncs. "
+                "Use 'failed' to find connections with recent failures."
+            ),
+            default=StatusFilter.ALL,
+        ),
+    ],
+    organization_id: Annotated[
+        str | OrganizationAliasEnum | None,
+        Field(
+            description=(
+                "Optional organization ID (UUID) or alias to filter results. "
+                "If provided, only syncs from this organization will be returned. "
+                "Accepts '@airbyte-internal' as an alias for the Airbyte internal org."
+            ),
+            default=None,
+        ),
+    ],
+    lookback_days: Annotated[
+        int,
+        Field(description="Number of days to look back (default: 7)", default=7),
+    ],
+    limit: Annotated[
+        int,
+        Field(description="Maximum number of results (default: 100)", default=100),
+    ],
+) -> list[dict[str, Any]]:
+    """List recent sync jobs for ALL actors using a connector type.
+
+    This tool finds all actors with the given connector definition and returns their
+    recent sync jobs, regardless of whether they have explicit version pins. It filters
+    out deleted actors, deleted workspaces, and deprecated connections.
+
+    Use this tool to:
+    - Find healthy connections with recent successful syncs (status_filter='succeeded')
+    - Investigate connector issues across all users (status_filter='failed')
+    - Get an overview of all recent sync activity (status_filter='all')
+
+    Supports both SOURCE and DESTINATION connectors. Provide exactly one of:
+    source_definition_id, source_canonical_name, destination_definition_id,
+    or destination_canonical_name.
+
+    Key fields in results:
+    - job_status: 'succeeded', 'failed', 'cancelled', etc.
+    - connection_id, connection_name: The connection that ran the sync
+    - actor_id, actor_name: The source or destination actor
+    - pin_origin_type, pin_origin, pinned_version_id: Version pin context (NULL if not pinned)
+    """
+    # Validate that exactly one connector parameter is provided
+    provided_params = [
+        source_definition_id,
+        source_canonical_name,
+        destination_definition_id,
+        destination_canonical_name,
+    ]
+    num_provided = sum(p is not None for p in provided_params)
+    if num_provided != 1:
+        raise PyAirbyteInputError(
+            message=(
+                "Exactly one of source_definition_id, source_canonical_name, "
+                "destination_definition_id, or destination_canonical_name must be provided."
+            ),
+        )
+
+    # Determine if this is a destination connector
+    is_destination = (
+        destination_definition_id is not None or destination_canonical_name is not None
+    )
+
+    # Resolve canonical name to definition ID if needed
+    resolved_definition_id: str
+    if source_canonical_name:
+        resolved_definition_id = _resolve_canonical_name_to_definition_id(
+            canonical_name=source_canonical_name,
+        )
+    elif destination_canonical_name:
+        resolved_definition_id = _resolve_canonical_name_to_definition_id(
+            canonical_name=destination_canonical_name,
+        )
+    elif source_definition_id:
+        resolved_definition_id = source_definition_id
+    else:
+        # We've validated exactly one param is provided, so this must be set
+        assert destination_definition_id is not None
+        resolved_definition_id = destination_definition_id
+
+    # Resolve organization ID alias
+    resolved_organization_id = OrganizationAliasEnum.resolve(organization_id)
+
+    return query_recent_syncs_for_connector(
+        connector_definition_id=resolved_definition_id,
+        is_destination=is_destination,
+        status_filter=status_filter,
+        organization_id=resolved_organization_id,
+        days=lookback_days,
+        limit=limit,
+    )
+
+
 @mcp_tool(
     read_only=True,
     idempotent=True,
@@ -678,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__)

airbyte_ops_mcp/mcp/regression_tests.py

@@ -37,6 +37,7 @@ from airbyte_ops_mcp.github_actions import (
     trigger_workflow_dispatch,
 )
 from airbyte_ops_mcp.mcp._mcp_utils import mcp_tool, register_mcp_tools
+from airbyte_ops_mcp.mcp.prerelease import ConnectorRepo
 
 logger = logging.getLogger(__name__)
 
@@ -309,7 +310,13 @@ def run_regression_tests(
     ],
     pr: Annotated[
         int,
-        "PR number from the airbyte monorepo to checkout and build from (e.g., 70847). Required.",
+        "PR number to checkout and build from (e.g., 70847). Required. "
+        "The PR must be from the repository specified by the 'repo' parameter.",
+    ],
+    repo: Annotated[
+        ConnectorRepo,
+        "Repository where the connector PR is located. "
+        "Use 'airbyte' for OSS connectors (default) or 'airbyte-enterprise' for enterprise connectors.",
     ],
     connection_id: Annotated[
         str | None,
@@ -347,6 +354,10 @@ def run_regression_tests(
     This tool triggers the regression test workflow which builds the connector
     from the specified PR and runs tests against it.
 
+    Supports both OSS connectors (from airbytehq/airbyte) and enterprise connectors
+    (from airbytehq/airbyte-enterprise). Use the 'repo' parameter to specify which
+    repository contains the connector PR.
+
     - skip_compare=False (default): Comparison mode - compares the PR version
       against the baseline (control) version.
     - skip_compare=True: Single-version mode - runs tests without comparison.
@@ -390,10 +401,11 @@ def run_regression_tests(
                 workflow_url=None,
             )
 
-    # Build workflow inputs - connector_name and pr are required
+    # Build workflow inputs - connector_name, pr, and repo are required
     workflow_inputs: dict[str, str] = {
         "connector_name": connector_name,
         "pr": str(pr),
+        "repo": repo,
     }
 
     # Add optional inputs
@@ -431,12 +443,13 @@ def run_regression_tests(
 
     view_url = dispatch_result.run_url or dispatch_result.workflow_url
     connection_info = f" for connection {connection_id}" if connection_id else ""
+    repo_info = f" from {repo}" if repo != ConnectorRepo.AIRBYTE else ""
     return RunRegressionTestsResponse(
         run_id=run_id,
         status=TestRunStatus.QUEUED,
         message=(
             f"{mode_description.capitalize()} regression test workflow triggered "
-            f"for {connector_name} (PR #{pr}){connection_info}. View progress at: {view_url}"
+            f"for {connector_name} (PR #{pr}{repo_info}){connection_info}. View progress at: {view_url}"
         ),
         workflow_url=dispatch_result.workflow_url,
         github_run_id=dispatch_result.run_id,