airbyte-internal-ops 0.9.1__py3-none-any.whl → 0.10.0__py3-none-any.whl

airbyte_ops_mcp/mcp/prod_db_queries.py

@@ -23,6 +23,7 @@ from airbyte_ops_mcp.prod_db_access.queries import (
     query_connections_by_connector,
     query_connections_by_destination_connector,
     query_connections_by_stream,
+    query_connector_rollouts,
     query_connector_versions,
     query_dataplanes_list,
     query_destination_connection_stats,
@@ -1272,6 +1273,232 @@ def query_prod_connector_connection_stats(
     )
 
 
+# =============================================================================
+# Connector Rollout Models and Tools
+# =============================================================================
+
+
+class ConnectorRolloutInfo(BaseModel):
+    """Information about a connector rollout."""
+
+    rollout_id: str = Field(description="The rollout UUID")
+    actor_definition_id: str = Field(description="The connector definition UUID")
+    state: str = Field(
+        description="Rollout state: initialized, workflow_started, in_progress, "
+        "paused, finalizing, succeeded, errored, failed_rolled_back, canceled"
+    )
+    initial_rollout_pct: int | None = Field(
+        default=None, description="Initial rollout percentage"
+    )
+    current_target_rollout_pct: int | None = Field(
+        default=None, description="Current target rollout percentage"
+    )
+    final_target_rollout_pct: int | None = Field(
+        default=None, description="Final target rollout percentage"
+    )
+    has_breaking_changes: bool = Field(
+        description="Whether the RC has breaking changes"
+    )
+    max_step_wait_time_mins: int | None = Field(
+        default=None, description="Maximum wait time between rollout steps in minutes"
+    )
+    rollout_strategy: str | None = Field(
+        default=None, description="Rollout strategy: manual, automated, overridden"
+    )
+    workflow_run_id: str | None = Field(
+        default=None, description="Temporal workflow run ID"
+    )
+    error_msg: str | None = Field(default=None, description="Error message if errored")
+    failed_reason: str | None = Field(
+        default=None, description="Reason for failure if failed"
+    )
+    paused_reason: str | None = Field(
+        default=None, description="Reason for pause if paused"
+    )
+    tag: str | None = Field(default=None, description="Optional tag for the rollout")
+    created_at: datetime | None = Field(
+        default=None, description="When the rollout was created"
+    )
+    updated_at: datetime | None = Field(
+        default=None, description="When the rollout was last updated"
+    )
+    completed_at: datetime | None = Field(
+        default=None, description="When the rollout completed (if terminal)"
+    )
+    expires_at: datetime | None = Field(
+        default=None, description="When the rollout expires"
+    )
+    rc_docker_image_tag: str | None = Field(
+        default=None, description="Docker image tag of the release candidate"
+    )
+    rc_docker_repository: str | None = Field(
+        default=None, description="Docker repository of the release candidate"
+    )
+    initial_docker_image_tag: str | None = Field(
+        default=None, description="Docker image tag of the initial version"
+    )
+    initial_docker_repository: str | None = Field(
+        default=None, description="Docker repository of the initial version"
+    )
+
+
+def _row_to_connector_rollout_info(row: dict[str, Any]) -> ConnectorRolloutInfo:
+    """Convert a database row to a ConnectorRolloutInfo model."""
+    return ConnectorRolloutInfo(
+        rollout_id=str(row["rollout_id"]),
+        actor_definition_id=str(row["actor_definition_id"]),
+        state=row["state"],
+        initial_rollout_pct=row.get("initial_rollout_pct"),
+        current_target_rollout_pct=row.get("current_target_rollout_pct"),
+        final_target_rollout_pct=row.get("final_target_rollout_pct"),
+        has_breaking_changes=row["has_breaking_changes"],
+        max_step_wait_time_mins=row.get("max_step_wait_time_mins"),
+        rollout_strategy=row.get("rollout_strategy"),
+        workflow_run_id=row.get("workflow_run_id"),
+        error_msg=row.get("error_msg"),
+        failed_reason=row.get("failed_reason"),
+        paused_reason=row.get("paused_reason"),
+        tag=row.get("tag"),
+        created_at=row.get("created_at"),
+        updated_at=row.get("updated_at"),
+        completed_at=row.get("completed_at"),
+        expires_at=row.get("expires_at"),
+        rc_docker_image_tag=row.get("rc_docker_image_tag"),
+        rc_docker_repository=row.get("rc_docker_repository"),
+        initial_docker_image_tag=row.get("initial_docker_image_tag"),
+        initial_docker_repository=row.get("initial_docker_repository"),
+    )
+
+
+@mcp_tool(
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_connector_rollouts(
+    actor_definition_id: Annotated[
+        str | None,
+        Field(description="Connector definition UUID to filter by (optional)"),
+    ] = None,
+    rollout_id: Annotated[
+        str | None,
+        Field(description="Specific rollout UUID to look up (optional)"),
+    ] = None,
+    active_only: Annotated[
+        bool,
+        Field(description="If true, only return active (non-terminal) rollouts"),
+    ] = False,
+    limit: Annotated[
+        int,
+        Field(description="Maximum number of results (default: 100)"),
+    ] = 100,
+) -> list[ConnectorRolloutInfo]:
+    """Query connector rollouts with flexible filtering.
+
+    Returns rollouts based on the provided filters. If no filters are specified,
+    returns all active rollouts. Useful for monitoring rollout status and history.
+
+    Filter behavior:
+    - rollout_id: Returns that specific rollout (ignores other filters)
+    - active_only: Returns only active (non-terminal) rollouts
+    - actor_definition_id: Returns rollouts for that specific connector
+    - No filters: Returns all active rollouts (same as active_only=True)
+    """
+    rows = query_connector_rollouts(
+        actor_definition_id=actor_definition_id,
+        rollout_id=rollout_id,
+        active_only=active_only,
+        limit=limit,
+    )
+    return [_row_to_connector_rollout_info(row) for row in rows]
+
+
+class RolloutMonitoringStats(BaseModel):
+    """Aggregate monitoring stats for a connector rollout."""
+
+    rollout_id: str = Field(description="Rollout UUID")
+    actor_definition_id: str = Field(description="Connector definition UUID")
+    num_actors: int = Field(description="Total actors using this connector")
+    num_actors_pinned: int = Field(
+        description="Actors with any version pin (eligible or already pinned)"
+    )
+    num_pinned_to_rollout: int = Field(
+        description="Actors specifically pinned to this rollout"
+    )
+
+
+class RolloutActorSyncStats(BaseModel):
+    """Per-actor sync stats for a rollout."""
+
+    actor_id: str = Field(description="Actor UUID")
+    num_succeeded: int = Field(description="Number of successful syncs")
+    num_failed: int = Field(description="Number of failed syncs")
+    num_connections: int = Field(description="Number of connections using this actor")
+
+
+class RolloutMonitoringResult(BaseModel):
+    """Complete monitoring result for a rollout."""
+
+    aggregate_stats: RolloutMonitoringStats = Field(
+        description="Aggregate counts for the rollout"
+    )
+    actor_stats: list[RolloutActorSyncStats] = Field(
+        description="Per-actor sync stats for actors pinned to the rollout"
+    )
+
+
+@mcp_tool(
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_rollout_monitoring_stats(
+    rollout_id: Annotated[
+        str,
+        Field(description="Rollout UUID to get monitoring stats for"),
+    ],
+    days_back: Annotated[
+        int,
+        Field(description="Number of days to look back for sync stats (default: 7)"),
+    ] = 7,
+) -> RolloutMonitoringResult:
+    """Get monitoring stats for a connector rollout.
+
+    Returns aggregate counts (total actors, pinned actors, rollout-pinned actors)
+    and per-actor sync stats (succeeded, failed, connections) for actors
+    participating in the rollout. This replaces the Retool rollout monitoring UI.
+    """
+    from airbyte_ops_mcp.prod_db_access.queries import (
+        query_rollout_actor_sync_stats,
+        query_rollout_monitoring_stats,
+    )
+
+    agg_stats = query_rollout_monitoring_stats(rollout_id=rollout_id)
+    if agg_stats is None:
+        raise ValueError(f"Rollout not found: {rollout_id}")
+
+    actor_stats = query_rollout_actor_sync_stats(
+        rollout_id=rollout_id, days_back=days_back
+    )
+
+    return RolloutMonitoringResult(
+        aggregate_stats=RolloutMonitoringStats(
+            rollout_id=str(agg_stats["rollout_id"]),
+            actor_definition_id=str(agg_stats["actor_definition_id"]),
+            num_actors=agg_stats["num_actors"] or 0,
+            num_actors_pinned=agg_stats["num_actors_pinned"] or 0,
+            num_pinned_to_rollout=agg_stats["num_pinned_to_rollout"] or 0,
+        ),
+        actor_stats=[
+            RolloutActorSyncStats(
+                actor_id=str(row["actor_id"]),
+                num_succeeded=row["num_succeeded"] or 0,
+                num_failed=row["num_failed"] or 0,
+                num_connections=row["num_connections"] or 0,
+            )
+            for row in actor_stats
+        ],
+    )
+
+
 def register_prod_db_query_tools(app: FastMCP) -> None:
     """Register prod DB query tools with the FastMCP app."""
     register_mcp_tools(app, mcp_module=__name__)

airbyte_ops_mcp/mcp/server.py

@@ -33,6 +33,7 @@ from airbyte_ops_mcp.mcp._guidance import MCP_SERVER_INSTRUCTIONS
 from airbyte_ops_mcp.mcp.cloud_connector_versions import (
     register_cloud_connector_version_tools,
 )
+from airbyte_ops_mcp.mcp.connector_rollout import register_connector_rollout_tools
 from airbyte_ops_mcp.mcp.gcp_logs import register_gcp_logs_tools
 from airbyte_ops_mcp.mcp.github_actions import register_github_actions_tools
 from airbyte_ops_mcp.mcp.github_repo_ops import register_github_repo_ops_tools
@@ -117,6 +118,7 @@ def register_server_assets(app: FastMCP) -> None:
     register_github_actions_tools(app)
     register_prerelease_tools(app)
     register_cloud_connector_version_tools(app)
+    register_connector_rollout_tools(app)
     register_prod_db_query_tools(app)
     register_gcp_logs_tools(app)
     register_prompts(app)

airbyte_ops_mcp/prod_db_access/queries.py

@@ -19,6 +19,8 @@ from google.cloud import secretmanager
 from airbyte_ops_mcp.gcp_auth import get_secret_manager_client
 from airbyte_ops_mcp.prod_db_access.db_engine import get_pool
 from airbyte_ops_mcp.prod_db_access.sql import (
+    SELECT_ACTIVE_CONNECTOR_ROLLOUTS,
+    SELECT_ACTIVE_CONNECTOR_ROLLOUTS_BY_DEFINITION,
     SELECT_ACTORS_PINNED_TO_VERSION,
     SELECT_CONNECTIONS_BY_CONNECTOR,
     SELECT_CONNECTIONS_BY_CONNECTOR_AND_ORG,
@@ -26,6 +28,8 @@ from airbyte_ops_mcp.prod_db_access.sql import (
     SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR_AND_ORG,
     SELECT_CONNECTIONS_BY_SOURCE_CONNECTOR_AND_STREAM,
     SELECT_CONNECTIONS_BY_SOURCE_CONNECTOR_AND_STREAM_AND_ORG,
+    SELECT_CONNECTOR_ROLLOUT_BY_ID,
+    SELECT_CONNECTOR_ROLLOUTS,
     SELECT_CONNECTOR_VERSIONS,
     SELECT_DATAPLANES_LIST,
     SELECT_DESTINATION_CONNECTION_STATS,
@@ -38,6 +42,8 @@ from airbyte_ops_mcp.prod_db_access.sql import (
     SELECT_RECENT_SUCCESSFUL_SYNCS_FOR_SOURCE_CONNECTOR,
     SELECT_RECENT_SYNCS_FOR_DESTINATION_CONNECTOR,
     SELECT_RECENT_SYNCS_FOR_SOURCE_CONNECTOR,
+    SELECT_ROLLOUT_ACTOR_SYNC_STATS,
+    SELECT_ROLLOUT_AGGREGATE_STATS,
     SELECT_SOURCE_CONNECTION_STATS,
     SELECT_SUCCESSFUL_SYNCS_FOR_VERSION,
     SELECT_SYNC_RESULTS_FOR_VERSION,
@@ -616,3 +622,141 @@ def query_connections_by_stream(
         query_name="SELECT_CONNECTIONS_BY_SOURCE_CONNECTOR_AND_STREAM_AND_ORG",
         gsm_client=gsm_client,
     )
+
+
+def query_connector_rollouts(
+    actor_definition_id: str | None = None,
+    rollout_id: str | None = None,
+    active_only: bool = False,
+    limit: int = 100,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query connector rollouts with flexible filtering.
+
+    This is the unified query function for connector rollouts. Based on the
+    arguments provided, it will:
+    - If rollout_id is provided: Return that specific rollout (as a single-item list)
+    - If active_only is True AND actor_definition_id is provided: Return active rollouts for that connector
+    - If active_only is True: Return all active (non-terminal) rollouts
+    - If actor_definition_id is provided: Return all rollouts for that connector (including terminal)
+    - Otherwise: Return all active rollouts (up to limit)
+
+    Args:
+        actor_definition_id: Optional connector definition UUID to filter by
+        rollout_id: Optional specific rollout UUID to look up
+        active_only: If True, only return active (non-terminal) rollouts
+        limit: Maximum number of results (default: 100)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of rollout records with version details
+    """
+    if rollout_id is not None:
+        rows = _run_sql_query(
+            SELECT_CONNECTOR_ROLLOUT_BY_ID,
+            parameters={"rollout_id": rollout_id},
+            query_name="SELECT_CONNECTOR_ROLLOUT_BY_ID",
+            gsm_client=gsm_client,
+        )
+        return rows
+
+    # Handle active_only with optional actor_definition_id filter
+    if active_only:
+        if actor_definition_id is not None:
+            # Filter by both active states AND actor_definition_id
+            return _run_sql_query(
+                SELECT_ACTIVE_CONNECTOR_ROLLOUTS_BY_DEFINITION,
+                parameters={
+                    "actor_definition_id": actor_definition_id,
+                    "limit": limit,
+                },
+                query_name="SELECT_ACTIVE_CONNECTOR_ROLLOUTS_BY_DEFINITION",
+                gsm_client=gsm_client,
+            )
+        # Only active filter, no actor_definition_id
+        return _run_sql_query(
+            SELECT_ACTIVE_CONNECTOR_ROLLOUTS,
+            parameters={"limit": limit},
+            query_name="SELECT_ACTIVE_CONNECTOR_ROLLOUTS",
+            gsm_client=gsm_client,
+        )
+
+    # Not active_only, but filter by actor_definition_id (all states)
+    if actor_definition_id is not None:
+        return _run_sql_query(
+            SELECT_CONNECTOR_ROLLOUTS,
+            parameters={
+                "actor_definition_id": actor_definition_id,
+                "limit": limit,
+            },
+            query_name="SELECT_CONNECTOR_ROLLOUTS",
+            gsm_client=gsm_client,
+        )
+
+    # Default: return active rollouts if no filters specified
+    return _run_sql_query(
+        SELECT_ACTIVE_CONNECTOR_ROLLOUTS,
+        parameters={"limit": limit},
+        query_name="SELECT_ACTIVE_CONNECTOR_ROLLOUTS",
+        gsm_client=gsm_client,
+    )
+
+
+def query_rollout_monitoring_stats(
+    rollout_id: str,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> dict[str, Any] | None:
+    """Query aggregate monitoring stats for a rollout.
+
+    Returns counts of total actors, actors with any pin, and actors pinned
+    specifically to this rollout.
+
+    Args:
+        rollout_id: Rollout UUID to get stats for
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        Dict with rollout_id, actor_definition_id, num_actors, num_actors_pinned,
+        num_pinned_to_rollout, or None if rollout not found
+    """
+    rows = _run_sql_query(
+        SELECT_ROLLOUT_AGGREGATE_STATS,
+        parameters={"rollout_id": rollout_id},
+        query_name="SELECT_ROLLOUT_AGGREGATE_STATS",
+        gsm_client=gsm_client,
+    )
+    # The SQL query uses scalar subqueries that always return one row,
+    # but with NULL values if the rollout doesn't exist
+    if rows and rows[0].get("rollout_id") is not None:
+        return rows[0]
+    return None
+
+
+def query_rollout_actor_sync_stats(
+    rollout_id: str,
+    days_back: int = 7,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query per-actor sync stats for actors pinned to a rollout.
+
+    Returns succeeded/failed job counts and connection counts for each actor
+    that is pinned to the specified rollout.
+
+    Args:
+        rollout_id: Rollout UUID to get actor stats for
+        days_back: Number of days to look back for job stats (default: 7)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of dicts with actor_id, num_succeeded, num_failed, num_connections
+    """
+    cutoff_date = datetime.now(timezone.utc) - timedelta(days=days_back)
+    return _run_sql_query(
+        SELECT_ROLLOUT_ACTOR_SYNC_STATS,
+        parameters={"rollout_id": rollout_id, "cutoff_date": cutoff_date},
+        query_name="SELECT_ROLLOUT_ACTOR_SYNC_STATS",
+        gsm_client=gsm_client,
+    )