Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
4 changes: 4 additions & 0 deletions docs/src/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
176 changes: 69 additions & 107 deletions singlestoredb/management/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
kesmit13 marked this conversation as resolved.
Outdated
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:
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 2 additions & 3 deletions singlestoredb/tests/test_management.py
Original file line number Diff line number Diff line change
Expand Up @@ -404,16 +404,15 @@ 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,
)

@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

Expand Down