free tier endpoints

Author: Pedro Rodrigues

singlestoredb/management/workspace.py

@@ -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:
@@ -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: 'name', 'cell_id')
+
+        Returns
+        -------
+        :class:`StarterWorkspace`
+        """
+        return manager.create_starter_workspace(name, database_name, workspace_group)
+
 
 class Billing(object):
     """Billing information."""
@@ -1717,6 +1838,131 @@ 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: 'name', '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: '
+                "'name', 'cell_id'",
+            )
+        if set(workspace_group.keys()) != {'name', 'cell_id'}:
+            raise ValueError("workspace_group must contain only 'name' and 'cell_id'")
+
+        payload = {
+            'name': name,
+            'databaseName': database_name,
+            'workspaceGroup': {
+                'name': workspace_group['name'],
+                '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 {
+            'userID': user_id,
+            'password': returned_password,
+        }
+
 
 def manage_workspaces(
     access_token: Optional[str] = None,