Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add reset all namespaces option #43

Merged
merged 4 commits into from
Sep 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 68 additions & 3 deletions src/commands/client/reset/index.test.ts
Original file line number Diff line number Diff line change
@@ -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();

Expand Down Expand Up @@ -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);
});
});
});
9 changes: 7 additions & 2 deletions src/commands/client/reset/index.ts
Original file line number Diff line number Diff line change
@@ -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);
}
}
4 changes: 3 additions & 1 deletion src/commands/command.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
Expand Down
31 changes: 28 additions & 3 deletions src/vector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import {
QueryManyCommand,
RangeCommand,
ResetCommand,
ResetCommandOptions,
UpdateCommand,
UpsertCommand,
} from "@commands/client";
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand Down
Loading