feat(documents): add abort signal to import and importStream

src/Typesense/Documents.ts

@@ -197,6 +197,10 @@ export interface SearchOptions {
   abortSignal?: AbortSignal | null;
 }
 
+export interface ImportOptions {
+  abortSignal?: AbortSignal | null;
+}
+
 const isNodeJSEnvironment =
   typeof process !== "undefined" &&
   process.versions != null &&
@@ -303,19 +307,23 @@ export default class Documents<T extends DocumentSchema = object>
    * Import a set of documents in a batch.
    * @param {string|Array} documents - Can be a JSONL string of documents or an array of document objects.
    * @param options
+   * @param importOptions - An `abortSignal` cancels the request, and its retries, when aborted.
    * @return {string|Array} Returns a JSONL string if the input was a JSONL string, otherwise it returns an array of results.
    */
   async import(
     documents: string,
     options?: DocumentImportParameters,
+    importOptions?: ImportOptions,
   ): Promise<string>;
   async import(
     documents: T[],
     options?: DocumentImportParameters,
+    importOptions?: ImportOptions,
   ): Promise<ImportResponse<T>[]>;
   async import(
     documents: T[] | string,
     options: DocumentImportParameters = {},
+    { abortSignal = null }: ImportOptions = {},
   ): Promise<string | ImportResponse<T>[]> {
     // Set default value for throwOnFail if not provided
     const finalOptions = { throwOnFail: true, ...options };
@@ -358,6 +366,7 @@ export default class Documents<T extends DocumentSchema = object>
         queryParameters: finalOptions,
         bodyParameters: documentsInJSONLFormat,
         additionalHeaders: { "Content-Type": "text/plain" },
+        abortSignal,
         skipConnectionTimeout: true, // We never want to client-side-timeout on an import and retry, since imports are syncronous and we want to let them take as long as it takes to complete fully
         enableKeepAlive: isNodeJSEnvironment ? true : false, // This is to prevent ECONNRESET socket hang up errors. Reference: https://github.com/axios/axios/issues/2936#issuecomment-779439991
       },
@@ -398,6 +407,7 @@ export default class Documents<T extends DocumentSchema = object>
   async importStream(
     readableStream: ReadStream,
     options: DocumentImportParameters = {},
+    { abortSignal = null }: ImportOptions = {},
   ): Promise<ImportResponse<T>[]> {
     const finalOptions = { throwOnFail: true, ...options };
 
@@ -408,6 +418,7 @@ export default class Documents<T extends DocumentSchema = object>
         queryParameters: finalOptions,
         bodyParameters: readableStream,
         additionalHeaders: { "Content-Type": "text/plain" },
+        abortSignal,
         skipConnectionTimeout: true, // We never want to client-side-timeout on an import and retry, since imports are syncronous and we want to let them take as long as it takes to complete fully
         enableKeepAlive: isNodeJSEnvironment ? true : false, // This is to prevent ECONNRESET socket hang up errors. Reference: https://github.com/axios/axios/issues/2936#issuecomment-779439991
       },

src/Typesense/Types.ts

@@ -9,6 +9,7 @@ import type {
   DeleteQuery,
   DeleteResponse,
   ImportResponse,
+  ImportOptions,
   DocumentsExportParameters,
 } from "./Documents";
 
@@ -222,6 +223,7 @@ export interface WriteableDocuments<T> {
   import(
     documents: T[] | string,
     options: DocumentWriteParameters,
+    importOptions?: ImportOptions,
   ): Promise<string | ImportResponse[]>;
   export(options: DocumentsExportParameters): Promise<string>;
 }

test/Typesense/Documents.spec.ts

@@ -348,6 +348,41 @@ describe("Documents", function () {
       expect(importResult[1].success).toBe(false);
       expect((importResult[1]).error).toBeDefined();
     });
+
+    it("imports the documents when given an abort signal", async function () {
+      const abortController = new AbortController();
+
+      const importResult = await typesense
+        .collections<typeof documentData>(testCollectionName)
+        .documents()
+        .import(
+          [documentData, anotherDocumentData],
+          {},
+          {
+            abortSignal: abortController.signal,
+          },
+        );
+
+      expect(importResult.map((r) => r.success)).toEqual([true, true]);
+    });
+
+    it("imports nothing when the abort signal is already aborted", async function () {
+      const abortController = new AbortController();
+      abortController.abort();
+
+      await expect(
+        typesense
+          .collections<typeof documentData>(testCollectionName)
+          .documents()
+          .import([documentData], {}, { abortSignal: abortController.signal }),
+      ).rejects.toThrow("Request aborted by caller.");
+
+      const exported = await typesense
+        .collections(testCollectionName)
+        .documents()
+        .export();
+      expect(exported).toBe("");
+    });
   });
 
   describe(".export", function () {

test/import-abort.spec.ts

@@ -0,0 +1,196 @@
+import { createReadStream } from "fs";
+import { fileURLToPath } from "url";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { Client as TypesenseClient } from "../src/Typesense";
+import type { WriteableDocuments } from "../src/Typesense/Types";
+
+const { mockAxios, mockCancelFn } = vi.hoisted(() => {
+  const mockCancelFn = vi.fn();
+  const mockCancelSource = {
+    token: "mock-cancel-token",
+    cancel: mockCancelFn,
+  };
+  const mockAxios = vi.fn() as any;
+  mockAxios.CancelToken = {
+    source: vi.fn(() => mockCancelSource),
+  };
+
+  return { mockAxios, mockCancelFn };
+});
+
+vi.mock("axios", () => ({
+  default: mockAxios,
+  __esModule: true,
+}));
+
+describe("Documents Import Abort Signal Tests", () => {
+  let typesense: TypesenseClient;
+
+  const documents = [
+    { id: "1", company_name: "Stark Industries" },
+    { id: "2", company_name: "Acme Corp" },
+  ];
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    typesense = new TypesenseClient({
+      nodes: [
+        {
+          host: "node0",
+          port: 8108,
+          protocol: "http",
+        },
+      ],
+      apiKey: "abcd",
+      randomizeNodes: false,
+      retryIntervalSeconds: 0.05,
+    });
+
+    mockAxios.mockResolvedValue({
+      status: 200,
+      data: '{"success":true}\n{"success":true}',
+    });
+  });
+
+  afterEach(() => {
+    vi.resetAllMocks();
+  });
+
+  it("should reject immediately if AbortSignal is already aborted", async () => {
+    const abortController = new AbortController();
+    abortController.abort();
+
+    await expect(
+      typesense
+        .collections("companies")
+        .documents()
+        .import(documents, {}, { abortSignal: abortController.signal }),
+    ).rejects.toThrow("Request aborted by caller.");
+
+    expect(mockAxios).not.toHaveBeenCalled();
+  });
+
+  it("should create cancel token and keep the signal out of the query parameters", async () => {
+    const abortController = new AbortController();
+
+    const result = await typesense
+      .collections("companies")
+      .documents()
+      .import(
+        documents,
+        { action: "upsert" },
+        { abortSignal: abortController.signal },
+      );
+
+    expect(mockAxios.CancelToken.source).toHaveBeenCalled();
+
+    expect(mockAxios).toHaveBeenCalledWith(
+      expect.objectContaining({
+        cancelToken: "mock-cancel-token",
+        params: expect.objectContaining({ action: "upsert" }),
+      }),
+    );
+    expect(mockAxios.mock.calls[0][0].params).not.toHaveProperty("abortSignal");
+
+    expect(result).toEqual([{ success: true }, { success: true }]);
+  });
+
+  it("should cancel a request in flight and not retry it", async () => {
+    mockAxios.mockImplementationOnce(
+      () =>
+        new Promise((_resolve, reject) => {
+          mockCancelFn.mockImplementationOnce(() =>
+            reject(new Error("canceled")),
+          );
+        }),
+    );
+    const abortController = new AbortController();
+
+    const importing = typesense
+      .collections("companies")
+      .documents()
+      .import(documents, {}, { abortSignal: abortController.signal });
+
+    await vi.waitFor(() => expect(mockAxios).toHaveBeenCalledTimes(1));
+    abortController.abort();
+
+    await expect(importing).rejects.toThrow("Request aborted by caller.");
+    expect(mockCancelFn).toHaveBeenCalledTimes(1);
+    expect(mockAxios).toHaveBeenCalledTimes(1);
+  });
+
+  it("should not retry a failed request once aborted", async () => {
+    const abortController = new AbortController();
+    mockAxios.mockImplementationOnce(async () => {
+      // Fires during the wait before the retry
+      setTimeout(() => abortController.abort(), 0);
+      return { status: 503, data: { message: "Not Ready or Lagging" } };
+    });
+
+    await expect(
+      typesense
+        .collections("companies")
+        .documents()
+        .import(documents, {}, { abortSignal: abortController.signal }),
+    ).rejects.toThrow("Request aborted by caller.");
+
+    expect(mockAxios).toHaveBeenCalledTimes(1);
+  });
+
+  it("should pass the abort signal for a stream import", async () => {
+    const abortController = new AbortController();
+    const stream = createReadStream(fileURLToPath(import.meta.url));
+
+    try {
+      await typesense
+        .collections("companies")
+        .documents()
+        .importStream(stream, {}, { abortSignal: abortController.signal });
+    } finally {
+      stream.destroy();
+    }
+
+    expect(mockAxios).toHaveBeenCalledWith(
+      expect.objectContaining({
+        cancelToken: "mock-cancel-token",
+      }),
+    );
+  });
+
+  it("should accept the abort signal through WriteableDocuments", async () => {
+    const abortController = new AbortController();
+    const writeable: WriteableDocuments<(typeof documents)[number]> = typesense
+      .collections<(typeof documents)[number]>("companies")
+      .documents();
+
+    await writeable.import(
+      documents,
+      {},
+      { abortSignal: abortController.signal },
+    );
+
+    expect(mockAxios).toHaveBeenCalledWith(
+      expect.objectContaining({
+        cancelToken: "mock-cancel-token",
+      }),
+    );
+  });
+
+  it("should work correctly without abort signal", async () => {
+    const result = await typesense
+      .collections("companies")
+      .documents()
+      .import(documents);
+
+    expect(mockAxios.CancelToken.source).not.toHaveBeenCalled();
+
+    expect(mockAxios).toHaveBeenCalledWith(
+      expect.not.objectContaining({
+        cancelToken: expect.anything(),
+      }),
+    );
+
+    expect(result).toEqual([{ success: true }, { success: true }]);
+  });
+});