Skip to content

Commit 89c8e93

Browse files
authored
feat(free-tier): add support for new free tier endpoints (#68)
* free tier endpoints * add management api key to starter workspace tests * remove name of workspacegroup from starter workspace creation * replace userID with user_id on create_starter_workspace_user return dictionary
1 parent 9620a71 commit 89c8e93

File tree

3 files changed

+361
-1
lines changed

3 files changed

+361
-1
lines changed

coverage-mysql.cov

116 KB
Binary file not shown.

singlestoredb/management/workspace.py

Lines changed: 246 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1301,24 +1301,38 @@ class StarterWorkspace(object):
13011301
See Also
13021302
--------
13031303
:meth:`WorkspaceManager.get_starter_workspace`
1304+
:meth:`WorkspaceManager.create_starter_workspace`
1305+
:meth:`WorkspaceManager.terminate_starter_workspace`
1306+
:meth:`WorkspaceManager.create_starter_workspace_user`
13041307
:attr:`WorkspaceManager.starter_workspaces`
13051308
13061309
"""
13071310

13081311
name: str
13091312
id: str
1313+
database_name: str
1314+
endpoint: Optional[str]
13101315

13111316
def __init__(
13121317
self,
13131318
name: str,
13141319
id: str,
1320+
database_name: str,
1321+
endpoint: Optional[str] = None,
13151322
):
13161323
#: Name of the starter workspace
13171324
self.name = name
13181325

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

1329+
#: Name of the database associated with the starter workspace
1330+
self.database_name = database_name
1331+
1332+
#: Endpoint to connect to the starter workspace. The endpoint is in the form
1333+
#: of ``hostname:port``
1334+
self.endpoint = endpoint
1335+
13221336
self._manager: Optional[WorkspaceManager] = None
13231337

13241338
def __str__(self) -> str:
@@ -1351,10 +1365,56 @@ def from_dict(
13511365
out = cls(
13521366
name=obj['name'],
13531367
id=obj['virtualWorkspaceID'],
1368+
database_name=obj['databaseName'],
1369+
endpoint=obj.get('endpoint'),
13541370
)
13551371
out._manager = manager
13561372
return out
13571373

1374+
def connect(self, **kwargs: Any) -> connection.Connection:
1375+
"""
1376+
Create a connection to the database server for this starter workspace.
1377+
1378+
Parameters
1379+
----------
1380+
**kwargs : keyword-arguments, optional
1381+
Parameters to the SingleStoreDB `connect` function except host
1382+
and port which are supplied by the starter workspace object
1383+
1384+
Returns
1385+
-------
1386+
:class:`Connection`
1387+
1388+
"""
1389+
if not self.endpoint:
1390+
raise ManagementError(
1391+
msg='An endpoint has not been set in this '
1392+
'starter workspace configuration',
1393+
)
1394+
# Parse endpoint as host:port
1395+
if ':' in self.endpoint:
1396+
host, port = self.endpoint.split(':', 1)
1397+
kwargs['host'] = host
1398+
kwargs['port'] = int(port)
1399+
else:
1400+
kwargs['host'] = self.endpoint
1401+
return connection.connect(**kwargs)
1402+
1403+
def terminate(self) -> None:
1404+
"""
1405+
Terminate the starter workspace.
1406+
1407+
Raises
1408+
------
1409+
ManagementError
1410+
If no workspace manager is associated with this object.
1411+
"""
1412+
if self._manager is None:
1413+
raise ManagementError(
1414+
msg='No workspace manager is associated with this object.',
1415+
)
1416+
self._manager.terminate_starter_workspace(self.id)
1417+
13581418
@property
13591419
def organization(self) -> Organization:
13601420
if self._manager is None:
@@ -1375,7 +1435,7 @@ def stage(self) -> Stage:
13751435
stages = stage
13761436

13771437
@property
1378-
def starter_workspaces(self) -> NamedList[StarterWorkspace]:
1438+
def starter_workspaces(self) -> NamedList['StarterWorkspace']:
13791439
"""Return a list of available starter workspaces."""
13801440
if self._manager is None:
13811441
raise ManagementError(
@@ -1386,6 +1446,67 @@ def starter_workspaces(self) -> NamedList[StarterWorkspace]:
13861446
[StarterWorkspace.from_dict(item, self._manager) for item in res.json()],
13871447
)
13881448

1449+
def create_user(
1450+
self,
1451+
user_name: str,
1452+
password: Optional[str] = None,
1453+
) -> Dict[str, str]:
1454+
"""
1455+
Create a new user for this starter workspace.
1456+
1457+
Parameters
1458+
----------
1459+
user_name : str
1460+
The starter workspace user name to connect the new user to the database
1461+
password : str, optional
1462+
Password for the new user. If not provided, a password will be
1463+
auto-generated by the system.
1464+
1465+
Returns
1466+
-------
1467+
Dict[str, str]
1468+
Dictionary containing 'userID' and 'password' of the created user
1469+
1470+
Raises
1471+
------
1472+
ManagementError
1473+
If no workspace manager is associated with this object.
1474+
"""
1475+
if self._manager is None:
1476+
raise ManagementError(
1477+
msg='No workspace manager is associated with this object.',
1478+
)
1479+
1480+
return self._manager.create_starter_workspace_user(self.id, user_name, password)
1481+
1482+
@classmethod
1483+
def create_starter_workspace(
1484+
cls,
1485+
manager: 'WorkspaceManager',
1486+
name: str,
1487+
database_name: str,
1488+
workspace_group: dict[str, str],
1489+
) -> 'StarterWorkspace':
1490+
"""
1491+
Create a new starter (shared tier) workspace.
1492+
1493+
Parameters
1494+
----------
1495+
manager : WorkspaceManager
1496+
The WorkspaceManager instance to use for the API call
1497+
name : str
1498+
Name of the starter workspace
1499+
database_name : str
1500+
Name of the database for the starter workspace
1501+
workspace_group : dict[str, str]
1502+
Workspace group input (dict with keys: 'cell_id')
1503+
1504+
Returns
1505+
-------
1506+
:class:`StarterWorkspace`
1507+
"""
1508+
return manager.create_starter_workspace(name, database_name, workspace_group)
1509+
13891510

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

1841+
def create_starter_workspace(
1842+
self,
1843+
name: str,
1844+
database_name: str,
1845+
workspace_group: dict[str, str],
1846+
) -> 'StarterWorkspace':
1847+
"""
1848+
Create a new starter (shared tier) workspace.
1849+
1850+
Parameters
1851+
----------
1852+
name : str
1853+
Name of the starter workspace
1854+
database_name : str
1855+
Name of the database for the starter workspace
1856+
workspace_group : dict[str, str]
1857+
Workspace group input (dict with keys: 'cell_id')
1858+
1859+
Returns
1860+
-------
1861+
:class:`StarterWorkspace`
1862+
"""
1863+
if not workspace_group or not isinstance(workspace_group, dict):
1864+
raise ValueError(
1865+
'workspace_group must be a dict with keys: '
1866+
"'cell_id'",
1867+
)
1868+
if set(workspace_group.keys()) != {'cell_id'}:
1869+
raise ValueError("workspace_group must contain only 'cell_id'")
1870+
1871+
payload = {
1872+
'name': name,
1873+
'databaseName': database_name,
1874+
'workspaceGroup': {
1875+
'cellID': workspace_group['cell_id'],
1876+
},
1877+
}
1878+
1879+
res = self._post('sharedtier/virtualWorkspaces', json=payload)
1880+
virtual_workspace_id = res.json().get('virtualWorkspaceID')
1881+
if not virtual_workspace_id:
1882+
raise ManagementError(msg='No virtualWorkspaceID returned from API')
1883+
1884+
res = self._get(f'sharedtier/virtualWorkspaces/{virtual_workspace_id}')
1885+
return StarterWorkspace.from_dict(res.json(), self)
1886+
1887+
def terminate_starter_workspace(
1888+
self,
1889+
id: str,
1890+
) -> None:
1891+
"""
1892+
Terminate a starter (shared tier) workspace.
1893+
1894+
Parameters
1895+
----------
1896+
id : str
1897+
ID of the starter workspace
1898+
wait_on_terminated : bool, optional
1899+
Wait for the starter workspace to go into 'Terminated' mode before returning
1900+
wait_interval : int, optional
1901+
Number of seconds between each server check
1902+
wait_timeout : int, optional
1903+
Total number of seconds to check server before giving up
1904+
1905+
Raises
1906+
------
1907+
ManagementError
1908+
If timeout is reached
1909+
1910+
"""
1911+
self._delete(f'sharedtier/virtualWorkspaces/{id}')
1912+
1913+
def create_starter_workspace_user(
1914+
self,
1915+
starter_workspace_id: str,
1916+
username: str,
1917+
password: Optional[str] = None,
1918+
) -> Dict[str, str]:
1919+
"""
1920+
Create a new user for a starter workspace.
1921+
1922+
Parameters
1923+
----------
1924+
starter_workspace_id : str
1925+
ID of the starter workspace
1926+
user_name : str
1927+
The starter workspace user name to connect the new user to the database
1928+
password : str, optional
1929+
Password for the new user. If not provided, a password will be
1930+
auto-generated by the system.
1931+
1932+
Returns
1933+
-------
1934+
Dict[str, str]
1935+
Dictionary containing 'userID' and 'password' of the created user
1936+
1937+
"""
1938+
payload = {
1939+
'userName': username,
1940+
}
1941+
if password is not None:
1942+
payload['password'] = password
1943+
1944+
res = self._post(
1945+
f'sharedtier/virtualWorkspaces/{starter_workspace_id}/users',
1946+
json=payload,
1947+
)
1948+
1949+
response_data = res.json()
1950+
user_id = response_data.get('userID')
1951+
if not user_id:
1952+
raise ManagementError(msg='No userID returned from API')
1953+
1954+
# Return the password provided by user or generated by API
1955+
returned_password = password if password is not None \
1956+
else response_data.get('password')
1957+
if not returned_password:
1958+
raise ManagementError(msg='No password available from API response')
1959+
1960+
return {
1961+
'user_id': user_id,
1962+
'password': returned_password,
1963+
}
1964+
17201965

17211966
def manage_workspaces(
17221967
access_token: Optional[str] = None,

0 commit comments

Comments
 (0)