lancedb/nodejs/__test__/remote.test.ts

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors

import * as http from "http";
import { RequestListener } from "http";
import packageJson = require("../package.json");
import {
  ClientAuthMethod,
  ClientConfig,
  Connection,
  ConnectionOptions,
  OAuthConfig,
  OAuthFlowType,
  TlsConfig,
  connect,
} from "../lancedb";
import {
  HeaderProvider,
  OAuthHeaderProvider,
  StaticHeaderProvider,
} from "../lancedb/header";
import { Index } from "../lancedb/indices";

// Test-only header providers
class CustomProvider extends HeaderProvider {
  getHeaders(): Record<string, string> {
    return { "X-Custom": "custom-value" };
  }
}

class ErrorProvider extends HeaderProvider {
  private errorMessage: string;
  public callCount: number = 0;

  constructor(errorMessage: string = "Test error") {
    super();
    this.errorMessage = errorMessage;
  }

  getHeaders(): Record<string, string> {
    this.callCount++;
    throw new Error(this.errorMessage);
  }
}

class ConcurrentProvider extends HeaderProvider {
  private counter: number = 0;

  getHeaders(): Record<string, string> {
    this.counter++;
    return { "X-Request-Id": String(this.counter) };
  }
}

async function withMockDatabase(
  listener: RequestListener,
  callback: (db: Connection) => void,
  connectionOptions?: ConnectionOptions,
) {
  const server = http.createServer(listener);
  server.listen(8000);

  const db = await connect(
    "db://dev",
    Object.assign(
      {
        apiKey: "fake",
        hostOverride: "http://localhost:8000",
      },
      connectionOptions,
    ),
  );

  try {
    await callback(db);
  } finally {
    // `close()` alone leaves the port bound until keep-alive sockets drain, so
    // a single failing test would cascade into EADDRINUSE for every test after
    // it. Destroy the connections and wait for the port to actually be free.
    await new Promise<void>((resolve) => {
      server.closeAllConnections();
      server.close(() => resolve());
    });
  }
}

describe("remote connection", () => {
  it("lists materialized views through the namespace route", async () => {
    await withMockDatabase(
      (req, res) => {
        expect(req.method).toBe("GET");
        expect(req.url).toBe("/v1/namespace/$/materialized_view/list");
        res
          .writeHead(200, { "content-type": "application/json" })
          .end(JSON.stringify({ views: ["daily_sales"] }));
      },
      async (db) => {
        expect(await db.listMaterializedViews()).toEqual(["daily_sales"]);
      },
    );
  });

  it("should accept partial connection options", async () => {
    await connect("db://test", {
      apiKey: "fake",
      clientConfig: {
        timeoutConfig: { readTimeout: 5 },
        retryConfig: { retries: 2 },
      },
    });
  });

  it("should accept overall timeout configuration", async () => {
    await connect("db://test", {
      apiKey: "fake",
      clientConfig: {
        timeoutConfig: { timeout: 30 },
      },
    });

    // Test with all timeout parameters
    await connect("db://test", {
      apiKey: "fake",
      clientConfig: {
        timeoutConfig: {
          timeout: 60,
          connectTimeout: 10,
          readTimeout: 20,
          poolIdleTimeout: 300,
        },
      },
    });
  });

  it("should pass down apiKey and userAgent", async () => {
    await withMockDatabase(
      (req, res) => {
        expect(req.headers["x-api-key"]).toEqual("fake");
        expect(req.headers["user-agent"]).toEqual(
          `LanceDB-Node-Client/${packageJson.version}`,
        );

        const body = JSON.stringify({ tables: [] });
        res.writeHead(200, { "Content-Type": "application/json" }).end(body);
      },
      async (db) => {
        const tableNames = await db.tableNames();
        expect(tableNames).toEqual([]);
      },
    );
  });

  it("allows customizing user agent", async () => {
    await withMockDatabase(
      (req, res) => {
        expect(req.headers["user-agent"]).toEqual("MyApp/1.0");

        const body = JSON.stringify({ tables: [] });
        res.writeHead(200, { "Content-Type": "application/json" }).end(body);
      },
      async (db) => {
        const tableNames = await db.tableNames();
        expect(tableNames).toEqual([]);
      },
      {
        clientConfig: {
          userAgent: "MyApp/1.0",
        },
      },
    );
  });

  it("shows the full error messages on retry errors", async () => {
    await withMockDatabase(
      (_req, res) => {
        // We retry on 500 errors, so we return 500s until the client gives up.
        res.writeHead(500).end("Internal Server Error");
      },
      async (db) => {
        try {
          await db.tableNames();
          fail("expected an error");
          // biome-ignore lint/suspicious/noExplicitAny: skip
        } catch (e: any) {
          expect(e.message).toContain("Hit retry limit for request_id=");
          expect(e.message).toContain("Caused by: Http error");
          expect(e.message).toContain("500 Internal Server Error");
        }
      },
      {
        clientConfig: {
          retryConfig: { retries: 2 },
        },
      },
    );
  });

  it("lists the rows a Function refresh skipped", async () => {
    const bodies: unknown[] = [];
    await withMockDatabase(
      (req, res) => {
        const path = req.url ?? "";
        if (path.endsWith("/describe/")) {
          res.writeHead(200, { "Content-Type": "application/json" }).end(
            JSON.stringify({
              name: "docs",
              version: 1,
              schema: { fields: [] },
            }),
          );
          return;
        }
        if (path === "/v1/table/docs/errors") {
          let body = "";
          req.on("data", (chunk) => {
            body += chunk;
          });
          req.on("end", () => {
            bodies.push(JSON.parse(body));
            res.writeHead(200, { "Content-Type": "application/json" }).end(
              `{"records": [{"job_id": "j-7", "fragment_id": 3, "row_offset": 9,
                "column": "embedding", "function": "embed", "function_version": "2",
                "table_version": 11, "error_type": "ValueError",
                "error_message": "bad input 'x'", "created_at_millis": 1700000000000}],
                "fragments": [{"job_id": "j-7", "fragment_id": 4, "rows_skipped": 500,
                "rows_recorded": 100}], "truncated": true}`,
            );
          });
          return;
        }
        res.writeHead(404).end();
      },
      async (db) => {
        const table = await db.openTable("docs");
        const errors = await table.functionErrors({
          jobId: "j-7",
          column: "embedding",
          limit: 2,
        });
        expect(errors.truncated).toBe(true);
        expect(errors.records.map((r) => r.errorMessage)).toEqual([
          "bad input 'x'",
        ]);
        expect(errors.records[0].rowOffset).toBe(9);
        expect(errors.fragments[0].rowsSkipped).toBe(500);
        await table.functionErrors();
        await expect(table.functionErrors({ limit: -1 })).rejects.toThrow(
          "limit must be a non-negative integer",
        );
      },
    );
    expect(bodies).toEqual([
      JSON.parse('{"job_id": "j-7", "column": "embedding", "limit": 2}'),
      {},
    ]);
  });

  it("surfaces JSON server errors from remote table operations", async () => {
    await withMockDatabase(
      (req, res) => {
        const path = req.url ?? "";
        if (path.endsWith("/describe/")) {
          res.writeHead(200, { "Content-Type": "application/json" }).end(
            JSON.stringify({
              name: "broken_table",
              version: 1,
              schema: { fields: [] },
            }),
          );
          return;
        }

        if (path.endsWith("/count_rows/")) {
          res
            .writeHead(400, { "Content-Type": "application/json" })
            .end(JSON.stringify({ error: "count rows failed" }));
          return;
        }

        res.writeHead(404).end();
      },
      async (db) => {
        const table = await db.openTable("broken_table");

        await expect(table.countRows()).rejects.toThrow("count rows failed");
      },
    );
  });

  it("should pass on requested extra headers", async () => {
    await withMockDatabase(
      (req, res) => {
        expect(req.headers["x-my-header"]).toEqual("my-value");

        const body = JSON.stringify({ tables: [] });
        res.writeHead(200, { "Content-Type": "application/json" }).end(body);
      },
      async (db) => {
        const tableNames = await db.tableNames();
        expect(tableNames).toEqual([]);
      },
      {
        clientConfig: {
          extraHeaders: {
            "x-my-header": "my-value",
          },
        },
      },
    );
  });

  it("supports version time-travel and branches on remote", async () => {
    await withMockDatabase(
      (req, res) => {
        const body = req.url?.includes("/branches/list")
          ? JSON.stringify({
              branches: {
                exp: { parentVersion: 1, createAt: 1, manifestSize: 1 },
              },
            })
          : JSON.stringify({ name: "t", version: 2, schema: { fields: [] } });
        res.writeHead(200, { "Content-Type": "application/json" }).end(body);
      },
      async (db) => {
        // version-only (and "main" + version) time-travel the main chain
        const v2 = await db.openTable("t", undefined, { version: 2 });
        expect(v2.currentBranch()).toBeNull();
        const mainV2 = await db.openTable("t", undefined, {
          branch: "main",
          version: 2,
        });
        expect(mainV2.currentBranch()).toBeNull();

        // a non-main branch opens a handle scoped to that branch
        const exp = await db.openTable("t", undefined, { branch: "exp" });
        expect(exp.currentBranch()).toBe("exp");
        const expV2 = await db.openTable("t", undefined, {
          branch: "exp",
          version: 2,
        });
        expect(expV2.currentBranch()).toBe("exp");
      },
    );
  });

  it("sends FTS options to remote tables", async () => {
    let createIndexBody: Record<string, unknown> | undefined;

    await withMockDatabase(
      (req, res) => {
        const path = req.url ?? "";
        if (path.endsWith("/describe/")) {
          res.writeHead(200, { "Content-Type": "application/json" }).end(
            JSON.stringify({
              name: "t",
              version: 1,
              schema: {
                fields: [
                  { name: "text", type: { type: "string" }, nullable: false },
                ],
              },
            }),
          );
          return;
        }

        if (path.endsWith("/create_index/")) {
          let raw = "";
          req.on("data", (chunk) => {
            raw += chunk;
          });
          req.on("end", () => {
            createIndexBody = JSON.parse(raw);
            res.writeHead(200).end();
          });
          return;
        }

        res.writeHead(404).end();
      },
      async (db) => {
        const table = await db.openTable("t");
        await table.createIndex("text", {
          config: Index.fts({
            blockSize: 256,
            removeStopWords: true,
            customStopWords: ["the"],
          }),
        });
      },
    );

    expect(createIndexBody?.["column"]).toBe("text");
    expect(createIndexBody?.["index_type"]).toBe("FTS");
    expect(createIndexBody?.["block_size"]).toBe(256);
    expect(createIndexBody?.["custom_stop_words"]).toEqual(["the"]);
  });

  it("diffs and cherry-picks remote branches", async () => {
    const sampleDiff = {
      fromBranch: "exp",
      parentVersion: 1,
      mainVersion: 2,
      branchVersion: 3,
      baseMoved: false,
      rowCountMain: 3,
      rowCountBranch: 3,
      rowSummary: {
        unchanged: 3,
        newOnBase: 0,
        newOnBranch: 0,
        staleRecompute: 0,
        inputsChanged: 0,
        deltaAvailable: false,
      },
      addedColumns: [{ name: "tag", dataType: "utf8", nullable: true }],
      removedColumns: [],
      changedColumns: [],
      addedIndexes: [],
      removedIndexes: [],
      errors: [],
    };
    const cherryPickBodies: Record<string, unknown>[] = [];

    await withMockDatabase(
      (req, res) => {
        const path = req.url ?? "";
        if (path.endsWith("/describe/")) {
          res.writeHead(200, { "Content-Type": "application/json" }).end(
            JSON.stringify({
              name: "t",
              version: 2,
              schema: { fields: [] },
            }),
          );
          return;
        }

        let raw = "";
        req.on("data", (chunk) => {
          raw += chunk;
        });
        req.on("end", () => {
          const body = raw ? JSON.parse(raw) : {};
          if (path.endsWith("/branches/diff/")) {
            // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
            expect(body).toEqual({ from_branch: "exp" });
            res
              .writeHead(200, { "Content-Type": "application/json" })
              .end(JSON.stringify(sampleDiff));
            return;
          }
          if (path.endsWith("/branches/cherry_pick/")) {
            cherryPickBodies.push(body);
            const dryRun = body["dry_run"] === true;
            const response = {
              status: dryRun ? "ready" : "failed",
              diff: dryRun
                ? sampleDiff
                : {
                    ...sampleDiff,
                    errors: [
                      { code: "baseMoved", message: "main has advanced" },
                    ],
                  },
              preview: { promotedColumns: dryRun ? ["tag"] : [] },
            };
            res
              .writeHead(dryRun ? 200 : 409, {
                "Content-Type": "application/json",
              })
              .end(JSON.stringify(response));
            return;
          }
          res.writeHead(404).end();
        });
      },
      async (db) => {
        const table = await db.openTable("t");
        const branches = await table.branches();

        await expect(branches.diff("exp")).resolves.toEqual(sampleDiff);

        const failed = await branches.cherryPick("exp");
        expect(failed.status).toBe("failed");
        expect(failed.diff.errors).toEqual([
          { code: "baseMoved", message: "main has advanced" },
        ]);

        const preview = await branches.cherryPick("exp", true);
        expect(preview.status).toBe("ready");
        expect(preview.preview.promotedColumns).toEqual(["tag"]);
      },
    );

    expect(cherryPickBodies).toEqual([
      // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
      { from_branch: "exp", dry_run: false },
      // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
      { from_branch: "exp", dry_run: true },
    ]);
  });

  describe("OAuthConfig", () => {
    it("should expose client auth method values", () => {
      expect(ClientAuthMethod.None).toBe("none");
      expect(ClientAuthMethod.ClientSecretBasic).toBe("client_secret_basic");
      expect(ClientAuthMethod.ClientSecretPost).toBe("client_secret_post");
    });

    it("should accept a confidential client with basic auth", () => {
      const config: OAuthConfig = {
        issuerUrl: "https://issuer.example.com",
        clientId: "client-id",
        clientSecret: "secret",
        scopes: ["openid"],
        flow: OAuthFlowType.AuthorizationCode,
        clientAuthMethod: ClientAuthMethod.ClientSecretBasic,
      };

      expect(config.clientAuthMethod).toBe(ClientAuthMethod.ClientSecretBasic);
    });

    it("should accept a public PKCE client without auth method or secret", () => {
      const config: OAuthConfig = {
        issuerUrl: "https://issuer.example.com",
        clientId: "client-id",
        scopes: ["openid"],
        flow: OAuthFlowType.AuthorizationCode,
        usePkce: true,
      };

      expect(config.clientSecret).toBeUndefined();
      expect(config.clientAuthMethod).toBeUndefined();
    });
  });

  describe("TlsConfig", () => {
    it("should create TlsConfig with all fields", () => {
      const tlsConfig: TlsConfig = {
        certFile: "/path/to/cert.pem",
        keyFile: "/path/to/key.pem",
        sslCaCert: "/path/to/ca.pem",
        assertHostname: false,
      };

      expect(tlsConfig.certFile).toBe("/path/to/cert.pem");
      expect(tlsConfig.keyFile).toBe("/path/to/key.pem");
      expect(tlsConfig.sslCaCert).toBe("/path/to/ca.pem");
      expect(tlsConfig.assertHostname).toBe(false);
    });

    it("should create TlsConfig with partial fields", () => {
      const tlsConfig: TlsConfig = {
        certFile: "/path/to/cert.pem",
        keyFile: "/path/to/key.pem",
      };

      expect(tlsConfig.certFile).toBe("/path/to/cert.pem");
      expect(tlsConfig.keyFile).toBe("/path/to/key.pem");
      expect(tlsConfig.sslCaCert).toBeUndefined();
      expect(tlsConfig.assertHostname).toBeUndefined();
    });

    it("should create ClientConfig with TlsConfig", () => {
      const tlsConfig: TlsConfig = {
        certFile: "/path/to/cert.pem",
        keyFile: "/path/to/key.pem",
        sslCaCert: "/path/to/ca.pem",
        assertHostname: true,
      };

      const clientConfig: ClientConfig = {
        userAgent: "test-agent",
        tlsConfig: tlsConfig,
      };

      expect(clientConfig.userAgent).toBe("test-agent");
      expect(clientConfig.tlsConfig).toBeDefined();
      expect(clientConfig.tlsConfig?.certFile).toBe("/path/to/cert.pem");
      expect(clientConfig.tlsConfig?.keyFile).toBe("/path/to/key.pem");
      expect(clientConfig.tlsConfig?.sslCaCert).toBe("/path/to/ca.pem");
      expect(clientConfig.tlsConfig?.assertHostname).toBe(true);
    });

    it("should handle empty TlsConfig", () => {
      const tlsConfig: TlsConfig = {};

      expect(tlsConfig.certFile).toBeUndefined();
      expect(tlsConfig.keyFile).toBeUndefined();
      expect(tlsConfig.sslCaCert).toBeUndefined();
      expect(tlsConfig.assertHostname).toBeUndefined();
    });

    it("should accept TlsConfig in connection options", () => {
      const tlsConfig: TlsConfig = {
        certFile: "/path/to/cert.pem",
        keyFile: "/path/to/key.pem",
        sslCaCert: "/path/to/ca.pem",
        assertHostname: false,
      };

      // Just verify that the ClientConfig accepts the TlsConfig
      const clientConfig: ClientConfig = {
        tlsConfig: tlsConfig,
      };

      const connectionOptions: ConnectionOptions = {
        apiKey: "fake",
        clientConfig: clientConfig,
      };

      // Verify the configuration structure is correct
      expect(connectionOptions.clientConfig).toBeDefined();
      expect(connectionOptions.clientConfig?.tlsConfig).toBeDefined();
      expect(connectionOptions.clientConfig?.tlsConfig?.certFile).toBe(
        "/path/to/cert.pem",
      );
    });
  });

  describe("header providers", () => {
    it("should work with StaticHeaderProvider", async () => {
      const provider = new StaticHeaderProvider({
        authorization: "Bearer test-token",
        "X-Custom": "value",
      });

      const headers = provider.getHeaders();
      expect(headers).toEqual({
        authorization: "Bearer test-token",
        "X-Custom": "value",
      });

      // Test that it returns a copy
      headers["X-Modified"] = "modified";
      const headers2 = provider.getHeaders();
      expect(headers2).not.toHaveProperty("X-Modified");
    });

    it("should pass headers from StaticHeaderProvider to requests", async () => {
      const provider = new StaticHeaderProvider({
        "X-Custom-Auth": "secret-token",
        "X-Request-Source": "test-suite",
      });

      await withMockDatabase(
        (req, res) => {
          expect(req.headers["x-custom-auth"]).toEqual("secret-token");
          expect(req.headers["x-request-source"]).toEqual("test-suite");

          const body = JSON.stringify({ tables: [] });
          res.writeHead(200, { "Content-Type": "application/json" }).end(body);
        },
        async () => {
          // Use actual header provider mechanism instead of extraHeaders
          const conn = await connect(
            "db://dev",
            {
              apiKey: "fake",
              hostOverride: "http://localhost:8000",
            },
            undefined, // session
            provider, // headerProvider
          );

          const tableNames = await conn.tableNames();
          expect(tableNames).toEqual([]);
        },
      );
    });

    it("should work with CustomProvider", () => {
      const provider = new CustomProvider();
      const headers = provider.getHeaders();
      expect(headers).toEqual({ "X-Custom": "custom-value" });
    });

    it("should handle ErrorProvider errors", () => {
      const provider = new ErrorProvider("Authentication failed");

      expect(() => provider.getHeaders()).toThrow("Authentication failed");
      expect(provider.callCount).toBe(1);

      // Test that error is thrown each time
      expect(() => provider.getHeaders()).toThrow("Authentication failed");
      expect(provider.callCount).toBe(2);
    });

    it("should work with ConcurrentProvider", () => {
      const provider = new ConcurrentProvider();

      const headers1 = provider.getHeaders();
      const headers2 = provider.getHeaders();
      const headers3 = provider.getHeaders();

      expect(headers1).toEqual({ "X-Request-Id": "1" });
      expect(headers2).toEqual({ "X-Request-Id": "2" });
      expect(headers3).toEqual({ "X-Request-Id": "3" });
    });

    describe("OAuthHeaderProvider", () => {
      it("should initialize correctly", () => {
        const fetcher = () => ({
          accessToken: "token123",
          expiresIn: 3600,
        });

        const provider = new OAuthHeaderProvider(fetcher);
        expect(provider).toBeInstanceOf(HeaderProvider);
      });

      it("should fetch token on first use", async () => {
        let callCount = 0;
        const fetcher = () => {
          callCount++;
          return {
            accessToken: "token123",
            expiresIn: 3600,
          };
        };

        const provider = new OAuthHeaderProvider(fetcher);

        // Need to manually refresh first due to sync limitation
        await provider.refreshToken();

        const headers = provider.getHeaders();
        expect(headers).toEqual({ authorization: "Bearer token123" });
        expect(callCount).toBe(1);

        // Second call should not fetch again
        const headers2 = provider.getHeaders();
        expect(headers2).toEqual({ authorization: "Bearer token123" });
        expect(callCount).toBe(1);
      });

      it("should handle tokens without expiry", async () => {
        const fetcher = () => ({
          accessToken: "permanent_token",
        });

        const provider = new OAuthHeaderProvider(fetcher);
        await provider.refreshToken();

        const headers = provider.getHeaders();
        expect(headers).toEqual({ authorization: "Bearer permanent_token" });
      });

      it("should throw error when access_token is missing", async () => {
        const fetcher = () =>
          ({
            expiresIn: 3600,
          }) as { accessToken?: string; expiresIn?: number };

        const provider = new OAuthHeaderProvider(
          fetcher as () => {
            accessToken: string;
            expiresIn?: number;
          },
        );

        await expect(provider.refreshToken()).rejects.toThrow(
          "Token fetcher did not return 'accessToken'",
        );
      });

      it("should handle async token fetchers", async () => {
        const fetcher = async () => {
          // Simulate async operation
          await new Promise((resolve) => setTimeout(resolve, 10));
          return {
            accessToken: "async_token",
            expiresIn: 3600,
          };
        };

        const provider = new OAuthHeaderProvider(fetcher);
        await provider.refreshToken();

        const headers = provider.getHeaders();
        expect(headers).toEqual({ authorization: "Bearer async_token" });
      });
    });

    it("should merge header provider headers with extra headers", async () => {
      const provider = new StaticHeaderProvider({
        "X-From-Provider": "provider-value",
      });

      await withMockDatabase(
        (req, res) => {
          expect(req.headers["x-from-provider"]).toEqual("provider-value");
          expect(req.headers["x-extra-header"]).toEqual("extra-value");

          const body = JSON.stringify({ tables: [] });
          res.writeHead(200, { "Content-Type": "application/json" }).end(body);
        },
        async () => {
          // Use header provider with additional extraHeaders
          const conn = await connect(
            "db://dev",
            {
              apiKey: "fake",
              hostOverride: "http://localhost:8000",
              clientConfig: {
                extraHeaders: {
                  "X-Extra-Header": "extra-value",
                },
              },
            },
            undefined, // session
            provider, // headerProvider
          );

          const tableNames = await conn.tableNames();
          expect(tableNames).toEqual([]);
        },
      );
    });
  });

  describe("header provider integration", () => {
    it("should work with TypeScript StaticHeaderProvider", async () => {
      let requestCount = 0;

      await withMockDatabase(
        (req, res) => {
          requestCount++;

          // Check headers are present on each request
          expect(req.headers["authorization"]).toEqual("Bearer test-token-123");
          expect(req.headers["x-custom"]).toEqual("custom-value");

          // Return different responses based on the endpoint
          if (req.url === "/v1/table/test_table/describe/") {
            const body = JSON.stringify({
              name: "test_table",
              schema: { fields: [] },
            });
            res
              .writeHead(200, { "Content-Type": "application/json" })
              .end(body);
          } else {
            const body = JSON.stringify({ tables: ["test_table"] });
            res
              .writeHead(200, { "Content-Type": "application/json" })
              .end(body);
          }
        },
        async () => {
          // Create provider with static headers
          const provider = new StaticHeaderProvider({
            authorization: "Bearer test-token-123",
            "X-Custom": "custom-value",
          });

          // Connect with the provider
          const conn = await connect(
            "db://dev",
            {
              apiKey: "fake",
              hostOverride: "http://localhost:8000",
            },
            undefined, // session
            provider, // headerProvider
          );

          // Make multiple requests to verify headers are sent each time
          const tables1 = await conn.tableNames();
          expect(tables1).toEqual(["test_table"]);

          const tables2 = await conn.tableNames();
          expect(tables2).toEqual(["test_table"]);

          // Verify headers were sent with each request
          expect(requestCount).toBeGreaterThanOrEqual(2);
        },
      );
    });

    it("should work with JavaScript function provider", async () => {
      let requestId = 0;

      await withMockDatabase(
        (req, res) => {
          // Check dynamic header is present
          expect(req.headers["x-request-id"]).toBeDefined();
          expect(req.headers["x-request-id"]).toMatch(/^req-\d+$/);

          const body = JSON.stringify({ tables: [] });
          res.writeHead(200, { "Content-Type": "application/json" }).end(body);
        },
        async () => {
          // Create a JavaScript function that returns dynamic headers
          const getHeaders = async () => {
            requestId++;
            return {
              "X-Request-Id": `req-${requestId}`,
              "X-Timestamp": new Date().toISOString(),
            };
          };

          // Connect with the function directly
          const conn = await connect(
            "db://dev",
            {
              apiKey: "fake",
              hostOverride: "http://localhost:8000",
            },
            undefined, // session
            getHeaders, // headerProvider
          );

          // Make requests - each should have different headers
          const tables = await conn.tableNames();
          expect(tables).toEqual([]);
        },
      );
    });

    it("should support OAuth-like token refresh pattern", async () => {
      let tokenVersion = 0;

      await withMockDatabase(
        (req, res) => {
          // Verify authorization header
          const authHeader = req.headers["authorization"];
          expect(authHeader).toBeDefined();
          expect(authHeader).toMatch(/^Bearer token-v\d+$/);

          const body = JSON.stringify({ tables: [] });
          res.writeHead(200, { "Content-Type": "application/json" }).end(body);
        },
        async () => {
          // Simulate OAuth token fetcher
          const fetchToken = async () => {
            tokenVersion++;
            return {
              authorization: `Bearer token-v${tokenVersion}`,
            };
          };

          // Connect with the function directly
          const conn = await connect(
            "db://dev",
            {
              apiKey: "fake",
              hostOverride: "http://localhost:8000",
            },
            undefined, // session
            fetchToken, // headerProvider
          );

          // Each request will fetch a new token
          await conn.tableNames();

          // Token should be different on next request
          await conn.tableNames();
        },
      );
    });
  });

  describe("renameTable", () => {
    async function captureRenameRequest(
      call: (db: Connection) => Promise<void>,
    ): Promise<{ url: string; body: Record<string, unknown> }> {
      let captured: { url: string; body: Record<string, unknown> } | undefined;
      await withMockDatabase((req, res) => {
        let raw = "";
        req.on("data", (chunk) => {
          raw += chunk;
        });
        req.on("end", () => {
          captured = {
            url: req.url ?? "",
            body: raw ? JSON.parse(raw) : {},
          };
          res.writeHead(200, { "Content-Type": "application/json" }).end("");
        });
      }, call);
      if (!captured) {
        throw new Error("mock server never saw a request");
      }
      return captured;
    }

    it("sends rename request for a table in the root namespace", async () => {
      const { url, body } = await captureRenameRequest(async (db) => {
        await db.renameTable("table1", "table2");
      });
      expect(url).toBe("/v1/table/table1/rename/");
      // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
      expect(body).toEqual({ new_table_name: "table2" });
    });

    it("omits new_namespace when only the current namespace is supplied", async () => {
      // Safe-default check: passing namespacePath alone must not send
      // `new_namespace`, so the server keeps the table in its current
      // namespace instead of silently moving it to root.
      const { url, body } = await captureRenameRequest(async (db) => {
        await db.renameTable("table1", "table2", {
          namespacePath: ["ns1"],
        });
      });
      expect(url).toBe("/v1/table/ns1$table1/rename/");
      // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
      expect(body).toEqual({ new_table_name: "table2" });
    });

    it("includes new_namespace in the body for a cross-namespace rename", async () => {
      const { url, body } = await captureRenameRequest(async (db) => {
        await db.renameTable("table1", "table2", {
          namespacePath: ["ns1"],
          newNamespacePath: ["ns2"],
        });
      });
      expect(url).toBe("/v1/table/ns1$table1/rename/");
      expect(body).toEqual({
        // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
        new_table_name: "table2",
        // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
        new_namespace: ["ns2"],
      });
    });
  });
});

describe("remote connection jobs surface", () => {
  it("lists, describes, cancels, and reads history", async () => {
    const { tableFromArrays, tableToIPC } = await import("apache-arrow");
    const eventsTable = tableFromArrays({ state: ["created", "succeeded"] });
    const eventsBody = Buffer.from(tableToIPC(eventsTable, "stream"));
    const queryEventsPayloads: Record<string, unknown>[] = [];

    await withMockDatabase(
      (req, res) => {
        let body = "";
        req.on("data", (chunk) => {
          body += chunk;
        });
        req.on("end", () => {
          const payload = body.length > 0 ? JSON.parse(body) : {};
          if (req.url === "/v1/jobs/list") {
            if (payload["page_token"] === undefined) {
              res
                .writeHead(200, { "Content-Type": "application/json" })
                .end(
                  '{"jobs": [{"job_id": "job-1", "table": "t1", ' +
                    '"job_type": "create_index", "state": "in_progress", ' +
                    '"created_at_millis": 1000}], "page_token": "next"}',
                );
            } else {
              res
                .writeHead(200, { "Content-Type": "application/json" })
                .end(
                  '{"jobs": [{"job_id": "job-2", "table": "t2", ' +
                    '"job_type": "create_index", "state": "succeeded", ' +
                    '"created_at_millis": 2000}]}',
                );
            }
          } else if (req.url === "/v1/jobs/describe") {
            if (payload["job_id"] === "job-2") {
              res
                .writeHead(200, { "Content-Type": "application/json" })
                .end(
                  '{"job_id": "job-2", "job_type": "refresh_column", ' +
                    '"job_state": "DONE", "creation_ms": 2000, ' +
                    '"result": {"rows_assigned": 1000000}}',
                );
              return;
            }
            if (payload["job_id"] !== "job-1") {
              res.writeHead(404).end("no such job");
              return;
            }
            res
              .writeHead(200, { "Content-Type": "application/json" })
              .end(
                '{"job_id": "job-1", "job_type": "create_index", ' +
                  '"job_state": "FAILED", "creation_ms": 1000, ' +
                  '"spec": {"column": "vec"}, "failure": {"phase": "execute", ' +
                  '"message": "worker died", "retryable": true}}',
              );
          } else if (req.url === "/v1/jobs/cancel") {
            if (payload["job_id"] !== "job-1") {
              res.writeHead(404).end("no such job");
              return;
            }
            res
              .writeHead(200, { "Content-Type": "application/json" })
              .end('{"job_id": "job-1"}');
          } else if (req.url === "/v1/jobs/query_events") {
            queryEventsPayloads.push(payload);
            res
              .writeHead(200, {
                "Content-Type": "application/vnd.apache.arrow.stream",
              })
              .end(eventsBody);
          } else {
            res.writeHead(404).end();
          }
        });
      },
      async (db) => {
        const jobs = await db.listJobs();
        expect(jobs.map((job) => job.jobId)).toEqual(["job-1", "job-2"]);
        expect(jobs[0].state).toEqual("running");
        expect(jobs[1].state).toEqual("finished");

        expect(await db.cancelJob("job-1")).toBe(true);
        expect(await db.cancelJob("missing")).toBe(false);

        // Opening a job hands back a populated handle; a missing one rejects.
        await expect(db.openJob("missing")).rejects.toThrow("not found");
        const finished = await db.openJob("job-2");
        expect(finished.state).toEqual("finished");
        expect(finished.result).toEqual({
          // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
          rows_assigned: 1000000,
        });

        const job = await db.openJob("job-1");
        expect(job.id).toEqual("job-1");

        // openJob already populated the handle; refresh() re-reads it.
        expect(job.state).toEqual("failed");
        await job.refresh();
        expect(job.state).toEqual("failed");
        expect(job.jobType).toEqual("create_index");
        expect(job.creationMs).toEqual(1000);
        expect(job.spec).toEqual({ column: "vec" });
        expect(job.result).toBeNull();
        expect(job.failure?.message).toEqual("worker died");

        // The handle reaches its own events, supplying its job id.
        const jobEvents = await job.events({
          limit: 500,
          filter: "state = 'claim_complete'",
        });
        expect(jobEvents.numRows).toEqual(2);
        expect(queryEventsPayloads.pop()).toEqual({
          // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
          job_id: "job-1",
          limit: 500,
          filter: "state = 'claim_complete'",
        });

        // Printing lays every known field out on its own line, with the JSON
        // payloads indented rather than crammed onto one line.
        expect(`${job}`).toEqual(
          [
            "Job(",
            '    id="job-1",',
            '    state="failed",',
            '    jobType="create_index",',
            "    creationMs=1000,",
            "    spec={",
            '        "column": "vec"',
            "    },",
            "    failure={",
            '        "phase": "execute",',
            '        "message": "worker died",',
            '        "retryable": true',
            "    },",
            ")",
          ].join("\n"),
        );

        expect(await job.status()).toEqual("failed");
        await expect(job.wait()).rejects.toThrow("worker died");
      },
    );
  });
});