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
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
158 changes: 46 additions & 112 deletions singlestoredb/management/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -1401,19 +1401,26 @@ def connect(self, **kwargs: Any) -> connection.Connection:
return connection.connect(**kwargs)

def terminate(self) -> None:
"""
Terminate the starter workspace.
"""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}')

Raises
------
ManagementError
If no workspace manager is associated with this object.
"""
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.',
)
self._manager.terminate_starter_workspace(self.id)
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 +1484,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 +1647,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 +1896,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