diff --git a/coverage-mysql.cov b/coverage-mysql.cov new file mode 100644 index 000000000..4d5f12c89 Binary files /dev/null and b/coverage-mysql.cov differ diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py index 8d61030d1..7b3c08cc2 100644 --- a/singlestoredb/management/workspace.py +++ b/singlestoredb/management/workspace.py @@ -1301,17 +1301,24 @@ 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 @@ -1319,6 +1326,13 @@ def __init__( #: 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: @@ -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: @@ -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( @@ -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.""" @@ -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, diff --git a/singlestoredb/tests/test_management.py b/singlestoredb/tests/test_management.py index 222c80001..7c51b34db 100755 --- a/singlestoredb/tests/test_management.py +++ b/singlestoredb/tests/test_management.py @@ -363,6 +363,121 @@ def test_connect(self): assert 'endpoint' in cm.exception.msg, cm.exception.msg +@pytest.mark.management +class TestStarterWorkspace(unittest.TestCase): + + manager = None + starter_workspace = None + starter_workspace_user = { + 'username': 'starter_user', + 'password': None, + } + + @property + def starter_username(self): + """Return the username for the starter workspace user.""" + return self.starter_workspace_user['username'] + + @property + def password(self): + """Return the password for the starter workspace user.""" + return self.starter_workspace_user['password'] + + @classmethod + def setUpClass(cls): + cls.manager = s2.manage_workspaces() + + us_regions = [x for x in cls.manager.regions if 'US' in x.name] + cls.password = secrets.token_urlsafe(20) + '-x&$' + + name = clean_name(secrets.token_urlsafe(20)[:20]) + + cls.starter_workspace = cls.manager.create_starter_workspace( + f'starter-ws-test-{name}', + database_name=f'starter_db_{name}', + workspace_group={ + 'cell_id': random.choice(us_regions).id, + }, + ) + + cls.manager.create_starter_workspace_user( + starter_workspace_id=cls.starter_workspace.id, + 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.manager = None + cls.password = None + + def test_str(self): + assert self.starter_workspace.name in str(self.starter_workspace.name) + + def test_repr(self): + assert repr(self.starter_workspace) == str(self.starter_workspace) + + def test_get_starter_workspace(self): + workspace = self.manager.get_starter_workspace(self.starter_workspace.id) + assert workspace.id == self.starter_workspace.id, workspace.id + + with self.assertRaises(s2.ManagementError) as cm: + workspace = self.manager.get_starter_workspace('bad id') + + assert 'UUID' in cm.exception.msg, cm.exception.msg + + def test_starter_workspaces(self): + workspaces = self.manager.starter_workspaces + ids = [x.id for x in workspaces] + names = [x.name for x in workspaces] + assert self.starter_workspace.id in ids + assert self.starter_workspace.name in names + + objs = {} + for item in workspaces: + objs[item.id] = item + objs[item.name] = item + + name = random.choice(names) + assert workspaces[name] == objs[name] + id = random.choice(ids) + assert workspaces[id] == objs[id] + + def test_no_manager(self): + workspace = self.manager.get_starter_workspace(self.starter_workspace.id) + workspace._manager = None + + with self.assertRaises(s2.ManagementError) as cm: + workspace.refresh() + + assert 'No workspace manager' in cm.exception.msg, cm.exception.msg + + with self.assertRaises(s2.ManagementError) as cm: + workspace.terminate() + + assert 'No workspace manager' in cm.exception.msg, cm.exception.msg + + def test_connect(self): + with self.starter_workspace.connect( + user=self.starter_username, + password=self.password, + ) as conn: + with conn.cursor() as cur: + cur.execute('show databases') + assert 'starter_db' in [x[0] for x in list(cur)] + + # Test missing endpoint + workspace = self.manager.get_starter_workspace(self.starter_workspace.id) + workspace.endpoint = None + + with self.assertRaises(s2.ManagementError) as cm: + workspace.connect(user='admin', password=self.password) + + assert 'endpoint' in cm.exception.msg, cm.exception.msg + + @pytest.mark.management class TestStage(unittest.TestCase):