Skip to content

Commit 6b2804e

Browse files
authored
Update starter workspace API (singlestore-labs#73)
* Update starter workspace API * Remove wait from terminate
1 parent 9c649cf commit 6b2804e

File tree

3 files changed

+52
-115
lines changed

3 files changed

+52
-115
lines changed

docs/src/api.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,11 +256,15 @@ create new ones.
256256
WorkspaceManager
257257
WorkspaceManager.organization
258258
WorkspaceManager.workspace_groups
259+
WorkspaceManager.starter_workspaces
259260
WorkspaceManager.regions
261+
WorkspaceManager.shared_tier_regions
260262
WorkspaceManager.create_workspace_group
261263
WorkspaceManager.create_workspace
264+
WorkspaceManager.create_starter_workspace
262265
WorkspaceManager.get_workspace_group
263266
WorkspaceManager.get_workspace
267+
WorkspaceManager.get_starter_workspace
264268

265269

266270
WorkspaceGroup

singlestoredb/management/workspace.py

Lines changed: 46 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -1401,19 +1401,26 @@ def connect(self, **kwargs: Any) -> connection.Connection:
14011401
return connection.connect(**kwargs)
14021402

14031403
def terminate(self) -> None:
1404-
"""
1405-
Terminate the starter workspace.
1404+
"""Terminate the starter workspace."""
1405+
if self._manager is None:
1406+
raise ManagementError(
1407+
msg='No workspace manager is associated with this object.',
1408+
)
1409+
self._manager._delete(f'sharedtier/virtualWorkspaces/{self.id}')
14061410

1407-
Raises
1408-
------
1409-
ManagementError
1410-
If no workspace manager is associated with this object.
1411-
"""
1411+
def refresh(self) -> StarterWorkspace:
1412+
"""Update the object to the current state."""
14121413
if self._manager is None:
14131414
raise ManagementError(
14141415
msg='No workspace manager is associated with this object.',
14151416
)
1416-
self._manager.terminate_starter_workspace(self.id)
1417+
new_obj = self._manager.get_starter_workspace(self.id)
1418+
for name, value in vars(new_obj).items():
1419+
if isinstance(value, Mapping):
1420+
setattr(self, name, snake_to_camel_dict(value))
1421+
else:
1422+
setattr(self, name, value)
1423+
return self
14171424

14181425
@property
14191426
def organization(self) -> Organization:
@@ -1477,35 +1484,32 @@ def create_user(
14771484
msg='No workspace manager is associated with this object.',
14781485
)
14791486

1480-
return self._manager.create_starter_workspace_user(self.id, user_name, password)
1487+
payload = {
1488+
'userName': user_name,
1489+
}
1490+
if password is not None:
1491+
payload['password'] = password
14811492

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.
1493+
res = self._manager._post(
1494+
f'sharedtier/virtualWorkspaces/{self.id}/users',
1495+
json=payload,
1496+
)
14921497

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')
1498+
response_data = res.json()
1499+
user_id = response_data.get('userID')
1500+
if not user_id:
1501+
raise ManagementError(msg='No userID returned from API')
15031502

1504-
Returns
1505-
-------
1506-
:class:`StarterWorkspace`
1507-
"""
1508-
return manager.create_starter_workspace(name, database_name, workspace_group)
1503+
# Return the password provided by user or generated by API
1504+
returned_password = password if password is not None \
1505+
else response_data.get('password')
1506+
if not returned_password:
1507+
raise ManagementError(msg='No password available from API response')
1508+
1509+
return {
1510+
'user_id': user_id,
1511+
'password': returned_password,
1512+
}
15091513

15101514

15111515
class Billing(object):
@@ -1643,6 +1647,14 @@ def regions(self) -> NamedList[Region]:
16431647
res = self._get('regions')
16441648
return NamedList([Region.from_dict(item, self) for item in res.json()])
16451649

1650+
@ttl_property(datetime.timedelta(hours=1))
1651+
def shared_tier_regions(self) -> NamedList[Region]:
1652+
"""Return a list of regions that support shared tier workspaces."""
1653+
res = self._get('regions/sharedtier')
1654+
return NamedList(
1655+
[Region.from_dict(item, self) for item in res.json()],
1656+
)
1657+
16461658
def create_workspace_group(
16471659
self,
16481660
name: str,
@@ -1884,84 +1896,6 @@ def create_starter_workspace(
18841896
res = self._get(f'sharedtier/virtualWorkspaces/{virtual_workspace_id}')
18851897
return StarterWorkspace.from_dict(res.json(), self)
18861898

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

19661900
def manage_workspaces(
19671901
access_token: Optional[str] = None,

singlestoredb/tests/test_management.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -404,16 +404,15 @@ def setUpClass(cls):
404404
},
405405
)
406406

407-
cls.manager.create_starter_workspace_user(
408-
starter_workspace_id=cls.starter_workspace.id,
407+
cls.starter_workspace.create_user(
409408
username=cls.starter_username,
410409
password=cls.password,
411410
)
412411

413412
@classmethod
414413
def tearDownClass(cls):
415414
if cls.starter_workspace is not None:
416-
cls.starter_workspace.terminate(force=True)
415+
cls.starter_workspace.terminate()
417416
cls.manager = None
418417
cls.password = None
419418

0 commit comments

Comments
 (0)