airbyte-internal-ops 0.1.11__py3-none-any.whl → 0.2.1__py3-none-any.whl

airbyte_ops_mcp/cloud_admin/api_client.py

@@ -48,21 +48,36 @@ class _OriginType(str, Enum):
 
 
 def _get_access_token(
-    client_id: str,
-    client_secret: str,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
 ) -> str:
     """Get an access token for Airbyte Cloud API.
 
+    If a bearer_token is provided, it is returned directly (already an access token).
+    Otherwise, exchanges client_id/client_secret for an access token.
+
     Args:
-        client_id: The Airbyte Cloud client ID
-        client_secret: The Airbyte Cloud client secret
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
 
     Returns:
         Access token string
 
     Raises:
-        PyAirbyteInputError: If authentication fails
+        PyAirbyteInputError: If authentication fails or no valid credentials provided
     """
+    # If bearer token provided, use it directly
+    if bearer_token:
+        return bearer_token
+
+    # Otherwise, exchange client credentials for access token
+    if not client_id or not client_secret:
+        raise PyAirbyteInputError(
+            message="Either bearer_token or both client_id and client_secret must be provided",
+        )
+
     # Always authenticate via the public API endpoint
     auth_url = f"{constants.CLOUD_API_ROOT}/applications/token"
     response = requests.post(
@@ -90,16 +105,18 @@ def _get_access_token(
 def get_user_id_by_email(
     email: str,
     api_root: str,
-    client_id: str,
-    client_secret: str,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
 ) -> str:
     """Get user ID from email address.
 
     Args:
         email: The user's email address
         api_root: The API root URL
-        client_id: The Airbyte Cloud client ID
-        client_secret: The Airbyte Cloud client secret
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
 
     Returns:
         User ID (UUID string)
@@ -107,7 +124,7 @@ def get_user_id_by_email(
     Raises:
         PyAirbyteInputError: If user not found or API request fails
     """
-    access_token = _get_access_token(client_id, client_secret)
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
 
     endpoint = f"{api_root}/users/list_instance_admin"
     response = requests.post(
@@ -152,8 +169,9 @@ def resolve_connector_version_id(
     connector_type: Literal["source", "destination"],
     version: str,
     api_root: str,
-    client_id: str,
-    client_secret: str,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
 ) -> str:
     """Resolve a version string to an actor_definition_version_id.
 
@@ -162,8 +180,9 @@ def resolve_connector_version_id(
         connector_type: Either "source" or "destination"
         version: The version string (e.g., "0.1.47-preview.abe7cb4")
         api_root: The API root URL
-        client_id: The Airbyte Cloud client ID
-        client_secret: The Airbyte Cloud client secret
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
 
     Returns:
         Version ID (UUID string)
@@ -171,7 +190,7 @@ def resolve_connector_version_id(
     Raises:
         PyAirbyteInputError: If version cannot be resolved or API request fails
     """
-    access_token = _get_access_token(client_id, client_secret)
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
 
     endpoint = f"{api_root}/actor_definition_versions/resolve"
     payload = {
@@ -223,8 +242,9 @@ def get_connector_version(
     connector_id: str,
     connector_type: Literal["source", "destination"],
     api_root: str,
-    client_id: str,
-    client_secret: str,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    bearer_token: str | None = None,
 ) -> dict[str, Any]:
     """Get version information for a deployed connector.
 
@@ -232,8 +252,9 @@ def get_connector_version(
         connector_id: The ID of the deployed connector (source or destination)
         connector_type: Either "source" or "destination"
         api_root: The API root URL
-        client_id: The Airbyte Cloud client ID
-        client_secret: The Airbyte Cloud client secret
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
 
     Returns:
         Dictionary containing:
@@ -243,7 +264,7 @@ def get_connector_version(
     Raises:
         PyAirbyteInputError: If the API request fails
     """
-    access_token = _get_access_token(client_id, client_secret)
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
 
     # Determine endpoint based on connector type
     # api_root already includes /v1
@@ -294,14 +315,15 @@ def set_connector_version_override(
     connector_id: str,
     connector_type: Literal["source", "destination"],
     api_root: str,
-    client_id: str,
-    client_secret: str,
-    workspace_id: str,
+    client_id: str | None = None,
+    client_secret: str | None = None,
+    workspace_id: str | None = None,
     version: str | None = None,
     unset: bool = False,
     override_reason: str | None = None,
     override_reason_reference_url: str | None = None,
     user_email: str | None = None,
+    bearer_token: str | None = None,
 ) -> bool:
     """Set or clear a version override for a deployed connector.
 
@@ -309,14 +331,15 @@ def set_connector_version_override(
         connector_id: The ID of the deployed connector
         connector_type: Either "source" or "destination"
         api_root: The API root URL
-        client_id: The Airbyte Cloud client ID
-        client_secret: The Airbyte Cloud client secret
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
         workspace_id: The workspace ID
         version: The version to pin to (e.g., "0.1.0"), or None to unset
         unset: If True, removes any existing override
         override_reason: Required when setting. Explanation for the override
         override_reason_reference_url: Optional URL with more context
         user_email: Email of user creating the override
+        bearer_token: Pre-existing bearer token (takes precedence over client credentials)
 
     Returns:
         True if operation succeeded, False if no override existed (unset only)
@@ -335,7 +358,7 @@ def set_connector_version_override(
             message="override_reason is required when setting a version and must be at least 10 characters",
         )
 
-    access_token = _get_access_token(client_id, client_secret)
+    access_token = _get_access_token(client_id, client_secret, bearer_token)
 
     # Build the scoped configuration
     scope_type = _ScopeType.ACTOR
@@ -479,6 +502,7 @@ def set_connector_version_override(
             api_root=api_root,
             client_id=client_id,
             client_secret=client_secret,
+            bearer_token=bearer_token,
         )
 
         # Get user ID from email if provided
@@ -489,6 +513,7 @@ def set_connector_version_override(
                 api_root=api_root,
                 client_id=client_id,
                 client_secret=client_secret,
+                bearer_token=bearer_token,
             )
 
         # Create the override with correct schema

airbyte_ops_mcp/cloud_admin/auth.py

@@ -41,6 +41,20 @@ def check_internal_admin_flag() -> bool:
     return bool(admin_user and EXPECTED_ADMIN_EMAIL_DOMAIN in admin_user)
 
 
+def check_internal_admin_flag_only() -> bool:
+    """Check if internal admin flag is set (without requiring user email env var).
+
+    This is a lighter check that only validates AIRBYTE_INTERNAL_ADMIN_FLAG,
+    allowing the admin user email to be provided as a parameter instead of
+    an environment variable.
+
+    Returns:
+        True if AIRBYTE_INTERNAL_ADMIN_FLAG is set correctly, False otherwise
+    """
+    admin_flag = os.environ.get(ENV_AIRBYTE_INTERNAL_ADMIN_FLAG)
+    return admin_flag == EXPECTED_ADMIN_FLAG_VALUE
+
+
 def require_internal_admin() -> None:
     """Require internal admin access for the current operation.
 
@@ -59,6 +73,24 @@ def require_internal_admin() -> None:
         )
 
 
+def require_internal_admin_flag_only() -> None:
+    """Require internal admin flag for the current operation.
+
+    This is a lighter check that only validates AIRBYTE_INTERNAL_ADMIN_FLAG,
+    allowing the admin user email to be provided as a parameter instead of
+    an environment variable.
+
+    Raises:
+        CloudAuthError: If AIRBYTE_INTERNAL_ADMIN_FLAG is not properly configured
+    """
+    if not check_internal_admin_flag_only():
+        raise CloudAuthError(
+            "This operation requires internal admin access. "
+            f"Set {ENV_AIRBYTE_INTERNAL_ADMIN_FLAG}={EXPECTED_ADMIN_FLAG_VALUE} "
+            "environment variable."
+        )
+
+
 def get_admin_user_email() -> str:
     """Get the admin user email from environment.
 

airbyte_ops_mcp/cloud_admin/connection_config.py

@@ -13,8 +13,8 @@ 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 (
+from airbyte_ops_mcp.regression_tests.connection_fetcher import fetch_connection_data
+from airbyte_ops_mcp.regression_tests.connection_secret_retriever import (
     retrieve_unmasked_config,
 )
 

airbyte_ops_mcp/constants.py

@@ -20,6 +20,24 @@ ENV_GCP_PROD_DB_ACCESS_CREDENTIALS = "GCP_PROD_DB_ACCESS_CREDENTIALS"
 EXPECTED_ADMIN_FLAG_VALUE = "airbyte.io"
 EXPECTED_ADMIN_EMAIL_DOMAIN = "@airbyte.io"
 
+# =============================================================================
+# HTTP Header Names for Airbyte Cloud Authentication
+# =============================================================================
+# These headers follow the PyAirbyte convention for passing credentials
+# via HTTP when running as an MCP HTTP server.
+
+HEADER_AIRBYTE_CLOUD_CLIENT_ID = "X-Airbyte-Cloud-Client-Id"
+"""HTTP header for OAuth client ID."""
+
+HEADER_AIRBYTE_CLOUD_CLIENT_SECRET = "X-Airbyte-Cloud-Client-Secret"
+"""HTTP header for OAuth client secret."""
+
+HEADER_AIRBYTE_CLOUD_WORKSPACE_ID = "X-Airbyte-Cloud-Workspace-Id"
+"""HTTP header for default workspace ID."""
+
+HEADER_AIRBYTE_CLOUD_API_URL = "X-Airbyte-Cloud-Api-Url"
+"""HTTP header for API root URL override."""
+
 # =============================================================================
 # GCP and Prod DB Constants (from connection-retriever)
 # =============================================================================

airbyte_ops_mcp/github_actions.py

@@ -9,6 +9,8 @@ are used by MCP tools but are not MCP-specific.
 from __future__ import annotations
 
 import os
+import shutil
+import subprocess
 import time
 from dataclasses import dataclass
 from datetime import datetime, timedelta
@@ -19,10 +21,11 @@ GITHUB_API_BASE = "https://api.github.com"
 
 
 def resolve_github_token(preferred_env_vars: list[str] | None = None) -> str:
-    """Resolve GitHub token from environment variables.
+    """Resolve GitHub token from environment variables or gh CLI.
 
     Checks environment variables in order of preference, returning the first
-    non-empty value found.
+    non-empty value found. If no environment variables are set, attempts to
+    get a token from the gh CLI tool using 'gh auth token'.
 
     Args:
         preferred_env_vars: List of environment variable names to check in order.
@@ -32,19 +35,37 @@ def resolve_github_token(preferred_env_vars: list[str] | None = None) -> str:
         GitHub token string.
 
     Raises:
-        ValueError: If no GitHub token is found in any of the specified env vars.
+        ValueError: If no GitHub token is found in env vars or gh CLI.
     """
     if preferred_env_vars is None:
         preferred_env_vars = ["GITHUB_CI_WORKFLOW_TRIGGER_PAT", "GITHUB_TOKEN"]
 
+    # Check environment variables first
     for env_var in preferred_env_vars:
         token = os.getenv(env_var)
         if token:
             return token
 
+    # Fall back to gh CLI if available
+    gh_path = shutil.which("gh")
+    if gh_path:
+        try:
+            result = subprocess.run(
+                [gh_path, "auth", "token"],
+                capture_output=True,
+                text=True,
+                timeout=5,
+                check=False,
+            )
+            if result.returncode == 0 and result.stdout.strip():
+                return result.stdout.strip()
+        except (subprocess.TimeoutExpired, subprocess.SubprocessError):
+            pass
+
     env_var_list = ", ".join(preferred_env_vars)
     raise ValueError(
-        f"No GitHub token found. Set one of: {env_var_list} environment variable."
+        f"No GitHub token found. Set one of: {env_var_list} environment variable, "
+        "or authenticate with 'gh auth login'."
     )
 
 
@@ -62,6 +83,74 @@ class WorkflowDispatchResult:
     """Direct URL to the workflow run, if discovered"""
 
 
+@dataclass
+class WorkflowJobInfo:
+    """Information about a single job in a workflow run."""
+
+    job_id: int
+    """GitHub job ID (use with git_ci_job_logs to retrieve logs)"""
+
+    name: str
+    """Job name as defined in the workflow"""
+
+    status: str
+    """Job status: queued, in_progress, completed"""
+
+    conclusion: str | None = None
+    """Job conclusion: success, failure, cancelled, skipped (only set when completed)"""
+
+    started_at: str | None = None
+    """ISO 8601 timestamp when the job started"""
+
+    completed_at: str | None = None
+    """ISO 8601 timestamp when the job completed"""
+
+
+def get_workflow_jobs(
+    owner: str,
+    repo: str,
+    run_id: int,
+) -> list[WorkflowJobInfo]:
+    """Get jobs for a workflow run from GitHub API.
+
+    Args:
+        owner: Repository owner (e.g., "airbytehq")
+        repo: Repository name (e.g., "airbyte")
+        run_id: Workflow run ID
+
+    Returns:
+        List of WorkflowJobInfo objects for each job in the workflow run.
+
+    Raises:
+        ValueError: If no GitHub token is found.
+        requests.HTTPError: If API request fails.
+    """
+    token = resolve_github_token()
+
+    url = f"{GITHUB_API_BASE}/repos/{owner}/{repo}/actions/runs/{run_id}/jobs"
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github+json",
+        "X-GitHub-Api-Version": "2022-11-28",
+    }
+
+    response = requests.get(url, headers=headers, timeout=30)
+    response.raise_for_status()
+
+    jobs_data = response.json()
+    return [
+        WorkflowJobInfo(
+            job_id=job["id"],
+            name=job["name"],
+            status=job["status"],
+            conclusion=job.get("conclusion"),
+            started_at=job.get("started_at"),
+            completed_at=job.get("completed_at"),
+        )
+        for job in jobs_data.get("jobs", [])
+    ]
+
+
 def find_workflow_run(
     owner: str,
     repo: str,
@@ -146,7 +235,7 @@ def trigger_workflow_dispatch(
     Args:
         owner: Repository owner (e.g., "airbytehq")
         repo: Repository name (e.g., "airbyte-ops-mcp")
-        workflow_file: Workflow file name (e.g., "connector-live-test.yml")
+        workflow_file: Workflow file name (e.g., "connector-regression-test.yml")
         ref: Git ref to run the workflow on (branch name)
         inputs: Workflow inputs dictionary
         token: GitHub API token

airbyte_ops_mcp/mcp/_http_headers.py

@@ -0,0 +1,254 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""HTTP header extraction for Airbyte Cloud credentials.
+
+This module provides internal helper functions for extracting Airbyte Cloud
+authentication credentials from HTTP headers when running as an MCP HTTP server.
+This enables per-request credential passing from upstream services like coral-agents.
+
+The resolution order for credentials is:
+1. HTTP headers (when running as MCP HTTP server)
+2. Environment variables (fallback)
+
+Note: This module is prefixed with "_" to indicate it is internal helper logic
+for the MCP module and should not be imported directly by external code.
+"""
+
+from __future__ import annotations
+
+import os
+
+from airbyte.cloud.auth import (
+    resolve_cloud_api_url,
+    resolve_cloud_client_id,
+    resolve_cloud_client_secret,
+    resolve_cloud_workspace_id,
+)
+from airbyte.secrets.base import SecretString
+from fastmcp.server.dependencies import get_http_headers
+
+from airbyte_ops_mcp.constants import (
+    HEADER_AIRBYTE_CLOUD_API_URL,
+    HEADER_AIRBYTE_CLOUD_CLIENT_ID,
+    HEADER_AIRBYTE_CLOUD_CLIENT_SECRET,
+    HEADER_AIRBYTE_CLOUD_WORKSPACE_ID,
+)
+
+
+def _get_header_value(headers: dict[str, str], header_name: str) -> str | None:
+    """Get a header value from a headers dict, case-insensitively.
+
+    Args:
+        headers: Dictionary of HTTP headers.
+        header_name: The header name to look for (case-insensitive).
+
+    Returns:
+        The header value if found, None otherwise.
+    """
+    header_name_lower = header_name.lower()
+    for key, value in headers.items():
+        if key.lower() == header_name_lower:
+            return value
+    return None
+
+
+def get_bearer_token_from_headers() -> SecretString | None:
+    """Extract bearer token from HTTP Authorization header.
+
+    This function extracts the bearer token from the standard HTTP
+    `Authorization: Bearer <token>` header when running as an MCP HTTP server.
+
+    Returns:
+        The bearer token as a SecretString, or None if not found or not in HTTP context.
+    """
+    headers = get_http_headers()
+    if not headers:
+        return None
+
+    auth_header = _get_header_value(headers, "Authorization")
+    if not auth_header:
+        return None
+
+    # Parse "Bearer <token>" format (case-insensitive prefix check)
+    bearer_prefix = "bearer "
+    if auth_header.lower().startswith(bearer_prefix):
+        token = auth_header[len(bearer_prefix) :].strip()
+        if token:
+            return SecretString(token)
+
+    return None
+
+
+def get_client_id_from_headers() -> SecretString | None:
+    """Extract client ID from HTTP headers.
+
+    Returns:
+        The client ID as a SecretString, or None if not found or not in HTTP context.
+    """
+    headers = get_http_headers()
+    if not headers:
+        return None
+
+    value = _get_header_value(headers, HEADER_AIRBYTE_CLOUD_CLIENT_ID)
+    if value:
+        return SecretString(value)
+    return None
+
+
+def get_client_secret_from_headers() -> SecretString | None:
+    """Extract client secret from HTTP headers.
+
+    Returns:
+        The client secret as a SecretString, or None if not found or not in HTTP context.
+    """
+    headers = get_http_headers()
+    if not headers:
+        return None
+
+    value = _get_header_value(headers, HEADER_AIRBYTE_CLOUD_CLIENT_SECRET)
+    if value:
+        return SecretString(value)
+    return None
+
+
+def get_workspace_id_from_headers() -> str | None:
+    """Extract workspace ID from HTTP headers.
+
+    Returns:
+        The workspace ID, or None if not found or not in HTTP context.
+    """
+    headers = get_http_headers()
+    if not headers:
+        return None
+
+    return _get_header_value(headers, HEADER_AIRBYTE_CLOUD_WORKSPACE_ID)
+
+
+def get_api_url_from_headers() -> str | None:
+    """Extract API URL from HTTP headers.
+
+    Returns:
+        The API URL, or None if not found or not in HTTP context.
+    """
+    headers = get_http_headers()
+    if not headers:
+        return None
+
+    return _get_header_value(headers, HEADER_AIRBYTE_CLOUD_API_URL)
+
+
+def resolve_client_id() -> SecretString:
+    """Resolve client ID from HTTP headers or environment variables.
+
+    Resolution order:
+    1. HTTP header X-Airbyte-Cloud-Client-Id
+    2. Environment variable AIRBYTE_CLOUD_CLIENT_ID (via PyAirbyte)
+
+    Returns:
+        The resolved client ID as a SecretString.
+
+    Raises:
+        PyAirbyteSecretNotFoundError: If no client ID can be resolved.
+    """
+    header_value = get_client_id_from_headers()
+    if header_value:
+        return header_value
+
+    return resolve_cloud_client_id()
+
+
+def resolve_client_secret() -> SecretString:
+    """Resolve client secret from HTTP headers or environment variables.
+
+    Resolution order:
+    1. HTTP header X-Airbyte-Cloud-Client-Secret
+    2. Environment variable AIRBYTE_CLOUD_CLIENT_SECRET (via PyAirbyte)
+
+    Returns:
+        The resolved client secret as a SecretString.
+
+    Raises:
+        PyAirbyteSecretNotFoundError: If no client secret can be resolved.
+    """
+    header_value = get_client_secret_from_headers()
+    if header_value:
+        return header_value
+
+    return resolve_cloud_client_secret()
+
+
+def resolve_workspace_id(workspace_id: str | None = None) -> str:
+    """Resolve workspace ID from multiple sources.
+
+    Resolution order:
+    1. Explicit workspace_id parameter (if provided)
+    2. HTTP header X-Airbyte-Cloud-Workspace-Id
+    3. Environment variable AIRBYTE_CLOUD_WORKSPACE_ID (via PyAirbyte)
+
+    Args:
+        workspace_id: Optional explicit workspace ID.
+
+    Returns:
+        The resolved workspace ID.
+
+    Raises:
+        PyAirbyteSecretNotFoundError: If no workspace ID can be resolved.
+    """
+    if workspace_id is not None:
+        return workspace_id
+
+    header_value = get_workspace_id_from_headers()
+    if header_value:
+        return header_value
+
+    return resolve_cloud_workspace_id()
+
+
+def resolve_api_url(api_url: str | None = None) -> str:
+    """Resolve API URL from multiple sources.
+
+    Resolution order:
+    1. Explicit api_url parameter (if provided)
+    2. HTTP header X-Airbyte-Cloud-Api-Url
+    3. Environment variable / default (via PyAirbyte)
+
+    Args:
+        api_url: Optional explicit API URL.
+
+    Returns:
+        The resolved API URL.
+    """
+    if api_url is not None:
+        return api_url
+
+    header_value = get_api_url_from_headers()
+    if header_value:
+        return header_value
+
+    return resolve_cloud_api_url()
+
+
+def resolve_bearer_token() -> SecretString | None:
+    """Resolve bearer token from HTTP headers or environment variables.
+
+    Resolution order:
+    1. HTTP Authorization header (Bearer <token>)
+    2. Environment variable AIRBYTE_CLOUD_BEARER_TOKEN
+
+    Returns:
+        The resolved bearer token as a SecretString, or None if not found.
+
+    Note:
+        Unlike resolve_client_id/resolve_client_secret, this function returns
+        None instead of raising an exception if no bearer token is found,
+        since bearer token auth is optional (can fall back to client credentials).
+    """
+    header_value = get_bearer_token_from_headers()
+    if header_value:
+        return header_value
+
+    # Try env var directly
+    env_value = os.environ.get("AIRBYTE_CLOUD_BEARER_TOKEN")
+    if env_value:
+        return SecretString(env_value)
+
+    return None

airbyte_ops_mcp/mcp/_mcp_utils.py

@@ -76,8 +76,8 @@ class ToolDomain(str, Enum):
     PROMPTS = "prompts"
     """Prompt templates for common workflows"""
 
-    LIVE_TESTS = "live_tests"
-    """Live tests for connector validation and regression testing"""
+    REGRESSION_TESTS = "regression_tests"
+    """Regression tests for connector validation and comparison testing"""
 
 
 _REGISTERED_TOOLS: list[tuple[Callable[..., Any], dict[str, Any]]] = []