Merge pull request #43 from upstash/DX-1225

Add reset all namespaces option

src/commands/client/reset/index.test.ts

@@ -1,7 +1,7 @@
-import { describe, expect, test } from "bun:test";
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
 import type { FetchResult } from "@commands/index";
 import { FetchCommand, ResetCommand, UpsertCommand } from "@commands/index";
-import { awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils";
+import { Index, awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -35,8 +35,73 @@ describe("RESET", () => {
         },
       ]).exec(client);
 
-      expect(resAfterReset).toEqual(Array.from({ length: 20 }).fill(null) as FetchResult[]);
+      expect(resAfterReset).toEqual(new Array(20).fill(null));
     },
     { timeout: 30_000 }
   );
+
+  describe("reset options", () => {
+    const namespaces = ["ns-1", "ns-2"];
+    const index = new Index({
+      token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+      url: process.env.UPSTASH_VECTOR_REST_URL!,
+    });
+
+    const vectorCount = 5;
+    const vectorIds = new Array(vectorCount).fill("").map((_, index) => `vector-${index}`);
+
+    beforeEach(async () => {
+      // insert vertors to the default namespace and to the two namespaces
+      await Promise.all(
+        [undefined, ...namespaces].map(async (ns) => {
+          const randomizedData = new Array(vectorCount)
+            .fill("")
+            .map((_, index) => ({ id: vectorIds[index], vector: range(0, 384) }));
+
+          await new UpsertCommand(randomizedData, { namespace: ns }).exec(client);
+          await awaitUntilIndexed(client);
+        })
+      );
+
+      const info = await index.info();
+      expect(info.namespaces[""].vectorCount).toBe(vectorCount);
+      expect(info.namespaces[namespaces[0]].vectorCount).toBe(vectorCount);
+      expect(info.namespaces[namespaces[1]].vectorCount).toBe(vectorCount);
+    });
+
+    afterEach(async () => {
+      await index.reset({ all: true });
+
+      // make sure that index is empty
+      const info = await index.info();
+      expect(info.vectorCount).toBe(0);
+    });
+
+    test("should reset default namespace", async () => {
+      await index.reset();
+      const info = await index.info();
+
+      expect(info.namespaces[""].vectorCount).toBe(0);
+      expect(info.namespaces[namespaces[0]].vectorCount).toBe(vectorCount);
+      expect(info.namespaces[namespaces[1]].vectorCount).toBe(vectorCount);
+    });
+
+    test("should reset given namespace", async () => {
+      await index.reset({ namespace: namespaces[0] });
+      const info = await index.info();
+
+      expect(info.namespaces[""].vectorCount).toBe(vectorCount);
+      expect(info.namespaces[namespaces[0]].vectorCount).toBe(0);
+      expect(info.namespaces[namespaces[1]].vectorCount).toBe(vectorCount);
+    });
+
+    test("should reset all namespaces", async () => {
+      await index.reset({ all: true });
+      const info = await index.info();
+
+      expect(info.namespaces[""].vectorCount).toBe(0);
+      expect(info.namespaces[namespaces[0]].vectorCount).toBe(0);
+      expect(info.namespaces[namespaces[1]].vectorCount).toBe(0);
+    });
+  });
 });

src/commands/client/reset/index.ts

@@ -1,16 +1,21 @@
 import type { NAMESPACE } from "@commands/client/types";
 import { Command } from "@commands/command";
 
-type ResetEndpointVariants = `reset` | `reset/${NAMESPACE}`;
+type ResetEndpointVariants = `reset` | `reset/${NAMESPACE}` | `reset?all`;
 
-type ResetCommandOptions = { namespace?: string };
+export type ResetCommandOptions =
+  | { namespace?: string; all?: never }
+  | { namespace?: never; all?: true };
 export class ResetCommand extends Command<string> {
   constructor(options?: ResetCommandOptions) {
     let endpoint: ResetEndpointVariants = "reset";
 
     if (options?.namespace) {
       endpoint = `${endpoint}/${options.namespace}`;
+    } else if (options?.all) {
+      endpoint = `${endpoint}?all`;
     }
+
     super([], endpoint);
   }
 }

src/commands/command.ts

@@ -24,7 +24,9 @@ const _ENDPOINTS = [
 
 export type EndpointVariants =
   | (typeof _ENDPOINTS)[number]
-  | `${(typeof _ENDPOINTS)[number]}/${NAMESPACE}`;
+  | `${(typeof _ENDPOINTS)[number]}/${NAMESPACE}`
+  | `reset?all`; // here to make sure that reset?all/<namespace> is invalid
+
 /**
  * TResult is the raw data returned from upstash, which may need to be transformed or parsed.
  */

src/vector.ts

@@ -7,6 +7,7 @@ import {
   QueryManyCommand,
   RangeCommand,
   ResetCommand,
+  ResetCommandOptions,
   UpdateCommand,
   UpsertCommand,
 } from "@commands/client";
@@ -239,17 +240,35 @@ export class Index<TIndexMetadata extends Dict = Dict> {
     new FetchCommand<TMetadata>(args).exec(this.client);
 
   /**
-   * It's used for wiping an entire index.
+   * It's used for wiping the index.
+   *
+   * By default, resets the default namespace:
    *
    * @example
    * ```js
    * await index.reset();
-   * console.log('Index has been reset');
+   * console.log('Default namespace has been reset');
+   * ```
+   *
+   * To reset a namespace, call reset like:
+   *
+   * @example
+   * ```js
+   * await index.reset({ namespace: "ns" });
+   * console.log('Namespace ns has been reset');
+   * ```
+   *
+   * If you want to reset all namespaces, call reset like:
+   *
+   * @example
+   * ```js
+   * await index.reset({ all: true });
+   * console.log('All namespaces have been reset');
    * ```
    *
    * @returns {Promise<string>} A promise that resolves with the result of the reset operation after the command is executed.
    */
-  reset = (options?: { namespace?: string }) => new ResetCommand(options).exec(this.client);
+  reset = (options?: ResetCommandOptions) => new ResetCommand(options).exec(this.client);
 
   /**
    * Retrieves a range of items from the index.
@@ -266,6 +285,12 @@ export class Index<TIndexMetadata extends Dict = Dict> {
    * console.log(rangeResults); // Outputs the result of the range operation
    * ```
    *
+   * You can also pass a namespace like:
+   *
+   * ```js
+   * const rangeResults = await index.range(rangeArgs, { namespace: "ns" });
+   * ```
+   *
    * @param {CommandArgs<typeof RangeCommand>} args - The arguments for the range command.
    * @param {number|string} args.cursor - The starting point (cursor) for the range query.
    * @param {number} args.limit - The maximum number of items to return in this range.