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

airbyte_ops_mcp/cli/cloud.py

@@ -23,7 +23,12 @@ from airbyte_protocol.models import ConfiguredAirbyteCatalog
 from cyclopts import App, Parameter
 
 from airbyte_ops_mcp.cli._base import app
-from airbyte_ops_mcp.cli._shared import print_error, print_json, print_success
+from airbyte_ops_mcp.cli._shared import (
+    exit_with_error,
+    print_error,
+    print_json,
+    print_success,
+)
 from airbyte_ops_mcp.cloud_admin.connection_config import fetch_connection_config
 from airbyte_ops_mcp.live_tests.ci_output import (
     generate_regression_report,
@@ -391,12 +396,11 @@ def live_test(
     # If connector_name is provided, build the image from source
     if connector_name:
         if connector_image:
-            print_error("Cannot specify both connector_image and connector_name")
             write_github_output("success", False)
             write_github_output(
                 "error", "Cannot specify both connector_image and connector_name"
             )
-            return
+            exit_with_error("Cannot specify both connector_image and connector_name")
 
         repo_root_path = Path(repo_root) if repo_root else None
         built_image = _build_connector_image_from_source(
@@ -407,19 +411,18 @@ def live_test(
         if not built_image:
             write_github_output("success", False)
             write_github_output("error", f"Failed to build image for {connector_name}")
-            return
+            exit_with_error(f"Failed to build image for {connector_name}")
         resolved_connector_image = built_image
 
     if connection_id:
         if config_path or catalog_path:
-            print_error(
-                "Cannot specify both connection_id and config_path/catalog_path"
-            )
             write_github_output("success", False)
             write_github_output(
                 "error", "Cannot specify both connection_id and file paths"
             )
-            return
+            exit_with_error(
+                "Cannot specify both connection_id and config_path/catalog_path"
+            )
 
         print_success(f"Fetching config/catalog from connection: {connection_id}")
         connection_data = fetch_connection_data(connection_id)
@@ -439,25 +442,23 @@ def live_test(
         catalog_file = Path(catalog_path) if catalog_path else None
 
     if not resolved_connector_image:
-        print_error(
+        write_github_output("success", False)
+        write_github_output("error", "Missing connector image")
+        exit_with_error(
             "You must provide one of the following: a connector_image, a connector_name, "
             "or a connection_id for a connection that has an associated connector image. "
             "If using connection_id, ensure the connection has a connector image configured."
         )
-        write_github_output("success", False)
-        write_github_output("error", "Missing connector image")
-        return
 
     # If connector_name was provided, we just built the image locally and it is already
     # available in Docker, so we skip the image availability check/pull. Only try to pull
     # if we didn't just build it (i.e., using a pre-built image from registry).
     if not connector_name and not ensure_image_available(resolved_connector_image):
-        print_error(f"Failed to pull connector image: {resolved_connector_image}")
         write_github_output("success", False)
         write_github_output(
             "error", f"Failed to pull image: {resolved_connector_image}"
         )
-        return
+        exit_with_error(f"Failed to pull connector image: {resolved_connector_image}")
 
     result = _run_connector_command(
         connector_image=resolved_connector_image,
@@ -494,7 +495,7 @@ def live_test(
     if result["success"]:
         print_success(f"Live test passed for {resolved_connector_image}")
     else:
-        print_error(f"Live test failed for {resolved_connector_image}")
+        exit_with_error(f"Live test failed for {resolved_connector_image}")
 
 
 def _run_with_optional_http_metrics(
@@ -673,12 +674,11 @@ def regression_test(
     # If connector_name is provided, build the target image from source
     if connector_name:
         if target_image:
-            print_error("Cannot specify both target_image and connector_name")
             write_github_output("success", False)
             write_github_output(
                 "error", "Cannot specify both target_image and connector_name"
             )
-            return
+            exit_with_error("Cannot specify both target_image and connector_name")
 
         repo_root_path = Path(repo_root) if repo_root else None
         built_image = _build_connector_image_from_source(
@@ -689,19 +689,18 @@ def regression_test(
         if not built_image:
             write_github_output("success", False)
             write_github_output("error", f"Failed to build image for {connector_name}")
-            return
+            exit_with_error(f"Failed to build image for {connector_name}")
         resolved_target_image = built_image
 
     if connection_id:
         if config_path or catalog_path:
-            print_error(
-                "Cannot specify both connection_id and config_path/catalog_path"
-            )
             write_github_output("success", False)
             write_github_output(
                 "error", "Cannot specify both connection_id and file paths"
             )
-            return
+            exit_with_error(
+                "Cannot specify both connection_id and config_path/catalog_path"
+            )
 
         print_success(f"Fetching config/catalog from connection: {connection_id}")
         connection_data = fetch_connection_data(connection_id)
@@ -723,36 +722,36 @@ def regression_test(
 
     # Validate that we have both images
     if not resolved_target_image:
-        print_error(
+        write_github_output("success", False)
+        write_github_output("error", "No target image specified")
+        exit_with_error(
             "You must provide one of the following: a target_image or a connector_name "
             "to build the target image from source."
         )
-        write_github_output("success", False)
-        write_github_output("error", "No target image specified")
-        return
 
     if not resolved_control_image:
-        print_error(
+        write_github_output("success", False)
+        write_github_output("error", "No control image specified")
+        exit_with_error(
             "You must provide one of the following: a control_image or a connection_id "
             "for a connection that has an associated connector image."
         )
-        write_github_output("success", False)
-        write_github_output("error", "No control image specified")
-        return
 
     # Pull images if they weren't just built locally
     # If connector_name was provided, we just built the target image locally
     if not connector_name and not ensure_image_available(resolved_target_image):
-        print_error(f"Failed to pull target connector image: {resolved_target_image}")
         write_github_output("success", False)
         write_github_output("error", f"Failed to pull image: {resolved_target_image}")
-        return
+        exit_with_error(
+            f"Failed to pull target connector image: {resolved_target_image}"
+        )
 
     if not ensure_image_available(resolved_control_image):
-        print_error(f"Failed to pull control connector image: {resolved_control_image}")
         write_github_output("success", False)
         write_github_output("error", f"Failed to pull image: {resolved_control_image}")
-        return
+        exit_with_error(
+            f"Failed to pull control connector image: {resolved_control_image}"
+        )
 
     target_output = output_path / "target"
     control_output = output_path / "control"
@@ -819,7 +818,7 @@ def regression_test(
     write_github_summary(summary)
 
     if regression_detected:
-        print_error(
+        exit_with_error(
             f"Regression detected between {resolved_target_image} and {resolved_control_image}"
         )
     elif both_succeeded:
@@ -827,7 +826,7 @@ def regression_test(
             f"Regression test passed for {resolved_target_image} vs {resolved_control_image}"
         )
     else:
-        print_error(
+        exit_with_error(
             f"Both versions failed for {resolved_target_image} vs {resolved_control_image}"
         )
 

airbyte_ops_mcp/cli/registry.py

@@ -1,21 +1,36 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
 """CLI commands for connector registry operations.
 
+This module provides CLI wrappers for registry operations. The core logic
+lives in the `airbyte_ops_mcp.registry` capability module.
+
 Commands:
     airbyte-ops registry connector publish-prerelease - Publish connector prerelease
+    airbyte-ops registry connector publish - Publish connector (apply/rollback version override)
     airbyte-ops registry image inspect - Inspect Docker image on DockerHub
 """
 
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Annotated
 
 from cyclopts import App, Parameter
 
 from airbyte_ops_mcp.cli._base import app
-from airbyte_ops_mcp.cli._shared import print_error, print_json, print_success
+from airbyte_ops_mcp.cli._shared import (
+    exit_with_error,
+    print_error,
+    print_json,
+    print_success,
+)
 from airbyte_ops_mcp.mcp.github import get_docker_image_info
 from airbyte_ops_mcp.mcp.prerelease import publish_connector_to_airbyte_registry
+from airbyte_ops_mcp.registry import (
+    ConnectorPublishResult,
+    PublishAction,
+    publish_connector,
+)
 
 # Create the registry sub-app
 registry_app = App(
@@ -65,6 +80,80 @@ def publish_prerelease(
     print_json(result.model_dump())
 
 
+@connector_app.command(name="publish")
+def publish(
+    name: Annotated[
+        str,
+        Parameter(help="Connector technical name (e.g., source-github)."),
+    ],
+    repo_path: Annotated[
+        Path,
+        Parameter(help="Path to the Airbyte monorepo. Defaults to current directory."),
+    ] = Path.cwd(),
+    apply_override: Annotated[
+        bool,
+        Parameter(
+            help="Apply a version override (promote RC to stable).",
+            negative="",  # Disable --no-apply-override
+        ),
+    ] = False,
+    rollback_override: Annotated[
+        bool,
+        Parameter(
+            help="Rollback a version override.",
+            negative="",  # Disable --no-rollback-override
+        ),
+    ] = False,
+    dry_run: Annotated[
+        bool,
+        Parameter(help="Show what would be published without making changes."),
+    ] = False,
+    prod: Annotated[
+        bool,
+        Parameter(
+            help="Target the production GCS bucket. Without this flag, operations target the dev bucket for safe testing.",
+            negative="",  # Disable --no-prod
+        ),
+    ] = False,
+) -> None:
+    """Publish a connector to the Airbyte registry.
+
+    This command handles connector publishing operations including applying
+    version overrides (promoting RC to stable) or rolling back version overrides.
+
+    By default, operations target the dev bucket (dev-airbyte-cloud-connector-metadata-service-2)
+    for safe testing. Use --prod to target the production bucket.
+    """
+    if apply_override and rollback_override:
+        exit_with_error("Cannot use both --apply-override and --rollback-override")
+
+    if not apply_override and not rollback_override:
+        exit_with_error("Must specify either --apply-override or --rollback-override")
+
+    # Map CLI flags to PublishAction
+    action: PublishAction = (
+        "apply-version-override" if apply_override else "rollback-version-override"
+    )
+
+    # Delegate to the capability module
+    if not repo_path.exists():
+        exit_with_error(f"Repository path not found: {repo_path}")
+
+    result: ConnectorPublishResult = publish_connector(
+        repo_path=repo_path,
+        connector_name=name,
+        action=action,
+        dry_run=dry_run,
+        use_prod=prod,
+    )
+
+    # Output result as JSON
+    print_json(result.model_dump())
+
+    if result.status == "failure":
+        exit_with_error(result.message or "Operation failed", code=1)
+
+
 @image_app.command(name="inspect")
 def inspect_image(
     image: Annotated[

airbyte_ops_mcp/cli/repo.py

@@ -27,6 +27,7 @@ from airbyte_ops_mcp.airbyte_repo.list_connectors import (
     CONNECTOR_PATH_PREFIX,
     METADATA_FILE_NAME,
     _detect_connector_language,
+    get_connectors_with_local_cdk,
 )
 from airbyte_ops_mcp.cli._base import app
 from airbyte_ops_mcp.cli._shared import exit_with_error, print_json
@@ -160,6 +161,15 @@ def list_connectors(
         bool,
         Parameter(help="Include only modified connectors (requires PR context)."),
     ] = False,
+    local_cdk: Annotated[
+        bool,
+        Parameter(
+            help=(
+                "Include connectors using local CDK reference. "
+                "When combined with --modified-only, adds local-CDK connectors to the modified set."
+            )
+        ),
+    ] = False,
     language: Annotated[
         list[str] | None,
         Parameter(help="Languages to include (python, java, low-code, manifest-only)."),
@@ -286,6 +296,11 @@ def list_connectors(
     connectors = list(result.connectors)
     repo_path_obj = Path(repo_path)
 
+    # Add connectors with local CDK reference if --local-cdk flag is set
+    if local_cdk:
+        local_cdk_connectors = get_connectors_with_local_cdk(repo_path)
+        connectors = sorted(set(connectors) | local_cdk_connectors)
+
     # Apply connector type filter
     if connector_type_filter:
         connectors = [

airbyte_ops_mcp/connection_config_retriever/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Connection config retriever module.
+
+This module provides functionality to retrieve unmasked connection configuration
+from Airbyte Cloud's internal database, including secret resolution from GCP
+Secret Manager and audit logging to GCP Cloud Logging.
+
+Refactored from: live_tests/_connection_retriever
+Original source: airbyte-platform-internal/tools/connection-retriever
+"""
+
+from airbyte_ops_mcp.connection_config_retriever.retrieval import (
+    ConnectionNotFoundError,
+    RetrievalMetadata,
+    TestingCandidate,
+    retrieve_objects,
+)
+from airbyte_ops_mcp.constants import ConnectionObject
+
+__all__ = [
+    "ConnectionNotFoundError",
+    "ConnectionObject",
+    "RetrievalMetadata",
+    "TestingCandidate",
+    "retrieve_objects",
+]

airbyte_ops_mcp/{live_tests/_connection_retriever → connection_config_retriever}/audit_logging.py

@@ -1,7 +1,8 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
-"""Audit logging for vendored connection-retriever.
+"""Audit logging for connection config retrieval.
 
-Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/audit_logging.py
+Refactored from: live_tests/_connection_retriever/audit_logging.py
+Original source: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/audit_logging.py
 """
 
 from __future__ import annotations
@@ -12,12 +13,10 @@ 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,
-)
+from airbyte_ops_mcp.constants import GCP_PROJECT_NAME
 
 if TYPE_CHECKING:
-    from airbyte_ops_mcp.live_tests._connection_retriever.retrieval import (
+    from airbyte_ops_mcp.connection_config_retriever.retrieval import (
         RetrievalMetadata,
     )
 

airbyte_ops_mcp/{live_tests/_connection_retriever → connection_config_retriever}/retrieval.py

@@ -1,7 +1,8 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
-"""Core retrieval logic for vendored connection-retriever.
+"""Core retrieval logic for connection config retrieval.
 
-Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/retrieval.py
+Refactored from: live_tests/_connection_retriever/retrieval.py
+Original source: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/retrieval.py
 
 This is a minimal subset focused on retrieving unmasked source config.
 For testing candidate discovery, see issue #91.
@@ -18,19 +19,12 @@ import requests
 import sqlalchemy
 from google.cloud import secretmanager
 
-from airbyte_ops_mcp.live_tests._connection_retriever.audit_logging import (
-    audit,
-)
-from airbyte_ops_mcp.live_tests._connection_retriever.consts import (
-    CLOUD_REGISTRY_URL,
-    ConnectionObject,
-)
-from airbyte_ops_mcp.live_tests._connection_retriever.db_access import (
-    get_pool,
-)
-from airbyte_ops_mcp.live_tests._connection_retriever.secrets_resolution import (
+from airbyte_ops_mcp.connection_config_retriever.audit_logging import audit
+from airbyte_ops_mcp.connection_config_retriever.secrets_resolution import (
     get_resolved_config,
 )
+from airbyte_ops_mcp.constants import CLOUD_REGISTRY_URL, ConnectionObject
+from airbyte_ops_mcp.prod_db_access.db_engine import get_pool
 
 LOGGER = logging.getLogger(__name__)
 
@@ -59,7 +53,7 @@ SELECT_ON_CONNECTION_NOT_EU = sqlalchemy.text(
 SELECT_ON_CONNECTION_DATAPLANE_GROUP_IS_EU = sqlalchemy.text(
     """
     SELECT
-        CASE WHEN dataplane_group.name = 'EU' THEN TRUE ELSE FALSE END as is_eu
+        CASE WHEN dataplane_group.name = 'EU' THEN TRUE ELSE FALSE END AS is_eu
     FROM
         connection
     JOIN
@@ -309,14 +303,6 @@ def retrieve_objects(
 
     This is a simplified version that only supports retrieval by connection_id.
     For testing candidate discovery by docker image, see issue #91.
-
-    Args:
-        connection_objects: List of ConnectionObject types to retrieve
-        retrieval_reason: Reason for retrieval (for audit logging)
-        connection_id: The connection ID to retrieve objects for
-
-    Returns:
-        List containing a single TestingCandidate with the requested objects
     """
     connection_candidates = [TestingCandidate(connection_id=connection_id)]
 

airbyte_ops_mcp/{live_tests/_connection_retriever → connection_config_retriever}/secrets_resolution.py

@@ -1,7 +1,8 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
-"""Secret resolution for vendored connection-retriever.
+"""Secret resolution for connection config retrieval.
 
-Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/secrets_resolution.py
+Refactored from: live_tests/_connection_retriever/secrets_resolution.py
+Original source: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/secrets_resolution.py
 """
 
 from __future__ import annotations
@@ -11,23 +12,13 @@ from typing import Any
 import dpath
 from google.cloud import secretmanager
 
-from airbyte_ops_mcp.live_tests._connection_retriever.consts import (
-    GCP_PROJECT_NAME,
-)
+from airbyte_ops_mcp.constants import GCP_PROJECT_NAME
 
 
 def get_secret_value(
     secret_manager_client: secretmanager.SecretManagerServiceClient, secret_id: str
 ) -> str:
-    """Get the value of the enabled version of a secret.
-
-    Args:
-        secret_manager_client: The secret manager client
-        secret_id: The id of the secret
-
-    Returns:
-        The value of the enabled version of the secret
-    """
+    """Get the value of the enabled version of a secret."""
     response = secret_manager_client.list_secret_versions(
         request={"parent": secret_id, "filter": "state:ENABLED"}
     )
@@ -39,14 +30,7 @@ def get_secret_value(
 
 
 def is_secret(value: Any) -> bool:
-    """Determine if a value is a secret.
-
-    Args:
-        value: The value to check
-
-    Returns:
-        True if the value is a secret, False otherwise
-    """
+    """Determine if a value is a secret."""
     return isinstance(value, dict) and value.get("_secret") is not None
 
 
@@ -54,15 +38,7 @@ def resolve_secrets_in_config(
     secret_manager_client: secretmanager.SecretManagerServiceClient,
     connector_config: dict,
 ) -> dict:
-    """Recursively resolve secrets in the connector_config.
-
-    Args:
-        secret_manager_client: The secret manager client
-        connector_config: The connector_config to resolve secrets in
-
-    Returns:
-        The connector_config with secrets resolved
-    """
+    """Recursively resolve secrets in the connector_config."""
     for key in connector_config:
         if is_secret(connector_config[key]):
             secret_id = f"projects/{GCP_PROJECT_NAME}/secrets/{connector_config[key]['_secret']}"
@@ -91,17 +67,7 @@ def get_resolved_config(
     actor_oauth_parameter: dict,
     spec: dict,
 ) -> dict:
-    """Get the resolved configuration, resolving secrets and merging OAuth params.
-
-    Args:
-        secret_manager_client: The secret manager client
-        actor_configuration: The actor configuration
-        actor_oauth_parameter: The actor oauth parameter
-        spec: The connector spec
-
-    Returns:
-        The resolved configuration
-    """
+    """Get the resolved configuration, resolving secrets and merging OAuth params."""
     resolved_configuration = resolve_secrets_in_config(
         secret_manager_client, actor_configuration
     )

airbyte_ops_mcp/constants.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+from enum import Enum
+
 MCP_SERVER_NAME = "airbyte-internal-ops"
 """The name of the MCP server."""
 
@@ -13,3 +15,36 @@ ENV_AIRBYTE_INTERNAL_ADMIN_USER = "AIRBYTE_INTERNAL_ADMIN_USER"
 # Expected values for internal admin authentication
 EXPECTED_ADMIN_FLAG_VALUE = "airbyte.io"
 EXPECTED_ADMIN_EMAIL_DOMAIN = "@airbyte.io"
+
+# =============================================================================
+# GCP and Prod DB Constants (from connection-retriever)
+# =============================================================================
+
+GCP_PROJECT_NAME = "prod-ab-cloud-proj"
+"""The GCP project name for Airbyte Cloud production."""
+
+CLOUD_REGISTRY_URL = (
+    "https://connectors.airbyte.com/files/registries/v0/cloud_registry.json"
+)
+"""URL for the Airbyte Cloud connector registry."""
+
+CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID = (
+    "projects/587336813068/secrets/CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS"
+)
+"""GCP Secret Manager ID for Prod DB 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_secret_retriever.py

@@ -35,7 +35,7 @@ import os
 from dataclasses import replace
 from typing import TYPE_CHECKING
 
-from airbyte_ops_mcp.live_tests._connection_retriever import (
+from airbyte_ops_mcp.connection_config_retriever import (
     ConnectionObject,
     retrieve_objects,
 )

airbyte_ops_mcp/mcp/github_repo_ops.py

@@ -59,6 +59,14 @@ def list_connectors_in_repo(
         set[str] | None,
         "Set of languages to exclude (mutually exclusive with language_filter)",
     ] = None,
+    connector_type: Annotated[
+        Literal["source", "destination"] | None,
+        "Filter by connector type: 'source' or 'destination', None=all",
+    ] = None,
+    connector_subtype: Annotated[
+        Literal["api", "database", "file", "custom"] | None,
+        "Filter by connector subtype: 'api', 'database', 'file', 'custom', None=all",
+    ] = None,
     pr_num_or_url: Annotated[
         str | None,
         "PR number (e.g., '123'), GitHub URL, or None to auto-detect from GITHUB_REF environment variable",
@@ -81,6 +89,8 @@ def list_connectors_in_repo(
         modified=modified,
         language_filter=language_filter,
         language_exclude=language_exclude,
+        connector_type=connector_type,
+        connector_subtype=connector_subtype,
         base_ref=base_ref,
         head_ref=head_ref,
     )

airbyte_ops_mcp/mcp/live_tests.py

@@ -17,6 +17,7 @@ from typing import Annotated, Any
 
 import requests
 from airbyte.cloud import CloudWorkspace
+from airbyte.cloud.auth import resolve_cloud_client_id, resolve_cloud_client_secret
 from fastmcp import FastMCP
 from pydantic import BaseModel, Field
 
@@ -388,15 +389,29 @@ def run_live_connection_tests(
 
     # Validate workspace membership if workspace_id is provided
     if workspace_id:
-        try:
-            workspace = CloudWorkspace(workspace_id=workspace_id)
-            # This will raise an exception if the connection doesn't belong to the workspace
-            workspace.get_connection(connection_id)
-        except Exception as e:
+        client_id = resolve_cloud_client_id()
+        client_secret = resolve_cloud_client_secret()
+        if not client_id or not client_secret:
+            return RunLiveConnectionTestsResponse(
+                run_id=run_id,
+                status=TestRunStatus.FAILED,
+                message=(
+                    "Missing Airbyte Cloud credentials. "
+                    "Set AIRBYTE_CLOUD_CLIENT_ID and AIRBYTE_CLOUD_CLIENT_SECRET env vars."
+                ),
+                workflow_url=None,
+            )
+        workspace = CloudWorkspace(
+            workspace_id=workspace_id,
+            client_id=client_id,
+            client_secret=client_secret,
+        )
+        connection = workspace.get_connection(connection_id)
+        if connection is None:
             return RunLiveConnectionTestsResponse(
                 run_id=run_id,
                 status=TestRunStatus.FAILED,
-                message=f"Connection {connection_id} validation failed for workspace {workspace_id}: {e}",
+                message=f"Connection {connection_id} not found in workspace {workspace_id}",
                 workflow_url=None,
             )