Skip to content

Commit 1b2de7c

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 1253368 commit 1b2de7c

File tree

1 file changed

+33
-35
lines changed

1 file changed

+33
-35
lines changed

src/io/vss_store.rs

Lines changed: 33 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -474,35 +474,32 @@ impl VssStoreInner {
474474
}
475475
}
476476

477-
async fn list_all_keys(
477+
async fn list_keys(
478478
&self, client: &VssClient<CustomRetryPolicy>, primary_namespace: &str,
479-
secondary_namespace: &str,
480-
) -> io::Result<Vec<String>> {
481-
let mut page_token = None;
482-
let mut keys = vec![];
479+
secondary_namespace: &str, page_token: Option<String>, page_size: Option<i32>,
480+
) -> io::Result<(Vec<String>, Option<String>)> {
483481
let key_prefix = self.build_obfuscated_prefix(primary_namespace, secondary_namespace);
484-
while page_token != Some("".to_string()) {
485-
let request = ListKeyVersionsRequest {
486-
store_id: self.store_id.clone(),
487-
key_prefix: Some(key_prefix.clone()),
488-
page_token,
489-
page_size: None,
490-
};
482+
let request = ListKeyVersionsRequest {
483+
store_id: self.store_id.clone(),
484+
key_prefix: Some(key_prefix),
485+
page_token,
486+
page_size,
487+
};
491488

492-
let response = client.list_key_versions(&request).await.map_err(|e| {
493-
let msg = format!(
494-
"Failed to list keys in {}/{}: {}",
495-
primary_namespace, secondary_namespace, e
496-
);
497-
Error::new(ErrorKind::Other, msg)
498-
})?;
489+
let response = client.list_key_versions(&request).await.map_err(|e| {
490+
let msg = format!(
491+
"Failed to list keys in {}/{}: {}",
492+
primary_namespace, secondary_namespace, e
493+
);
494+
Error::new(ErrorKind::Other, msg)
495+
})?;
499496

500-
for kv in response.key_versions {
501-
keys.push(self.extract_key(&kv.key)?);
502-
}
503-
page_token = response.next_page_token;
497+
let mut keys = Vec::with_capacity(response.key_versions.len());
498+
for kv in response.key_versions {
499+
keys.push(self.extract_key(&kv.key)?);
504500
}
505-
Ok(keys)
501+
502+
Ok((keys, response.next_page_token))
506503
}
507504

508505
async fn read_internal(
@@ -624,17 +621,18 @@ impl VssStoreInner {
624621
) -> io::Result<Vec<String>> {
625622
check_namespace_key_validity(&primary_namespace, &secondary_namespace, None, "list")?;
626623

627-
let keys = self
628-
.list_all_keys(client, &primary_namespace, &secondary_namespace)
629-
.await
630-
.map_err(|e| {
631-
let msg = format!(
632-
"Failed to retrieve keys in namespace: {}/{} : {}",
633-
primary_namespace, secondary_namespace, e
634-
);
635-
Error::new(ErrorKind::Other, msg)
636-
})?;
637-
624+
let mut page_token: Option<String> = None;
625+
let mut keys = vec![];
626+
loop {
627+
let (page_keys, next_page_token) = self
628+
.list_keys(client, &primary_namespace, &secondary_namespace, page_token, None)
629+
.await?;
630+
keys.extend(page_keys);
631+
match next_page_token {
632+
Some(t) if !t.is_empty() => page_token = Some(t),
633+
_ => break,
634+
}
635+
}
638636
Ok(keys)
639637
}
640638

0 commit comments

Comments
 (0)