Skip to content

Commit 81e9f9f

Browse files
benthecarmanclaude
andcommitted
Refactor list_all_keys into reusable list_keys
Extract the single-page VSS listing logic into a list_keys method that accepts page_token and page_size parameters. list_internal now drives the pagination loop itself, calling list_keys per page. This prepares for PaginatedKVStore support which will reuse list_keys for single-page queries. This also fixes a potential issue where if the VSS server returned None for the page token we could enter into an infinite loop. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent d8c98a3 commit 81e9f9f

1 file changed

Lines changed: 33 additions & 35 deletions

File tree

src/io/vss_store.rs

Lines changed: 33 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -391,35 +391,32 @@ impl VssStoreInner {
391391
}
392392
}
393393

394-
async fn list_all_keys(
394+
async fn list_keys(
395395
&self, client: &VssClient<CustomRetryPolicy>, primary_namespace: &str,
396-
secondary_namespace: &str,
397-
) -> io::Result<Vec<String>> {
398-
let mut page_token = None;
399-
let mut keys = vec![];
396+
secondary_namespace: &str, page_token: Option<String>, page_size: Option<i32>,
397+
) -> io::Result<(Vec<String>, Option<String>)> {
400398
let key_prefix = self.build_obfuscated_prefix(primary_namespace, secondary_namespace);
401-
while page_token != Some("".to_string()) {
402-
let request = ListKeyVersionsRequest {
403-
store_id: self.store_id.clone(),
404-
key_prefix: Some(key_prefix.clone()),
405-
page_token,
406-
page_size: None,
407-
};
399+
let request = ListKeyVersionsRequest {
400+
store_id: self.store_id.clone(),
401+
key_prefix: Some(key_prefix),
402+
page_token,
403+
page_size,
404+
};
408405

409-
let response = client.list_key_versions(&request).await.map_err(|e| {
410-
let msg = format!(
411-
"Failed to list keys in {}/{}: {}",
412-
primary_namespace, secondary_namespace, e
413-
);
414-
Error::new(ErrorKind::Other, msg)
415-
})?;
406+
let response = client.list_key_versions(&request).await.map_err(|e| {
407+
let msg = format!(
408+
"Failed to list keys in {}/{}: {}",
409+
primary_namespace, secondary_namespace, e
410+
);
411+
Error::new(ErrorKind::Other, msg)
412+
})?;
416413

417-
for kv in response.key_versions {
418-
keys.push(self.extract_key(&kv.key)?);
419-
}
420-
page_token = response.next_page_token;
414+
let mut keys = Vec::with_capacity(response.key_versions.len());
415+
for kv in response.key_versions {
416+
keys.push(self.extract_key(&kv.key)?);
421417
}
422-
Ok(keys)
418+
419+
Ok((keys, response.next_page_token))
423420
}
424421

425422
async fn read_internal(
@@ -543,17 +540,18 @@ impl VssStoreInner {
543540
) -> io::Result<Vec<String>> {
544541
check_namespace_key_validity(&primary_namespace, &secondary_namespace, None, "list")?;
545542

546-
let keys = self
547-
.list_all_keys(client, &primary_namespace, &secondary_namespace)
548-
.await
549-
.map_err(|e| {
550-
let msg = format!(
551-
"Failed to retrieve keys in namespace: {}/{} : {}",
552-
primary_namespace, secondary_namespace, e
553-
);
554-
Error::new(ErrorKind::Other, msg)
555-
})?;
556-
543+
let mut page_token: Option<String> = None;
544+
let mut keys = vec![];
545+
loop {
546+
let (page_keys, next_page_token) = self
547+
.list_keys(client, &primary_namespace, &secondary_namespace, page_token, None)
548+
.await?;
549+
keys.extend(page_keys);
550+
match next_page_token {
551+
Some(t) if !t.is_empty() => page_token = Some(t),
552+
_ => break,
553+
}
554+
}
557555
Ok(keys)
558556
}
559557

0 commit comments

Comments
 (0)