Skip to content

Commit 2d495ba

Browse files
aaronsteersdevin-ai-integration[bot]
aaronsteers
and
devin-ai-integration[bot]
authored
feat(mcp): add list_cloud_workspaces and describe_cloud_organization tools (#883)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1 parent aac9cae commit 2d495ba

File tree

2 files changed

+355
-1
lines changed

2 files changed

+355
-1
lines changed

‎airbyte/_util/api_util.py‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1562,3 +1562,113 @@ def get_connector_builder_project_for_definition_id(
15621562
client_secret=client_secret,
15631563
)
15641564
return json_result.get("builderProjectId")
1565+
1566+
1567+
# Organization and workspace listing
1568+
1569+
1570+
def list_organizations_for_user(
1571+
*,
1572+
api_root: str,
1573+
client_id: SecretString,
1574+
client_secret: SecretString,
1575+
) -> list[models.OrganizationResponse]:
1576+
"""List all organizations accessible to the current user.
1577+
1578+
Uses the public API endpoint: GET /organizations
1579+
1580+
Args:
1581+
api_root: The API root URL
1582+
client_id: OAuth client ID
1583+
client_secret: OAuth client secret
1584+
1585+
Returns:
1586+
List of OrganizationResponse objects containing organization_id, organization_name, email
1587+
"""
1588+
airbyte_instance = get_airbyte_server_instance(
1589+
api_root=api_root,
1590+
client_id=client_id,
1591+
client_secret=client_secret,
1592+
)
1593+
response = airbyte_instance.organizations.list_organizations_for_user()
1594+
1595+
if status_ok(response.status_code) and response.organizations_response:
1596+
return response.organizations_response.data
1597+
1598+
raise AirbyteError(
1599+
message="Failed to list organizations for user.",
1600+
context={
1601+
"status_code": response.status_code,
1602+
"response": response,
1603+
},
1604+
)
1605+
1606+
1607+
def list_workspaces_in_organization(
1608+
organization_id: str,
1609+
*,
1610+
api_root: str,
1611+
client_id: SecretString,
1612+
client_secret: SecretString,
1613+
name_contains: str | None = None,
1614+
max_items_limit: int | None = None,
1615+
) -> list[dict[str, Any]]:
1616+
"""List workspaces within a specific organization.
1617+
1618+
Uses the Config API endpoint: POST /v1/workspaces/list_by_organization_id
1619+
1620+
Args:
1621+
organization_id: The organization ID to list workspaces for
1622+
api_root: The API root URL
1623+
client_id: OAuth client ID
1624+
client_secret: OAuth client secret
1625+
name_contains: Optional substring filter for workspace names (server-side)
1626+
max_items_limit: Optional maximum number of workspaces to return
1627+
1628+
Returns:
1629+
List of workspace dictionaries containing workspaceId, organizationId, name, slug, etc.
1630+
"""
1631+
result: list[dict[str, Any]] = []
1632+
page_size = 100
1633+
1634+
# Build base payload
1635+
payload: dict[str, Any] = {
1636+
"organizationId": organization_id,
1637+
"pagination": {
1638+
"pageSize": page_size,
1639+
"rowOffset": 0,
1640+
},
1641+
}
1642+
if name_contains:
1643+
payload["nameContains"] = name_contains
1644+
1645+
# Fetch pages until we have all results or reach the limit
1646+
while True:
1647+
json_result = _make_config_api_request(
1648+
path="/workspaces/list_by_organization_id",
1649+
json=payload,
1650+
api_root=api_root,
1651+
client_id=client_id,
1652+
client_secret=client_secret,
1653+
)
1654+
1655+
workspaces = json_result.get("workspaces", [])
1656+
1657+
# If no results returned, we've exhausted all pages
1658+
if not workspaces:
1659+
break
1660+
1661+
result.extend(workspaces)
1662+
1663+
# Check if we've reached the limit
1664+
if max_items_limit is not None and len(result) >= max_items_limit:
1665+
return result[:max_items_limit]
1666+
1667+
# If we got fewer results than page_size, this was the last page
1668+
if len(workspaces) < page_size:
1669+
break
1670+
1671+
# Bump offset for next iteration
1672+
payload["pagination"]["rowOffset"] += page_size
1673+
1674+
return result

‎airbyte/mcp/cloud_ops.py‎

Lines changed: 245 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from pydantic import BaseModel, Field
99

1010
from airbyte import cloud, get_destination, get_source
11-
from airbyte._util.api_util import PyAirbyteInputError
11+
from airbyte._util import api_util
1212
from airbyte.cloud.auth import (
1313
resolve_cloud_api_url,
1414
resolve_cloud_client_id,
@@ -18,13 +18,15 @@
1818
from airbyte.cloud.connectors import CustomCloudSourceDefinition
1919
from airbyte.cloud.workspaces import CloudWorkspace
2020
from airbyte.destinations.util import get_noop_destination
21+
from airbyte.exceptions import AirbyteMissingResourceError, PyAirbyteInputError
2122
from airbyte.mcp._tool_utils import (
2223
check_guid_created_in_session,
2324
mcp_tool,
2425
register_guid_created_in_session,
2526
register_tools,
2627
)
2728
from airbyte.mcp._util import resolve_config, resolve_list_of_strings
29+
from airbyte.secrets import SecretString
2830

2931

3032
CLOUD_AUTH_TIP_TEXT = (
@@ -121,6 +123,28 @@ class CloudConnectionDetails(BaseModel):
121123
"""Table prefix applied when syncing to the destination."""
122124

123125

126+
class CloudOrganizationResult(BaseModel):
127+
"""Information about an organization in Airbyte Cloud."""
128+
129+
id: str
130+
"""The organization ID."""
131+
name: str
132+
"""Display name of the organization."""
133+
email: str
134+
"""Email associated with the organization."""
135+
136+
137+
class CloudWorkspaceResult(BaseModel):
138+
"""Information about a workspace in Airbyte Cloud."""
139+
140+
id: str
141+
"""The workspace ID."""
142+
name: str
143+
"""Display name of the workspace."""
144+
organization_id: str
145+
"""ID of the organization this workspace belongs to."""
146+
147+
124148
def _get_cloud_workspace(workspace_id: str | None = None) -> CloudWorkspace:
125149
"""Get an authenticated CloudWorkspace.
126150
@@ -852,6 +876,226 @@ def list_deployed_cloud_connections(
852876
]
853877

854878

879+
def _resolve_organization(
880+
organization_id: str | None,
881+
organization_name: str | None,
882+
*,
883+
api_root: str,
884+
client_id: SecretString,
885+
client_secret: SecretString,
886+
) -> api_util.models.OrganizationResponse:
887+
"""Resolve organization from either ID or exact name match.
888+
889+
Args:
890+
organization_id: The organization ID (if provided directly)
891+
organization_name: The organization name (exact match required)
892+
api_root: The API root URL
893+
client_id: OAuth client ID
894+
client_secret: OAuth client secret
895+
896+
Returns:
897+
The resolved OrganizationResponse object
898+
899+
Raises:
900+
PyAirbyteInputError: If neither or both parameters are provided,
901+
or if no organization matches the exact name
902+
AirbyteMissingResourceError: If the organization is not found
903+
"""
904+
if organization_id and organization_name:
905+
raise PyAirbyteInputError(
906+
message="Provide either 'organization_id' or 'organization_name', not both."
907+
)
908+
if not organization_id and not organization_name:
909+
raise PyAirbyteInputError(
910+
message="Either 'organization_id' or 'organization_name' must be provided."
911+
)
912+
913+
# Get all organizations for the user
914+
orgs = api_util.list_organizations_for_user(
915+
api_root=api_root,
916+
client_id=client_id,
917+
client_secret=client_secret,
918+
)
919+
920+
if organization_id:
921+
# Find by ID
922+
matching_orgs = [org for org in orgs if org.organization_id == organization_id]
923+
if not matching_orgs:
924+
raise AirbyteMissingResourceError(
925+
resource_type="organization",
926+
context={
927+
"organization_id": organization_id,
928+
"message": f"No organization found with ID '{organization_id}' "
929+
"for the current user.",
930+
},
931+
)
932+
return matching_orgs[0]
933+
934+
# Find by exact name match (case-sensitive)
935+
matching_orgs = [org for org in orgs if org.organization_name == organization_name]
936+
937+
if not matching_orgs:
938+
raise AirbyteMissingResourceError(
939+
resource_type="organization",
940+
context={
941+
"organization_name": organization_name,
942+
"message": f"No organization found with exact name '{organization_name}' "
943+
"for the current user.",
944+
},
945+
)
946+
947+
if len(matching_orgs) > 1:
948+
raise PyAirbyteInputError(
949+
message=f"Multiple organizations found with name '{organization_name}'. "
950+
"Please use 'organization_id' instead to specify the exact organization."
951+
)
952+
953+
return matching_orgs[0]
954+
955+
956+
def _resolve_organization_id(
957+
organization_id: str | None,
958+
organization_name: str | None,
959+
*,
960+
api_root: str,
961+
client_id: SecretString,
962+
client_secret: SecretString,
963+
) -> str:
964+
"""Resolve organization ID from either ID or exact name match.
965+
966+
This is a convenience wrapper around _resolve_organization that returns just the ID.
967+
"""
968+
org = _resolve_organization(
969+
organization_id=organization_id,
970+
organization_name=organization_name,
971+
api_root=api_root,
972+
client_id=client_id,
973+
client_secret=client_secret,
974+
)
975+
return org.organization_id
976+
977+
978+
@mcp_tool(
979+
domain="cloud",
980+
read_only=True,
981+
idempotent=True,
982+
open_world=True,
983+
extra_help_text=CLOUD_AUTH_TIP_TEXT,
984+
)
985+
def list_cloud_workspaces(
986+
*,
987+
organization_id: Annotated[
988+
str | None,
989+
Field(
990+
description="Organization ID. Required if organization_name is not provided.",
991+
default=None,
992+
),
993+
],
994+
organization_name: Annotated[
995+
str | None,
996+
Field(
997+
description=(
998+
"Organization name (exact match). " "Required if organization_id is not provided."
999+
),
1000+
default=None,
1001+
),
1002+
],
1003+
name_contains: Annotated[
1004+
str | None,
1005+
"Optional substring to filter workspaces by name (server-side filtering)",
1006+
] = None,
1007+
max_items_limit: Annotated[
1008+
int | None,
1009+
"Optional maximum number of items to return (default: no limit)",
1010+
] = None,
1011+
) -> list[CloudWorkspaceResult]:
1012+
"""List all workspaces in a specific organization.
1013+
1014+
Requires either organization_id OR organization_name (exact match) to be provided.
1015+
This tool will NOT list workspaces across all organizations - you must specify
1016+
which organization to list workspaces from.
1017+
"""
1018+
api_root = resolve_cloud_api_url()
1019+
client_id = resolve_cloud_client_id()
1020+
client_secret = resolve_cloud_client_secret()
1021+
1022+
resolved_org_id = _resolve_organization_id(
1023+
organization_id=organization_id,
1024+
organization_name=organization_name,
1025+
api_root=api_root,
1026+
client_id=client_id,
1027+
client_secret=client_secret,
1028+
)
1029+
1030+
workspaces = api_util.list_workspaces_in_organization(
1031+
organization_id=resolved_org_id,
1032+
api_root=api_root,
1033+
client_id=client_id,
1034+
client_secret=client_secret,
1035+
name_contains=name_contains,
1036+
max_items_limit=max_items_limit,
1037+
)
1038+
1039+
return [
1040+
CloudWorkspaceResult(
1041+
id=ws.get("workspaceId", ""),
1042+
name=ws.get("name", ""),
1043+
organization_id=ws.get("organizationId", ""),
1044+
)
1045+
for ws in workspaces
1046+
]
1047+
1048+
1049+
@mcp_tool(
1050+
domain="cloud",
1051+
read_only=True,
1052+
idempotent=True,
1053+
open_world=True,
1054+
extra_help_text=CLOUD_AUTH_TIP_TEXT,
1055+
)
1056+
def describe_cloud_organization(
1057+
*,
1058+
organization_id: Annotated[
1059+
str | None,
1060+
Field(
1061+
description="Organization ID. Required if organization_name is not provided.",
1062+
default=None,
1063+
),
1064+
],
1065+
organization_name: Annotated[
1066+
str | None,
1067+
Field(
1068+
description=(
1069+
"Organization name (exact match). " "Required if organization_id is not provided."
1070+
),
1071+
default=None,
1072+
),
1073+
],
1074+
) -> CloudOrganizationResult:
1075+
"""Get details about a specific organization.
1076+
1077+
Requires either organization_id OR organization_name (exact match) to be provided.
1078+
This tool is useful for looking up an organization's ID from its name, or vice versa.
1079+
"""
1080+
api_root = resolve_cloud_api_url()
1081+
client_id = resolve_cloud_client_id()
1082+
client_secret = resolve_cloud_client_secret()
1083+
1084+
org = _resolve_organization(
1085+
organization_id=organization_id,
1086+
organization_name=organization_name,
1087+
api_root=api_root,
1088+
client_id=client_id,
1089+
client_secret=client_secret,
1090+
)
1091+
1092+
return CloudOrganizationResult(
1093+
id=org.organization_id,
1094+
name=org.organization_name,
1095+
email=org.email,
1096+
)
1097+
1098+
8551099
def _get_custom_source_definition_description(
8561100
custom_source: CustomCloudSourceDefinition,
8571101
) -> str:

0 commit comments

Comments
 (0)