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

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,
+    )

airbyte_ops_mcp/prod_db_access/sql.py

@@ -0,0 +1,353 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""SQL query templates and schema documentation for Airbyte Cloud Prod DB Replica.
+
+Prod DB Replica Schema Reference
+================================
+
+Database: prod-configapi
+Instance: prod-ab-cloud-proj:us-west3:prod-pgsql-replica
+
+connection
+----------
+id, namespace_definition, namespace_format, prefix, source_id, destination_id, name,
+catalog, status, schedule, manual, resource_requirements, created_at, updated_at,
+source_catalog_id, schedule_type, schedule_data, non_breaking_change_preference,
+breaking_change, field_selection_data, destination_catalog_id, status_reason
+
+actor
+-----
+id, workspace_id, actor_definition_id, name, configuration, actor_type, tombstone,
+created_at, updated_at, resource_requirements
+
+workspace
+---------
+id, customer_id, name, slug, email, initial_setup_complete, anonymous_data_collection,
+send_newsletter, send_security_updates, display_setup_wizard, tombstone, notifications,
+first_sync_complete, feedback_complete, created_at, updated_at, webhook_operation_configs,
+notification_settings, organization_id, dataplane_group_id
+
+dataplane_group
+---------------
+id, organization_id, name, enabled, created_at, updated_at, tombstone
+
+Note: Main dataplane groups are:
+- 645a183f-b12b-4c6e-8ad3-99e165603450 = US (default, ~133K workspaces)
+- 153996d3-208e-4887-b8b1-e5fe48104450 = US-Central (~12K workspaces)
+- b9e48d61-f082-4a14-a8d0-799a907938cb = EU (~3K workspaces)
+
+actor_definition_version
+------------------------
+id, actor_definition_id, created_at, updated_at, documentation_url, docker_repository,
+docker_image_tag, spec, protocol_version, release_date, normalization_repository,
+normalization_tag, supports_dbt, normalization_integration_type, allowed_hosts,
+suggested_streams, release_stage, support_state, support_level, supports_refreshes,
+cdk_version, last_published, internal_support_level, language, supports_file_transfer,
+supports_data_activation, connector_ipc_options
+
+scoped_configuration
+--------------------
+id, key, resource_type, resource_id, scope_type, scope_id, value, description,
+reference_url, origin_type, origin, expires_at, created_at, updated_at
+
+Note: Version overrides are stored with key='connector_version', resource_type='actor_definition',
+scope_type='actor', and value=actor_definition_version.id (UUID).
+
+jobs
+----
+id, config_type, scope (connection_id), config, status, started_at, created_at,
+updated_at, metadata, is_scheduled
+
+Note: status values: 'succeeded', 'failed', 'cancelled', 'running', 'incomplete'
+      config_type values: 'sync', 'reset_connection', 'refresh'
+
+attempts
+--------
+id, job_id, attempt_number, log_path, output, status, created_at, updated_at,
+ended_at, failure_summary, processing_task_queue, attempt_sync_config
+"""
+
+from __future__ import annotations
+
+import sqlalchemy
+
+# =============================================================================
+# Connection Queries
+# =============================================================================
+
+# Query connections by connector type, optionally filtered by organization
+SELECT_CONNECTIONS_BY_CONNECTOR = sqlalchemy.text(
+    """
+    SELECT
+         connection.id AS connection_id,
+         connection.name AS connection_name,
+         connection.source_id,
+         workspace.id AS workspace_id,
+         workspace.name AS workspace_name,
+         workspace.organization_id,
+         workspace.dataplane_group_id,
+         dataplane_group.name AS dataplane_name,
+         source_actor.actor_definition_id AS source_definition_id,
+         source_actor.name AS source_name
+    FROM connection
+    JOIN actor AS source_actor
+      ON connection.source_id = source_actor.id
+    JOIN workspace
+      ON source_actor.workspace_id = workspace.id
+    LEFT JOIN dataplane_group
+      ON workspace.dataplane_group_id = dataplane_group.id
+    WHERE
+         source_actor.actor_definition_id = :connector_definition_id
+     AND (:organization_id IS NULL OR workspace.organization_id = :organization_id)
+    LIMIT :limit
+    """
+)
+
+# =============================================================================
+# Connector Version Queries
+# =============================================================================
+
+SELECT_CONNECTOR_VERSIONS = sqlalchemy.text(
+    """
+    SELECT
+         actor_definition_version.id AS version_id,
+         actor_definition_version.docker_image_tag,
+         actor_definition_version.docker_repository,
+         actor_definition_version.release_stage,
+         actor_definition_version.support_level,
+         actor_definition_version.cdk_version,
+         actor_definition_version.language,
+         actor_definition_version.last_published,
+         actor_definition_version.release_date
+    FROM actor_definition_version
+    WHERE
+         actor_definition_version.actor_definition_id = :actor_definition_id
+    ORDER BY
+         actor_definition_version.last_published DESC NULLS LAST,
+         actor_definition_version.created_at DESC
+    """
+)
+
+# List new connector releases within the last N days
+# Uses last_published (timestamp) rather than release_date (date only, often NULL)
+# Note: No index on last_published, but table is small (~39K rows)
+SELECT_NEW_CONNECTOR_RELEASES = sqlalchemy.text(
+    """
+    SELECT
+         actor_definition_version.id AS version_id,
+         actor_definition_version.actor_definition_id,
+         actor_definition_version.docker_repository,
+         actor_definition_version.docker_image_tag,
+         actor_definition_version.last_published,
+         actor_definition_version.release_date,
+         actor_definition_version.release_stage,
+         actor_definition_version.support_level,
+         actor_definition_version.cdk_version,
+         actor_definition_version.language,
+         actor_definition_version.created_at
+    FROM actor_definition_version
+    WHERE
+         actor_definition_version.last_published >= :cutoff_date
+    ORDER BY
+         actor_definition_version.last_published DESC
+    LIMIT :limit
+    """
+)
+
+SELECT_ACTORS_PINNED_TO_VERSION = sqlalchemy.text(
+    """
+    SELECT
+         scoped_configuration.scope_id AS actor_id,
+         scoped_configuration.resource_id AS actor_definition_id,
+         scoped_configuration.origin_type,
+         scoped_configuration.origin,
+         scoped_configuration.description,
+         scoped_configuration.created_at,
+         scoped_configuration.expires_at,
+         actor.name AS actor_name,
+         actor.workspace_id,
+         workspace.name AS workspace_name,
+         workspace.organization_id,
+         workspace.dataplane_group_id,
+         dataplane_group.name AS dataplane_name
+    FROM scoped_configuration
+    JOIN actor
+      ON scoped_configuration.scope_id = actor.id
+    JOIN workspace
+      ON actor.workspace_id = workspace.id
+    LEFT JOIN dataplane_group
+      ON workspace.dataplane_group_id = dataplane_group.id
+    WHERE
+         scoped_configuration.key = 'connector_version'
+     AND scoped_configuration.scope_type = 'actor'
+     AND scoped_configuration.value = :actor_definition_version_id
+    ORDER BY
+         scoped_configuration.created_at DESC
+    """
+)
+
+# =============================================================================
+# Sync Results Queries
+# =============================================================================
+
+# Get sync results for actors pinned to a specific connector definition VERSION ID
+# This joins through scoped_configuration to find actors with version overrides
+SELECT_SYNC_RESULTS_FOR_VERSION = sqlalchemy.text(
+    """
+    SELECT
+         jobs.id AS job_id,
+         jobs.scope AS connection_id,
+         jobs.status AS job_status,
+         jobs.started_at,
+         jobs.updated_at AS job_updated_at,
+         connection.name AS connection_name,
+         actor.id AS actor_id,
+         actor.name AS actor_name,
+         actor.actor_definition_id,
+         scoped_configuration.origin_type AS pin_origin_type,
+         scoped_configuration.origin AS pin_origin,
+         workspace.id AS workspace_id,
+         workspace.name AS workspace_name,
+         workspace.organization_id,
+         workspace.dataplane_group_id,
+         dataplane_group.name AS dataplane_name
+    FROM jobs
+    JOIN connection
+      ON jobs.scope = connection.id::text
+    JOIN actor
+      ON connection.source_id = actor.id
+    JOIN scoped_configuration
+      ON scoped_configuration.scope_id = actor.id
+     AND scoped_configuration.key = 'connector_version'
+     AND scoped_configuration.scope_type = 'actor'
+    JOIN workspace
+      ON actor.workspace_id = workspace.id
+    LEFT JOIN dataplane_group
+      ON workspace.dataplane_group_id = dataplane_group.id
+    WHERE
+         jobs.config_type = 'sync'
+     AND scoped_configuration.value = :actor_definition_version_id
+     AND jobs.updated_at >= :cutoff_date
+    ORDER BY
+         jobs.updated_at DESC
+    LIMIT :limit
+    """
+)
+
+# Get successful sync results for actors pinned to a specific connector definition VERSION ID
+SELECT_SUCCESSFUL_SYNCS_FOR_VERSION = sqlalchemy.text(
+    """
+    SELECT
+         jobs.id AS job_id,
+         jobs.scope AS connection_id,
+         jobs.started_at,
+         jobs.updated_at AS job_updated_at,
+         connection.name AS connection_name,
+         actor.id AS actor_id,
+         actor.name AS actor_name,
+         actor.actor_definition_id,
+         scoped_configuration.origin_type AS pin_origin_type,
+         scoped_configuration.origin AS pin_origin,
+         workspace.id AS workspace_id,
+         workspace.name AS workspace_name,
+         workspace.organization_id,
+         workspace.dataplane_group_id,
+         dataplane_group.name AS dataplane_name
+    FROM jobs
+    JOIN connection
+      ON jobs.scope = connection.id::text
+    JOIN actor
+      ON connection.source_id = actor.id
+    JOIN scoped_configuration
+      ON scoped_configuration.scope_id = actor.id
+     AND scoped_configuration.key = 'connector_version'
+     AND scoped_configuration.scope_type = 'actor'
+    JOIN workspace
+      ON actor.workspace_id = workspace.id
+    LEFT JOIN dataplane_group
+      ON workspace.dataplane_group_id = dataplane_group.id
+    WHERE
+         jobs.config_type = 'sync'
+     AND jobs.status = 'succeeded'
+     AND scoped_configuration.value = :actor_definition_version_id
+     AND jobs.updated_at >= :cutoff_date
+    ORDER BY
+         jobs.updated_at DESC
+    LIMIT :limit
+    """
+)
+
+# =============================================================================
+# Dataplane and Workspace Queries
+# =============================================================================
+
+# List all dataplane groups with workspace counts
+SELECT_DATAPLANES_LIST = sqlalchemy.text(
+    """
+    SELECT
+         dataplane_group.id AS dataplane_group_id,
+         dataplane_group.name AS dataplane_name,
+         dataplane_group.organization_id,
+         dataplane_group.enabled,
+         dataplane_group.tombstone,
+         dataplane_group.created_at,
+         COUNT(workspace.id) AS workspace_count
+    FROM dataplane_group
+    LEFT JOIN workspace
+      ON workspace.dataplane_group_id = dataplane_group.id
+     AND workspace.tombstone = false
+    WHERE
+         dataplane_group.tombstone = false
+    GROUP BY
+         dataplane_group.id,
+         dataplane_group.name,
+         dataplane_group.organization_id,
+         dataplane_group.enabled,
+         dataplane_group.tombstone,
+         dataplane_group.created_at
+    ORDER BY
+         workspace_count DESC
+    """
+)
+
+# Get workspace info including dataplane group for EU filtering
+SELECT_WORKSPACE_INFO = sqlalchemy.text(
+    """
+    SELECT
+         workspace.id AS workspace_id,
+         workspace.name AS workspace_name,
+         workspace.slug,
+         workspace.organization_id,
+         workspace.dataplane_group_id,
+         dataplane_group.name AS dataplane_name,
+         workspace.created_at,
+         workspace.tombstone
+    FROM workspace
+    LEFT JOIN dataplane_group
+      ON workspace.dataplane_group_id = dataplane_group.id
+    WHERE
+         workspace.id = :workspace_id
+    """
+)
+
+# Get all workspaces in an organization with dataplane info
+SELECT_ORG_WORKSPACES = sqlalchemy.text(
+    """
+    SELECT
+         workspace.id AS workspace_id,
+         workspace.name AS workspace_name,
+         workspace.slug,
+         workspace.organization_id,
+         workspace.dataplane_group_id,
+         dataplane_group.name AS dataplane_name,
+         workspace.created_at,
+         workspace.tombstone
+    FROM workspace
+    LEFT JOIN dataplane_group
+      ON workspace.dataplane_group_id = dataplane_group.id
+    WHERE
+         workspace.organization_id = :organization_id
+     AND workspace.tombstone = false
+    ORDER BY
+         workspace.name
+    """
+)

airbyte_ops_mcp/registry/__init__.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Registry operations for Airbyte connectors.
+
+This package provides functionality for publishing connectors to the Airbyte
+registry, including promoting and rolling back release candidates.
+"""
+
+from __future__ import annotations
+
+from airbyte_ops_mcp.registry.models import (
+    ConnectorMetadata,
+    ConnectorPublishResult,
+    PublishAction,
+)
+from airbyte_ops_mcp.registry.publish import (
+    CONNECTOR_PATH_PREFIX,
+    METADATA_FILE_NAME,
+    get_connector_metadata,
+    is_release_candidate,
+    publish_connector,
+    strip_rc_suffix,
+)
+
+__all__ = [
+    "CONNECTOR_PATH_PREFIX",
+    "METADATA_FILE_NAME",
+    "ConnectorMetadata",
+    "ConnectorPublishResult",
+    "PublishAction",
+    "get_connector_metadata",
+    "is_release_candidate",
+    "publish_connector",
+    "strip_rc_suffix",
+]