airbyte-internal-ops 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

airbyte_ops_mcp/mcp/{live_tests.py → regression_tests.py}

@@ -1,9 +1,16 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
-"""MCP tools for live connection tests.
+"""MCP tools for connector regression tests.
 
-This module provides MCP tools for triggering live validation and regression tests
-on Airbyte Cloud connections via GitHub Actions workflows. Tests run asynchronously
-in GitHub Actions and results can be polled via workflow status.
+This module provides MCP tools for triggering regression tests on Airbyte Cloud
+connections via GitHub Actions workflows. Regression tests can run in two modes:
+- Single version mode: Tests a connector version against a connection config
+- Comparison mode: Compares a target version against a control (baseline) version
+
+Tests run asynchronously in GitHub Actions and results can be polled via workflow status.
+
+Note: The term "regression tests" encompasses all connector validation testing.
+The term "live tests" is reserved for scenarios where actual Cloud connections
+are pinned to pre-release versions for real-world validation.
 """
 
 from __future__ import annotations
@@ -17,6 +24,10 @@ 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 airbyte.exceptions import (
+    AirbyteMissingResourceError,
+    AirbyteWorkspaceMismatchError,
+)
 from fastmcp import FastMCP
 from pydantic import BaseModel, Field
 
@@ -33,13 +44,49 @@ logger = logging.getLogger(__name__)
 # GitHub Workflow Configuration
 # =============================================================================
 
-LIVE_TEST_REPO_OWNER = "airbytehq"
-LIVE_TEST_REPO_NAME = "airbyte-ops-mcp"
-LIVE_TEST_DEFAULT_BRANCH = "main"
-LIVE_TEST_WORKFLOW_FILE = "connector-live-test.yml"
+REGRESSION_TEST_REPO_OWNER = "airbytehq"
+REGRESSION_TEST_REPO_NAME = "airbyte-ops-mcp"
+REGRESSION_TEST_DEFAULT_BRANCH = "main"
+# Unified regression test workflow (handles both single-version and comparison modes)
 REGRESSION_TEST_WORKFLOW_FILE = "connector-regression-test.yml"
 
 
+# =============================================================================
+# Workspace Validation Helpers
+# =============================================================================
+
+
+def validate_connection_workspace(
+    connection_id: str,
+    workspace_id: str,
+) -> None:
+    """Validate that a connection belongs to the expected workspace.
+
+    Uses PyAirbyte's CloudConnection.check_is_valid() method to verify that
+    the connection exists and belongs to the specified workspace.
+
+    Raises:
+        ValueError: If Airbyte Cloud credentials are missing.
+        AirbyteWorkspaceMismatchError: If connection belongs to a different workspace.
+        AirbyteMissingResourceError: If connection is not found.
+    """
+    client_id = resolve_cloud_client_id()
+    client_secret = resolve_cloud_client_secret()
+    if not client_id or not client_secret:
+        raise ValueError(
+            "Missing Airbyte Cloud credentials. "
+            "Set AIRBYTE_CLOUD_CLIENT_ID and AIRBYTE_CLOUD_CLIENT_SECRET env vars."
+        )
+
+    workspace = CloudWorkspace(
+        workspace_id=workspace_id,
+        client_id=client_id,
+        client_secret=client_secret,
+    )
+    connection = workspace.get_connection(connection_id)
+    connection.check_is_valid()
+
+
 def _get_workflow_run_status(
     owner: str,
     repo: str,
@@ -90,8 +137,8 @@ class TestRunStatus(str, Enum):
     FAILED = "failed"
 
 
-class TestPhaseStatus(str, Enum):
-    """Status of a test phase (live or regression)."""
+class TestOutcome(str, Enum):
+    """Outcome of a test (execution or comparison)."""
 
     PENDING = "pending"
     RUNNING = "running"
@@ -130,10 +177,10 @@ class StreamComparisonResultModel(BaseModel):
     message: str = Field(description="Human-readable comparison summary")
 
 
-class LivePhaseResult(BaseModel):
-    """Results from the live test phase."""
+class RegressionTestExecutionResult(BaseModel):
+    """Results from executing the connector (validations and record counts)."""
 
-    status: TestPhaseStatus = Field(description="Status of the live phase")
+    outcome: TestOutcome = Field(description="Outcome of the execution")
     catalog_validations: list[ValidationResultModel] = Field(
         default_factory=list,
         description="Results of catalog validation checks",
@@ -148,14 +195,14 @@ class LivePhaseResult(BaseModel):
     )
     error_message: str | None = Field(
         default=None,
-        description="Error message if the phase failed",
+        description="Error message if the execution failed",
     )
 
 
-class RegressionPhaseResult(BaseModel):
-    """Results from the regression test phase."""
+class RegressionTestComparisonResult(BaseModel):
+    """Results from comparing target vs control connector versions."""
 
-    status: TestPhaseStatus = Field(description="Status of the regression phase")
+    outcome: TestOutcome = Field(description="Outcome of the comparison")
     baseline_version: str | None = Field(
         default=None,
         description="Version of the baseline (control) connector",
@@ -166,12 +213,12 @@ class RegressionPhaseResult(BaseModel):
     )
     error_message: str | None = Field(
         default=None,
-        description="Error message if the phase failed",
+        description="Error message if the comparison failed",
     )
 
 
-class LiveConnectionTestResult(BaseModel):
-    """Complete result of a live connection test run."""
+class RegressionTestResult(BaseModel):
+    """Complete result of a regression test run."""
 
     run_id: str = Field(description="Unique identifier for this test run")
     connection_id: str = Field(description="The connection being tested")
@@ -183,23 +230,23 @@ class LiveConnectionTestResult(BaseModel):
     )
     baseline_version: str | None = Field(
         default=None,
-        description="Version of the baseline connector (if regression ran)",
+        description="Version of the baseline connector (if comparison mode)",
     )
     evaluation_mode: str = Field(
         default="diagnostic",
         description="Evaluation mode used (diagnostic or strict)",
     )
-    skip_regression_tests: bool = Field(
+    compare_versions: bool = Field(
         default=False,
-        description="Whether regression tests were skipped by request",
+        description="Whether comparison mode was used (target vs control)",
     )
-    live_phase: LivePhaseResult | None = Field(
+    execution_result: RegressionTestExecutionResult | None = Field(
         default=None,
-        description="Results from the live test phase",
+        description="Results from executing the connector (validations and record counts)",
     )
-    regression_phase: RegressionPhaseResult | None = Field(
+    comparison_result: RegressionTestComparisonResult | None = Field(
         default=None,
-        description="Results from the regression test phase",
+        description="Results from comparing target vs control connector versions",
     )
     artifacts: dict[str, str] = Field(
         default_factory=dict,
@@ -223,8 +270,8 @@ class LiveConnectionTestResult(BaseModel):
     )
 
 
-class RunLiveConnectionTestsResponse(BaseModel):
-    """Response from starting a live connection test via GitHub Actions workflow."""
+class RunRegressionTestsResponse(BaseModel):
+    """Response from starting a regression test via GitHub Actions workflow."""
 
     run_id: str = Field(
         description="Unique identifier for the test run (internal tracking ID)"
@@ -255,66 +302,57 @@ class RunLiveConnectionTestsResponse(BaseModel):
     idempotent=False,
     open_world=True,
 )
-def run_live_connection_tests(
-    connection_id: Annotated[str, "The Airbyte Cloud connection ID to test"],
-    command: Annotated[
+def run_regression_tests(
+    connector_name: Annotated[
         str,
-        "Airbyte command to run: 'spec', 'check', 'discover', or 'read'",
-    ] = "check",
-    workspace_id: Annotated[
+        "Connector name to build from source (e.g., 'source-pokeapi'). Required.",
+    ],
+    pr: Annotated[
+        int,
+        "PR number from the airbyte monorepo to checkout and build from (e.g., 70847). Required.",
+    ],
+    connection_id: Annotated[
         str | None,
-        "Optional Airbyte Cloud workspace ID. If provided, validates that the connection "
-        "belongs to this workspace before triggering tests. If omitted, no validation is done.",
+        "Airbyte Cloud connection ID to fetch config/catalog from. "
+        "If not provided, uses GSM integration test secrets.",
     ] = None,
-    skip_regression_tests: Annotated[
+    skip_compare: Annotated[
         bool,
-        "If True, run only live tests (connector-live-test workflow). "
-        "If False, run regression tests comparing target vs control versions "
-        "(connector-regression-test workflow).",
-    ] = True,
-    connector_image: Annotated[
-        str | None,
-        "Optional connector image with tag for live tests (e.g., 'airbyte/source-github:1.0.0'). "
-        "If not provided, auto-detected from connection. Only used when skip_regression_tests=True.",
-    ] = None,
-    target_image: Annotated[
+        "If True, skip comparison and run single-version tests only. "
+        "If False (default), run comparison tests (target vs control versions).",
+    ] = False,
+    skip_read_action: Annotated[
+        bool,
+        "If True, skip the read action (run only spec, check, discover). "
+        "If False (default), run all verbs including read.",
+    ] = False,
+    override_test_image: Annotated[
         str | None,
-        "Target connector image (new version) with tag for regression tests "
-        "(e.g., 'airbyte/source-github:2.0.0'). Optional if connector_name is provided. "
-        "Only used when skip_regression_tests=False.",
+        "Override test connector image with tag (e.g., 'airbyte/source-github:1.0.0'). "
+        "Ignored if skip_compare=False.",
     ] = None,
-    control_image: Annotated[
+    override_control_image: Annotated[
         str | None,
-        "Control connector image (baseline version) with tag for regression tests "
-        "(e.g., 'airbyte/source-github:1.0.0'). Optional if connection_id is provided "
-        "(auto-detected from connection). Only used when skip_regression_tests=False.",
+        "Override control connector image (baseline version) with tag. "
+        "Ignored if skip_compare=True.",
     ] = None,
-    connector_name: Annotated[
+    workspace_id: Annotated[
         str | None,
-        "Connector name to build the connector image from source "
-        "(e.g., 'source-pokeapi'). If provided, builds the image locally with tag 'dev'. "
-        "For live tests, this builds the test image. For regression tests, this builds "
-        "the target image while control is auto-detected from the connection.",
+        "Optional Airbyte Cloud workspace ID. If provided with connection_id, validates "
+        "that the connection belongs to this workspace before triggering tests.",
     ] = None,
-    pr: Annotated[
-        int | None,
-        "PR number from the airbyte monorepo to checkout and build from "
-        "(e.g., 70847). Only used when connector_name is provided. "
-        "If not specified, builds from the default branch (master).",
-    ] = None,
-) -> RunLiveConnectionTestsResponse:
-    """Start a live connection test run via GitHub Actions workflow.
+) -> RunRegressionTestsResponse:
+    """Start a regression test run via GitHub Actions workflow.
 
-    This tool triggers either the live-test or regression-test workflow depending
-    on the skip_regression_tests parameter:
+    This tool triggers the regression test workflow which builds the connector
+    from the specified PR and runs tests against it.
 
-    - skip_regression_tests=True (default): Triggers connector-live-test workflow.
-      Runs the specified command against the connection and validates the output.
+    - skip_compare=False (default): Comparison mode - compares the PR version
+      against the baseline (control) version.
+    - skip_compare=True: Single-version mode - runs tests without comparison.
 
-    - skip_regression_tests=False: Triggers connector-regression-test workflow.
-      Compares the target connector version against a control (baseline) version.
-      For regression tests, provide either target_image or connector_name to specify
-      the target version.
+    If connection_id is provided, config/catalog are fetched from Airbyte Cloud.
+    Otherwise, GSM integration test secrets are used.
 
     Returns immediately with a run_id and workflow URL. Check the workflow URL
     to monitor progress and view results.
@@ -329,133 +367,77 @@ def run_live_connection_tests(
     try:
         token = resolve_github_token()
     except ValueError as e:
-        return RunLiveConnectionTestsResponse(
+        return RunRegressionTestsResponse(
             run_id=run_id,
             status=TestRunStatus.FAILED,
             message=str(e),
             workflow_url=None,
         )
 
-    # Validate workspace membership if workspace_id is provided
-    if workspace_id:
-        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} not found in workspace {workspace_id}",
-                workflow_url=None,
-            )
-
-    if skip_regression_tests:
-        # Live test workflow
-        workflow_inputs: dict[str, str] = {
-            "connection_id": connection_id,
-            "command": command,
-        }
-        if connector_image:
-            workflow_inputs["connector_image"] = connector_image
-        if connector_name:
-            workflow_inputs["connector_name"] = connector_name
-            if pr:
-                workflow_inputs["pr"] = str(pr)
-
+    # Validate workspace membership if workspace_id and connection_id are provided
+    if workspace_id and connection_id:
         try:
-            dispatch_result = trigger_workflow_dispatch(
-                owner=LIVE_TEST_REPO_OWNER,
-                repo=LIVE_TEST_REPO_NAME,
-                workflow_file=LIVE_TEST_WORKFLOW_FILE,
-                ref=LIVE_TEST_DEFAULT_BRANCH,
-                inputs=workflow_inputs,
-                token=token,
-            )
-        except Exception as e:
-            logger.exception("Failed to trigger live test workflow")
-            return RunLiveConnectionTestsResponse(
+            validate_connection_workspace(connection_id, workspace_id)
+        except (
+            ValueError,
+            AirbyteWorkspaceMismatchError,
+            AirbyteMissingResourceError,
+        ) as e:
+            return RunRegressionTestsResponse(
                 run_id=run_id,
                 status=TestRunStatus.FAILED,
-                message=f"Failed to trigger live-test workflow: {e}",
+                message=str(e),
                 workflow_url=None,
             )
 
-        view_url = dispatch_result.run_url or dispatch_result.workflow_url
-        return RunLiveConnectionTestsResponse(
-            run_id=run_id,
-            status=TestRunStatus.QUEUED,
-            message=f"Live-test workflow triggered for connection {connection_id}. "
-            f"View progress at: {view_url}",
-            workflow_url=dispatch_result.workflow_url,
-            github_run_id=dispatch_result.run_id,
-            github_run_url=dispatch_result.run_url,
-        )
-
-    # Regression test workflow (skip_regression_tests=False)
-    # Validate that we have enough info to run regression tests
-    if not target_image and not connector_name:
-        return RunLiveConnectionTestsResponse(
-            run_id=run_id,
-            status=TestRunStatus.FAILED,
-            message=(
-                "For regression tests (skip_regression_tests=False), provide either "
-                "target_image or connector_name so the workflow can determine the target image."
-            ),
-            workflow_url=None,
-        )
-
-    workflow_inputs = {
-        "connection_id": connection_id,
-        "command": command,
+    # Build workflow inputs - connector_name and pr are required
+    workflow_inputs: dict[str, str] = {
+        "connector_name": connector_name,
+        "pr": str(pr),
     }
-    if target_image:
-        workflow_inputs["target_image"] = target_image
-    if control_image:
-        workflow_inputs["control_image"] = control_image
-    if connector_name:
-        workflow_inputs["connector_name"] = connector_name
-        if pr:
-            workflow_inputs["pr"] = str(pr)
 
+    # Add optional inputs
+    if connection_id:
+        workflow_inputs["connection_id"] = connection_id
+    if skip_compare:
+        workflow_inputs["skip_compare"] = "true"
+    if skip_read_action:
+        workflow_inputs["skip_read_action"] = "true"
+    if override_test_image:
+        workflow_inputs["override_test_image"] = override_test_image
+    if override_control_image:
+        workflow_inputs["override_control_image"] = override_control_image
+
+    mode_description = "single-version" if skip_compare else "comparison"
     try:
         dispatch_result = trigger_workflow_dispatch(
-            owner=LIVE_TEST_REPO_OWNER,
-            repo=LIVE_TEST_REPO_NAME,
+            owner=REGRESSION_TEST_REPO_OWNER,
+            repo=REGRESSION_TEST_REPO_NAME,
             workflow_file=REGRESSION_TEST_WORKFLOW_FILE,
-            ref=LIVE_TEST_DEFAULT_BRANCH,
+            ref=REGRESSION_TEST_DEFAULT_BRANCH,
             inputs=workflow_inputs,
             token=token,
         )
     except Exception as e:
-        logger.exception("Failed to trigger regression test workflow")
-        return RunLiveConnectionTestsResponse(
+        logger.exception(
+            f"Failed to trigger {mode_description} regression test workflow"
+        )
+        return RunRegressionTestsResponse(
             run_id=run_id,
             status=TestRunStatus.FAILED,
-            message=f"Failed to trigger regression-test workflow: {e}",
+            message=f"Failed to trigger {mode_description} regression test workflow: {e}",
             workflow_url=None,
         )
 
     view_url = dispatch_result.run_url or dispatch_result.workflow_url
-    return RunLiveConnectionTestsResponse(
+    connection_info = f" for connection {connection_id}" if connection_id else ""
+    return RunRegressionTestsResponse(
         run_id=run_id,
         status=TestRunStatus.QUEUED,
-        message=f"Regression-test workflow triggered for connection {connection_id}. "
-        f"View progress at: {view_url}",
+        message=(
+            f"{mode_description.capitalize()} regression test workflow triggered "
+            f"for {connector_name} (PR #{pr}){connection_info}. View progress at: {view_url}"
+        ),
         workflow_url=dispatch_result.workflow_url,
         github_run_id=dispatch_result.run_id,
         github_run_url=dispatch_result.run_url,
@@ -467,8 +449,8 @@ def run_live_connection_tests(
 # =============================================================================
 
 
-def register_live_tests_tools(app: FastMCP) -> None:
-    """Register live tests tools with the FastMCP app.
+def register_regression_tests_tools(app: FastMCP) -> None:
+    """Register regression tests tools with the FastMCP app.
 
     Args:
         app: FastMCP application instance

airbyte_ops_mcp/mcp/server.py

@@ -26,10 +26,10 @@ from airbyte_ops_mcp.mcp.cloud_connector_versions import (
 )
 from airbyte_ops_mcp.mcp.github import register_github_tools
 from airbyte_ops_mcp.mcp.github_repo_ops import register_github_repo_ops_tools
-from airbyte_ops_mcp.mcp.live_tests import register_live_tests_tools
 from airbyte_ops_mcp.mcp.prerelease import register_prerelease_tools
 from airbyte_ops_mcp.mcp.prod_db_queries import register_prod_db_query_tools
 from airbyte_ops_mcp.mcp.prompts import register_prompts
+from airbyte_ops_mcp.mcp.regression_tests import register_regression_tests_tools
 from airbyte_ops_mcp.mcp.server_info import register_server_info_resources
 
 # Default HTTP server configuration
@@ -47,7 +47,7 @@ def register_server_assets(app: FastMCP) -> None:
     - REPO: GitHub repository operations
     - CLOUD: Cloud connector version management
     - PROMPTS: Prompt templates for common workflows
-    - LIVE_TESTS: Live connection validation and regression tests
+    - REGRESSION_TESTS: Connector regression tests (single-version and comparison)
     - REGISTRY: Connector registry operations (future)
     - METADATA: Connector metadata operations (future)
     - QA: Connector quality assurance (future)
@@ -63,7 +63,7 @@ def register_server_assets(app: FastMCP) -> None:
     register_cloud_connector_version_tools(app)
     register_prod_db_query_tools(app)
     register_prompts(app)
-    register_live_tests_tools(app)
+    register_regression_tests_tools(app)
 
 
 register_server_assets(app)

airbyte_ops_mcp/prod_db_access/db_engine.py

@@ -166,24 +166,20 @@ def _get_connector() -> Connector:
 
 
 def _get_secret_value(
-    gsm_client: secretmanager.SecretManagerServiceClient, secret_id: str
+    gsm_client: secretmanager.SecretManagerServiceClient,
+    secret_id: str,
 ) -> str:
-    """Get the value of the enabled version of a secret.
+    """Get the value of the latest version of a secret.
 
     Args:
         gsm_client: GCP Secret Manager client
-        secret_id: The id of the secret
+        secret_id: The full resource ID of the secret
+            (e.g., "projects/123/secrets/my-secret")
 
     Returns:
-        The value of the enabled version of the secret
+        The value of the latest version of the secret
     """
-    response = gsm_client.list_secret_versions(
-        request={"parent": secret_id, "filter": "state:ENABLED"}
-    )
-    if len(response.versions) == 0:
-        raise ValueError(f"No enabled version of secret {secret_id} found")
-    enabled_version = response.versions[0]
-    response = gsm_client.access_secret_version(name=enabled_version.name)
+    response = gsm_client.access_secret_version(name=f"{secret_id}/versions/latest")
     return response.payload.data.decode("UTF-8")
 
 

airbyte_ops_mcp/prod_db_access/queries.py

@@ -22,6 +22,8 @@ from airbyte_ops_mcp.prod_db_access.sql import (
     SELECT_ACTORS_PINNED_TO_VERSION,
     SELECT_CONNECTIONS_BY_CONNECTOR,
     SELECT_CONNECTIONS_BY_CONNECTOR_AND_ORG,
+    SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR,
+    SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR_AND_ORG,
     SELECT_CONNECTOR_VERSIONS,
     SELECT_DATAPLANES_LIST,
     SELECT_FAILED_SYNC_ATTEMPTS_FOR_CONNECTOR,
@@ -30,6 +32,7 @@ from airbyte_ops_mcp.prod_db_access.sql import (
     SELECT_SUCCESSFUL_SYNCS_FOR_VERSION,
     SELECT_SYNC_RESULTS_FOR_VERSION,
     SELECT_WORKSPACE_INFO,
+    SELECT_WORKSPACES_BY_EMAIL_DOMAIN,
 )
 
 logger = logging.getLogger(__name__)
@@ -113,6 +116,48 @@ def query_connections_by_connector(
     )
 
 
+def query_connections_by_destination_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 destination connector type, optionally filtered by organization.
+
+    Args:
+        connector_definition_id: Destination 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
+    """
+    # Use separate queries to avoid pg8000 NULL parameter type issues
+    if organization_id is None:
+        return _run_sql_query(
+            SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR,
+            parameters={
+                "connector_definition_id": connector_definition_id,
+                "limit": limit,
+            },
+            query_name="SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR",
+            gsm_client=gsm_client,
+        )
+
+    return _run_sql_query(
+        SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR_AND_ORG,
+        parameters={
+            "connector_definition_id": connector_definition_id,
+            "organization_id": organization_id,
+            "limit": limit,
+        },
+        query_name="SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR_AND_ORG",
+        gsm_client=gsm_client,
+    )
+
+
 def query_connector_versions(
     connector_definition_id: str,
     *,
@@ -337,3 +382,37 @@ def query_org_workspaces(
         query_name="SELECT_ORG_WORKSPACES",
         gsm_client=gsm_client,
     )
+
+
+def query_workspaces_by_email_domain(
+    email_domain: str,
+    limit: int = 100,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query workspaces by email domain.
+
+    This is useful for identifying workspaces based on user email domains.
+    For example, searching for "motherduck.com" will find workspaces where users have
+    @motherduck.com email addresses, which may belong to partner accounts.
+
+    Args:
+        email_domain: Email domain to search for (e.g., "motherduck.com", "fivetran.com").
+            Do not include the "@" symbol.
+        limit: Maximum number of results (default: 100)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of workspace records with organization_id, workspace_id, workspace_name,
+        slug, email, dataplane_group_id, dataplane_name, and created_at.
+        Results are ordered by organization_id and workspace_name.
+    """
+    # Strip leading @ if provided
+    clean_domain = email_domain.lstrip("@")
+
+    return _run_sql_query(
+        SELECT_WORKSPACES_BY_EMAIL_DOMAIN,
+        parameters={"email_domain": clean_domain, "limit": limit},
+        query_name="SELECT_WORKSPACES_BY_EMAIL_DOMAIN",
+        gsm_client=gsm_client,
+    )