feat: support KVS.listKeys() prefix and collection parameters (#3001)

Adds `prefix` and `collection` parameters to the `KVS.forEachKey()`
interface and implements the `prefix` filter in the `MemoryStorage`
implementation of KVS. Note that `MemoryStorage` doesn't currently
support the `collection` parameter (it's noop).

Closes #2974

Author: barjin

packages/core/src/storages/key_value_store.ts

@@ -460,22 +460,24 @@ export class KeyValueStore {
         options: KeyValueStoreIteratorOptions = {},
         index = 0,
     ): Promise<void> {
-        const { exclusiveStartKey } = options;
+        const { exclusiveStartKey, prefix, collection } = options;
         ow(iteratee, ow.function);
         ow(
             options,
             ow.object.exactShape({
                 exclusiveStartKey: ow.optional.string,
+                prefix: ow.optional.string,
+                collection: ow.optional.string,
             }),
         );
 
-        const response = await this.client.listKeys({ exclusiveStartKey });
+        const response = await this.client.listKeys({ exclusiveStartKey, prefix, collection });
         const { nextExclusiveStartKey, isTruncated, items } = response;
         for (const item of items) {
             await iteratee(item.key, index++, { size: item.size });
         }
         return isTruncated
-            ? this._forEachKey(iteratee, { exclusiveStartKey: nextExclusiveStartKey }, index)
+            ? this._forEachKey(iteratee, { exclusiveStartKey: nextExclusiveStartKey, prefix, collection }, index)
             : undefined; // [].forEach() returns undefined.
     }
 
@@ -758,4 +760,12 @@ export interface KeyValueStoreIteratorOptions {
      * All keys up to this one (including) are skipped from the result.
      */
     exclusiveStartKey?: string;
+    /**
+     * If set, only keys that start with this prefix are returned.
+     */
+    prefix?: string;
+    /**
+     * Collection name to use for listing keys.
+     */
+    collection?: string;
 }

packages/memory-storage/src/resource-clients/key-value-store.ts

@@ -119,10 +119,16 @@ export class KeyValueStoreClient extends BaseClient {
     }
 
     async listKeys(options: storage.KeyValueStoreClientListOptions = {}): Promise<storage.KeyValueStoreClientListData> {
-        const { limit = DEFAULT_API_PARAM_LIMIT, exclusiveStartKey } = s
+        const {
+            limit = DEFAULT_API_PARAM_LIMIT,
+            exclusiveStartKey,
+            prefix,
+        } = s
             .object({
                 limit: s.number.greaterThan(0).optional,
                 exclusiveStartKey: s.string.optional,
+                collection: s.string.optional, // This is ignored, but kept for validation consistency with API client.
+                prefix: s.string.optional,
             })
             .parse(options);
 
@@ -151,23 +157,25 @@ export class KeyValueStoreClient extends BaseClient {
             return a.key.localeCompare(b.key);
         });
 
-        let truncatedItems = items;
+        const filteredItems = items.filter((item) => !prefix || item.key.startsWith(prefix));
+
+        let truncatedItems = filteredItems;
         if (exclusiveStartKey) {
-            const keyPos = items.findIndex((item) => item.key === exclusiveStartKey);
-            if (keyPos !== -1) truncatedItems = items.slice(keyPos + 1);
+            const keyPos = filteredItems.findIndex((item) => item.key === exclusiveStartKey);
+            if (keyPos !== -1) truncatedItems = filteredItems.slice(keyPos + 1);
         }
 
         const limitedItems = truncatedItems.slice(0, limit);
 
-        const lastItemInStore = items[items.length - 1];
-        const lastSelectedItem = limitedItems[limitedItems.length - 1];
+        const lastItemInStore = filteredItems.at(-1);
+        const lastSelectedItem = limitedItems.at(-1);
         const isLastSelectedItemAbsolutelyLast = lastItemInStore === lastSelectedItem;
-        const nextExclusiveStartKey = isLastSelectedItemAbsolutelyLast ? undefined : lastSelectedItem.key;
+        const nextExclusiveStartKey = isLastSelectedItemAbsolutelyLast ? undefined : lastSelectedItem?.key;
 
         existingStoreById.updateTimestamps(false);
 
         return {
-            count: items.length,
+            count: limitedItems.length,
             limit,
             exclusiveStartKey,
             isTruncated: !isLastSelectedItemAbsolutelyLast,

packages/types/src/storages.ts

@@ -136,6 +136,8 @@ export interface KeyValueStoreClientUpdateOptions {
 export interface KeyValueStoreClientListOptions {
     limit?: number;
     exclusiveStartKey?: string;
+    collection?: string;
+    prefix?: string;
 }
 
 export interface KeyValueStoreItemData {

test/core/storages/key_value_store.test.ts

@@ -509,6 +509,39 @@ describe('KeyValueStore', () => {
     });
 
     describe('forEachKey', () => {
+        test('should work with prefixes', async () => {
+            const store = await KeyValueStore.open();
+
+            for (const [key, value] of Object.entries({
+                'img-key1': 'PAYLOAD',
+                'img-key2': 'PAYLOAD',
+                'txt-key1': 'PAYLOAD',
+                'txt-key2': 'PAYLOAD',
+            })) {
+                await store.setValue(key, value);
+            }
+
+            const imgKeys: string[] = [];
+            const txtKeys: string[] = [];
+
+            await store.forEachKey(
+                (key) => {
+                    imgKeys.push(key);
+                },
+                { prefix: 'img-' },
+            );
+
+            await store.forEachKey(
+                (key) => {
+                    txtKeys.push(key);
+                },
+                { prefix: 'txt-' },
+            );
+
+            expect(imgKeys).toEqual(['img-key1', 'img-key2']);
+            expect(txtKeys).toEqual(['txt-key1', 'txt-key2']);
+        });
+
         test('should work remotely', async () => {
             const store = new KeyValueStore({
                 id: 'my-store-id-1',
@@ -555,13 +588,13 @@ describe('KeyValueStore', () => {
                 async (key, index, info) => {
                     results.push([key, index, info]);
                 },
-                { exclusiveStartKey: 'key0' },
+                { exclusiveStartKey: 'key0', prefix: 'img/' },
             );
 
             expect(mockListKeys).toBeCalledTimes(3);
-            expect(mockListKeys).toHaveBeenNthCalledWith(1, { exclusiveStartKey: 'key0' });
-            expect(mockListKeys).toHaveBeenNthCalledWith(2, { exclusiveStartKey: 'key2' });
-            expect(mockListKeys).toHaveBeenNthCalledWith(3, { exclusiveStartKey: 'key4' });
+            expect(mockListKeys).toHaveBeenNthCalledWith(1, { exclusiveStartKey: 'key0', prefix: 'img/' });
+            expect(mockListKeys).toHaveBeenNthCalledWith(2, { exclusiveStartKey: 'key2', prefix: 'img/' });
+            expect(mockListKeys).toHaveBeenNthCalledWith(3, { exclusiveStartKey: 'key4', prefix: 'img/' });
 
             expect(results).toHaveLength(5);
             results.forEach((r, i) => {