Skip to content

Commit

Permalink
feat: add blob support to uniqueCount
Browse files Browse the repository at this point in the history
  • Loading branch information
kitsonk committed Mar 28, 2024
1 parent cffeaf5 commit 2a0c155
Show file tree
Hide file tree
Showing 2 changed files with 158 additions and 15 deletions.
101 changes: 101 additions & 0 deletions keys.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { assert, assertEquals, setup, teardown } from "./_test_util.ts";
import { batchedAtomic } from "./batched_atomic.ts";

import {
equals,
Expand Down Expand Up @@ -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() {
Expand Down
72 changes: 57 additions & 15 deletions keys.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<Deno.KvKeyPart>, item: Uint8Array) {
for (const i of set) {
if (ArrayBuffer.isView(i) && timingSafeEqual(i, item)) {
Expand All @@ -156,17 +158,29 @@ function addIfUnique(set: Set<Deno.KvKeyPart>, item: Uint8Array) {
}

function addOrIncrement(
map: Map<Deno.KvKeyPart, number>,
map: Map<Deno.KvKeyPart, { count: number; isBlob?: boolean }>,
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
Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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"]
Expand All @@ -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";
Expand All @@ -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<UniqueCountElement[]> {
const list = kv.list({ prefix }, options);
const prefixLength = prefix.length;
const prefixCounts = new Map<Deno.KvKeyPart, number>();
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,
}));
}

Expand Down

0 comments on commit 2a0c155

Please sign in to comment.