feat: implement type safety for JSON.GET and JSON.SET

Author: aleclarson

package.json

@@ -58,6 +58,7 @@
     "@cloudflare/workers-types": "^4.20241011.0",
     "@sinclair/typebox": "^0.34.21",
     "@types/node": "^22.7.5",
+    "jsonpath-ts": "^0.1.1",
     "prettier": "^3.3.3",
     "prettier-plugin-organize-imports": "^4.1.0",
     "radashi": "13.0.0-beta.ffa4778",

pnpm-lock.yaml

@@ -1,35 +1,96 @@
 import { StaticEncode, TSchema } from "@sinclair/typebox";
+import { Parse } from "jsonpath-ts";
 import { RedisCommand } from "../command";
-import { RedisKey } from "../key";
+import { RedisJSONEntity } from "../key";
 import { encodeModifiers, Modifiers } from "../modifier";
 import { NX, XX } from "../modifiers";
+import { JSONPath } from "../utils/json-path";
 
 /**
  * Set the JSON value at path in key
  * @see https://redis.io/commands/json.set/
  */
-export function SET<T extends TSchema>(
-  key: RedisKey<T>,
-  path: string,
-  value: StaticEncode<T>,
+export function SET<T extends TSchema, TPath extends JSONPath>(
+  key: RedisJSONEntity<T>,
+  path: TPath,
+  value: Parse<TPath, StaticEncode<T>>,
   ...modifiers: Modifiers<[NX | XX]>
 ) {
   return new RedisCommand<"OK" | null>([
     "JSON.SET",
     key.name,
     path,
-    key.encode(value),
+    key.encode(value, path),
     ...encodeModifiers(modifiers),
   ]);
 }
 
+type ParseMultiGet<TPaths extends JSONPath[], T> = (
+  TPaths extends [
+    infer TPath extends JSONPath,
+    ...infer TPaths extends JSONPath[],
+  ]
+    ? { [P in TPath]: Parse<TPath, T>[] } & ParseMultiGet<TPaths, T>
+    : unknown
+) extends infer TResult
+  ? { [K in keyof TResult]: TResult[K] }
+  : never;
+
+/**
+ * Retrieve the root (`$`) JSON value for the given key.
+ *
+ * @see https://redis.io/commands/json.get/
+ */
+export function GET<T extends TSchema>(
+  key: RedisJSONEntity<T>,
+): RedisCommand<StaticEncode<T> | undefined>;
+
 /**
- * Get the JSON value at path in key
+ * Retrieve the JSON value at the given path in key.
+ *
  * @see https://redis.io/commands/json.get/
  */
-export function GET<T extends TSchema>(key: RedisKey<T>, paths: string[]) {
-  return new RedisCommand<StaticEncode<T> | undefined>(
-    ["JSON.GET", key.name, ...paths],
-    (result) => (result !== null ? key.decode(result) : undefined),
-  );
+export function GET<T extends TSchema, TPath extends JSONPath>(
+  key: RedisJSONEntity<T>,
+  path: TPath,
+): RedisCommand<Parse<TPath, StaticEncode<T>> | undefined>;
+
+/**
+ * Retrieve the JSON values at the given paths in key.
+ *
+ * Returns an object with the paths as keys. Each key points to an array of
+ * parsed JSON values.
+ *
+ * @see https://redis.io/commands/json.get/
+ */
+export function GET<
+  T extends TSchema,
+  TPaths extends [JSONPath, ...JSONPath[]],
+>(
+  key: RedisJSONEntity<T>,
+  paths: TPaths,
+): RedisCommand<ParseMultiGet<TPaths, StaticEncode<T>> | undefined>;
+
+export function GET(key: RedisJSONEntity, path?: JSONPath | JSONPath[]) {
+  if (!path) {
+    return new RedisCommand(["JSON.GET", key.name], (result) =>
+      result !== null ? JSON.parse(result) : undefined,
+    );
+  }
+  if (!Array.isArray(path)) {
+    return new RedisCommand(
+      ["JSON.GET", key.name],
+      (result) => JSON.parse(result)[0],
+    );
+  }
+  const paths: JSONPath[] = path;
+  if (!paths.length) {
+    throw new Error("Expected at least one path");
+  }
+  return new RedisCommand(["JSON.GET", key.name, ...paths], (results) => {
+    // Normalize the results to an object with the paths as keys.
+    return paths.length > 1
+      ? JSON.parse(results)
+      : { [paths[0]]: JSON.parse(results) };
+  });
 }

src/commands/json.ts

@@ -3,6 +3,7 @@ import { StaticEncode, TObject, TSchema } from "@sinclair/typebox/type";
 import { Decode, Encode } from "@sinclair/typebox/value";
 import { RedisValue } from "./command";
 import { RedisTransform } from "./transform";
+import { JSONPath, resolveSchemaForJSONPath } from "./utils/json-path";
 
 /**
  * Represents a namespace of keys in a Redis database.
@@ -81,6 +82,40 @@ export class RedisEntity<
   };
 }
 
+/**
+ * Represents a key in a Redis database with a JSON value.
+ *
+ * You must use `JSON.*` commands to read or edit this key's value. Of
+ * course, you may use key-focused commands (e.g. `DEL` may delete this
+ * key).
+ *
+ * @see https://redis.io/docs/latest/develop/data-types/json/
+ */
+export class RedisJSONEntity<
+  T extends TSchema = TSchema,
+  K extends string | RedisKeyspace = string,
+> extends RedisKey<T, K> {
+  declare $$typeof: "RedisKey" & { subtype: "RedisJSONEntity" };
+
+  encode(value: unknown, path?: JSONPath): RedisValue {
+    // The schema is defined for JS, not Redis, so a "decoded" value
+    // represents a Redis value.
+    return Decode(
+      path ? resolveSchemaForJSONPath(this.schema, path) : this.schema,
+      value,
+    );
+  }
+
+  decode(value: unknown, path?: JSONPath): StaticEncode<T> {
+    // The schema is defined for JS, not Redis, so an "encoded" value
+    // represents a JS value.
+    return Encode(
+      path ? resolveSchemaForJSONPath(this.schema, path) : this.schema,
+      value,
+    );
+  }
+}
+
 export class RedisSet<
   T extends TSchema = TSchema,
   K extends string | RedisKeyspace = string,

src/key.ts

@@ -0,0 +1,46 @@
+import { KindGuard, TSchema } from "@sinclair/typebox";
+import * as Type from "@sinclair/typebox/type";
+
+export type JSONPath = `$${string}`;
+
+/**
+ * Given a TypeBox schema and a JSON path, return the TypeBox schema that
+ * corresponds to the JSON path.
+ */
+export function resolveSchemaForJSONPath(
+  schema: TSchema,
+  path: JSONPath,
+  state = parseJSONPath(path),
+) {
+  const part = state.parts[state.index++];
+  if (!part) {
+    return schema;
+  }
+  if (part === "$") {
+    return resolveSchemaForJSONPath(schema, path, state);
+  }
+  let key: string | undefined;
+  if (part.charCodeAt(0) === 46 /* . */) {
+    key = part.slice(1);
+  } else if (part.charCodeAt(0) === 91 /* [ */) {
+    if (part.charCodeAt(1) === 39 /* ' */) {
+      key = part.slice(2, -2);
+    }
+    // If no single quote is present, then this is either [*] or a specific
+    // array index, like [0] or [1].
+    else if (KindGuard.IsArray(schema)) {
+      return resolveSchemaForJSONPath(schema.items, path, state);
+    }
+  }
+  if (key != null && KindGuard.IsObject(schema) && key in schema.properties) {
+    return resolveSchemaForJSONPath(schema.properties[key], path, state);
+  }
+  return Type.Never();
+}
+
+function parseJSONPath(path: JSONPath) {
+  return {
+    parts: path.match(/(^\$|\.[^.[\]*]+|\[(?:\*|\d+|'[^']+')\])/gi) ?? [],
+    index: 0,
+  };
+}