Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
"functions",
"storage",
"branching",
"compute",
)
DEFAULT_FEATURE_GROUPS = frozenset({
"account",
Expand All @@ -22,6 +23,7 @@
"docs",
"functions",
"branching",
"compute",
})

ALL_TOOL_NAMES = frozenset(tool.name for tool in TOOL_DEFINITIONS)
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
import os
import asyncio
import random
from collections.abc import Callable
Expand Down Expand Up @@ -28,6 +29,14 @@
ComputeSettingsForCreateWorkspaceInput,
StartWorkspaceRequest,
StopWorkspaceRequest,
DescribeWorkspaceDetailRequest,
DescribeComputesRequest,
ModifyComputeSettingsRequest,
CreateDatabaseRequest,
DescribeDatabasesRequest,
CreateDBAccountRequest,
DescribeDBAccountsRequest,
DescribeDBAccountConnectionRequest,
)
except ImportError:
logger.error("volcenginesdkaidap client dependencies not installed")
Expand Down Expand Up @@ -254,6 +263,10 @@ async def create_workspace(
workspace_name: str,
engine_type: str = "Supabase",
engine_version: str = "Supabase_1_24",
agent_plan_api_key: Optional[str] = None,
min_cu: float = 0.25,
max_cu: float = 1,
suspend_timeout_seconds: int = 300,
) -> dict:
try:
request = CreateWorkspaceRequest(
Expand All @@ -262,15 +275,18 @@ async def create_workspace(
engine_version=engine_version,
branch_settings=BranchSettingsForCreateWorkspaceInput(branch_name="main"),
compute_settings=ComputeSettingsForCreateWorkspaceInput(
auto_scaling_limit_min_cu=0.25,
auto_scaling_limit_max_cu=1,
suspend_timeout_seconds=300
auto_scaling_limit_min_cu=min_cu,
auto_scaling_limit_max_cu=max_cu,
suspend_timeout_seconds=suspend_timeout_seconds
),
workspace_settings=WorkspaceSettingsForCreateWorkspaceInput(
public_connection="Disabled",
deletion_protection="Disabled"
),
)
agent_plan_api_key = agent_plan_api_key or os.getenv("ARK_AGENT_PLAN_API_KEY")
if agent_plan_api_key:
request.agent_plan_api_key = agent_plan_api_key
response = self.client.create_workspace(request)

workspace_id = getattr(response, 'workspace_id', None)
Expand All @@ -291,6 +307,207 @@ async def create_workspace(
"error": str(e),
}

async def get_workspace_detail(self, workspace_id: str) -> Optional[Any]:
try:
request = DescribeWorkspaceDetailRequest(workspace_id=workspace_id)
response = self.client.describe_workspace_detail(request)
return getattr(response, "workspace", None)
except Exception as e:
logger.error(f"Error getting workspace detail: {e}")
raise RuntimeError(str(e))

async def describe_computes(
self,
workspace_id: str,
branch_id: Optional[str] = None,
service_type: Optional[str] = None,
) -> list[dict]:
if not branch_id:
branch_id = await self.get_default_branch_id(workspace_id)
if not branch_id:
raise RuntimeError(f"Could not resolve default branch for workspace {workspace_id}")
request_kwargs = {"workspace_id": workspace_id, "branch_id": branch_id}
if service_type:
request_kwargs["service_type"] = service_type
request = DescribeComputesRequest(**request_kwargs)
response = self.client.describe_computes(request)
computes = []
for compute in getattr(response, "computes", []) or []:
computes.append(compute.to_dict() if hasattr(compute, "to_dict") else compute)
return computes

async def modify_compute_settings(
self,
workspace_id: str,
min_cu: float,
max_cu: float,
suspend_timeout_seconds: Optional[int] = None,
service_type: Optional[str] = None,
) -> dict:
try:
request_kwargs = {
"workspace_id": workspace_id,
"auto_scaling_limit_min_cu": min_cu,
"auto_scaling_limit_max_cu": max_cu,
}
if suspend_timeout_seconds is not None:
request_kwargs["suspend_timeout_seconds"] = suspend_timeout_seconds
if service_type:
request_kwargs["service_type"] = service_type
request = ModifyComputeSettingsRequest(**request_kwargs)
self.client.modify_compute_settings(request)
return {"success": True, "workspace_id": workspace_id}
except Exception as e:
logger.error(f"Error modifying compute settings: {e}")
return {"success": False, "error": str(e)}

async def _ensure_branch_id(self, workspace_id: str, branch_id: Optional[str] = None) -> str:
if branch_id:
return branch_id
resolved = await self.get_default_branch_id(workspace_id)
if not resolved:
raise RuntimeError(f"Could not resolve default branch for workspace {workspace_id}")
return resolved

async def _resolve_primary_compute_id(self, workspace_id: str, branch_id: str) -> Optional[str]:
computes = await self.describe_computes(workspace_id, branch_id=branch_id)
for compute in computes:
if str(compute.get("compute_role") or "").lower() == "primary":
return compute.get("compute_id")
return computes[0].get("compute_id") if computes else None

async def describe_databases(
self,
workspace_id: str,
branch_id: Optional[str] = None,
search: Optional[str] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
) -> list[dict]:
branch_id = await self._ensure_branch_id(workspace_id, branch_id)
request_kwargs = {"workspace_id": workspace_id, "branch_id": branch_id}
if search:
request_kwargs["search"] = search
if limit is not None:
request_kwargs["limit"] = limit
if offset is not None:
request_kwargs["offset"] = offset
request = DescribeDatabasesRequest(**request_kwargs)
response = self.client.describe_databases(request)
databases = []
for database in getattr(response, "databases", []) or []:
databases.append(database.to_dict() if hasattr(database, "to_dict") else database)
return databases

async def create_database(
self,
workspace_id: str,
database_name: str,
branch_id: Optional[str] = None,
database_owner: Optional[str] = None,
database_desc: Optional[str] = None,
) -> dict:
try:
branch_id = await self._ensure_branch_id(workspace_id, branch_id)
request_kwargs = {
"workspace_id": workspace_id,
"branch_id": branch_id,
"database_name": database_name,
}
if database_owner:
request_kwargs["database_owner"] = database_owner
if database_desc:
request_kwargs["database_desc"] = database_desc
request = CreateDatabaseRequest(**request_kwargs)
response = self.client.create_database(request)
database = getattr(response, "database", None)
return {
"success": True,
"database": database.to_dict() if hasattr(database, "to_dict") else database,
}
except Exception as e:
logger.error(f"Error creating database: {e}")
return {"success": False, "error": str(e)}

async def describe_db_accounts(
self,
workspace_id: str,
branch_id: Optional[str] = None,
search: Optional[str] = None,
limit: Optional[int] = None,
offset: Optional[int] = None,
) -> list[dict]:
branch_id = await self._ensure_branch_id(workspace_id, branch_id)
request_kwargs = {"workspace_id": workspace_id, "branch_id": branch_id}
if search:
request_kwargs["search"] = search
if limit is not None:
request_kwargs["limit"] = limit
if offset is not None:
request_kwargs["offset"] = offset
request = DescribeDBAccountsRequest(**request_kwargs)
response = self.client.describe_db_accounts(request)
accounts = []
for account in getattr(response, "accounts", []) or []:
accounts.append(account.to_dict() if hasattr(account, "to_dict") else account)
return accounts

async def create_db_account(
self,
workspace_id: str,
account_name: str,
account_password: str,
branch_id: Optional[str] = None,
account_desc: Optional[str] = None,
) -> dict:
try:
branch_id = await self._ensure_branch_id(workspace_id, branch_id)
request_kwargs = {
"workspace_id": workspace_id,
"branch_id": branch_id,
"account_name": account_name,
"account_password": account_password,
}
if account_desc:
request_kwargs["account_desc"] = account_desc
request = CreateDBAccountRequest(**request_kwargs)
response = self.client.create_db_account(request)
account = getattr(response, "account", None)
return {
"success": True,
"account": account.to_dict() if hasattr(account, "to_dict") else account,
}
except Exception as e:
logger.error(f"Error creating db account: {e}")
return {"success": False, "error": str(e)}

async def describe_db_account_connection(
self,
workspace_id: str,
account_name: str,
database_name: str,
branch_id: Optional[str] = None,
compute_id: Optional[str] = None,
address_id: Optional[str] = None,
) -> dict:
branch_id = await self._ensure_branch_id(workspace_id, branch_id)
if not compute_id:
compute_id = await self._resolve_primary_compute_id(workspace_id, branch_id)
if not compute_id:
raise RuntimeError(f"Could not resolve compute for workspace {workspace_id}")
request_kwargs = {
"workspace_id": workspace_id,
"branch_id": branch_id,
"account_name": account_name,
"database_name": database_name,
"compute_id": compute_id,
}
if address_id:
request_kwargs["address_id"] = address_id
request = DescribeDBAccountConnectionRequest(**request_kwargs)
response = self.client.describe_db_account_connection(request)
return response.to_dict() if hasattr(response, "to_dict") else response

async def start_workspace(self, workspace_id: str) -> dict:
try:
request = StartWorkspaceRequest(workspace_id=workspace_id)
Expand Down
13 changes: 12 additions & 1 deletion server/mcp_server_supabase/src/mcp_server_supabase/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,14 @@
from typing import Any, Callable

from .platform import AidapClient
from .tools import DatabaseTools, EdgeFunctionTools, StorageTools, WorkspaceTools
from .tools import (
ComputeTools,
DatabaseAdminTools,
DatabaseTools,
EdgeFunctionTools,
StorageTools,
WorkspaceTools,
)


@dataclass(slots=True)
Expand All @@ -11,7 +18,9 @@ class SupabaseRuntime:
edge_tools: EdgeFunctionTools
storage_tools: StorageTools
database_tools: DatabaseTools
database_admin_tools: DatabaseAdminTools
workspace_tools: WorkspaceTools
compute_tools: ComputeTools


def create_runtime(
Expand All @@ -24,5 +33,7 @@ def create_runtime(
edge_tools=EdgeFunctionTools(client),
storage_tools=StorageTools(client),
database_tools=DatabaseTools(client),
database_admin_tools=DatabaseAdminTools(client),
workspace_tools=WorkspaceTools(client),
compute_tools=ComputeTools(client),
)
Loading
Loading