mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-02 07:34:45 +08:00
refactor(persistence): adopt continuation-token pagination (#3249)
* refactor(persistence): adopt continuation-token pagination Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(pagination): address continuation review findings Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(tui): recover completed list refreshes Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(pagination): address review scalability findings Signed-off-by: Drew Newberry <anewberry@nvidia.com> * fix(pagination): repair branch validation Signed-off-by: Drew Newberry <anewberry@nvidia.com> * docs(go): use page size in template example Signed-off-by: Drew Newberry <anewberry@nvidia.com> --------- Signed-off-by: Drew Newberry <anewberry@nvidia.com>
This commit is contained in:
+88
-63
@@ -819,49 +819,61 @@ class SandboxClient:
|
||||
self,
|
||||
*,
|
||||
workspace: str,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
page_size: int = 100,
|
||||
label_selector: str | None = None,
|
||||
) -> builtins.list[SandboxRef]:
|
||||
request = openshell_pb2.ListSandboxesRequest(
|
||||
workspace_scope=_workspace_scope(workspace),
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
label_selector=label_selector or "",
|
||||
)
|
||||
response = self._stub.ListSandboxes(request, timeout=self._timeout)
|
||||
return [_sandbox_ref(item) for item in response.sandboxes]
|
||||
sandboxes: builtins.list[SandboxRef] = []
|
||||
page_token = ""
|
||||
while True:
|
||||
response = self._stub.ListSandboxes(
|
||||
openshell_pb2.ListSandboxesRequest(
|
||||
workspace_scope=_workspace_scope(workspace),
|
||||
page_size=page_size,
|
||||
page_token=page_token,
|
||||
label_selector=label_selector or "",
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
sandboxes.extend(_sandbox_ref(item) for item in response.sandboxes)
|
||||
if not getattr(response, "next_page_token", ""):
|
||||
return sandboxes
|
||||
page_token = response.next_page_token
|
||||
|
||||
def list_for_all_workspaces(
|
||||
self,
|
||||
*,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
page_size: int = 100,
|
||||
label_selector: str | None = None,
|
||||
) -> builtins.list[SandboxRef]:
|
||||
request = openshell_pb2.ListSandboxesRequest(
|
||||
workspace_scope=_all_workspaces_scope(),
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
label_selector=label_selector or "",
|
||||
)
|
||||
response = self._stub.ListSandboxes(request, timeout=self._timeout)
|
||||
return [_sandbox_ref(item) for item in response.sandboxes]
|
||||
sandboxes: builtins.list[SandboxRef] = []
|
||||
page_token = ""
|
||||
while True:
|
||||
response = self._stub.ListSandboxes(
|
||||
openshell_pb2.ListSandboxesRequest(
|
||||
workspace_scope=_all_workspaces_scope(),
|
||||
page_size=page_size,
|
||||
page_token=page_token,
|
||||
label_selector=label_selector or "",
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
sandboxes.extend(_sandbox_ref(item) for item in response.sandboxes)
|
||||
if not getattr(response, "next_page_token", ""):
|
||||
return sandboxes
|
||||
page_token = response.next_page_token
|
||||
|
||||
def list_ids(
|
||||
self,
|
||||
*,
|
||||
workspace: str,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
page_size: int = 100,
|
||||
label_selector: str | None = None,
|
||||
) -> builtins.list[str]:
|
||||
return [
|
||||
item.id
|
||||
for item in self.list(
|
||||
workspace=workspace,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
page_size=page_size,
|
||||
label_selector=label_selector,
|
||||
)
|
||||
]
|
||||
@@ -869,15 +881,13 @@ class SandboxClient:
|
||||
def list_ids_for_all_workspaces(
|
||||
self,
|
||||
*,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
page_size: int = 100,
|
||||
label_selector: str | None = None,
|
||||
) -> builtins.list[str]:
|
||||
return [
|
||||
item.id
|
||||
for item in self.list_for_all_workspaces(
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
page_size=page_size,
|
||||
label_selector=label_selector,
|
||||
)
|
||||
]
|
||||
@@ -1181,38 +1191,48 @@ class SandboxTemplateClient:
|
||||
self,
|
||||
*,
|
||||
workspace: str,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
page_size: int = 100,
|
||||
label_selector: str = "",
|
||||
) -> builtins.list[openshell_pb2.SandboxWorkloadTemplate]:
|
||||
response = self._stub.ListSandboxTemplates(
|
||||
openshell_pb2.ListSandboxTemplatesRequest(
|
||||
workspace_scope=_workspace_scope(workspace),
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
label_selector=label_selector,
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return list(response.templates)
|
||||
templates: builtins.list[openshell_pb2.SandboxWorkloadTemplate] = []
|
||||
page_token = ""
|
||||
while True:
|
||||
response = self._stub.ListSandboxTemplates(
|
||||
openshell_pb2.ListSandboxTemplatesRequest(
|
||||
workspace_scope=_workspace_scope(workspace),
|
||||
page_size=page_size,
|
||||
page_token=page_token,
|
||||
label_selector=label_selector,
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
templates.extend(response.templates)
|
||||
if not getattr(response, "next_page_token", ""):
|
||||
return templates
|
||||
page_token = response.next_page_token
|
||||
|
||||
def list_for_all_workspaces(
|
||||
self,
|
||||
*,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
page_size: int = 100,
|
||||
label_selector: str = "",
|
||||
) -> builtins.list[openshell_pb2.SandboxWorkloadTemplate]:
|
||||
response = self._stub.ListSandboxTemplates(
|
||||
openshell_pb2.ListSandboxTemplatesRequest(
|
||||
workspace_scope=_all_workspaces_scope(),
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
label_selector=label_selector,
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return list(response.templates)
|
||||
templates: builtins.list[openshell_pb2.SandboxWorkloadTemplate] = []
|
||||
page_token = ""
|
||||
while True:
|
||||
response = self._stub.ListSandboxTemplates(
|
||||
openshell_pb2.ListSandboxTemplatesRequest(
|
||||
workspace_scope=_all_workspaces_scope(),
|
||||
page_size=page_size,
|
||||
page_token=page_token,
|
||||
label_selector=label_selector,
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
templates.extend(response.templates)
|
||||
if not getattr(response, "next_page_token", ""):
|
||||
return templates
|
||||
page_token = response.next_page_token
|
||||
|
||||
def delete(self, name: str, *, workspace: str) -> bool:
|
||||
response = self._stub.DeleteSandboxTemplate(
|
||||
@@ -1276,19 +1296,24 @@ class WorkspaceClient:
|
||||
def list(
|
||||
self,
|
||||
*,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
page_size: int = 100,
|
||||
label_selector: str | None = None,
|
||||
) -> builtins.list[WorkspaceRef]:
|
||||
response = self._stub.ListWorkspaces(
|
||||
openshell_pb2.ListWorkspacesRequest(
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
label_selector=label_selector or "",
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
return [_workspace_ref(ws) for ws in response.workspaces]
|
||||
workspaces: builtins.list[WorkspaceRef] = []
|
||||
page_token = ""
|
||||
while True:
|
||||
response = self._stub.ListWorkspaces(
|
||||
openshell_pb2.ListWorkspacesRequest(
|
||||
page_size=page_size,
|
||||
page_token=page_token,
|
||||
label_selector=label_selector or "",
|
||||
),
|
||||
timeout=self._timeout,
|
||||
)
|
||||
workspaces.extend(_workspace_ref(ws) for ws in response.workspaces)
|
||||
if not getattr(response, "next_page_token", ""):
|
||||
return workspaces
|
||||
page_token = response.next_page_token
|
||||
|
||||
def delete(self, name: str) -> bool:
|
||||
response = self._stub.DeleteWorkspace(
|
||||
|
||||
@@ -1963,7 +1963,11 @@ def _make_workload_template_proto(
|
||||
|
||||
|
||||
class _FakeSandboxStub:
|
||||
def __init__(self, listed: list[openshell_pb2.Sandbox] | None = None) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
listed: list[openshell_pb2.Sandbox] | None = None,
|
||||
listed_pages: list[list[openshell_pb2.Sandbox]] | None = None,
|
||||
) -> None:
|
||||
self.create_request: openshell_pb2.CreateSandboxRequest | None = None
|
||||
self.list_request: openshell_pb2.ListSandboxesRequest | None = None
|
||||
self.get_request: openshell_pb2.GetSandboxRequest | None = None
|
||||
@@ -1981,6 +1985,8 @@ class _FakeSandboxStub:
|
||||
openshell_pb2.DeleteSandboxTemplateRequest | None
|
||||
) = None
|
||||
self._listed = listed or []
|
||||
self._listed_pages = listed_pages
|
||||
self.list_requests: list[openshell_pb2.ListSandboxesRequest] = []
|
||||
self._templates: list[openshell_pb2.SandboxWorkloadTemplate] = []
|
||||
|
||||
def GetSandbox(
|
||||
@@ -2061,7 +2067,17 @@ class _FakeSandboxStub:
|
||||
timeout: float | None = None,
|
||||
) -> Any:
|
||||
self.list_request = request
|
||||
self.list_requests.append(deepcopy(request))
|
||||
_ = timeout
|
||||
if self._listed_pages is not None:
|
||||
page = int(request.page_token or "0")
|
||||
next_page_token = (
|
||||
str(page + 1) if page + 1 < len(self._listed_pages) else ""
|
||||
)
|
||||
return SimpleNamespace(
|
||||
sandboxes=list(self._listed_pages[page]),
|
||||
next_page_token=next_page_token,
|
||||
)
|
||||
return SimpleNamespace(sandboxes=list(self._listed))
|
||||
|
||||
def CreateSandboxTemplate(
|
||||
@@ -2393,13 +2409,13 @@ def test_sandbox_template_client_crud_forwards_requests() -> None:
|
||||
assert _request_workspace(stub.get_template_request) == "default"
|
||||
|
||||
listed = client.list(
|
||||
workspace="default", limit=50, offset=10, label_selector="team=runtime"
|
||||
workspace="default", page_size=50, label_selector="team=runtime"
|
||||
)
|
||||
assert len(listed) == 1
|
||||
assert stub.list_template_request is not None
|
||||
assert _request_workspace(stub.list_template_request) == "default"
|
||||
assert stub.list_template_request.limit == 50
|
||||
assert stub.list_template_request.offset == 10
|
||||
assert stub.list_template_request.page_size == 50
|
||||
assert stub.list_template_request.page_token == ""
|
||||
assert stub.list_template_request.label_selector == "team=runtime"
|
||||
assert not _request_selects_all_workspaces(stub.list_template_request)
|
||||
|
||||
@@ -2413,13 +2429,13 @@ def test_sandbox_template_list_for_all_workspaces_selects_all() -> None:
|
||||
stub = _FakeSandboxStub()
|
||||
client = _template_client_with_fake_stub(stub)
|
||||
|
||||
client.list_for_all_workspaces(limit=100, offset=5, label_selector="team=runtime")
|
||||
client.list_for_all_workspaces(page_size=100, label_selector="team=runtime")
|
||||
|
||||
assert stub.list_template_request is not None
|
||||
assert _request_selects_all_workspaces(stub.list_template_request)
|
||||
assert _request_workspace(stub.list_template_request) is None
|
||||
assert stub.list_template_request.limit == 100
|
||||
assert stub.list_template_request.offset == 5
|
||||
assert stub.list_template_request.page_size == 100
|
||||
assert stub.list_template_request.page_token == ""
|
||||
assert stub.list_template_request.label_selector == "team=runtime"
|
||||
|
||||
|
||||
@@ -2534,6 +2550,26 @@ def test_list_without_selector_sends_empty_string() -> None:
|
||||
assert stub.list_request.label_selector == ""
|
||||
|
||||
|
||||
def test_list_follows_continuation_tokens() -> None:
|
||||
stub = _FakeSandboxStub(
|
||||
listed_pages=[
|
||||
[_make_sandbox_proto("sandbox-1", "job-1")],
|
||||
[_make_sandbox_proto("sandbox-2", "job-2")],
|
||||
]
|
||||
)
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
sandboxes = client.list(
|
||||
workspace="default", page_size=1, label_selector="team=core"
|
||||
)
|
||||
|
||||
assert [sandbox.name for sandbox in sandboxes] == ["job-1", "job-2"]
|
||||
assert len(stub.list_requests) == 2
|
||||
assert stub.list_requests[0].page_token == ""
|
||||
assert stub.list_requests[1].page_token == "1"
|
||||
assert stub.list_requests[1].label_selector == "team=core"
|
||||
|
||||
|
||||
def test_list_ids_forwards_label_selector() -> None:
|
||||
stub = _FakeSandboxStub(listed=[_make_sandbox_proto("sandbox-1", "job-1")])
|
||||
client = _client_with_fake_stub(stub)
|
||||
|
||||
Reference in New Issue
Block a user