@@ -2048,21 +2048,43 @@ def __get_workspace_by_id(workspace_id: str) -> WorkspaceTarget:
20482048 Raises:
20492049 ValueError: If workspace cannot be found
20502050 """
2051- workspace_manager = s2 .manage_workspaces (
2052- access_token = get_access_token (), organization_id = get_org_id ()
2053- )
20542051
20552052 target = None
20562053 is_shared = False
20572054
20582055 try :
2059- target = workspace_manager .get_workspace (workspace_id )
2056+ # Try as dedicated workspace first
2057+ workspace_data = build_request ("GET" , f"workspaces/{ workspace_id } " )
2058+
2059+ # Create a simple object to match the SDK interface
2060+ class SimpleWorkspace :
2061+ def __init__ (self , data ):
2062+ self .name = data .get ("name" , "" )
2063+ self .id = data .get ("workspaceID" , workspace_id )
2064+ self .endpoint = data .get ("endpoint" )
2065+
2066+ target = SimpleWorkspace (workspace_data )
20602067 is_shared = False # Dedicated workspace
20612068 except Exception as e :
20622069 if "404" in str (e ):
20632070 # Try as virtual workspace
2064- target = workspace_manager .get_starter_workspace (workspace_id )
2065- is_shared = True # Shared/virtual workspace
2071+ try :
2072+ virtual_workspace_data = build_request (
2073+ "GET" , f"sharedtier/virtualWorkspaces/{ workspace_id } "
2074+ )
2075+
2076+ # Create a simple object to match the SDK interface
2077+ class SimpleVirtualWorkspace :
2078+ def __init__ (self , data ):
2079+ self .name = data .get ("name" , "" )
2080+ self .id = data .get ("virtualWorkspaceID" , workspace_id )
2081+ self .endpoint = data .get ("endpoint" )
2082+ self .database_name = data .get ("databaseName" , "" )
2083+
2084+ target = SimpleVirtualWorkspace (virtual_workspace_data )
2085+ is_shared = True # Shared/virtual workspace
2086+ except Exception :
2087+ raise ValueError (f"Cannot find workspace { workspace_id } " )
20662088 else :
20672089 raise e
20682090
@@ -2120,31 +2142,29 @@ def create_starter_workspace(
21202142 )
21212143
21222144 try :
2123- # Initialize workspace manager using the SDK
2124- workspace_manager = s2 .manage_workspaces (
2125- access_token = get_access_token (),
2126- base_url = settings .s2_api_base_url ,
2127- organization_id = get_org_id (),
2128- )
2129-
2130- # Create the starter workspace
2131- starter_workspace = workspace_manager .create_starter_workspace (
2132- name = name ,
2133- database_name = database_name ,
2145+ # Create the starter workspace using the API
2146+ payload = {
2147+ "name" : name ,
2148+ "databaseName" : database_name ,
21342149 # TODO: Dinamically set region_id if needed
2135- workspace_group = {"cell_id" : "3482219c-a389-4079-b18b-d50662524e8a" },
2150+ "workspaceGroup" : {"cellID" : "3482219c-a389-4079-b18b-d50662524e8a" },
2151+ }
2152+
2153+ starter_workspace_data = build_request (
2154+ "POST" , "sharedtier/virtualWorkspaces" , data = payload
21362155 )
21372156
21382157 ctx .info (
2139- f"Starter workspace '{ name } ' created successfully with ID: { starter_workspace . workspace_id } "
2158+ f"Starter workspace '{ name } ' created successfully with ID: { starter_workspace_data . get ( 'virtualWorkspaceID' ) } "
21402159 )
21412160
21422161 return {
21432162 "status" : "success" ,
21442163 "message" : f"Starter workspace '{ name } ' created successfully" ,
2145- "name" : starter_workspace .name ,
2146- "endpoint" : starter_workspace .endpoint ,
2147- "database_name" : starter_workspace .database_name ,
2164+ "workspace_id" : starter_workspace_data .get ("virtualWorkspaceID" ),
2165+ "name" : starter_workspace_data .get ("name" ),
2166+ "endpoint" : starter_workspace_data .get ("endpoint" ),
2167+ "database_name" : starter_workspace_data .get ("databaseName" ),
21482168 }
21492169
21502170 except Exception as e :
@@ -2221,20 +2241,13 @@ def terminate_virtual_workspace(
22212241 )
22222242
22232243 try :
2224- # Initialize workspace manager using the SDK
2225- workspace_manager = s2 .manage_workspaces (
2226- access_token = get_access_token (),
2227- base_url = settings .s2_api_base_url ,
2228- organization_id = get_org_id (),
2229- )
2230-
22312244 # First, try to get the workspace details before termination
22322245 workspace_name = None
22332246 try :
2234- starter_workspace = workspace_manager . get_starter_workspace (
2235- validated_workspace_id
2247+ starter_workspace_data = build_request (
2248+ "GET" , f"sharedtier/virtualWorkspaces/ { validated_workspace_id } "
22362249 )
2237- workspace_name = starter_workspace . name
2250+ workspace_name = starter_workspace_data . get ( " name" )
22382251 ctx .info (
22392252 f"Found virtual workspace '{ workspace_name } ' (ID: { validated_workspace_id } )"
22402253 )
@@ -2246,7 +2259,9 @@ def terminate_virtual_workspace(
22462259 )
22472260
22482261 # Terminate the virtual workspace
2249- workspace_manager .terminate_starter_workspace (validated_workspace_id )
2262+ build_request (
2263+ "DELETE" , f"sharedtier/virtualWorkspaces/{ validated_workspace_id } "
2264+ )
22502265
22512266 termination_time = datetime .now ().isoformat ()
22522267
@@ -2297,7 +2312,7 @@ def terminate_virtual_workspace(
22972312 {"func" : complete_database_migration , "internal" : True },
22982313 # This tool is under development and not yet available for public use
22992314 {"func" : create_starter_workspace , "internal" : True },
2300- {"func" : terminate_virtual_workspace },
2315+ {"func" : terminate_virtual_workspace , "internal" : True },
23012316]
23022317
23032318# Export the tools
0 commit comments