mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-07 13:59:08 +00:00
Jack Ye
21f11b4463
A 1M-row column refresh over 200 fragments produced no visible result,
and the client could only ever say `"running"`. Everything needed to
diagnose it already existed server-side — the job registry records a
`claim`/`claim_complete` pair per fragment carrying `rows_processed` —
but none of it was reachable.
## Before
Four ways to ask about a job, none of which told you much.
```python
job = table.refresh_column_async("embedding")
job.status() # "running". That was the entire debug surface.
db.get_job(job_id) # state, and a spec. No result, no progress.
db.job_history(job_id) # raw record batches, no limit, no filter
db.job(job_id) # a handle that knew nothing
```
## After
Open a job the way you open a table; the handle answers everything.
```python
job = db.open_job(job_id) # raises JobNotFoundError if there is no such job
```
```python
>>> print(job)
Job(
id='job-1',
state='failed',
job_type='refresh_column',
creation_ms=1757000000000,
spec={
"column": "embedding",
"num_workers": 4
},
failure=JobFailureInfo(phase='execute', message='worker died', retryable=True),
)
```
Individual fields are there too — `job.state`, `job.job_type`,
`job.creation_ms`, `job.spec`, `job.result`, `job.failure` — and
`job.result` carries `rows_assigned` / `rows_failed` as soon as the job
succeeds, with no `wait()` required.
Per-fragment progress *while it is still running*:
```python
done = job.events(filter="state = 'claim_complete'", limit=10_000)
done.column("rows_processed").to_pylist() # [5000, 5000, ...]
```
The handle an async action returns is the same object, one `refresh()`
away:
```python
job = table.refresh_column_async("embedding")
job.refresh()
job.state, job.result
```
TypeScript is the same experience, down to `console.log`:
```ts
const job = await db.openJob(jobId); // rejects if there is no such job
console.log(job); // same multi-line layout
job.state; job.jobType; job.spec; job.result; job.failure;
const done = await job.events({ filter: "state = 'claim_complete'", limit: 10_000 });
```
## Why each piece matters
- **A result without waiting.** `rows_assigned` / `rows_failed` used to
live only on the terminal result, so a job that never terminated
reported nothing at all.
- **`limit`.** The server caps event rows at 1000 and truncates without
saying so, which silently hid most of a 200-fragment job's history.
- **`filter`.** `claim_complete` rows carry per-claim `rows_processed` —
the only progress signal that exists mid-flight.
- **Events outlive the worker.** They live in the job registry, not in
pod logs that vanish with the pod.
- **One place to ask.** `open_job` replaces `describe_job`,
`query_job_events` and `job`, so a question about a job has one answer
instead of one per calling location.
- **A missing job is an error, not a `None`.** The common case is a job
id copied out of a log, where absence is the surprise worth raising —
and it matches `open_table`.
- **Printing is the debug surface.** Every field on its own line, JSON
payloads keeping their structure. An unrefreshed handle stays on one
line, because there is nothing to lay out.
- **In-process jobs say so.** A local refresh reports `state` and leaves
the rest null rather than inventing fields it has no record for.
`list_jobs` and `cancel_job` stay as they were: one lists, the other is
a one-shot action that should not need a describe first.
## Breaking
All shipped in 0.38.0. No deprecated aliases.
| Was | Now |
| --- | --- |
| `Connection.get_job` → `describe_job` | `Connection.open_job` returns
a populated `Job`, or raises |
| `Connection.job_history` → `query_job_events` | `job.events(...)` |
| `Connection.job` | `Connection.open_job` |
| Python events → `List[pa.RecordBatch]` | `pa.Table` |
| `JobDescription.spec_json` / `.result_json` | internal; use `job.spec`
/ `job.result` |
Node's `Job` is now a TypeScript class wrapping the native handle, so it
returns an Arrow table and parsed values like Python does. New
`Error::JobNotFound` / `JobNotFoundError`; the three job exceptions are
now in the Python API reference.
1084 lines
34 KiB
TypeScript
1084 lines
34 KiB
TypeScript
// 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 {
|
|
ClientConfig,
|
|
Connection,
|
|
ConnectionOptions,
|
|
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("refuses materialized views before issuing any request", async () => {
|
|
const paths: string[] = [];
|
|
await withMockDatabase(
|
|
(req, res) => {
|
|
paths.push(req.url ?? "");
|
|
res.writeHead(404).end();
|
|
},
|
|
async (db) => {
|
|
await expect(db.openMaterializedView("secret_table")).rejects.toThrow(
|
|
/only on local databases/,
|
|
);
|
|
await expect(db.listMaterializedViews()).rejects.toThrow(
|
|
/only on local databases/,
|
|
);
|
|
expect(paths).toEqual([]);
|
|
},
|
|
);
|
|
});
|
|
|
|
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("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("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");
|
|
},
|
|
);
|
|
});
|
|
});
|