feat: add blob support to uniqueCount

Author: kitsonk

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() {

keys.ts

@@ -146,6 +146,8 @@
 
 import { timingSafeEqual } from "jsr:@std/[email protected]/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)) {
@@ -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
@@ -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<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,
   }));
 }