airbyte-internal-ops 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

airbyte_ops_mcp/prod_db_access/db_engine.py

@@ -0,0 +1,126 @@
+# 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
+
+from airbyte_ops_mcp.constants import (
+    CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID,
+)
+
+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),
+    )

airbyte_ops_mcp/prod_db_access/queries.py

@@ -0,0 +1,272 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Query execution functions for Airbyte Cloud Prod DB Replica.
+
+This module provides functions that execute SQL queries against the Prod DB Replica
+and return structured results. Each function wraps a SQL template from sql.py.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from datetime import datetime, timedelta, timezone
+from time import perf_counter
+from typing import Any
+
+import sqlalchemy
+from google.cloud import secretmanager
+
+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,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _run_sql_query(
+    statement: sqlalchemy.sql.elements.TextClause,
+    parameters: Mapping[str, Any] | None = None,
+    *,
+    query_name: str | None = None,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Execute a SQL text statement and return rows as list[dict], logging elapsed time.
+
+    Args:
+        statement: SQLAlchemy text clause to execute
+        parameters: Query parameters to bind
+        query_name: Optional name for logging (defaults to first line of SQL)
+        gsm_client: GCP Secret Manager client for retrieving credentials.
+            If None, a new client will be instantiated.
+
+    Returns:
+        List of row dicts from the query result
+    """
+    if gsm_client is None:
+        gsm_client = secretmanager.SecretManagerServiceClient()
+    pool = get_pool(gsm_client)
+    start = perf_counter()
+    with pool.connect() as conn:
+        result = conn.execute(statement, parameters or {})
+        rows = [dict(row._mapping) for row in result]
+    elapsed = perf_counter() - start
+
+    name = query_name or "SQL query"
+    logger.info("Prod DB query %s returned %d rows in %.3f s", name, len(rows), elapsed)
+
+    return rows
+
+
+def query_connections_by_connector(
+    connector_definition_id: str,
+    organization_id: str | None = None,
+    limit: int = 1000,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query connections by source connector type, optionally filtered by organization.
+
+    Args:
+        connector_definition_id: Connector definition UUID to filter by
+        organization_id: Optional organization UUID to search within
+        limit: Maximum number of results (default: 1000)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of connection records with workspace and dataplane info
+    """
+    return _run_sql_query(
+        SELECT_CONNECTIONS_BY_CONNECTOR,
+        parameters={
+            "connector_definition_id": connector_definition_id,
+            "organization_id": organization_id,
+            "limit": limit,
+        },
+        query_name="SELECT_CONNECTIONS_BY_CONNECTOR",
+        gsm_client=gsm_client,
+    )
+
+
+def query_connector_versions(
+    connector_definition_id: str,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query all versions for a connector definition.
+
+    Args:
+        connector_definition_id: Connector definition UUID
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of version records ordered by last_published DESC
+    """
+    return _run_sql_query(
+        SELECT_CONNECTOR_VERSIONS,
+        parameters={"actor_definition_id": connector_definition_id},
+        query_name="SELECT_CONNECTOR_VERSIONS",
+        gsm_client=gsm_client,
+    )
+
+
+def query_new_connector_releases(
+    days: int = 7,
+    limit: int = 100,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query recently published connector versions.
+
+    Args:
+        days: Number of days to look back (default: 7)
+        limit: Maximum number of results (default: 100)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of recently published connector versions
+    """
+    cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
+    return _run_sql_query(
+        SELECT_NEW_CONNECTOR_RELEASES,
+        parameters={"cutoff_date": cutoff_date, "limit": limit},
+        query_name="SELECT_NEW_CONNECTOR_RELEASES",
+        gsm_client=gsm_client,
+    )
+
+
+def query_actors_pinned_to_version(
+    connector_version_id: str,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query actors (sources/destinations) pinned to a specific connector version.
+
+    Args:
+        connector_version_id: Connector version UUID to search for
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of actors pinned to the specified version
+    """
+    return _run_sql_query(
+        SELECT_ACTORS_PINNED_TO_VERSION,
+        parameters={"actor_definition_version_id": connector_version_id},
+        query_name="SELECT_ACTORS_PINNED_TO_VERSION",
+        gsm_client=gsm_client,
+    )
+
+
+def query_sync_results_for_version(
+    connector_version_id: str,
+    days: int = 7,
+    limit: int = 100,
+    successful_only: bool = False,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query sync job results for actors pinned to a specific connector version.
+
+    Args:
+        connector_version_id: Connector version UUID to filter by
+        days: Number of days to look back (default: 7)
+        limit: Maximum number of results (default: 100)
+        successful_only: If True, only return successful syncs (default: False)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of sync job results
+    """
+    cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
+    query = (
+        SELECT_SUCCESSFUL_SYNCS_FOR_VERSION
+        if successful_only
+        else SELECT_SYNC_RESULTS_FOR_VERSION
+    )
+    query_name = (
+        "SELECT_SUCCESSFUL_SYNCS_FOR_VERSION"
+        if successful_only
+        else "SELECT_SYNC_RESULTS_FOR_VERSION"
+    )
+    return _run_sql_query(
+        query,
+        parameters={
+            "actor_definition_version_id": connector_version_id,
+            "cutoff_date": cutoff_date,
+            "limit": limit,
+        },
+        query_name=query_name,
+        gsm_client=gsm_client,
+    )
+
+
+def query_dataplanes_list(
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query all dataplane groups with workspace counts.
+
+    Args:
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of dataplane groups ordered by workspace count DESC
+    """
+    return _run_sql_query(
+        SELECT_DATAPLANES_LIST,
+        query_name="SELECT_DATAPLANES_LIST",
+        gsm_client=gsm_client,
+    )
+
+
+def query_workspace_info(
+    workspace_id: str,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> dict[str, Any] | None:
+    """Query workspace info including dataplane group.
+
+    Args:
+        workspace_id: Workspace UUID
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        Workspace info dict, or None if not found
+    """
+    rows = _run_sql_query(
+        SELECT_WORKSPACE_INFO,
+        parameters={"workspace_id": workspace_id},
+        query_name="SELECT_WORKSPACE_INFO",
+        gsm_client=gsm_client,
+    )
+    return rows[0] if rows else None
+
+
+def query_org_workspaces(
+    organization_id: str,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query all workspaces in an organization with dataplane info.
+
+    Args:
+        organization_id: Organization UUID
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of workspaces in the organization
+    """
+    return _run_sql_query(
+        SELECT_ORG_WORKSPACES,
+        parameters={"organization_id": organization_id},
+        query_name="SELECT_ORG_WORKSPACES",
+        gsm_client=gsm_client,
+    )