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

airbyte_ops_mcp/cloud_admin/connection_config.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Fetch connection configuration from Airbyte Cloud.
+
+This module provides utilities for fetching connection configuration
+from Airbyte Cloud, optionally including unmasked secrets.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from airbyte_ops_mcp.live_tests.connection_fetcher import fetch_connection_data
+from airbyte_ops_mcp.live_tests.connection_secret_retriever import (
+    retrieve_unmasked_config,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FetchConnectionConfigResult(BaseModel):
+    """Result of fetching connection configuration."""
+
+    success: bool = Field(description="Whether the operation succeeded")
+    message: str = Field(description="Human-readable status message")
+    connection_id: str = Field(description="The connection ID that was fetched")
+    source_id: str | None = Field(
+        default=None, description="The source ID for this connection"
+    )
+    source_name: str | None = Field(default=None, description="The name of the source")
+    output_path: str | None = Field(
+        default=None, description="Path where the config was written"
+    )
+    with_secrets: bool = Field(
+        default=False, description="Whether secrets were included"
+    )
+
+
+def fetch_connection_config(
+    connection_id: str,
+    output_path: Path | None = None,
+    with_secrets: bool = False,
+    oc_issue_url: str | None = None,
+) -> FetchConnectionConfigResult:
+    """Fetch connection configuration from Airbyte Cloud.
+
+    This function retrieves the source configuration for a given connection ID
+    and writes it to a local JSON file. When with_secrets is True, it uses the
+    internal connection-retriever to fetch unmasked secrets from the database.
+
+    Args:
+        connection_id: The UUID of the Airbyte Cloud connection.
+        output_path: Path to output file or directory. If directory, writes
+            connection-<id>-config.json. Default: ./connection-<id>-config.json
+        with_secrets: If True, fetch unmasked secrets from the internal database.
+            Requires appropriate GCP credentials and Cloud SQL Proxy access.
+        oc_issue_url: Required when with_secrets is True. The OC issue URL for
+            audit logging purposes.
+
+    Returns:
+        FetchConnectionConfigResult with operation status and details.
+
+    Raises:
+        ValueError: If with_secrets is True but oc_issue_url is not provided.
+    """
+    if with_secrets and not oc_issue_url:
+        return FetchConnectionConfigResult(
+            success=False,
+            message="--oc-issue-url is required when using --with-secrets for audit logging",
+            connection_id=connection_id,
+            with_secrets=with_secrets,
+        )
+
+    # Resolve output path
+    if output_path is None:
+        resolved_path = Path(f"./connection-{connection_id}-config.json")
+    elif output_path.is_dir():
+        resolved_path = output_path / f"connection-{connection_id}-config.json"
+    else:
+        resolved_path = output_path
+
+    # Fetch connection data via public Cloud API
+    connection_data = fetch_connection_data(connection_id)
+
+    # Get the config - either masked or unmasked
+    if with_secrets:
+        # Use the internal connection-retriever to get unmasked secrets
+        retrieval_reason = f"CLI fetch-connection-config: {oc_issue_url}"
+        unmasked_config = retrieve_unmasked_config(
+            connection_id=connection_id,
+            retrieval_reason=retrieval_reason,
+        )
+
+        if unmasked_config is None:
+            return FetchConnectionConfigResult(
+                success=False,
+                message=(
+                    f"Failed to retrieve unmasked secrets for connection {connection_id}. "
+                    "Ensure you have GCP credentials and Cloud SQL Proxy access."
+                ),
+                connection_id=connection_id,
+                source_id=connection_data.source_id,
+                source_name=connection_data.source_name,
+                with_secrets=True,
+            )
+
+        config = unmasked_config
+        logger.info(
+            f"Retrieved unmasked config for connection {connection_id} "
+            f"(reason: {retrieval_reason})"
+        )
+    else:
+        # Use the masked config from the public API
+        config = connection_data.config
+
+    # Write config to file
+    resolved_path.parent.mkdir(parents=True, exist_ok=True)
+    resolved_path.write_text(json.dumps(config, indent=2))
+
+    return FetchConnectionConfigResult(
+        success=True,
+        message=f"Successfully wrote config to {resolved_path}",
+        connection_id=connection_id,
+        source_id=connection_data.source_id,
+        source_name=connection_data.source_name,
+        output_path=str(resolved_path),
+        with_secrets=with_secrets,
+    )

airbyte_ops_mcp/live_tests/__init__.py

@@ -5,6 +5,16 @@ This module provides tools for testing Airbyte connectors against live data
 without using Dagger. It uses Docker SDK directly for container orchestration.
 """
 
+from airbyte_ops_mcp.live_tests.connection_fetcher import (
+    ConnectionData,
+    fetch_connection_data,
+)
+from airbyte_ops_mcp.live_tests.connection_secret_retriever import (
+    enrich_config_with_secrets,
+    is_secret_retriever_enabled,
+    retrieve_unmasked_config,
+    should_use_secret_retriever,
+)
 from airbyte_ops_mcp.live_tests.models import (
     Command,
     ConnectorUnderTest,
@@ -14,7 +24,13 @@ from airbyte_ops_mcp.live_tests.models import (
 
 __all__ = [
     "Command",
+    "ConnectionData",
     "ConnectorUnderTest",
     "ExecutionResult",
     "TargetOrControl",
+    "enrich_config_with_secrets",
+    "fetch_connection_data",
+    "is_secret_retriever_enabled",
+    "retrieve_unmasked_config",
+    "should_use_secret_retriever",
 ]

airbyte_ops_mcp/live_tests/_connection_retriever/__init__.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Vendored subset of connection-retriever from airbyte-platform-internal.
+
+This module contains a minimal subset of the connection-retriever tool,
+vendored to avoid depending on an unpublished internal package. It provides
+functionality to retrieve unmasked source configuration from Airbyte Cloud's
+internal database.
+
+Original source: airbyte-platform-internal/tools/connection-retriever
+Vendored: 2025-01-XX
+
+Only the following functionality is included:
+- Retrieve unmasked source config for a given connection ID
+- Secret resolution from GCP Secret Manager
+- Audit logging to GCP Cloud Logging
+
+NOT included (see issue #91 for future work):
+- retrieve_testing_candidates() - BigQuery-based candidate discovery
+- Destination config retrieval
+- CLI interface
+"""
+
+from airbyte_ops_mcp.live_tests._connection_retriever.consts import (
+    ConnectionObject,
+)
+from airbyte_ops_mcp.live_tests._connection_retriever.retrieval import (
+    TestingCandidate,
+    retrieve_objects,
+)
+
+__all__ = [
+    "ConnectionObject",
+    "TestingCandidate",
+    "retrieve_objects",
+]

airbyte_ops_mcp/live_tests/_connection_retriever/audit_logging.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Audit logging for vendored connection-retriever.
+
+Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/audit_logging.py
+"""
+
+from __future__ import annotations
+
+import logging
+import subprocess
+from typing import TYPE_CHECKING, Any, Callable
+
+from google.cloud import logging as gcloud_logging
+
+from airbyte_ops_mcp.live_tests._connection_retriever.consts import (
+    GCP_PROJECT_NAME,
+)
+
+if TYPE_CHECKING:
+    from airbyte_ops_mcp.live_tests._connection_retriever.retrieval import (
+        RetrievalMetadata,
+    )
+
+LOGGER = logging.getLogger(__name__)
+
+# Lazy-initialized to avoid import-time GCP calls
+_logging_client: gcloud_logging.Client | None = None
+_airbyte_gcloud_logger: Any = None
+
+
+def _get_logger() -> Any:
+    """Get the GCP Cloud Logger, initializing lazily on first use."""
+    global _logging_client, _airbyte_gcloud_logger
+
+    if _airbyte_gcloud_logger is not None:
+        return _airbyte_gcloud_logger
+
+    _logging_client = gcloud_logging.Client(project=GCP_PROJECT_NAME)
+    _airbyte_gcloud_logger = _logging_client.logger(
+        "airbyte-cloud-connection-retriever"
+    )
+    return _airbyte_gcloud_logger
+
+
+def get_user_email() -> str:
+    """Get the email of the currently authenticated GCP user."""
+    # This is a bit hacky - should use service account impersonation
+    # https://cloud.google.com/iam/docs/impersonating-service-accounts
+    command = [
+        "gcloud",
+        "auth",
+        "list",
+        "--filter=status:ACTIVE",
+        "--format=value(account)",
+    ]
+    output = subprocess.run(command, capture_output=True, text=True, check=True)
+    return output.stdout.strip()
+
+
+def get_audit_log_message(
+    retrieval_metadata: RetrievalMetadata,
+) -> dict:
+    """Build an audit log message for a retrieval operation."""
+    user_email = get_user_email()
+    return {
+        "message": (
+            f"{user_email} is accessing {retrieval_metadata.connection_object.value} "
+            f"for connection_id {retrieval_metadata.connection_id} "
+            f"for {retrieval_metadata.retrieval_reason}"
+        ),
+        "retrieval_reason": retrieval_metadata.retrieval_reason,
+        "connection_object_type": retrieval_metadata.connection_object.value,
+        "user": user_email,
+        "connection_id": retrieval_metadata.connection_id,
+    }
+
+
+def audit(function_to_audit: Callable) -> Callable:
+    """Decorator to audit function calls to GCP Cloud Logging."""
+
+    def wrapper(
+        retrieval_metadata: RetrievalMetadata, *args: Any, **kwargs: Any
+    ) -> Callable:
+        audit_log_message = get_audit_log_message(retrieval_metadata)
+        _get_logger().log_struct(audit_log_message)
+        return function_to_audit(*args, **kwargs)
+
+    return wrapper

airbyte_ops_mcp/live_tests/_connection_retriever/consts.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Constants for vendored connection-retriever.
+
+Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/consts.py
+"""
+
+from enum import Enum
+
+GCP_PROJECT_NAME = "prod-ab-cloud-proj"
+
+CLOUD_REGISTRY_URL = (
+    "https://connectors.airbyte.com/files/registries/v0/cloud_registry.json"
+)
+
+CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID = (
+    "projects/587336813068/secrets/CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS"
+)
+
+
+class ConnectionObject(Enum):
+    """Types of connection objects that can be retrieved."""
+
+    CONNECTION = "connection"
+    SOURCE_ID = "source-id"
+    DESTINATION_ID = "destination-id"
+    DESTINATION_CONFIG = "destination-config"
+    SOURCE_CONFIG = "source-config"
+    CATALOG = "catalog"
+    CONFIGURED_CATALOG = "configured-catalog"
+    STATE = "state"
+    WORKSPACE_ID = "workspace-id"
+    DESTINATION_DOCKER_IMAGE = "destination-docker-image"
+    SOURCE_DOCKER_IMAGE = "source-docker-image"

airbyte_ops_mcp/live_tests/_connection_retriever/db_access.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Database access for vendored connection-retriever.
+
+Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/db_access.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.live_tests._connection_retriever.consts import (
+    CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID,
+)
+from airbyte_ops_mcp.live_tests._connection_retriever.secrets_resolution import (
+    get_secret_value,
+)
+
+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_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(
+    secret_manager_client: secretmanager.SecretManagerServiceClient,
+) -> sqlalchemy.Engine:
+    """Get a SQLAlchemy connection pool for the Airbyte Cloud database."""
+    pg_connection_details = json.loads(
+        get_secret_value(
+            secret_manager_client, CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID
+        )
+    )
+
+    if os.getenv("CI"):
+        # In CI we connect via Cloud SQL Auth Proxy, running on localhost
+        host = "127.0.0.1"
+        try:
+            return sqlalchemy.create_engine(
+                f"postgresql+{PG_DRIVER}://{pg_connection_details['pg_user']}:{pg_connection_details['pg_password']}@127.0.0.1/{pg_connection_details['database_name']}",
+            )
+        except Exception as e:
+            raise AssertionError(
+                f"sqlalchemy.create_engine exception; could not connect to the proxy at {host}. "
+                f"Error: {traceback.format_exception(e)}"
+            ) from e
+    else:
+        return sqlalchemy.create_engine(
+            f"postgresql+{PG_DRIVER}://",
+            creator=get_database_creator(pg_connection_details),
+        )