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
Binary file added coverage-mysql.cov
Binary file not shown.
247 changes: 246 additions & 1 deletion singlestoredb/management/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -1301,24 +1301,38 @@ class StarterWorkspace(object):
See Also
--------
:meth:`WorkspaceManager.get_starter_workspace`
:meth:`WorkspaceManager.create_starter_workspace`
:meth:`WorkspaceManager.terminate_starter_workspace`
:meth:`WorkspaceManager.create_starter_workspace_user`
:attr:`WorkspaceManager.starter_workspaces`

"""

name: str
id: str
database_name: str
endpoint: Optional[str]

def __init__(
self,
name: str,
id: str,
database_name: str,
endpoint: Optional[str] = None,
):
#: Name of the starter workspace
self.name = name

#: Unique ID of the starter workspace
self.id = id

#: Name of the database associated with the starter workspace
self.database_name = database_name

#: Endpoint to connect to the starter workspace. The endpoint is in the form
#: of ``hostname:port``
self.endpoint = endpoint

self._manager: Optional[WorkspaceManager] = None

def __str__(self) -> str:
Expand Down Expand Up @@ -1351,10 +1365,56 @@ def from_dict(
out = cls(
name=obj['name'],
id=obj['virtualWorkspaceID'],
database_name=obj['databaseName'],
endpoint=obj.get('endpoint'),
)
out._manager = manager
return out

def connect(self, **kwargs: Any) -> connection.Connection:
"""
Create a connection to the database server for this starter workspace.

Parameters
----------
**kwargs : keyword-arguments, optional
Parameters to the SingleStoreDB `connect` function except host
and port which are supplied by the starter workspace object

Returns
-------
:class:`Connection`

"""
if not self.endpoint:
raise ManagementError(
msg='An endpoint has not been set in this '
'starter workspace configuration',
)
# Parse endpoint as host:port
if ':' in self.endpoint:
host, port = self.endpoint.split(':', 1)
kwargs['host'] = host
kwargs['port'] = int(port)
else:
kwargs['host'] = self.endpoint
return connection.connect(**kwargs)

def terminate(self) -> None:
"""
Terminate the starter workspace.

Raises
------
ManagementError
If no workspace manager is associated with this object.
"""
if self._manager is None:
raise ManagementError(
msg='No workspace manager is associated with this object.',
)
self._manager.terminate_starter_workspace(self.id)

@property
def organization(self) -> Organization:
if self._manager is None:
Expand All @@ -1375,7 +1435,7 @@ def stage(self) -> Stage:
stages = stage

@property
def starter_workspaces(self) -> NamedList[StarterWorkspace]:
def starter_workspaces(self) -> NamedList['StarterWorkspace']:
"""Return a list of available starter workspaces."""
if self._manager is None:
raise ManagementError(
Expand All @@ -1386,6 +1446,67 @@ def starter_workspaces(self) -> NamedList[StarterWorkspace]:
[StarterWorkspace.from_dict(item, self._manager) for item in res.json()],
)

def create_user(
self,
user_name: str,
password: Optional[str] = None,
) -> Dict[str, str]:
"""
Create a new user for this starter workspace.

Parameters
----------
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

Raises
------
ManagementError
If no workspace manager is associated with this object.
"""
if self._manager is None:
raise ManagementError(
msg='No workspace manager is associated with this object.',
)

return self._manager.create_starter_workspace_user(self.id, user_name, 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.

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')

Returns
-------
:class:`StarterWorkspace`
"""
return manager.create_starter_workspace(name, database_name, workspace_group)


class Billing(object):
"""Billing information."""
Expand Down Expand Up @@ -1717,6 +1838,130 @@ def get_starter_workspace(self, id: str) -> StarterWorkspace:
res = self._get(f'sharedtier/virtualWorkspaces/{id}')
return StarterWorkspace.from_dict(res.json(), manager=self)

def create_starter_workspace(
self,
name: str,
database_name: str,
workspace_group: dict[str, str],
) -> 'StarterWorkspace':
"""
Create a new starter (shared tier) workspace.

Parameters
----------
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')

Returns
-------
:class:`StarterWorkspace`
"""
if not workspace_group or not isinstance(workspace_group, dict):
raise ValueError(
'workspace_group must be a dict with keys: '
"'cell_id'",
)
if set(workspace_group.keys()) != {'cell_id'}:
raise ValueError("workspace_group must contain only 'cell_id'")

payload = {
'name': name,
'databaseName': database_name,
'workspaceGroup': {
'cellID': workspace_group['cell_id'],
},
}

res = self._post('sharedtier/virtualWorkspaces', json=payload)
virtual_workspace_id = res.json().get('virtualWorkspaceID')
if not virtual_workspace_id:
raise ManagementError(msg='No virtualWorkspaceID returned from API')

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
Loading