From 2a0c155e4153f39f15fd671b5557cfe1294597e5 Mon Sep 17 00:00:00 2001 From: Kitson Kelly Date: Thu, 28 Mar 2024 13:55:35 +1100 Subject: [PATCH] feat: add blob support to uniqueCount --- keys.test.ts | 101 +++++++++++++++++++++++++++++++++++++++++++++++++++ keys.ts | 72 ++++++++++++++++++++++++++++-------- 2 files changed, 158 insertions(+), 15 deletions(-) diff --git a/keys.test.ts b/keys.test.ts index e2f4c06..1640e3a 100644 --- a/keys.test.ts +++ b/keys.test.ts @@ -1,4 +1,5 @@ import { assert, assertEquals, setup, teardown } from "./_test_util.ts"; +import { batchedAtomic } from "./batched_atomic.ts"; import { equals, @@ -120,6 +121,106 @@ Deno.test({ }, }); +Deno.test({ + name: "uniqueCount - detects blobs", + async fn() { + const kv = await setup(); + const res = await batchedAtomic(kv) + .set(["a"], "a") + .set(["a", "b"], "b") + .set(["a", "b", "c"], "c") + .set(["a", "d", "f", "g"], "g") + .set(["a", "h"], "h") + .setBlob(["a", "i"], new Uint8Array([1, 2, 3])) + .set(["a", "i", "j"], "j") + .set(["e"], "e") + .commit(); + assert(res.every(({ ok }) => ok)); + + const actual = await uniqueCount(kv, ["a"]); + + assertEquals(actual, [ + { key: ["a", "b"], count: 1 }, + { key: ["a", "d"], count: 1 }, + { key: ["a", "h"], count: 0 }, + { key: ["a", "i"], count: 1, isBlob: true }, + ]); + + return teardown(); + }, +}); + +Deno.test({ + name: "uniqueCount - ignores blob keys", + async fn() { + const kv = await setup(); + const res = await batchedAtomic(kv) + .set(["a"], "a") + .set(["a", "b"], "b") + .set(["a", "b", "c"], "c") + .set(["a", "d", "f", "g"], "g") + .set(["a", "h"], "h") + .setBlob(["a", "i"], new Uint8Array([1, 2, 3])) + .set(["a", "i", "j"], "j") + .set(["e"], "e") + .commit(); + assert(res.every(({ ok }) => ok)); + + const actual = await uniqueCount(kv, ["a", "i"]); + + assertEquals(actual, [{ key: ["a", "i", "j"], count: 0 }]); + + return teardown(); + }, +}); + +Deno.test({ + name: "uniqueCount - handles Uint8Array equality with blobs", + async fn() { + const kv = await setup(); + const res = await batchedAtomic(kv) + .set(["a"], "a") + .setBlob(["a", new Uint8Array([2, 3, 4])], new Uint8Array([1, 2, 3])) + .set(["a", new Uint8Array([2, 3, 4]), "c"], "c") + .set(["a", new Uint8Array([4, 5, 6]), "c"], "c") + .set(["e"], "e") + .commit(); + assert(res.every(({ ok }) => ok)); + + const actual = await uniqueCount(kv, ["a"]); + + assertEquals(actual, [ + { key: ["a", new Uint8Array([2, 3, 4])], count: 1, isBlob: true }, + { key: ["a", new Uint8Array([4, 5, 6])], count: 1 }, + ]); + return teardown(); + }, +}); + +Deno.test({ + name: "uniqueCount - ignores blob keys with Uint8Array key parts", + async fn() { + const kv = await setup(); + const res = await batchedAtomic(kv) + .set(["a"], "a") + .setBlob(["a", new Uint8Array([2, 3, 4])], new Uint8Array([1, 2, 3])) + .set(["a", new Uint8Array([2, 3, 4]), "c"], "c") + .set(["a", new Uint8Array([4, 5, 6]), "c"], "c") + .set(["e"], "e") + .commit(); + assert(res.every(({ ok }) => ok)); + + const actual = await uniqueCount(kv, ["a", new Uint8Array([2, 3, 4])]); + + assertEquals(actual, [{ + key: ["a", new Uint8Array([2, 3, 4]), "c"], + count: 0, + }]); + + return teardown(); + }, +}); + Deno.test({ name: "equals", fn() { diff --git a/keys.ts b/keys.ts index f5bff61..f22c86b 100644 --- a/keys.ts +++ b/keys.ts @@ -146,6 +146,8 @@ import { timingSafeEqual } from "jsr:@std/crypto@0.220/timing_safe_equal"; +import { BLOB_KEY, BLOB_META_KEY } from "./blob_util.ts"; + function addIfUnique(set: Set, item: Uint8Array) { for (const i of set) { if (ArrayBuffer.isView(i) && timingSafeEqual(i, item)) { @@ -156,17 +158,29 @@ function addIfUnique(set: Set, item: Uint8Array) { } function addOrIncrement( - map: Map, + map: Map, item: Uint8Array, - increment: boolean, + next: Deno.KvKeyPart | undefined, ) { + let count = 0; + let isBlob = false; + if (next) { + if (next === BLOB_KEY) { + isBlob = true; + } else if (next !== BLOB_META_KEY) { + count = 1; + } + } for (const [k, v] of map) { if (ArrayBuffer.isView(k) && timingSafeEqual(k, item)) { - map.set(k, increment ? v + 1 : v); + if (isBlob) { + v.isBlob = true; + } + v.count = count; return; } } - map.set(item, increment ? 1 : 0); + map.set(item, isBlob ? { count, isBlob } : { count }); } /** Determines if one {@linkcode Deno.KvKeyPart} equals another. This is more @@ -318,6 +332,9 @@ export async function unique( throw new TypeError(`Unexpected key length of ${key.length}.`); } const part = key[prefixLength]; + if (part === BLOB_KEY || part === BLOB_META_KEY) { + continue; + } if (ArrayBuffer.isView(part)) { addIfUnique(prefixes, part); } else { @@ -327,6 +344,17 @@ export async function unique( return [...prefixes].map((part) => [...prefix, part]); } +/** Elements of an array that gets resolved when calling + * {@linkcode uniqueCount}. */ +export interface UniqueCountElement { + /** The key of the element. */ + key: Deno.KvKey; + /** The number of sub-keys the key has. */ + count: number; + /** Indicates if the value of the key is a kv-toolbox blob value. */ + isBlob?: boolean; +} + /** Resolves with an array of unique sub keys/prefixes for the provided prefix * along with the number of sub keys that match that prefix. The `count` * represents the number of sub keys, a value of `0` indicates that only the @@ -336,7 +364,11 @@ export async function unique( * where you are retrieving a list including counts and you want to know all the * unique _descendants_ of a key in order to be able to enumerate them. * - * For example if you had the following keys stored in a datastore: + * If you omit a `prefix`, all unique root keys are resolved. + * + * @example + * + * If you had the following keys stored in a datastore: * * ```ts * ["a", "b"] @@ -345,7 +377,7 @@ export async function unique( * ["a", "d", "f"] * ``` * - * And you would get the following results when using `unique()`: + * And you would get the following results when using `uniqueCount()`: * * ```ts * import { uniqueCount } from "jsr:@kitsonk/kv-toolbox/keys"; @@ -356,36 +388,46 @@ export async function unique( * // { key: ["a", "d"], count: 2 } * await kv.close(); * ``` - * - * If you omit a `prefix`, all unique root keys are resolved. */ export async function uniqueCount( kv: Deno.Kv, prefix: Deno.KvKey = [], options?: Deno.KvListOptions, -): Promise<{ key: Deno.KvKey; count: number }[]> { +): Promise { const list = kv.list({ prefix }, options); const prefixLength = prefix.length; - const prefixCounts = new Map(); + const prefixCounts = new Map< + Deno.KvKeyPart, + { count: number; isBlob?: boolean } + >(); for await (const { key } of list) { if (key.length <= prefixLength) { throw new TypeError(`Unexpected key length of ${key.length}.`); } const part = key[prefixLength]; + if (part === BLOB_KEY || part === BLOB_META_KEY) { + continue; + } + const next = key[prefixLength + 1]; if (ArrayBuffer.isView(part)) { - addOrIncrement(prefixCounts, part, key.length > (prefixLength + 1)); + addOrIncrement(prefixCounts, part, next); } else { if (!prefixCounts.has(part)) { - prefixCounts.set(part, 0); + prefixCounts.set(part, { count: 0 }); } - if (key.length > (prefixLength + 1)) { - prefixCounts.set(part, prefixCounts.get(part)! + 1); + if (next) { + const count = prefixCounts.get(part)!; + if (next === BLOB_KEY) { + count.isBlob = true; + } else if (next !== BLOB_META_KEY) { + count.count++; + } } } } return [...prefixCounts].map(([part, count]) => ({ key: [...prefix, part], - count, + ...count, })); }