airbyte-internal-ops 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

airbyte_ops_mcp/cli/registry.py

@@ -8,6 +8,8 @@ Commands:
     airbyte-ops registry connector compute-prerelease-tag - Compute prerelease version tag
     airbyte-ops registry connector publish-prerelease - Publish connector prerelease
     airbyte-ops registry connector publish - Publish connector (apply/rollback version override)
+    airbyte-ops registry enterprise-stubs sync --bucket prod|dev - Sync connector_stubs.json to GCS
+    airbyte-ops registry enterprise-stubs check --bucket prod|dev - Compare local file with GCS
     airbyte-ops registry image inspect - Inspect Docker image on DockerHub
 """
 
@@ -16,13 +18,14 @@ from __future__ import annotations
 import contextlib
 import sys
 from pathlib import Path
-from typing import Annotated
+from typing import Annotated, Literal
 
 import yaml
 from cyclopts import App, Parameter
 
 from airbyte_ops_mcp.cli._base import app
 from airbyte_ops_mcp.cli._shared import (
+    error_console,
     exit_with_error,
     print_error,
     print_json,
@@ -42,6 +45,15 @@ from airbyte_ops_mcp.registry import (
     PublishAction,
     publish_connector,
 )
+from airbyte_ops_mcp.registry._gcs_util import get_bucket_name
+from airbyte_ops_mcp.registry.connector_stubs import (
+    CONNECTOR_STUBS_FILE,
+    CONNECTOR_STUBS_PATH,
+    ConnectorStub,
+    load_local_stubs,
+    read_connector_stubs,
+    write_connector_stubs,
+)
 
 # Create the registry sub-app
 registry_app = App(
@@ -57,11 +69,39 @@ registry_app.command(connector_app)
 image_app = App(name="image", help="Docker image operations.")
 registry_app.command(image_app)
 
+# Create the enterprise-stubs sub-app under registry (for whole-file GCS operations)
+enterprise_stubs_app = App(
+    name="enterprise-stubs",
+    help="Enterprise connector stubs GCS operations (whole-file sync).",
+)
+registry_app.command(enterprise_stubs_app)
+
 
 AIRBYTE_REPO_OWNER = "airbytehq"
+AIRBYTE_ENTERPRISE_REPO_NAME = "airbyte-enterprise"
 AIRBYTE_REPO_NAME = "airbyte"
 CONNECTOR_PATH_PREFIX = "airbyte-integrations/connectors"
 
+# Type alias for bucket argument
+BucketArg = Literal["dev", "prod"]
+
+
+def _validate_bucket_arg(bucket: str) -> BucketArg:
+    """Validate and return the bucket argument.
+
+    Args:
+        bucket: The bucket argument from CLI.
+
+    Returns:
+        The validated bucket value.
+
+    Raises:
+        SystemExit: If the bucket value is invalid.
+    """
+    if bucket not in ("prod", "dev"):
+        exit_with_error(f"Invalid bucket '{bucket}'. Must be 'prod' or 'dev'.")
+    return bucket  # type: ignore[return-value]
+
 
 def _get_connector_version_from_github(
     connector_name: str,
@@ -295,3 +335,177 @@ def inspect_image(
     else:
         print_error(f"Image {result.full_name} not found.")
     print_json(result.model_dump())
+
+
+@enterprise_stubs_app.command(name="check")
+def enterprise_stubs_check(
+    bucket: Annotated[
+        BucketArg,
+        Parameter(
+            help="Target GCS bucket: 'prod' or 'dev'.",
+        ),
+    ],
+    repo_root: Annotated[
+        Path,
+        Parameter(
+            help="Path to the airbyte-enterprise repository root. Defaults to current directory."
+        ),
+    ] = Path.cwd(),
+) -> None:
+    """Compare local connector_stubs.json with the version in GCS.
+
+    This command reads the entire local connector_stubs.json file and compares it
+    with the version currently published in GCS.
+
+    Exit codes:
+        0: Local file matches GCS (check passed)
+        1: Differences found (check failed)
+
+    Output:
+        STDOUT: JSON representation of the comparison result
+        STDERR: Informational messages and comparison details
+
+    Example:
+        airbyte-ops registry enterprise-stubs check --bucket prod --repo-root /path/to/airbyte-enterprise
+        airbyte-ops registry enterprise-stubs check --bucket dev
+    """
+    bucket = _validate_bucket_arg(bucket)
+
+    # Load local stubs
+    try:
+        local_stubs = load_local_stubs(repo_root)
+    except FileNotFoundError as e:
+        exit_with_error(str(e))
+    except ValueError as e:
+        exit_with_error(str(e))
+
+    # Load published stubs from GCS
+    bucket_name = get_bucket_name(bucket)
+    published_stubs = read_connector_stubs(bucket_name)
+
+    error_console.print(
+        f"Comparing local {CONNECTOR_STUBS_FILE} with {bucket_name}/{CONNECTOR_STUBS_PATH}"
+    )
+
+    # Build lookup dicts by stub ID (filter out stubs without IDs)
+    local_by_id = {stub["id"]: stub for stub in local_stubs if stub.get("id")}
+    published_by_id = {stub["id"]: stub for stub in published_stubs if stub.get("id")}
+
+    all_ids = set(local_by_id.keys()) | set(published_by_id.keys())
+    differences: list[dict[str, str]] = []
+
+    for stub_id in sorted(all_ids):
+        local_stub = local_by_id.get(stub_id)
+        published_stub = published_by_id.get(stub_id)
+
+        if local_stub is None:
+            differences.append({"id": stub_id, "status": "only_in_gcs"})
+        elif published_stub is None:
+            differences.append({"id": stub_id, "status": "only_in_local"})
+        elif local_stub != published_stub:
+            differences.append({"id": stub_id, "status": "modified"})
+
+    result = {
+        "local_count": len(local_stubs),
+        "published_count": len(published_stubs),
+        "in_sync": len(differences) == 0,
+        "differences": differences,
+    }
+
+    if differences:
+        error_console.print(
+            f"[yellow]Warning:[/yellow] {len(differences)} difference(s) found:"
+        )
+        for diff in differences:
+            error_console.print(f"  {diff['id']}: {diff['status']}")
+        print_json(result)
+        sys.exit(1)
+
+    error_console.print(
+        f"[green]Local file is in sync with GCS ({len(local_stubs)} stubs)[/green]"
+    )
+    print_json(result)
+
+
+@enterprise_stubs_app.command(name="sync")
+def enterprise_stubs_sync(
+    bucket: Annotated[
+        BucketArg,
+        Parameter(
+            help="Target GCS bucket: 'prod' or 'dev'.",
+        ),
+    ],
+    repo_root: Annotated[
+        Path,
+        Parameter(
+            help="Path to the airbyte-enterprise repository root. Defaults to current directory."
+        ),
+    ] = Path.cwd(),
+    dry_run: Annotated[
+        bool,
+        Parameter(help="Show what would be uploaded without making changes."),
+    ] = False,
+) -> None:
+    """Sync local connector_stubs.json to GCS.
+
+    This command uploads the entire local connector_stubs.json file to GCS,
+    replacing the existing file. Use this after merging changes to master
+    in the airbyte-enterprise repository.
+
+    Exit codes:
+        0: Sync successful (or dry-run completed)
+        1: Error (file not found, validation failed, etc.)
+
+    Output:
+        STDOUT: JSON representation of the sync result
+        STDERR: Informational messages and status updates
+
+    Example:
+        airbyte-ops registry enterprise-stubs sync --bucket prod --repo-root /path/to/airbyte-enterprise
+        airbyte-ops registry enterprise-stubs sync --bucket dev
+        airbyte-ops registry enterprise-stubs sync --bucket dev --dry-run
+    """
+    bucket = _validate_bucket_arg(bucket)
+
+    # Load local stubs
+    try:
+        local_stubs = load_local_stubs(repo_root)
+    except FileNotFoundError as e:
+        exit_with_error(str(e))
+    except ValueError as e:
+        exit_with_error(str(e))
+
+    # Validate all stubs
+    for stub in local_stubs:
+        ConnectorStub(**stub)
+
+    bucket_name = get_bucket_name(bucket)
+
+    if dry_run:
+        error_console.print(
+            f"[DRY RUN] Would upload {len(local_stubs)} stubs to "
+            f"{bucket_name}/{CONNECTOR_STUBS_PATH}"
+        )
+        result = {
+            "dry_run": True,
+            "stub_count": len(local_stubs),
+            "bucket": bucket_name,
+            "path": CONNECTOR_STUBS_PATH,
+        }
+        print_json(result)
+        return
+
+    # Write to GCS (replaces entire file)
+    write_connector_stubs(bucket_name, local_stubs)
+
+    error_console.print(
+        f"[green]Synced {len(local_stubs)} stubs to {bucket_name}/{CONNECTOR_STUBS_PATH}[/green]"
+    )
+    result = {
+        "dry_run": False,
+        "stub_count": len(local_stubs),
+        "bucket": bucket_name,
+        "path": CONNECTOR_STUBS_PATH,
+        "stub_ids": [stub.get("id") for stub in local_stubs],
+    }
+    print_json(result)

airbyte_ops_mcp/github_api.py

@@ -9,6 +9,8 @@ modules but are not MCP-specific.
 
 from __future__ import annotations
 
+import datetime
+import functools
 import os
 import re
 import shutil
@@ -17,6 +19,8 @@ from dataclasses import dataclass
 from urllib.parse import urlparse
 
 import requests
+from github import Auth, Github, GithubException
+from github.Repository import Repository
 
 GITHUB_API_BASE = "https://api.github.com"
 
@@ -130,6 +134,58 @@ def get_pr_head_ref(
     )
 
 
+@functools.lru_cache(maxsize=32)
+def _get_github_repo(owner: str, repo: str, token: str) -> Repository:
+    """Get a cached GitHub repository object.
+
+    This function caches repository objects to avoid redundant API calls
+    when fetching multiple PRs from the same repository. Uses lazy=True
+    to avoid making an API call just to create the repo object.
+    """
+    auth = Auth.Token(token)
+    gh = Github(auth=auth)
+    return gh.get_repo(f"{owner}/{repo}", lazy=True)
+
+
+def get_pr_merge_date(
+    owner: str,
+    repo: str,
+    pr_number: int,
+    token: str | None = None,
+) -> datetime.date | None:
+    """Get the merge date for a PR.
+
+    Args:
+        owner: Repository owner (e.g., "airbytehq")
+        repo: Repository name (e.g., "airbyte")
+        pr_number: Pull request number
+        token: GitHub API token. If None, will be resolved from environment.
+
+    Returns:
+        The date the PR was merged, or None if not merged.
+
+    Raises:
+        GitHubAPIError: If the API request fails.
+    """
+    if token is None:
+        token = resolve_github_token()
+
+    try:
+        gh_repo = _get_github_repo(owner, repo, token)
+        pr = gh_repo.get_pull(pr_number)
+    except GithubException as e:
+        if e.status == 404:
+            raise GitHubAPIError(f"PR {owner}/{repo}#{pr_number} not found") from e
+        raise GitHubAPIError(
+            f"Failed to fetch PR {owner}/{repo}#{pr_number}: {e.status} {e.data}"
+        ) from e
+
+    if pr.merged_at is None:
+        return None
+
+    return pr.merged_at.date()
+
+
 def get_file_contents_at_ref(
     owner: str,
     repo: str,

airbyte_ops_mcp/registry/_gcs_util.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Shared GCS utilities for registry operations.
+
+This module provides common GCS helper functions used across registry
+operations, including connector stubs and metadata management.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from google.cloud import storage
+
+from airbyte_ops_mcp.gcp_auth import get_gcp_credentials
+
+# GCS bucket names for the metadata service
+PROD_METADATA_SERVICE_BUCKET_NAME = "prod-airbyte-cloud-connector-metadata-service"
+DEV_METADATA_SERVICE_BUCKET_NAME = "dev-airbyte-cloud-connector-metadata-service-2"
+
+# Type alias for environment ID
+EnvId = Literal["dev", "prod"]
+
+
+def get_gcs_client() -> storage.Client:
+    """Get a GCS storage client using centralized credentials.
+
+    Uses the centralized GCP authentication from gcp_auth module,
+    which supports both GCP_PROD_DB_ACCESS_CREDENTIALS env var
+    and standard ADC discovery.
+
+    Returns:
+        A configured GCS storage client.
+    """
+    credentials = get_gcp_credentials()
+    return storage.Client(credentials=credentials)
+
+
+def get_bucket_name(env: EnvId) -> str:
+    """Get the metadata service bucket name.
+
+    Args:
+        env: The environment ID ('dev' or 'prod').
+
+    Returns:
+        The bucket name to use for GCS operations.
+    """
+    if env == "prod":
+        return PROD_METADATA_SERVICE_BUCKET_NAME
+    return DEV_METADATA_SERVICE_BUCKET_NAME
+
+
+def get_gcs_file_text(
+    bucket_name: str,
+    path: str,
+    client: storage.Client | None = None,
+) -> str | None:
+    """Read a text file from GCS.
+
+    Args:
+        bucket_name: The GCS bucket name.
+        path: The path to the file within the bucket.
+        client: Optional GCS client. If not provided, creates one using get_gcs_client().
+
+    Returns:
+        The file contents as a string, or None if the file doesn't exist.
+    """
+    if client is None:
+        client = get_gcs_client()
+
+    bucket = client.bucket(bucket_name)
+    blob = bucket.blob(path)
+
+    if not blob.exists():
+        return None
+
+    return blob.download_as_string().decode("utf-8")
+
+
+def upload_gcs_file_text(
+    bucket_name: str,
+    path: str,
+    content: str,
+    content_type: str = "text/plain",
+    client: storage.Client | None = None,
+) -> None:
+    """Upload a text file to GCS.
+
+    Args:
+        bucket_name: The GCS bucket name.
+        path: The path to the file within the bucket.
+        content: The text content to upload.
+        content_type: The MIME type of the content. Defaults to "text/plain".
+        client: Optional GCS client. If not provided, creates one using get_gcs_client().
+    """
+    if client is None:
+        client = get_gcs_client()
+
+    bucket = client.bucket(bucket_name)
+    blob = bucket.blob(path)
+    blob.upload_from_string(content, content_type=content_type)

airbyte_ops_mcp/registry/connector_stubs.py

@@ -0,0 +1,186 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Core logic for connector stub operations.
+
+This module provides the core functionality for managing enterprise connector
+stubs in the Airbyte Cloud catalog. The stubs are stored in GCS and appear
+in the catalog as placeholder entries that direct users to a sales funnel.
+
+The stubs are stored at:
+    gs://prod-airbyte-cloud-connector-metadata-service/resources/connector_stubs/v1/connector_stubs.json
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from airbyte_ops_mcp.registry._gcs_util import (
+    get_bucket_name,
+    get_gcs_file_text,
+    upload_gcs_file_text,
+)
+
+# Path to connector stubs file in GCS
+CONNECTOR_STUBS_PATH = "resources/connector_stubs/v1/connector_stubs.json"
+
+# Local file name for connector stubs in airbyte-enterprise repo
+CONNECTOR_STUBS_FILE = "connector_stubs.json"
+
+
+class ConnectorStub(BaseModel):
+    """A connector stub entry for the enterprise connector catalog.
+
+    Stubs are placeholder entries that appear in the Airbyte Cloud catalog
+    for enterprise connectors. When users click on them, they are directed
+    to a sales funnel rather than being able to configure the connector directly.
+    """
+
+    model_config = {"populate_by_name": True}
+
+    id: str = Field(
+        description="Unique identifier for the stub (e.g., 'source-oracle-enterprise')"
+    )
+    name: str = Field(description="Display name of the connector (e.g., 'Oracle')")
+    url: str = Field(
+        description="URL to the connector's documentation page (must be publicly accessible)"
+    )
+    icon: str = Field(
+        description="URL to the connector's icon (typically stored in the same GCS bucket)"
+    )
+    definition_id: str | None = Field(
+        default=None,
+        description="UUID of the connector definition (if it exists in the registry)",
+        alias="definitionId",
+    )
+    label: str | None = Field(
+        default=None, description="Label for the connector (typically 'enterprise')"
+    )
+    type: str | None = Field(
+        default=None,
+        description="Type of connector (e.g., 'enterprise_source', 'enterprise_destination')",
+    )
+    codename: str | None = Field(
+        default=None, description="Internal codename for the connector (optional)"
+    )
+
+
+def read_connector_stubs(bucket_name: str) -> list[dict]:
+    """Read connector stubs from GCS.
+
+    Args:
+        bucket_name: The GCS bucket name.
+
+    Returns:
+        List of connector stub dictionaries.
+
+    Raises:
+        ValueError: If the file exists but contains invalid data.
+    """
+    content = get_gcs_file_text(bucket_name, CONNECTOR_STUBS_PATH)
+
+    if content is None:
+        return []
+
+    stubs = json.loads(content)
+
+    if not isinstance(stubs, list):
+        raise ValueError(
+            f"Expected connector_stubs.json to contain a list, got {type(stubs).__name__}"
+        )
+
+    return stubs
+
+
+def write_connector_stubs(bucket_name: str, stubs: list[dict]) -> None:
+    """Write connector stubs to GCS.
+
+    Args:
+        bucket_name: The GCS bucket name.
+        stubs: List of connector stub dictionaries to write.
+    """
+    content = json.dumps(stubs, indent=2)
+    upload_gcs_file_text(
+        bucket_name, CONNECTOR_STUBS_PATH, content, content_type="application/json"
+    )
+
+
+def find_stub_by_connector(stubs: list[dict], connector: str) -> dict | None:
+    """Find a stub by connector name or ID.
+
+    Matches by:
+    - Exact ID match (e.g., 'source-oracle-enterprise')
+    - ID with '-enterprise' suffix (e.g., 'source-oracle' matches 'source-oracle-enterprise')
+    - Name match (case-insensitive, spaces converted to hyphens)
+
+    Args:
+        stubs: List of connector stub dictionaries.
+        connector: Connector name or stub ID to find.
+
+    Returns:
+        The matching stub dictionary, or None if not found.
+    """
+    for stub in stubs:
+        stub_id = stub.get("id", "")
+        # Match by exact ID or by connector name pattern
+        if stub_id == connector or stub_id == f"{connector}-enterprise":
+            return stub
+        # Also check if the connector name matches the stub name
+        if stub.get("name", "").lower().replace(" ", "-") == connector.lower():
+            return stub
+    return None
+
+
+def load_local_stubs(repo_root: Path) -> list[dict]:
+    """Load connector stubs from local repository.
+
+    Args:
+        repo_root: Path to the airbyte-enterprise repository root.
+
+    Returns:
+        List of connector stub dictionaries.
+
+    Raises:
+        FileNotFoundError: If the connector stubs file doesn't exist.
+        ValueError: If the file contains invalid data.
+    """
+    stubs_file = repo_root / CONNECTOR_STUBS_FILE
+    if not stubs_file.exists():
+        raise FileNotFoundError(f"Connector stubs file not found: {stubs_file}")
+
+    content = stubs_file.read_text()
+    stubs = json.loads(content)
+
+    if not isinstance(stubs, list):
+        raise ValueError(
+            f"Expected {CONNECTOR_STUBS_FILE} to contain a list, got {type(stubs).__name__}"
+        )
+
+    return stubs
+
+
+def save_local_stubs(repo_root: Path, stubs: list[dict]) -> None:
+    """Save connector stubs to local repository.
+
+    Args:
+        repo_root: Path to the airbyte-enterprise repository root.
+        stubs: List of connector stub dictionaries to save.
+    """
+    stubs_file = repo_root / CONNECTOR_STUBS_FILE
+    content = json.dumps(stubs, indent=2) + "\n"
+    stubs_file.write_text(content)
+
+
+# Re-export get_bucket_name for convenience
+__all__ = [
+    "CONNECTOR_STUBS_FILE",
+    "CONNECTOR_STUBS_PATH",
+    "ConnectorStub",
+    "find_stub_by_connector",
+    "get_bucket_name",
+    "load_local_stubs",
+    "read_connector_stubs",
+    "save_local_stubs",
+    "write_connector_stubs",
+]

airbyte_ops_mcp/registry/publish.py

@@ -15,6 +15,10 @@ import yaml
 from google.cloud import storage
 from google.oauth2 import service_account
 
+from airbyte_ops_mcp.registry._gcs_util import (
+    DEV_METADATA_SERVICE_BUCKET_NAME,
+    PROD_METADATA_SERVICE_BUCKET_NAME,
+)
 from airbyte_ops_mcp.registry.models import (
     ConnectorMetadata,
     ConnectorPublishResult,
@@ -27,10 +31,6 @@ METADATA_FOLDER = "metadata"
 LATEST_GCS_FOLDER_NAME = "latest"
 RELEASE_CANDIDATE_GCS_FOLDER_NAME = "release_candidate"
 
-# Bucket names for metadata service
-PROD_METADATA_SERVICE_BUCKET_NAME = "prod-airbyte-cloud-connector-metadata-service"
-DEV_METADATA_SERVICE_BUCKET_NAME = "dev-airbyte-cloud-connector-metadata-service-2"
-
 # Default to dev bucket for safety - use --prod flag to target production
 DEFAULT_METADATA_SERVICE_BUCKET_NAME = DEV_METADATA_SERVICE_BUCKET_NAME