airbyte-internal-ops 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

airbyte_ops_mcp/mcp/prod_db_queries.py

@@ -0,0 +1,357 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""MCP tools for querying the Airbyte Cloud Prod DB Replica.
+
+This module provides MCP tools that wrap the query functions from
+airbyte_ops_mcp.prod_db_access.queries for use by AI agents.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Any
+
+import requests
+from airbyte.exceptions import PyAirbyteInputError
+from fastmcp import FastMCP
+from pydantic import Field
+
+from airbyte_ops_mcp.mcp._mcp_utils import ToolDomain, mcp_tool, register_mcp_tools
+from airbyte_ops_mcp.prod_db_access.queries import (
+    query_actors_pinned_to_version,
+    query_connections_by_connector,
+    query_connector_versions,
+    query_dataplanes_list,
+    query_new_connector_releases,
+    query_sync_results_for_version,
+    query_workspace_info,
+)
+
+# Cloud UI base URL for building connection URLs
+CLOUD_UI_BASE_URL = "https://cloud.airbyte.com"
+
+# Cloud registry URL for resolving canonical names
+CLOUD_REGISTRY_URL = (
+    "https://connectors.airbyte.com/files/registries/v0/cloud_registry.json"
+)
+
+
+def _resolve_canonical_name_to_definition_id(canonical_name: str) -> str:
+    """Resolve a canonical source name to a definition ID.
+
+    Args:
+        canonical_name: Canonical source name (e.g., 'source-youtube-analytics').
+
+    Returns:
+        The source definition ID (UUID).
+
+    Raises:
+        PyAirbyteInputError: If the canonical name cannot be resolved.
+    """
+    response = requests.get(CLOUD_REGISTRY_URL, timeout=60)
+
+    if response.status_code != 200:
+        raise PyAirbyteInputError(
+            message=f"Failed to fetch connector registry: {response.status_code}",
+            context={"response": response.text},
+        )
+
+    data = response.json()
+    sources = data.get("sources", [])
+
+    # Normalize the canonical name for matching
+    normalized_input = canonical_name.lower().strip()
+
+    # Try exact match on name field
+    for source in sources:
+        source_name = source.get("name", "").lower()
+        # The registry returns names like "YouTube Analytics"
+        # So we need to handle both formats
+        if source_name == normalized_input:
+            return source["sourceDefinitionId"]
+
+        # Also try matching against a slugified version
+        # e.g., "YouTube Analytics" -> "youtube-analytics"
+        slugified = source_name.replace(" ", "-")
+        if slugified == normalized_input or f"source-{slugified}" == normalized_input:
+            return source["sourceDefinitionId"]
+
+    raise PyAirbyteInputError(
+        message=f"Could not find source definition for canonical name: {canonical_name}",
+        context={
+            "hint": "Use the exact canonical name (e.g., 'source-youtube-analytics') "
+            "or display name (e.g., 'YouTube Analytics'). "
+            "You can list available sources using the connector registry tools.",
+            "searched_for": canonical_name,
+        },
+    )
+
+
+@mcp_tool(
+    ToolDomain.CLOUD_ADMIN,
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_dataplanes() -> list[dict[str, Any]]:
+    """List all dataplane groups with workspace counts.
+
+    Returns information about all active dataplane groups in Airbyte Cloud,
+    including the number of workspaces in each. Useful for understanding
+    the distribution of workspaces across regions (US, US-Central, EU).
+
+    Returns list of dicts with keys: dataplane_group_id, dataplane_name,
+    organization_id, enabled, tombstone, created_at, workspace_count
+    """
+    return query_dataplanes_list()
+
+
+@mcp_tool(
+    ToolDomain.CLOUD_ADMIN,
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_workspace_info(
+    workspace_id: Annotated[
+        str,
+        Field(description="Workspace UUID to look up"),
+    ],
+) -> dict[str, Any] | None:
+    """Get workspace information including dataplane group.
+
+    Returns details about a specific workspace, including which dataplane
+    (region) it belongs to. Useful for determining if a workspace is in
+    the EU region for filtering purposes.
+
+    Returns dict with keys: workspace_id, workspace_name, slug, organization_id,
+    dataplane_group_id, dataplane_name, created_at, tombstone
+    Or None if workspace not found.
+    """
+    return query_workspace_info(workspace_id)
+
+
+@mcp_tool(
+    ToolDomain.CLOUD_ADMIN,
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_connector_versions(
+    connector_definition_id: Annotated[
+        str,
+        Field(description="Connector definition UUID to list versions for"),
+    ],
+) -> list[dict[str, Any]]:
+    """List all versions for a connector definition.
+
+    Returns all published versions of a connector, ordered by last_published
+    date descending. Useful for understanding version history and finding
+    specific version IDs for pinning or rollout monitoring.
+
+    Returns list of dicts with keys: version_id, docker_image_tag, docker_repository,
+    release_stage, support_level, cdk_version, language, last_published, release_date
+    """
+    return query_connector_versions(connector_definition_id)
+
+
+@mcp_tool(
+    ToolDomain.CLOUD_ADMIN,
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_new_connector_releases(
+    days: Annotated[
+        int,
+        Field(description="Number of days to look back (default: 7)", default=7),
+    ] = 7,
+    limit: Annotated[
+        int,
+        Field(description="Maximum number of results (default: 100)", default=100),
+    ] = 100,
+) -> list[dict[str, Any]]:
+    """List recently published connector versions.
+
+    Returns connector versions published within the specified number of days.
+    Uses last_published timestamp which reflects when the version was actually
+    deployed to the registry (not the changelog date).
+
+    Returns list of dicts with keys: version_id, connector_definition_id, docker_repository,
+    docker_image_tag, last_published, release_date, release_stage, support_level,
+    cdk_version, language, created_at
+    """
+    return query_new_connector_releases(days=days, limit=limit)
+
+
+@mcp_tool(
+    ToolDomain.CLOUD_ADMIN,
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_actors_by_connector_version(
+    connector_version_id: Annotated[
+        str,
+        Field(description="Connector version UUID to find pinned instances for"),
+    ],
+) -> list[dict[str, Any]]:
+    """List actors (sources/destinations) pinned to a specific connector version.
+
+    Returns all actors that have been explicitly pinned to a specific
+    connector version via scoped_configuration. Useful for monitoring
+    rollouts and understanding which customers are affected by version pins.
+
+    The actor_id field is the actor ID (superset of source_id/destination_id).
+
+    Returns list of dicts with keys: actor_id, connector_definition_id, origin_type,
+    origin, description, created_at, expires_at, actor_name, workspace_id,
+    workspace_name, organization_id, dataplane_group_id, dataplane_name
+    """
+    return query_actors_pinned_to_version(connector_version_id)
+
+
+@mcp_tool(
+    ToolDomain.CLOUD_ADMIN,
+    read_only=True,
+    idempotent=True,
+)
+def query_prod_connector_version_sync_results(
+    connector_version_id: Annotated[
+        str,
+        Field(description="Connector version UUID to find sync results for"),
+    ],
+    days: Annotated[
+        int,
+        Field(description="Number of days to look back (default: 7)", default=7),
+    ] = 7,
+    limit: Annotated[
+        int,
+        Field(description="Maximum number of results (default: 100)", default=100),
+    ] = 100,
+    successful_only: Annotated[
+        bool,
+        Field(
+            description="If True, only return successful syncs (default: False)",
+            default=False,
+        ),
+    ] = False,
+) -> list[dict[str, Any]]:
+    """List sync job results for actors pinned to a specific connector version.
+
+    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.
+
+    The actor_id field is the actor ID (superset of source_id/destination_id).
+
+    Returns list of dicts with keys: job_id, connection_id, job_status, started_at,
+    job_updated_at, connection_name, actor_id, actor_name, connector_definition_id,
+    pin_origin_type, pin_origin, workspace_id, workspace_name, organization_id,
+    dataplane_group_id, dataplane_name
+    """
+    return query_sync_results_for_version(
+        connector_version_id,
+        days=days,
+        limit=limit,
+        successful_only=successful_only,
+    )
+
+
+@mcp_tool(
+    ToolDomain.CLOUD_ADMIN,
+    read_only=True,
+    idempotent=True,
+    open_world=True,
+)
+def query_prod_connections_by_connector(
+    source_definition_id: Annotated[
+        str | None,
+        Field(
+            description=(
+                "Source connector definition ID (UUID) to search for. "
+                "Exactly one of this or source_canonical_name is required. "
+                "Example: 'afa734e4-3571-11ec-991a-1e0031268139' for YouTube Analytics."
+            ),
+            default=None,
+        ),
+    ] = None,
+    source_canonical_name: Annotated[
+        str | None,
+        Field(
+            description=(
+                "Canonical source connector name to search for. "
+                "Exactly one of this or source_definition_id is required. "
+                "Examples: 'source-youtube-analytics', 'YouTube Analytics'."
+            ),
+            default=None,
+        ),
+    ] = None,
+    organization_id: Annotated[
+        str | None,
+        Field(
+            description=(
+                "Optional organization ID (UUID) to filter results. "
+                "If provided, only connections in this organization will be returned."
+            ),
+            default=None,
+        ),
+    ] = None,
+    limit: Annotated[
+        int,
+        Field(description="Maximum number of results (default: 1000)", default=1000),
+    ] = 1000,
+) -> list[dict[str, Any]]:
+    """Search for all connections using a specific source connector type.
+
+    This tool queries the Airbyte Cloud Prod DB Replica directly for fast results.
+    It finds all connections where the source connector matches the specified type,
+    regardless of how the source is named by users.
+
+    Optionally filter by organization_id to limit results to a specific organization.
+
+    Returns a list of connection dicts with workspace context and clickable Cloud UI URLs.
+    Each dict contains: connection_id, connection_name, connection_url, source_id,
+    source_name, source_definition_id, workspace_id, workspace_name, organization_id,
+    dataplane_group_id, dataplane_name.
+    """
+    # Validate that exactly one of the two parameters is provided
+    if (source_definition_id is None) == (source_canonical_name is None):
+        raise PyAirbyteInputError(
+            message=(
+                "Exactly one of source_definition_id or source_canonical_name "
+                "must be provided, but not both."
+            ),
+        )
+
+    # 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,
+        )
+    else:
+        resolved_definition_id = source_definition_id  # type: ignore[assignment]
+
+    # Query the database and transform rows to include connection URLs
+    return [
+        {
+            "organization_id": str(row.get("organization_id", "")),
+            "workspace_id": str(row["workspace_id"]),
+            "workspace_name": row.get("workspace_name", ""),
+            "connection_id": str(row["connection_id"]),
+            "connection_name": row.get("connection_name", ""),
+            "connection_url": (
+                f"{CLOUD_UI_BASE_URL}/workspaces/{row['workspace_id']}"
+                f"/connections/{row['connection_id']}/status"
+            ),
+            "source_id": str(row["source_id"]),
+            "source_name": row.get("source_name", ""),
+            "source_definition_id": str(row["source_definition_id"]),
+            "dataplane_group_id": str(row.get("dataplane_group_id", "")),
+            "dataplane_name": row.get("dataplane_name", ""),
+        }
+        for row in query_connections_by_connector(
+            connector_definition_id=resolved_definition_id,
+            organization_id=organization_id,
+            limit=limit,
+        )
+    ]
+
+
+def register_prod_db_query_tools(app: FastMCP) -> None:
+    """Register prod DB query tools with the FastMCP app."""
+    register_mcp_tools(app, ToolDomain.CLOUD_ADMIN)

airbyte_ops_mcp/mcp/server.py

@@ -19,6 +19,7 @@ from airbyte_ops_mcp.mcp.github import register_github_tools
 from airbyte_ops_mcp.mcp.github_repo_ops import register_github_repo_ops_tools
 from airbyte_ops_mcp.mcp.live_tests import register_live_tests_tools
 from airbyte_ops_mcp.mcp.prerelease import register_prerelease_tools
+from airbyte_ops_mcp.mcp.prod_db_queries import register_prod_db_query_tools
 from airbyte_ops_mcp.mcp.prompts import register_prompts
 from airbyte_ops_mcp.mcp.server_info import register_server_info_resources
 
@@ -47,6 +48,7 @@ def register_server_assets(app: FastMCP) -> None:
     register_github_tools(app)
     register_prerelease_tools(app)
     register_cloud_connector_version_tools(app)
+    register_prod_db_query_tools(app)
     register_prompts(app)
     register_live_tests_tools(app)
 

airbyte_ops_mcp/mcp/server_info.py

@@ -19,7 +19,7 @@ from airbyte_ops_mcp.mcp._mcp_utils import (
 
 
 @lru_cache(maxsize=1)
-def _get_version_info() -> dict[str, str | None]:
+def _get_version_info() -> dict[str, str | list[str] | None]:
     """Get version information for the MCP server.
 
     Returns:
@@ -66,7 +66,7 @@ def _get_version_info() -> dict[str, str | None]:
     mime_type="application/json",
     domain=ToolDomain.SERVER_INFO,
 )
-def mcp_server_info() -> dict[str, str | None]:
+def mcp_server_info() -> dict[str, str | list[str] | None]:
     """Resource that returns information for the MCP server.
 
     This includes package version, release history, help URLs, as well as other information.

airbyte_ops_mcp/prod_db_access/__init__.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Prod DB Access module for querying Airbyte Cloud Prod DB Replica.
+
+This module provides:
+- sql.py: SQL query templates and schema documentation
+- db_engine.py: Database connection and engine management
+- queries.py: Query execution functions
+"""
+
+from airbyte_ops_mcp.prod_db_access.db_engine import get_pool
+from airbyte_ops_mcp.prod_db_access.sql import (
+    SELECT_ACTORS_PINNED_TO_VERSION,
+    SELECT_CONNECTIONS_BY_CONNECTOR,
+    SELECT_CONNECTOR_VERSIONS,
+    SELECT_DATAPLANES_LIST,
+    SELECT_NEW_CONNECTOR_RELEASES,
+    SELECT_ORG_WORKSPACES,
+    SELECT_SUCCESSFUL_SYNCS_FOR_VERSION,
+    SELECT_SYNC_RESULTS_FOR_VERSION,
+    SELECT_WORKSPACE_INFO,
+)
+
+__all__ = [
+    "SELECT_ACTORS_PINNED_TO_VERSION",
+    "SELECT_CONNECTIONS_BY_CONNECTOR",
+    "SELECT_CONNECTOR_VERSIONS",
+    "SELECT_DATAPLANES_LIST",
+    "SELECT_NEW_CONNECTOR_RELEASES",
+    "SELECT_ORG_WORKSPACES",
+    "SELECT_SUCCESSFUL_SYNCS_FOR_VERSION",
+    "SELECT_SYNC_RESULTS_FOR_VERSION",
+    "SELECT_WORKSPACE_INFO",
+    "get_pool",
+]

airbyte_ops_mcp/prod_db_access/db_engine.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Database engine and connection management for Airbyte Cloud Prod DB Replica.
+
+This module provides connection pooling and engine management for querying
+the Airbyte Cloud production database replica.
+
+For SQL query templates and schema documentation, see sql.py.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import traceback
+from typing import Any, Callable
+
+import sqlalchemy
+from google.cloud import secretmanager
+from google.cloud.sql.connector import Connector
+from google.cloud.sql.connector.enums import IPTypes
+
+# Secret ID for database connection details
+CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID = (
+    "projects/587336813068/secrets/CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS"
+)
+
+PG_DRIVER = "pg8000"
+
+# Lazy-initialized to avoid import-time GCP auth
+_connector: Connector | None = None
+
+
+def _get_connector() -> Connector:
+    """Get the Cloud SQL connector, initializing lazily on first use."""
+    global _connector
+    if _connector is None:
+        _connector = Connector()
+    return _connector
+
+
+def _get_secret_value(
+    gsm_client: secretmanager.SecretManagerServiceClient, secret_id: str
+) -> str:
+    """Get the value of the enabled version of a secret.
+
+    Args:
+        gsm_client: GCP Secret Manager client
+        secret_id: The id of the secret
+
+    Returns:
+        The value of the enabled version of the secret
+    """
+    response = gsm_client.list_secret_versions(
+        request={"parent": secret_id, "filter": "state:ENABLED"}
+    )
+    if len(response.versions) == 0:
+        raise ValueError(f"No enabled version of secret {secret_id} found")
+    enabled_version = response.versions[0]
+    response = gsm_client.access_secret_version(name=enabled_version.name)
+    return response.payload.data.decode("UTF-8")
+
+
+def get_database_creator(pg_connection_details: dict) -> Callable:
+    """Create a database connection creator function."""
+
+    def creator() -> Any:
+        return _get_connector().connect(
+            pg_connection_details["database_address"],
+            PG_DRIVER,
+            user=pg_connection_details["pg_user"],
+            password=pg_connection_details["pg_password"],
+            db=pg_connection_details["database_name"],
+            ip_type=IPTypes.PRIVATE,
+        )
+
+    return creator
+
+
+def get_pool(
+    gsm_client: secretmanager.SecretManagerServiceClient,
+) -> sqlalchemy.Engine:
+    """Get a SQLAlchemy connection pool for the Airbyte Cloud database.
+
+    This function supports two connection modes:
+    1. Direct connection via Cloud SQL Python Connector (default, requires VPC access)
+    2. Connection via Cloud SQL Auth Proxy (when CI or USE_CLOUD_SQL_PROXY env var is set)
+
+    For proxy mode, start the proxy with:
+        cloud-sql-proxy prod-ab-cloud-proj:us-west3:prod-pgsql-replica --port=<port>
+
+    Environment variables:
+        CI: If set, uses proxy connection mode
+        USE_CLOUD_SQL_PROXY: If set, uses proxy connection mode
+        DB_PORT: Port for proxy connection (default: 5432)
+
+    Args:
+        gsm_client: GCP Secret Manager client for retrieving credentials
+
+    Returns:
+        SQLAlchemy Engine connected to the Prod DB Replica
+    """
+    pg_connection_details = json.loads(
+        _get_secret_value(
+            gsm_client, CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID
+        )
+    )
+
+    if os.getenv("CI") or os.getenv("USE_CLOUD_SQL_PROXY"):
+        # Connect via Cloud SQL Auth Proxy, running on localhost
+        # Port can be configured via DB_PORT env var (default: 5432)
+        host = "127.0.0.1"
+        port = os.getenv("DB_PORT", "5432")
+        try:
+            return sqlalchemy.create_engine(
+                f"postgresql+{PG_DRIVER}://{pg_connection_details['pg_user']}:{pg_connection_details['pg_password']}@{host}:{port}/{pg_connection_details['database_name']}",
+            )
+        except Exception as e:
+            raise AssertionError(
+                f"sqlalchemy.create_engine exception; could not connect to the proxy at {host}:{port}. "
+                f"Error: {traceback.format_exception(e)}"
+            ) from e
+
+    # Default: Connect via Cloud SQL Python Connector (requires VPC access)
+    return sqlalchemy.create_engine(
+        f"postgresql+{PG_DRIVER}://",
+        creator=get_database_creator(pg_connection_details),
+    )