airbyte-internal-ops 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

airbyte_ops_mcp/cloud_admin/api_client.py

@@ -15,6 +15,9 @@ from airbyte import constants
 from airbyte.exceptions import PyAirbyteInputError
 
 from airbyte_ops_mcp import constants as ops_constants
+from airbyte_ops_mcp.mcp.prod_db_queries import (
+    _resolve_canonical_name_to_definition_id,
+)
 
 # Internal enums for scoped configuration API "magic strings"
 # These values caused issues during development and are now centralized here
@@ -959,3 +962,473 @@ def set_connector_version_override(
             )
 
         return True
+
+
+def set_workspace_connector_version_override(
+    workspace_id: str,
+    connector_name: str,
+    connector_type: Literal["source", "destination"],
+    api_root: str,
+    client_id: str | None = None,
+    client_secret: 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 workspace-level version override for a connector type.
+
+    This pins ALL instances of a connector type within a workspace to a specific version.
+    For example, pinning 'source-github' at workspace level means all GitHub sources
+    in that workspace will use the pinned version.
+
+    Args:
+        workspace_id: The workspace ID
+        connector_name: The connector name (e.g., 'source-github')
+        connector_type: Either "source" or "destination"
+        api_root: The API root URL
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        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)
+
+    Raises:
+        PyAirbyteInputError: If the API request fails or parameters are invalid
+    """
+    # Input validation
+    if (version is None) == (not unset):
+        raise PyAirbyteInputError(
+            message="Must specify EXACTLY ONE of version (to set) OR unset=True (to clear), but not both",
+        )
+
+    if not unset and (not override_reason or len(override_reason.strip()) < 10):
+        raise PyAirbyteInputError(
+            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, bearer_token)
+
+    # Resolve connector name to actor_definition_id using the shared registry lookup
+    actor_definition_id = _resolve_canonical_name_to_definition_id(connector_name)
+
+    if unset:
+        # Get the existing workspace-level configuration
+        active_config = _get_scoped_configuration_context(
+            actor_definition_id=actor_definition_id,
+            scope_type=_ScopeType.WORKSPACE,
+            scope_id=workspace_id,
+            api_root=api_root,
+            access_token=access_token,
+        )
+
+        if not active_config:
+            return False
+
+        # Verify this is actually a workspace-scoped config (not inherited from org)
+        api_scope_type = active_config.get("scope_type", "").lower()
+        api_scope_id = active_config.get("scope_id", "")
+        if api_scope_type != _ScopeType.WORKSPACE.value or api_scope_id != workspace_id:
+            raise PyAirbyteInputError(
+                message=f"Cannot delete: the active config is not workspace-scoped. "
+                f"Expected scope_type='{_ScopeType.WORKSPACE.value}' and scope_id='{workspace_id}', "
+                f"but got scope_type='{api_scope_type}' and scope_id='{api_scope_id}'. "
+                f"This may be an inherited config from organization level.",
+                context={
+                    "workspace_id": workspace_id,
+                    "expected_scope_type": _ScopeType.WORKSPACE.value,
+                    "actual_scope_type": api_scope_type,
+                    "actual_scope_id": api_scope_id,
+                    "full_config": active_config,
+                },
+            )
+
+        # Delete the configuration
+        delete_endpoint = f"{api_root}/scoped_configuration/delete"
+        delete_payload = {"scopedConfigurationId": active_config["id"]}
+
+        response = requests.post(
+            delete_endpoint,
+            json=delete_payload,
+            headers={
+                "Authorization": f"Bearer {access_token}",
+                "User-Agent": ops_constants.USER_AGENT,
+                "Content-Type": "application/json",
+            },
+            timeout=30,
+        )
+
+        if response.status_code not in (200, 204):
+            raise PyAirbyteInputError(
+                message=f"Failed to delete workspace version override: {response.status_code} {response.text}",
+                context={
+                    "delete_endpoint": delete_endpoint,
+                    "config_id": active_config["id"],
+                    "status_code": response.status_code,
+                    "response": response.text,
+                },
+            )
+
+        return True
+
+    # Set a new workspace-level override
+    # Resolve version string to version ID
+    version_id = resolve_connector_version_id(
+        actor_definition_id=actor_definition_id,
+        connector_type=connector_type,
+        version=version,
+        api_root=api_root,
+        bearer_token=access_token,
+    )
+
+    # Check for existing workspace-level configuration
+    existing_config = _get_scoped_configuration_context(
+        actor_definition_id=actor_definition_id,
+        scope_type=_ScopeType.WORKSPACE,
+        scope_id=workspace_id,
+        api_root=api_root,
+        access_token=access_token,
+    )
+
+    if existing_config:
+        existing_version_id = existing_config.get("value")
+        existing_version_name = existing_config.get("value_name", "unknown")
+
+        # If already pinned to the same version, no action needed
+        if existing_version_id == version_id:
+            raise PyAirbyteInputError(
+                message=f"Workspace is already pinned to version {existing_version_name} for {connector_name}. "
+                f"Use unset=True first if you want to re-pin to a different version.",
+                context={
+                    "workspace_id": workspace_id,
+                    "connector_name": connector_name,
+                    "existing_version": existing_version_name,
+                    "requested_version": version,
+                },
+            )
+
+        # Verify this is a workspace-scoped config before deleting
+        api_scope_type = existing_config.get("scope_type", "").lower()
+        api_scope_id = existing_config.get("scope_id", "")
+        if (
+            api_scope_type == _ScopeType.WORKSPACE.value
+            and api_scope_id == workspace_id
+        ):
+            # Delete existing workspace-level config before creating new one
+            delete_endpoint = f"{api_root}/scoped_configuration/delete"
+            delete_payload = {"scopedConfigurationId": existing_config["id"]}
+
+            delete_response = requests.post(
+                delete_endpoint,
+                json=delete_payload,
+                headers={
+                    "Authorization": f"Bearer {access_token}",
+                    "User-Agent": ops_constants.USER_AGENT,
+                    "Content-Type": "application/json",
+                },
+                timeout=30,
+            )
+
+            if delete_response.status_code not in (200, 204):
+                raise PyAirbyteInputError(
+                    message=f"Failed to delete existing workspace version override: "
+                    f"{delete_response.status_code} {delete_response.text}",
+                )
+
+    # Get user ID from email
+    if not user_email:
+        raise PyAirbyteInputError(
+            message="user_email is required to set a version override",
+        )
+    origin = get_user_id_by_email(
+        email=user_email,
+        api_root=api_root,
+        bearer_token=access_token,
+    )
+
+    # Create the override
+    endpoint = f"{api_root}/scoped_configuration/create"
+    payload: dict[str, Any] = {
+        "config_key": _ScopedConfigKey.CONNECTOR_VERSION.value,
+        "resource_type": _ResourceType.ACTOR_DEFINITION.value,
+        "resource_id": actor_definition_id,
+        "scope_type": _ScopeType.WORKSPACE.value,
+        "scope_id": workspace_id,
+        "value": version_id,
+        "description": override_reason,
+        "origin_type": _OriginType.USER.value,
+        "origin": origin,
+    }
+
+    if override_reason_reference_url:
+        payload["reference_url"] = override_reason_reference_url
+
+    response = requests.post(
+        endpoint,
+        json=payload,
+        headers={
+            "Authorization": f"Bearer {access_token}",
+            "User-Agent": ops_constants.USER_AGENT,
+            "Content-Type": "application/json",
+        },
+        timeout=30,
+    )
+
+    if response.status_code not in (200, 201):
+        raise PyAirbyteInputError(
+            message=f"Failed to set workspace version override: {response.status_code} {response.text}",
+            context={
+                "workspace_id": workspace_id,
+                "connector_name": connector_name,
+                "version": version,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    return True
+
+
+def set_organization_connector_version_override(
+    organization_id: str,
+    connector_name: str,
+    connector_type: Literal["source", "destination"],
+    api_root: str,
+    client_id: str | None = None,
+    client_secret: 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 an organization-level version override for a connector type.
+
+    This pins ALL instances of a connector type across an entire organization to a
+    specific version. For example, pinning 'source-github' at organization level means
+    all GitHub sources in all workspaces within that organization will use the pinned version.
+
+    Args:
+        organization_id: The organization ID
+        connector_name: The connector name (e.g., 'source-github')
+        connector_type: Either "source" or "destination"
+        api_root: The API root URL
+        client_id: The Airbyte Cloud client ID (required if no bearer_token)
+        client_secret: The Airbyte Cloud client secret (required if no bearer_token)
+        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)
+
+    Raises:
+        PyAirbyteInputError: If the API request fails or parameters are invalid
+    """
+    # Input validation
+    if (version is None) == (not unset):
+        raise PyAirbyteInputError(
+            message="Must specify EXACTLY ONE of version (to set) OR unset=True (to clear), but not both",
+        )
+
+    if not unset and (not override_reason or len(override_reason.strip()) < 10):
+        raise PyAirbyteInputError(
+            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, bearer_token)
+
+    # Resolve connector name to actor_definition_id using the shared registry lookup
+    actor_definition_id = _resolve_canonical_name_to_definition_id(connector_name)
+
+    if unset:
+        # Get the existing organization-level configuration
+        active_config = _get_scoped_configuration_context(
+            actor_definition_id=actor_definition_id,
+            scope_type=_ScopeType.ORGANIZATION,
+            scope_id=organization_id,
+            api_root=api_root,
+            access_token=access_token,
+        )
+
+        if not active_config:
+            return False
+
+        # Verify this is actually an organization-scoped config
+        api_scope_type = active_config.get("scope_type", "").lower()
+        api_scope_id = active_config.get("scope_id", "")
+        if (
+            api_scope_type != _ScopeType.ORGANIZATION.value
+            or api_scope_id != organization_id
+        ):
+            raise PyAirbyteInputError(
+                message=f"Cannot delete: the active config is not organization-scoped. "
+                f"Expected scope_type='{_ScopeType.ORGANIZATION.value}' and scope_id='{organization_id}', "
+                f"but got scope_type='{api_scope_type}' and scope_id='{api_scope_id}'.",
+                context={
+                    "organization_id": organization_id,
+                    "expected_scope_type": _ScopeType.ORGANIZATION.value,
+                    "actual_scope_type": api_scope_type,
+                    "actual_scope_id": api_scope_id,
+                    "full_config": active_config,
+                },
+            )
+
+        # Delete the configuration
+        delete_endpoint = f"{api_root}/scoped_configuration/delete"
+        delete_payload = {"scopedConfigurationId": active_config["id"]}
+
+        response = requests.post(
+            delete_endpoint,
+            json=delete_payload,
+            headers={
+                "Authorization": f"Bearer {access_token}",
+                "User-Agent": ops_constants.USER_AGENT,
+                "Content-Type": "application/json",
+            },
+            timeout=30,
+        )
+
+        if response.status_code not in (200, 204):
+            raise PyAirbyteInputError(
+                message=f"Failed to delete organization version override: {response.status_code} {response.text}",
+                context={
+                    "delete_endpoint": delete_endpoint,
+                    "config_id": active_config["id"],
+                    "status_code": response.status_code,
+                    "response": response.text,
+                },
+            )
+
+        return True
+
+    # Set a new organization-level override
+    # Resolve version string to version ID
+    version_id = resolve_connector_version_id(
+        actor_definition_id=actor_definition_id,
+        connector_type=connector_type,
+        version=version,
+        api_root=api_root,
+        bearer_token=access_token,
+    )
+
+    # Check for existing organization-level configuration
+    existing_config = _get_scoped_configuration_context(
+        actor_definition_id=actor_definition_id,
+        scope_type=_ScopeType.ORGANIZATION,
+        scope_id=organization_id,
+        api_root=api_root,
+        access_token=access_token,
+    )
+
+    if existing_config:
+        existing_version_id = existing_config.get("value")
+        existing_version_name = existing_config.get("value_name", "unknown")
+
+        # If already pinned to the same version, no action needed
+        if existing_version_id == version_id:
+            raise PyAirbyteInputError(
+                message=f"Organization is already pinned to version {existing_version_name} for {connector_name}. "
+                f"Use unset=True first if you want to re-pin to a different version.",
+                context={
+                    "organization_id": organization_id,
+                    "connector_name": connector_name,
+                    "existing_version": existing_version_name,
+                    "requested_version": version,
+                },
+            )
+
+        # Verify this is an organization-scoped config before deleting
+        api_scope_type = existing_config.get("scope_type", "").lower()
+        api_scope_id = existing_config.get("scope_id", "")
+        if (
+            api_scope_type == _ScopeType.ORGANIZATION.value
+            and api_scope_id == organization_id
+        ):
+            # Delete existing organization-level config before creating new one
+            delete_endpoint = f"{api_root}/scoped_configuration/delete"
+            delete_payload = {"scopedConfigurationId": existing_config["id"]}
+
+            delete_response = requests.post(
+                delete_endpoint,
+                json=delete_payload,
+                headers={
+                    "Authorization": f"Bearer {access_token}",
+                    "User-Agent": ops_constants.USER_AGENT,
+                    "Content-Type": "application/json",
+                },
+                timeout=30,
+            )
+
+            if delete_response.status_code not in (200, 204):
+                raise PyAirbyteInputError(
+                    message=f"Failed to delete existing organization version override: "
+                    f"{delete_response.status_code} {delete_response.text}",
+                )
+
+    # Get user ID from email
+    if not user_email:
+        raise PyAirbyteInputError(
+            message="user_email is required to set a version override",
+        )
+    origin = get_user_id_by_email(
+        email=user_email,
+        api_root=api_root,
+        bearer_token=access_token,
+    )
+
+    # Create the override
+    endpoint = f"{api_root}/scoped_configuration/create"
+    payload: dict[str, Any] = {
+        "config_key": _ScopedConfigKey.CONNECTOR_VERSION.value,
+        "resource_type": _ResourceType.ACTOR_DEFINITION.value,
+        "resource_id": actor_definition_id,
+        "scope_type": _ScopeType.ORGANIZATION.value,
+        "scope_id": organization_id,
+        "value": version_id,
+        "description": override_reason,
+        "origin_type": _OriginType.USER.value,
+        "origin": origin,
+    }
+
+    if override_reason_reference_url:
+        payload["reference_url"] = override_reason_reference_url
+
+    response = requests.post(
+        endpoint,
+        json=payload,
+        headers={
+            "Authorization": f"Bearer {access_token}",
+            "User-Agent": ops_constants.USER_AGENT,
+            "Content-Type": "application/json",
+        },
+        timeout=30,
+    )
+
+    if response.status_code not in (200, 201):
+        raise PyAirbyteInputError(
+            message=f"Failed to set organization version override: {response.status_code} {response.text}",
+            context={
+                "organization_id": organization_id,
+                "connector_name": connector_name,
+                "version": version,
+                "endpoint": endpoint,
+                "status_code": response.status_code,
+                "response": response.text,
+            },
+        )
+
+    return True

airbyte_ops_mcp/cloud_admin/models.py

@@ -71,3 +71,59 @@ class VersionOverrideOperationResult(BaseModel):
         if self.success:
             return f"✓ {self.message}"
         return f"✗ {self.message}"
+
+
+class WorkspaceVersionOverrideResult(BaseModel):
+    """Result of a workspace-level version override operation.
+
+    This model provides detailed information about the outcome of a workspace-level
+    version pinning or unpinning operation.
+    """
+
+    success: bool = Field(description="Whether the operation succeeded")
+    message: str = Field(description="Human-readable message describing the result")
+    workspace_id: str = Field(description="The workspace ID")
+    connector_name: str = Field(
+        description="The connector name (e.g., 'source-github')"
+    )
+    connector_type: Literal["source", "destination"] = Field(
+        description="The type of connector (source or destination)"
+    )
+    version: str | None = Field(
+        default=None,
+        description="The version that was pinned (None if cleared or failed)",
+    )
+
+    def __str__(self) -> str:
+        """Return a string representation of the operation result."""
+        if self.success:
+            return f"✓ {self.message}"
+        return f"✗ {self.message}"
+
+
+class OrganizationVersionOverrideResult(BaseModel):
+    """Result of an organization-level version override operation.
+
+    This model provides detailed information about the outcome of an organization-level
+    version pinning or unpinning operation.
+    """
+
+    success: bool = Field(description="Whether the operation succeeded")
+    message: str = Field(description="Human-readable message describing the result")
+    organization_id: str = Field(description="The organization ID")
+    connector_name: str = Field(
+        description="The connector name (e.g., 'source-github')"
+    )
+    connector_type: Literal["source", "destination"] = Field(
+        description="The type of connector (source or destination)"
+    )
+    version: str | None = Field(
+        default=None,
+        description="The version that was pinned (None if cleared or failed)",
+    )
+
+    def __str__(self) -> str:
+        """Return a string representation of the operation result."""
+        if self.success:
+            return f"✓ {self.message}"
+        return f"✗ {self.message}"