Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions airbyte/_util/api_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -1672,3 +1672,39 @@ def list_workspaces_in_organization(
payload["pagination"]["rowOffset"] += page_size

return result


def get_workspace_organization_info(
workspace_id: str,
*,
api_root: str,
client_id: SecretString,
client_secret: SecretString,
) -> dict[str, Any]:
"""Get organization info for a workspace.

Uses the Config API endpoint: POST /v1/workspaces/get_organization_info

This is an efficient O(1) lookup that directly retrieves the organization
info for a workspace without needing to iterate through all organizations.

Args:
workspace_id: The workspace ID to look up
api_root: The API root URL
client_id: OAuth client ID
client_secret: OAuth client secret

Returns:
Dictionary containing organization info:
- organizationId: The organization ID
- organizationName: The organization name
- sso: Whether SSO is enabled
- billing: Billing information (optional)
"""
return _make_config_api_request(
path="/workspaces/get_organization_info",
json={"workspaceId": workspace_id},
api_root=api_root,
client_id=client_id,
client_secret=client_secret,
)
30 changes: 30 additions & 0 deletions airbyte/cloud/workspaces.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from __future__ import annotations

from dataclasses import dataclass
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

Expand Down Expand Up @@ -83,6 +84,35 @@ def workspace_url(self) -> str | None:
"""The web URL of the workspace."""
return f"{get_web_url_root(self.api_root)}/workspaces/{self.workspace_id}"

@cached_property
def _organization_info(self) -> dict[str, Any]:
"""Fetch and cache organization info for this workspace.

Uses the Config API endpoint for an efficient O(1) lookup.
"""
return api_util.get_workspace_organization_info(
workspace_id=self.workspace_id,
api_root=self.api_root,
client_id=self.client_id,
client_secret=self.client_secret,
)

@property
def organization_id(self) -> str | None:
"""The ID of the organization this workspace belongs to.

This value is cached after the first lookup.
"""
return self._organization_info.get("organizationId")

@property
def organization_name(self) -> str | None:
"""The name of the organization this workspace belongs to.

This value is cached after the first lookup.
"""
return self._organization_info.get("organizationName")

# Test connection and creds

def connect(self) -> None:
Expand Down
34 changes: 25 additions & 9 deletions airbyte/mcp/cloud_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,12 +152,16 @@ class CloudOrganizationResult(BaseModel):
class CloudWorkspaceResult(BaseModel):
"""Information about a workspace in Airbyte Cloud."""

id: str
workspace_id: str
"""The workspace ID."""
name: str
workspace_name: str
"""Display name of the workspace."""
workspace_url: str | None = None
"""URL to access the workspace in Airbyte Cloud."""
organization_id: str
"""ID of the organization this workspace belongs to."""
organization_name: str | None = None
"""Name of the organization this workspace belongs to."""


class LogReadResult(BaseModel):
Expand Down Expand Up @@ -463,18 +467,30 @@ def check_airbyte_cloud_workspace(
default=None,
),
],
) -> str:
) -> CloudWorkspaceResult:
"""Check if we have a valid Airbyte Cloud connection and return workspace info.

Returns workspace ID and workspace URL for verification.
Returns workspace details including workspace ID, name, and organization info.
"""
workspace: CloudWorkspace = _get_cloud_workspace(workspace_id)
workspace.connect()
api_root = resolve_cloud_api_url()
client_id = resolve_cloud_client_id()
client_secret = resolve_cloud_client_secret()

return (
f"✅ Successfully connected to Airbyte Cloud workspace.\n"
f"Workspace ID: {workspace.workspace_id}\n"
f"Workspace URL: {workspace.workspace_url}"
# Get workspace details from the public API
workspace_response = api_util.get_workspace(
workspace_id=workspace.workspace_id,
api_root=api_root,
client_id=client_id,
client_secret=client_secret,
)

return CloudWorkspaceResult(
workspace_id=workspace_response.workspace_id,
workspace_name=workspace_response.name,
workspace_url=workspace.workspace_url,
organization_id=workspace.organization_id or "[error: organization ID not discovered]",
organization_name=workspace.organization_name,
)


Expand Down