From a998dad3788148d89d28ac968663e98b80dfbc32 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 17 Jul 2025 10:33:36 -0500 Subject: [PATCH 1/2] Update starter workspace API --- docs/src/api.rst | 4 + singlestoredb/management/workspace.py | 176 ++++++++++--------------- singlestoredb/tests/test_management.py | 5 +- 3 files changed, 75 insertions(+), 110 deletions(-) diff --git a/docs/src/api.rst b/docs/src/api.rst index 3dc3a7e21..7af62fc98 100644 --- a/docs/src/api.rst +++ b/docs/src/api.rst @@ -256,11 +256,15 @@ create new ones. WorkspaceManager WorkspaceManager.organization WorkspaceManager.workspace_groups + WorkspaceManager.starter_workspaces WorkspaceManager.regions + WorkspaceManager.shared_tier_regions WorkspaceManager.create_workspace_group WorkspaceManager.create_workspace + WorkspaceManager.create_starter_workspace WorkspaceManager.get_workspace_group WorkspaceManager.get_workspace + WorkspaceManager.get_starter_workspace WorkspaceGroup diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py index 7b3c08cc2..aaf32aec9 100644 --- a/singlestoredb/management/workspace.py +++ b/singlestoredb/management/workspace.py @@ -1400,20 +1400,55 @@ def connect(self, **kwargs: Any) -> connection.Connection: kwargs['host'] = self.endpoint return connection.connect(**kwargs) - def terminate(self) -> None: + def terminate( + self, + wait_on_terminated: bool = False, + wait_interval: int = 10, + wait_timeout: int = 600, + ) -> None: """ Terminate the starter workspace. + Parameters + ---------- + wait_on_terminated : bool, optional + Wait for the workspace to go into 'Terminated' mode before returning + wait_interval : int, optional + Number of seconds between each server check + wait_timeout : int, optional + Total number of seconds to check server before giving up + Raises ------ ManagementError - If no workspace manager is associated with this object. + If timeout is reached + """ if self._manager is None: raise ManagementError( msg='No workspace manager is associated with this object.', ) - self._manager.terminate_starter_workspace(self.id) + self._manager._delete(f'sharedtier/virtualWorkspaces/{self.id}') + if wait_on_terminated: + self._manager._wait_on_state( + self._manager.get_starter_workspace(self.id), + 'Terminated', interval=wait_interval, timeout=wait_timeout, + ) + self.refresh() + + def refresh(self) -> StarterWorkspace: + """Update the object to the current state.""" + if self._manager is None: + raise ManagementError( + msg='No workspace manager is associated with this object.', + ) + new_obj = self._manager.get_starter_workspace(self.id) + for name, value in vars(new_obj).items(): + if isinstance(value, Mapping): + setattr(self, name, snake_to_camel_dict(value)) + else: + setattr(self, name, value) + return self @property def organization(self) -> Organization: @@ -1477,35 +1512,32 @@ def create_user( msg='No workspace manager is associated with this object.', ) - return self._manager.create_starter_workspace_user(self.id, user_name, password) + payload = { + 'userName': user_name, + } + if password is not None: + payload['password'] = password - @classmethod - def create_starter_workspace( - cls, - manager: 'WorkspaceManager', - name: str, - database_name: str, - workspace_group: dict[str, str], - ) -> 'StarterWorkspace': - """ - Create a new starter (shared tier) workspace. + res = self._manager._post( + f'sharedtier/virtualWorkspaces/{self.id}/users', + json=payload, + ) - Parameters - ---------- - manager : WorkspaceManager - The WorkspaceManager instance to use for the API call - name : str - Name of the starter workspace - database_name : str - Name of the database for the starter workspace - workspace_group : dict[str, str] - Workspace group input (dict with keys: 'cell_id') + response_data = res.json() + user_id = response_data.get('userID') + if not user_id: + raise ManagementError(msg='No userID returned from API') - Returns - ------- - :class:`StarterWorkspace` - """ - return manager.create_starter_workspace(name, database_name, workspace_group) + # Return the password provided by user or generated by API + returned_password = password if password is not None \ + else response_data.get('password') + if not returned_password: + raise ManagementError(msg='No password available from API response') + + return { + 'user_id': user_id, + 'password': returned_password, + } class Billing(object): @@ -1643,6 +1675,14 @@ def regions(self) -> NamedList[Region]: res = self._get('regions') return NamedList([Region.from_dict(item, self) for item in res.json()]) + @ttl_property(datetime.timedelta(hours=1)) + def shared_tier_regions(self) -> NamedList[Region]: + """Return a list of regions that support shared tier workspaces.""" + res = self._get('regions/sharedtier') + return NamedList( + [Region.from_dict(item, self) for item in res.json()], + ) + def create_workspace_group( self, name: str, @@ -1884,84 +1924,6 @@ def create_starter_workspace( res = self._get(f'sharedtier/virtualWorkspaces/{virtual_workspace_id}') return StarterWorkspace.from_dict(res.json(), self) - def terminate_starter_workspace( - self, - id: str, - ) -> None: - """ - Terminate a starter (shared tier) workspace. - - Parameters - ---------- - id : str - ID of the starter workspace - wait_on_terminated : bool, optional - Wait for the starter workspace to go into 'Terminated' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ - self._delete(f'sharedtier/virtualWorkspaces/{id}') - - def create_starter_workspace_user( - self, - starter_workspace_id: str, - username: str, - password: Optional[str] = None, - ) -> Dict[str, str]: - """ - Create a new user for a starter workspace. - - Parameters - ---------- - starter_workspace_id : str - ID of the starter workspace - user_name : str - The starter workspace user name to connect the new user to the database - password : str, optional - Password for the new user. If not provided, a password will be - auto-generated by the system. - - Returns - ------- - Dict[str, str] - Dictionary containing 'userID' and 'password' of the created user - - """ - payload = { - 'userName': username, - } - if password is not None: - payload['password'] = password - - res = self._post( - f'sharedtier/virtualWorkspaces/{starter_workspace_id}/users', - json=payload, - ) - - response_data = res.json() - user_id = response_data.get('userID') - if not user_id: - raise ManagementError(msg='No userID returned from API') - - # Return the password provided by user or generated by API - returned_password = password if password is not None \ - else response_data.get('password') - if not returned_password: - raise ManagementError(msg='No password available from API response') - - return { - 'user_id': user_id, - 'password': returned_password, - } - def manage_workspaces( access_token: Optional[str] = None, diff --git a/singlestoredb/tests/test_management.py b/singlestoredb/tests/test_management.py index 2cbe12b0a..11d263ee0 100755 --- a/singlestoredb/tests/test_management.py +++ b/singlestoredb/tests/test_management.py @@ -404,8 +404,7 @@ def setUpClass(cls): }, ) - cls.manager.create_starter_workspace_user( - starter_workspace_id=cls.starter_workspace.id, + cls.starter_workspace.create_user( username=cls.starter_username, password=cls.password, ) @@ -413,7 +412,7 @@ def setUpClass(cls): @classmethod def tearDownClass(cls): if cls.starter_workspace is not None: - cls.starter_workspace.terminate(force=True) + cls.starter_workspace.terminate() cls.manager = None cls.password = None From 78816167792f3e77701250336774bbba0e3c2bc2 Mon Sep 17 00:00:00 2001 From: Kevin Smith Date: Thu, 17 Jul 2025 10:38:25 -0500 Subject: [PATCH 2/2] Remove wait from terminate --- singlestoredb/management/workspace.py | 32 ++------------------------- 1 file changed, 2 insertions(+), 30 deletions(-) diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py index aaf32aec9..3b6777bb8 100644 --- a/singlestoredb/management/workspace.py +++ b/singlestoredb/management/workspace.py @@ -1400,41 +1400,13 @@ def connect(self, **kwargs: Any) -> connection.Connection: kwargs['host'] = self.endpoint return connection.connect(**kwargs) - def terminate( - self, - wait_on_terminated: bool = False, - wait_interval: int = 10, - wait_timeout: int = 600, - ) -> None: - """ - Terminate the starter workspace. - - Parameters - ---------- - wait_on_terminated : bool, optional - Wait for the workspace to go into 'Terminated' mode before returning - wait_interval : int, optional - Number of seconds between each server check - wait_timeout : int, optional - Total number of seconds to check server before giving up - - Raises - ------ - ManagementError - If timeout is reached - - """ + def terminate(self) -> None: + """Terminate the starter workspace.""" if self._manager is None: raise ManagementError( msg='No workspace manager is associated with this object.', ) self._manager._delete(f'sharedtier/virtualWorkspaces/{self.id}') - if wait_on_terminated: - self._manager._wait_on_state( - self._manager.get_starter_workspace(self.id), - 'Terminated', interval=wait_interval, timeout=wait_timeout, - ) - self.refresh() def refresh(self) -> StarterWorkspace: """Update the object to the current state."""