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:
Drew Newberry
2026-09-11 00:02:07 +00:00
committed by GitHub
parent ddc8bba967
commit 33bbda3d33
84 changed files with 3545 additions and 1574 deletions
+88 -63
View File
@@ -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(
+43 -7
View File
@@ -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)