upstash · ogzhanolguncu · Oct 26, 2023 · Oct 23, 2023 · Oct 26, 2023
diff --git a/pkg/commands/geo_search_store.test.ts b/pkg/commands/geo_search_store.test.ts
@@ -0,0 +1,119 @@
+import { expect, test, describe, afterAll } from "bun:test";
+import { keygen, newHttpClient } from "../test-utils.ts";
+
+import { GeoAddCommand } from "./geo_add.ts";
+import { GeoSearchStoreCommand } from "./geo_search_store.ts";
+import { ZRangeCommand } from "./zrange.ts";
+
+const client = newHttpClient();
+const { newKey, cleanup } = keygen();
+afterAll(cleanup);
+
+describe("GEOSSEARCHSTORE tests", () => {
+  test("should return members within the radius and store them in sorted set", async () => {
+    const key = newKey();
+    const destination = newKey();
+
+    await new GeoAddCommand([
+      key,
+      { longitude: -73.9857, latitude: 40.7488, member: "Empire State Building" },
+      { longitude: -74.0445, latitude: 40.6892, member: "Statue of Liberty" },
+      { longitude: -73.9632, latitude: 40.7789, member: "Central Park" },
+      { longitude: -73.873, latitude: 40.7769, member: "LaGuardia Airport" },
+      { longitude: -74.177, latitude: 40.6413, member: "JFK Airport" },
+      { longitude: -73.9772, latitude: 40.7527, member: "Grand Central Terminal" },
+    ]).exec(client);
+
+    const res = await new GeoSearchStoreCommand([
+      destination,
+      key,
+      { type: "FROMMEMBER", member: "Empire State Building" },
+      { type: "BYRADIUS", radius: 5, radiusType: "KM" },
+      "ASC",
+    ]).exec(client);
+    const zrangeRes = await new ZRangeCommand([destination, 0, -1, { withScores: true }]).exec(
+      client
+    );
+    expect(zrangeRes).toEqual([
+      "Empire State Building",
+      1791875672666387,
+      "Grand Central Terminal",
+      1791875708058440,
+      "Central Park",
+      1791875790048608,
+    ]);
+    expect(res).toEqual(zrangeRes.length / 2);
+  });
+
+  test("should store geosearch in sorted set with distances", async () => {
+    const key = newKey();
+    const destination = newKey();
+
+    await new GeoAddCommand([
+      key,
+      { longitude: -73.9857, latitude: 40.7488, member: "Empire State Building" },
+      { longitude: -74.0445, latitude: 40.6892, member: "Statue of Liberty" },
+      { longitude: -73.9632, latitude: 40.7789, member: "Central Park" },
+      { longitude: -73.873, latitude: 40.7769, member: "LaGuardia Airport" },
+      { longitude: -74.177, latitude: 40.6413, member: "JFK Airport" },
+      { longitude: -73.9772, latitude: 40.7527, member: "Grand Central Terminal" },
+    ]).exec(client);
+
+    const res = await new GeoSearchStoreCommand([
+      destination,
+      key,
+      { type: "FROMMEMBER", member: "Empire State Building" },
+      { type: "BYRADIUS", radius: 5, radiusType: "KM" },
+      "ASC",
+      { storeDist: true },
+    ]).exec(client);
+    const zrangeRes = await new ZRangeCommand([destination, 0, -1, { withScores: true }]).exec(
+      client
+    );
+    expect(zrangeRes).toEqual([
+      "Empire State Building",
+      0,
+      "Grand Central Terminal",
+      "0.83757447438393129",
+      "Central Park",
+      "3.8473905221815641",
+    ]);
+    expect(res).toEqual(zrangeRes.length / 2);
+  });
+
+  test("should return object members within the radius and store them in sorted set with distance and members", async () => {
+    const key = newKey();
+    const destination = newKey();
+
+    await new GeoAddCommand<{ name: string }>([
+      key,
+      { longitude: -73.9857, latitude: 40.7488, member: { name: "Empire State Building" } },
+      { longitude: -74.0445, latitude: 40.6892, member: { name: "Statue of Liberty" } },
+      { longitude: -73.9632, latitude: 40.7789, member: { name: "Central Park" } },
+      { longitude: -73.873, latitude: 40.7769, member: { name: "LaGuardia Airport" } },
+      { longitude: -74.177, latitude: 40.6413, member: { name: "JFK Airport" } },
+      { longitude: -73.9772, latitude: 40.7527, member: { name: "Grand Central Terminal" } },
+    ]).exec(client);
+
+    const res = await new GeoSearchStoreCommand([
+      destination,
+      key,
+      { type: "FROMMEMBER", member: { name: "Empire State Building" } },
+      { type: "BYRADIUS", radius: 5, radiusType: "KM" },
+      "DESC",
+      { storeDist: true },
+    ]).exec(client);
+    const zrangeRes = await new ZRangeCommand([destination, 0, -1, { withScores: true }]).exec(
+      client
+    );
+    expect(zrangeRes).toEqual([
+      { name: "Empire State Building" },
+      0,
+      { name: "Grand Central Terminal" },
+      "0.83757447438393129",
+      { name: "Central Park" },
+      "3.8473905221815641",
+    ]);
+    expect(res).toEqual(zrangeRes.length / 2);
+  });
+});
diff --git a/pkg/commands/geo_search_store.ts b/pkg/commands/geo_search_store.ts
@@ -0,0 +1,68 @@
+import { Command, CommandOptions } from "./command.ts";
+
+type RadiusOptions = "M" | "KM" | "FT" | "MI";
+type CenterPoint<TMemberType> =
+  | {
+      type: "FROMMEMBER" | "frommember";
+      member: TMemberType;
+    }
+  | {
+      type: "FROMLONLAT" | "fromlonlat";
+      coordinate: { lon: number; lat: number };
+    };
+
+type Shape =
+  | { type: "BYRADIUS" | "byradius"; radius: number; radiusType: RadiusOptions }
+  | {
+      type: "BYBOX" | "bybox";
+      rect: { width: number; height: number };
+      rectType: RadiusOptions;
+    };
+
+type GeoSearchCommandOptions = {
+  count?: { limit: number; any?: boolean };
+  storeDist?: boolean;
+};
+
+/**
+ * @see https://redis.io/commands/geosearchstore
+ */
+export class GeoSearchStoreCommand<
+  TMemberType = string,
+  TOptions extends GeoSearchCommandOptions = GeoSearchCommandOptions
+> extends Command<any[] | any[][], number> {
+  constructor(
+    [destination, key, centerPoint, shape, order, opts]: [
+      destination: string,
+      key: string,
+      centerPoint: CenterPoint<TMemberType>,
+      shape: Shape,
+      order: "ASC" | "DESC" | "asc" | "desc",
+      opts?: TOptions
+    ],
+    commandOptions?: CommandOptions<any[] | any[][], number>
+  ) {
+    const command: unknown[] = ["GEOSEARCHSTORE", destination, key];
+
+    if (centerPoint.type === "FROMMEMBER" || centerPoint.type === "frommember") {
+      command.push(centerPoint.type, centerPoint.member);
+    }
+    if (centerPoint.type === "FROMLONLAT" || centerPoint.type === "fromlonlat") {
+      command.push(centerPoint.type, centerPoint.coordinate.lon, centerPoint.coordinate.lat);
+    }
+
+    if (shape.type === "BYRADIUS" || shape.type === "byradius") {
+      command.push(shape.type, shape.radius, shape.radiusType);
+    }
+    if (shape.type === "BYBOX" || shape.type === "bybox") {
+      command.push(shape.type, shape.rect.width, shape.rect.height, shape.rectType);
+    }
+    command.push(order);
+
+    if (opts?.count) {
+      command.push(opts.count.limit, ...(opts.count.any ? ["ANY"] : []));
+    }
+
+    super([...command, ...(opts?.storeDist ? ["STOREDIST"] : [])], commandOptions);
+  }
+}
diff --git a/pkg/commands/mod.ts b/pkg/commands/mod.ts
@@ -20,6 +20,7 @@ export * from "./geo_dist";
 export * from "./geo_pos";
 export * from "./geo_hash";
 export * from "./geo_search";
+export * from "./geo_search_store";
 export * from "./get";
 export * from "./getbit";
 export * from "./getdel";

diff --git a/pkg/pipeline.ts b/pkg/pipeline.ts
@@ -20,6 +20,7 @@ import {
   GeoHashCommand,
   GeoAddCommand,
   GeoDistCommand,
+  GeoSearchStoreCommand,
   GeoSearchCommand,
   GeoPosCommand,
   GetBitCommand,
@@ -1138,6 +1139,12 @@ export class Pipeline<TCommands extends Command<any, any>[] = []> {
       geosearch: (...args: CommandArgs<typeof GeoSearchCommand>) =>
         new GeoSearchCommand(args, this.commandOptions).exec(this.client),
 
+      /**
+       * @see https://redis.io/commands/geosearchstore
+       */
+      geosearchstore: (...args: CommandArgs<typeof GeoSearchStoreCommand>) =>
+        new GeoSearchStoreCommand(args, this.commandOptions).exec(this.client),
+
       /**
        * @see https://redis.io/commands/json.get
        */

diff --git a/pkg/redis.ts b/pkg/redis.ts
@@ -18,6 +18,7 @@ import {
   FlushDBCommand,
   GeoAddCommand,
   GeoDistCommand,
+  GeoSearchStoreCommand,
   GeoSearchCommand,
   GeoHashCommand,
   GeoPosCommand,
@@ -276,6 +277,12 @@ export class Redis {
       geosearch: (...args: CommandArgs<typeof GeoSearchCommand>) =>
         new GeoSearchCommand(args, this.opts).exec(this.client),
 
+      /**
+       * @see https://redis.io/commands/geosearchstore
+       */
+      geosearchstore: (...args: CommandArgs<typeof GeoSearchStoreCommand>) =>
+        new GeoSearchStoreCommand(args, this.opts).exec(this.client),
+
       /**
        * @see https://redis.io/commands/json.get
        */