mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-25 13:12:39 +00:00
Merge remote-tracking branch 'origin/main' into remote-set-unenforced-primary-key
Keep the remote set_unenforced_primary_key implementation over main's NotSupported stub, and drop the remote test case that expected it to fail. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Daniel Rammer
Claude Opus 5.5
+1487
-1378
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,39 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// Every Apache Arrow release the package accepts as a peer dependency, as
|
||||
// separately installed copies so the tests can feed objects built by one
|
||||
// Arrow version into a LanceDB running on another. When a new Arrow major
|
||||
// is released, add the alias to `package.json` and extend this list; the
|
||||
// test files pick the change up from here.
|
||||
import * as arrow15 from "apache-arrow-15";
|
||||
import * as arrow16 from "apache-arrow-16";
|
||||
import * as arrow17 from "apache-arrow-17";
|
||||
import * as arrow18 from "apache-arrow-18";
|
||||
import * as arrow19 from "apache-arrow-19";
|
||||
import * as arrow20 from "apache-arrow-20";
|
||||
import * as arrow21 from "apache-arrow-21";
|
||||
|
||||
export { arrow15, arrow16, arrow17, arrow18, arrow19, arrow20, arrow21 };
|
||||
|
||||
export type ApacheArrow =
|
||||
| typeof arrow15
|
||||
| typeof arrow16
|
||||
| typeof arrow17
|
||||
| typeof arrow18
|
||||
| typeof arrow19
|
||||
| typeof arrow20
|
||||
| typeof arrow21;
|
||||
|
||||
export const arrowVersions: ApacheArrow[] = [
|
||||
arrow15,
|
||||
arrow16,
|
||||
arrow17,
|
||||
arrow18,
|
||||
arrow19,
|
||||
arrow20,
|
||||
arrow21,
|
||||
];
|
||||
|
||||
/** The newest supported Arrow release, for tests that only need one. */
|
||||
export const latestArrow = arrow21;
|
||||
@@ -0,0 +1,509 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
Field,
|
||||
Float,
|
||||
Float32,
|
||||
Int32,
|
||||
Int64,
|
||||
List,
|
||||
Schema,
|
||||
Struct,
|
||||
Utf8,
|
||||
} from "apache-arrow";
|
||||
import {
|
||||
convertToTable,
|
||||
makeArrowTable,
|
||||
resolveBlobInputs,
|
||||
} from "../lancedb/arrow";
|
||||
import { BlobFile, blob, coerceBlobValue, isBlobField } from "../lancedb/blob";
|
||||
import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
|
||||
import { getRegistry } from "../lancedb/embedding/registry";
|
||||
|
||||
describe("blob()", () => {
|
||||
it("marks the field as lance.blob.v2", () => {
|
||||
const field = blob("image", { nullable: false });
|
||||
expect(field.nullable).toBe(false);
|
||||
expect(isBlobField(field)).toBe(true);
|
||||
expect(field.metadata.get("ARROW:extension:name")).toBe("lance.blob.v2");
|
||||
});
|
||||
|
||||
it("writes encoding thresholds as field metadata", () => {
|
||||
const field = blob("video", {
|
||||
inlineSizeThreshold: 1024,
|
||||
dedicatedSizeThreshold: 2 * 1024 * 1024,
|
||||
packFileSizeThreshold: 64 * 1024 * 1024,
|
||||
});
|
||||
expect(
|
||||
field.metadata.get("lance-encoding:blob-inline-size-threshold"),
|
||||
).toBe("1024");
|
||||
expect(
|
||||
field.metadata.get("lance-encoding:blob-dedicated-size-threshold"),
|
||||
).toBe(String(2 * 1024 * 1024));
|
||||
expect(
|
||||
field.metadata.get("lance-encoding:blob-pack-file-size-threshold"),
|
||||
).toBe(String(64 * 1024 * 1024));
|
||||
});
|
||||
|
||||
it("rejects invalid thresholds", () => {
|
||||
expect(() => blob("image", { inlineSizeThreshold: -1 })).toThrow(
|
||||
/inlineSizeThreshold must be non-negative/,
|
||||
);
|
||||
expect(() => blob("image", { dedicatedSizeThreshold: 0 })).toThrow(
|
||||
/dedicatedSizeThreshold must be positive/,
|
||||
);
|
||||
expect(() => blob("image", { packFileSizeThreshold: 1.5 })).toThrow(
|
||||
/packFileSizeThreshold must be a safe integer/,
|
||||
);
|
||||
expect(() =>
|
||||
blob("image", { dedicatedSizeThreshold: Number.MAX_SAFE_INTEGER + 1 }),
|
||||
).toThrow(/dedicatedSizeThreshold must be a safe integer/);
|
||||
});
|
||||
});
|
||||
|
||||
class BlobStruct {
|
||||
constructor(
|
||||
public data?: unknown,
|
||||
public uri?: unknown,
|
||||
) {}
|
||||
}
|
||||
|
||||
describe("coerceBlobValue", () => {
|
||||
it.each([
|
||||
["Buffer", Buffer.from("x"), { data: Buffer.from("x"), uri: null }],
|
||||
[
|
||||
"Uint8Array",
|
||||
new Uint8Array([120]),
|
||||
{ data: new Uint8Array([120]), uri: null },
|
||||
],
|
||||
["URI string", "s3://bucket/key", { data: null, uri: "s3://bucket/key" }],
|
||||
[
|
||||
"data struct",
|
||||
{ data: Buffer.from("y") },
|
||||
{ data: Buffer.from("y"), uri: null },
|
||||
],
|
||||
[
|
||||
"uri struct",
|
||||
{ uri: "s3://bucket/key" },
|
||||
{ data: null, uri: "s3://bucket/key" },
|
||||
],
|
||||
[
|
||||
"ArrayBuffer",
|
||||
new Uint8Array([120]).buffer,
|
||||
{ data: new Uint8Array([120]), uri: null },
|
||||
],
|
||||
[
|
||||
"URL",
|
||||
new URL("https://example.com/blob"),
|
||||
{ data: null, uri: "https://example.com/blob" },
|
||||
],
|
||||
[
|
||||
"data struct with ArrayBuffer",
|
||||
{ data: new Uint8Array([121]).buffer },
|
||||
{ data: new Uint8Array([121]), uri: null },
|
||||
],
|
||||
[
|
||||
"uri struct with URL",
|
||||
{ uri: new URL("file:///tmp/a.png") },
|
||||
{ data: null, uri: "file:///tmp/a.png" },
|
||||
],
|
||||
[
|
||||
"class instance with Uint8Array data",
|
||||
new BlobStruct(new Uint8Array([122])),
|
||||
{ data: new Uint8Array([122]), uri: null },
|
||||
],
|
||||
[
|
||||
"class instance with ArrayBuffer data",
|
||||
new BlobStruct(new Uint8Array([122]).buffer),
|
||||
{ data: new Uint8Array([122]), uri: null },
|
||||
],
|
||||
[
|
||||
"class instance with URL uri",
|
||||
new BlobStruct(undefined, new URL("s3://bucket/key")),
|
||||
{ data: null, uri: "s3://bucket/key" },
|
||||
],
|
||||
["null", null, null],
|
||||
])("accepts %s", (_name, input, expected) => {
|
||||
expect(coerceBlobValue(input)).toEqual(expected);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["empty URI", "", /uri cannot be empty/],
|
||||
["object without data or uri", { position: 0 }, /data' or 'uri/],
|
||||
[
|
||||
"Int16Array",
|
||||
new Int16Array([1]),
|
||||
/Blob data must be Buffer, Uint8Array, ArrayBuffer, or Blob/,
|
||||
],
|
||||
[
|
||||
"both data and uri",
|
||||
{ data: Buffer.from("y"), uri: "s3://bucket/key" },
|
||||
/exactly one of 'data' or 'uri'/,
|
||||
],
|
||||
[
|
||||
"neither data nor uri",
|
||||
{ data: null, uri: null },
|
||||
/exactly one of 'data' or 'uri'/,
|
||||
],
|
||||
["Blob", new Blob(["x"]), /must be read asynchronously/],
|
||||
["data struct with Blob", { data: new Blob(["x"]) }, /asynchronously/],
|
||||
[
|
||||
"class instance with Blob data",
|
||||
new BlobStruct(new Blob(["x"])),
|
||||
/asynchronously/,
|
||||
],
|
||||
[
|
||||
"unknown struct field",
|
||||
{ data: Buffer.from("y"), mime: "image/png" },
|
||||
/Blob struct values only support 'data' and 'uri', got 'mime'/,
|
||||
],
|
||||
[
|
||||
"misspelled struct field",
|
||||
{ uri: "s3://bucket/key", dat: Buffer.from("y") },
|
||||
/got 'dat'/,
|
||||
],
|
||||
[
|
||||
"SharedArrayBuffer",
|
||||
new SharedArrayBuffer(2),
|
||||
/Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob/,
|
||||
],
|
||||
["Date", new Date(0), /Blob column values must be/],
|
||||
])("rejects %s", (_name, input, message) => {
|
||||
expect(() => coerceBlobValue(input)).toThrow(message);
|
||||
});
|
||||
});
|
||||
|
||||
describe("resolveBlobInputs", () => {
|
||||
it("reads Blob and File values in blob columns only", async () => {
|
||||
const schema = new Schema([
|
||||
new Field("name", new Utf8(), true),
|
||||
blob("image"),
|
||||
new Field("images", new List(blob("image")), true),
|
||||
]);
|
||||
const shared = new Blob(["same"]);
|
||||
const notABlobColumn = new Blob(["other"]);
|
||||
const data = [
|
||||
{
|
||||
name: notABlobColumn,
|
||||
image: new File(["file"], "a.txt"),
|
||||
images: [shared, { data: shared }, null],
|
||||
},
|
||||
{ name: "x", image: Buffer.from("buf"), images: null },
|
||||
];
|
||||
const resolved = await resolveBlobInputs(data, schema);
|
||||
|
||||
expect(resolved[0].name).toBe(notABlobColumn);
|
||||
expect(Buffer.from(resolved[0].image as Uint8Array).toString()).toBe(
|
||||
"file",
|
||||
);
|
||||
const images = resolved[0].images as [Uint8Array, { data: Uint8Array }];
|
||||
expect(Buffer.from(images[0]).toString()).toBe("same");
|
||||
expect(Buffer.from(images[1].data).toString()).toBe("same");
|
||||
// Rows without Blob values and the caller's records are left alone.
|
||||
expect(resolved[1]).toBe(data[1]);
|
||||
expect(data[0].images?.[0]).toBe(shared);
|
||||
});
|
||||
|
||||
it("reads each distinct Blob once, bare or inside data structs", async () => {
|
||||
class CountingBlob extends Blob {
|
||||
calls = 0;
|
||||
async arrayBuffer(): Promise<ArrayBuffer> {
|
||||
this.calls++;
|
||||
return super.arrayBuffer();
|
||||
}
|
||||
}
|
||||
const shared = new CountingBlob(["same"]);
|
||||
const schema = new Schema([blob("image")]);
|
||||
const resolved = await resolveBlobInputs(
|
||||
[
|
||||
{ image: { data: shared } },
|
||||
{ image: { data: shared } },
|
||||
{ image: shared },
|
||||
],
|
||||
schema,
|
||||
);
|
||||
expect(shared.calls).toBe(1);
|
||||
const first = (resolved[0].image as { data: Uint8Array }).data;
|
||||
expect(Buffer.from(first).toString()).toBe("same");
|
||||
expect((resolved[1].image as { data: Uint8Array }).data).toBe(first);
|
||||
expect(resolved[2].image).toBe(first);
|
||||
});
|
||||
|
||||
it("reads a Blob held by a class instance", async () => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const [row] = await resolveBlobInputs(
|
||||
[{ image: new BlobStruct(new Blob(["hi"])) }],
|
||||
schema,
|
||||
);
|
||||
expect(
|
||||
Buffer.from((row.image as { data: Uint8Array }).data).toString(),
|
||||
).toBe("hi");
|
||||
});
|
||||
|
||||
it("returns the input when there is nothing to read", async () => {
|
||||
const data = [{ image: Buffer.from("x") }];
|
||||
const schema = new Schema([blob("image")]);
|
||||
await expect(resolveBlobInputs(data, schema)).resolves.toBe(data);
|
||||
await expect(resolveBlobInputs(data)).resolves.toBe(data);
|
||||
});
|
||||
});
|
||||
|
||||
describe("convertToTable blob columns with embedding functions", () => {
|
||||
// Embedding metadata routes rows through ensureNestedFieldsExist, which must
|
||||
// leave blob values for the blob coercion instead of treating them as structs.
|
||||
class LengthEmbedding extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 2;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float32();
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map((text) => [text.length, 1]);
|
||||
}
|
||||
}
|
||||
getRegistry().register("blob-test-length")(LengthEmbedding);
|
||||
const func = getRegistry().get<LengthEmbedding>("blob-test-length")!.create();
|
||||
const embeddingSchema = LanceSchema({
|
||||
text: func.sourceField(new Utf8()),
|
||||
vector: func.vectorField(),
|
||||
});
|
||||
const schema = new Schema(
|
||||
[...embeddingSchema.fields, blob("image")],
|
||||
embeddingSchema.metadata,
|
||||
);
|
||||
|
||||
it.each([
|
||||
["Uint8Array", new Uint8Array([104, 105]), "hi", null],
|
||||
["ArrayBuffer", new Uint8Array([104, 105]).buffer, "hi", null],
|
||||
["Blob", new Blob(["hi"]), "hi", null],
|
||||
["URI string", "s3://b/k", null, "s3://b/k"],
|
||||
["URL", new URL("s3://b/k"), null, "s3://b/k"],
|
||||
["data struct", { data: new Uint8Array([104, 105]) }, "hi", null],
|
||||
])("accepts %s input", async (_name, image, data, uri) => {
|
||||
const table = await convertToTable([{ text: "a", image }], undefined, {
|
||||
schema,
|
||||
});
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(0);
|
||||
const bytes = column.getChild("data")!.get(0);
|
||||
expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data);
|
||||
expect(column.getChild("uri")!.get(0)).toBe(uri);
|
||||
});
|
||||
|
||||
it("writes a missing blob as null", async () => {
|
||||
const table = await convertToTable([{ text: "a" }], undefined, { schema });
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(1);
|
||||
expect(column.get(0)).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("BlobFile", () => {
|
||||
it("rejects constructing BlobFile without a native handle", () => {
|
||||
expect(() => new (BlobFile as unknown as { new (): BlobFile })()).toThrow(
|
||||
/fetchBlobFiles/,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("makeArrowTable blob columns", () => {
|
||||
it.each([
|
||||
["ReadableStream", new ReadableStream()],
|
||||
["Int16Array", new Int16Array([1])],
|
||||
])("rejects %s input rather than writing a null blob", (_name, image) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
expect(() => makeArrowTable([{ image }], { schema })).toThrow(/image/);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["Blob", new Blob(["hello"])],
|
||||
["File", new File(["hello"], "hello.txt")],
|
||||
["data struct with Blob", { data: new Blob(["hello"]) }],
|
||||
[
|
||||
"Blob with enumerable fields",
|
||||
Object.assign(new Blob(["hello"]), { position: 0 }),
|
||||
],
|
||||
])("points %s input at the async APIs", (_name, image) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
expect(() => makeArrowTable([{ id: 1, image }], { schema })).toThrow(
|
||||
/blob field image at row 0: Blob and File values must be read asynchronously/,
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["ArrayBuffer", new Uint8Array([104, 105]).buffer, "hi", null],
|
||||
[
|
||||
"data struct with ArrayBuffer",
|
||||
{ data: new Uint8Array([104, 105]).buffer },
|
||||
"hi",
|
||||
null,
|
||||
],
|
||||
["URL", new URL("https://example.com/a"), null, "https://example.com/a"],
|
||||
[
|
||||
"uri struct with URL",
|
||||
{ uri: new URL("file:///tmp/a.png") },
|
||||
null,
|
||||
"file:///tmp/a.png",
|
||||
],
|
||||
])("accepts %s input", (_name, image, data, uri) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const input = { image };
|
||||
const table = makeArrowTable([input], { schema });
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(0);
|
||||
const bytes = column.getChild("data")!.get(0);
|
||||
expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data);
|
||||
expect(column.getChild("uri")!.get(0)).toBe(uri);
|
||||
// The caller's record is left untouched.
|
||||
expect(input.image).toBe(image);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["Uint8Array", new Uint8Array([104, 105]), "hi", null],
|
||||
["URI string", "s3://b/k", null, "s3://b/k"],
|
||||
["data struct", { data: new Uint8Array([104, 105]) }, "hi", null],
|
||||
])("accepts %s input after a null row", (_name, image, data, uri) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const table = makeArrowTable([{ image: null }, { image }], { schema });
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(1);
|
||||
expect(column.get(0)).toBeNull();
|
||||
const bytes = column.getChild("data")!.get(1);
|
||||
expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data);
|
||||
expect(column.getChild("uri")!.get(1)).toBe(uri);
|
||||
});
|
||||
|
||||
it("accepts a null struct before one holding a blob", () => {
|
||||
const schema = new Schema([
|
||||
new Field("info", new Struct([blob("image")]), true),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[{ info: null }, { info: { image: new Uint8Array([97]) } }],
|
||||
{ schema },
|
||||
);
|
||||
const info = table.getChild("info")!;
|
||||
expect(info.get(0)).toBeNull();
|
||||
const image = (info.get(1) as { image: { data: Uint8Array } }).image;
|
||||
expect(Buffer.from(image.data).toString()).toBe("a");
|
||||
});
|
||||
|
||||
it("skips a missing blob column named like an Object.prototype member", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int32(), true),
|
||||
blob("constructor"),
|
||||
]);
|
||||
const table = makeArrowTable([{ id: 1 }], { schema });
|
||||
expect(table.schema.fields.map((f) => f.name)).toEqual(["id"]);
|
||||
});
|
||||
|
||||
it("accepts ArrayBuffer and URL inside lists and structs", () => {
|
||||
const schema = new Schema([
|
||||
new Field("images", new List(blob("image")), true),
|
||||
new Field("info", new Struct([blob("image")]), true),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{
|
||||
images: [new Uint8Array([97]).buffer, new URL("s3://b/k")],
|
||||
info: { image: new Uint8Array([98]).buffer },
|
||||
},
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const images = Array.from(
|
||||
table.getChild("images")!.get(0) as Iterable<{
|
||||
data: Uint8Array | null;
|
||||
uri: string | null;
|
||||
}>,
|
||||
);
|
||||
expect(Buffer.from(images[0].data!).toString()).toBe("a");
|
||||
expect(images[1].uri).toBe("s3://b/k");
|
||||
const info = table.getChild("info")!.get(0) as {
|
||||
image: { data: Uint8Array };
|
||||
};
|
||||
expect(Buffer.from(info.image.data).toString()).toBe("b");
|
||||
});
|
||||
|
||||
it("coerces Buffer input onto a blob field", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
blob("image"),
|
||||
]);
|
||||
const table = makeArrowTable([{ id: 1n, image: Buffer.from("hello") }], {
|
||||
schema,
|
||||
});
|
||||
expect(isBlobField(table.schema.fields[1])).toBe(true);
|
||||
const image = table.getChild("image")!;
|
||||
expect(image.nullCount).toBe(0);
|
||||
expect(image.getChild("uri")!.get(0)).toBeNull();
|
||||
expect(image.getChild("data")!.nullCount).toBe(0);
|
||||
expect(Buffer.from(image.getChild("data")!.get(0)!).toString()).toBe(
|
||||
"hello",
|
||||
);
|
||||
});
|
||||
|
||||
it("coerces Buffer elements inside a list and keeps null slots", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field("images", new List(blob("image")), true),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{ id: 1n, images: [Buffer.from("a"), Buffer.from("bb")] },
|
||||
{ id: 2n, images: null },
|
||||
{ id: 3n, images: [Buffer.from("c"), null] },
|
||||
{ id: 4n, images: [] },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const images = table.getChild("images")!;
|
||||
expect(images.nullCount).toBe(1);
|
||||
const rows = images.toArray();
|
||||
expect(rows[1]).toBeNull();
|
||||
expect(Array.from(rows[3] as Iterable<unknown>)).toHaveLength(0);
|
||||
const first = Array.from(rows[0] as Iterable<{ data: Uint8Array | null }>);
|
||||
expect(Buffer.from(first[0].data!).toString()).toBe("a");
|
||||
expect(Buffer.from(first[1].data!).toString()).toBe("bb");
|
||||
const third = Array.from(
|
||||
rows[2] as Iterable<{ data: Uint8Array | null } | null>,
|
||||
);
|
||||
expect(Buffer.from(third[0]!.data!).toString()).toBe("c");
|
||||
expect(third[1]).toBeNull();
|
||||
});
|
||||
|
||||
it("coerces Buffer fields inside list structs", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field(
|
||||
"items",
|
||||
new List(
|
||||
new Field(
|
||||
"item",
|
||||
new Struct([new Field("name", new Utf8(), true), blob("image")]),
|
||||
true,
|
||||
),
|
||||
),
|
||||
true,
|
||||
),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{
|
||||
id: 1n,
|
||||
items: [{ name: "one", image: Buffer.from("alpha") }],
|
||||
},
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const items = Array.from(
|
||||
table.getChild("items")!.toArray()[0] as Iterable<{
|
||||
name: string;
|
||||
image: { data: Uint8Array | null };
|
||||
}>,
|
||||
);
|
||||
expect(items[0].name).toBe("one");
|
||||
expect(Buffer.from(items[0].image.data!).toString()).toBe("alpha");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,306 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as http from "http";
|
||||
import { Catalog, connectCatalog } from "../lancedb";
|
||||
|
||||
type RecordedRequest = {
|
||||
url: string;
|
||||
headers: http.IncomingHttpHeaders;
|
||||
body: Record<string, unknown>;
|
||||
};
|
||||
|
||||
async function withCatalog(
|
||||
responses: [number, unknown][],
|
||||
callback: (catalog: Catalog, requests: RecordedRequest[]) => Promise<void>,
|
||||
) {
|
||||
const requests: RecordedRequest[] = [];
|
||||
const server = http.createServer(async (req, res) => {
|
||||
const chunks: Buffer[] = [];
|
||||
for await (const chunk of req) chunks.push(Buffer.from(chunk));
|
||||
const body = Buffer.concat(chunks).toString();
|
||||
requests.push({
|
||||
url: req.url ?? "",
|
||||
headers: req.headers,
|
||||
body: body ? JSON.parse(body) : {},
|
||||
});
|
||||
const [status, response] = responses.shift() ?? [
|
||||
500,
|
||||
{ error: "Unexpected request" },
|
||||
];
|
||||
res.writeHead(status, { "content-type": "application/json" });
|
||||
res.end(status === 204 ? undefined : JSON.stringify(response));
|
||||
});
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
||||
const address = server.address();
|
||||
if (!address || typeof address === "string")
|
||||
throw new Error("Missing server address");
|
||||
try {
|
||||
const catalog = await connectCatalog(`http://127.0.0.1:${address.port}`, {
|
||||
apiKey: "secret",
|
||||
sqlHostOverride: "grpc+tls://sql.example.com:10026",
|
||||
clientConfig: {
|
||||
extraHeaders: {
|
||||
"X-LanceDB-Database": "wrong-static",
|
||||
"X-LanceDB-Database-Prefix": "wrong",
|
||||
},
|
||||
},
|
||||
headerProvider: () => ({
|
||||
"X-LanceDB-Database": "wrong-dynamic",
|
||||
"X-LanceDB-Database-Prefix": "wrong",
|
||||
authorization: "Bearer refreshed",
|
||||
}),
|
||||
});
|
||||
await callback(catalog, requests);
|
||||
expect(responses).toHaveLength(0);
|
||||
} finally {
|
||||
server.closeAllConnections();
|
||||
await new Promise<void>((resolve) => server.close(() => resolve()));
|
||||
}
|
||||
}
|
||||
|
||||
describe("remote catalog", () => {
|
||||
it("uses root namespace routes and preserves independent database scope", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[204, null],
|
||||
[200, { tables: [] }],
|
||||
[200, {}],
|
||||
[200, { tables: [] }],
|
||||
[200, { tables: [] }],
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["team/search"], page_token: "next" }],
|
||||
[204, null],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const first = await catalog.createDatabase("team/search", {
|
||||
existOk: true,
|
||||
});
|
||||
expect(await first.tableNames()).toEqual([]);
|
||||
const second = await catalog.connectDatabase("other");
|
||||
expect(await second.tableNames()).toEqual([]);
|
||||
expect(await first.tableNames()).toEqual([]);
|
||||
const names = catalog.listDatabases({ pageLimit: 1, pageToken: "a/b" });
|
||||
expect(await names.next()).toEqual({
|
||||
done: false,
|
||||
value: "team/search",
|
||||
});
|
||||
expect(names.pageToken()).toBe("next");
|
||||
await catalog.dropDatabase("team/search", { ignoreMissing: true });
|
||||
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
|
||||
expect(requests[0].body).toEqual({ mode: "ExistOk" });
|
||||
expect(requests[5].url).toBe(
|
||||
"/v1/namespace/$/list?limit=1&page_token=a%2Fb",
|
||||
);
|
||||
expect(requests[6].body).toEqual({
|
||||
mode: "Skip",
|
||||
behavior: "Restrict",
|
||||
});
|
||||
for (const [i, request] of requests.entries()) {
|
||||
expect(request.headers["x-lancedb-database"]).toBe(
|
||||
i === 1 || i === 4 ? "team/search" : i === 3 ? "other" : undefined,
|
||||
);
|
||||
expect(request.headers["x-lancedb-database-prefix"]).toBeUndefined();
|
||||
expect(request.headers.authorization).toBe("Bearer refreshed");
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates lifecycle errors and sends restricted drops", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[404, {}],
|
||||
[409, {}],
|
||||
[400, {}],
|
||||
[404, {}],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
await expect(catalog.connectDatabase("missing")).rejects.toThrow(
|
||||
"missing",
|
||||
);
|
||||
await expect(catalog.createDatabase("exists")).rejects.toThrow(
|
||||
"exists",
|
||||
);
|
||||
await expect(catalog.dropDatabase("full")).rejects.toThrow();
|
||||
await catalog.dropDatabase("missing", { ignoreMissing: true });
|
||||
expect(requests[2].body).toEqual({
|
||||
mode: "Fail",
|
||||
behavior: "Restrict",
|
||||
});
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("lazily traverses pages, including empty pages, and exposes cached state", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["a", "b"], page_token: "empty" }],
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: [], page_token: "last" }],
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["c", "d"], page_token: "" }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const names = catalog.listDatabases({ pageLimit: 2 });
|
||||
expect(names[Symbol.asyncIterator]()).toBe(names);
|
||||
expect(requests).toHaveLength(0);
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(names.pageToken()).toBeUndefined();
|
||||
expect(await names.next()).toEqual({ done: false, value: "a" });
|
||||
expect(names.numPageResults()).toBe(1);
|
||||
expect(names.pageToken()).toBe("empty");
|
||||
expect(await names.next()).toEqual({ done: false, value: "b" });
|
||||
expect(requests).toHaveLength(1);
|
||||
expect(await names.next()).toEqual({ done: false, value: "c" });
|
||||
expect(requests).toHaveLength(3);
|
||||
expect(names.pageToken()).toBeUndefined();
|
||||
expect(names.numPageResults()).toBe(1);
|
||||
const remaining = [];
|
||||
for await (const name of names) remaining.push(name);
|
||||
expect(remaining).toEqual(["d"]);
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(await names.next()).toEqual({ done: true, value: undefined });
|
||||
expect(requests.map((request) => request.url)).toEqual([
|
||||
"/v1/namespace/$/list?limit=2",
|
||||
"/v1/namespace/$/list?limit=2&page_token=empty",
|
||||
"/v1/namespace/$/list?limit=2&page_token=last",
|
||||
]);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("resumes from a saved token after draining the cache", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["a", "b"], page_token: "next/token" }],
|
||||
[200, { namespaces: ["c"] }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const names = catalog.listDatabases({
|
||||
pageToken: "start/token",
|
||||
pageLimit: 2,
|
||||
});
|
||||
expect(names.pageToken()).toBe("start/token");
|
||||
expect(await names.next()).toEqual({ done: false, value: "a" });
|
||||
while (names.numPageResults() > 0) await names.next();
|
||||
const resumed = catalog.listDatabases({
|
||||
pageToken: names.pageToken(),
|
||||
pageLimit: 2,
|
||||
});
|
||||
const results = [];
|
||||
for await (const name of resumed) results.push(name);
|
||||
expect(results).toEqual(["c"]);
|
||||
expect(requests.map((request) => request.url)).toEqual([
|
||||
"/v1/namespace/$/list?limit=2&page_token=start%2Ftoken",
|
||||
"/v1/namespace/$/list?limit=2&page_token=next%2Ftoken",
|
||||
]);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("terminates on errors and retains the failed request token", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["a"], page_token: "retry" }],
|
||||
[400, {}],
|
||||
[200, { namespaces: ["b"] }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const names = catalog.listDatabases();
|
||||
await names.next();
|
||||
await expect(names.next()).rejects.toThrow();
|
||||
expect(names.pageToken()).toBe("retry");
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(await names.next()).toEqual({ done: true, value: undefined });
|
||||
expect(requests).toHaveLength(2);
|
||||
const resumed = catalog.listDatabases({ pageToken: names.pageToken() });
|
||||
expect(await resumed.next()).toEqual({ done: false, value: "b" });
|
||||
expect(await resumed.next()).toEqual({ done: true, value: undefined });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("handles an empty listing and concurrent advances", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[200, { namespaces: [] }],
|
||||
[200, { namespaces: ["a", "b"] }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const empty = catalog.listDatabases();
|
||||
expect(await empty.next()).toEqual({ done: true, value: undefined });
|
||||
expect(await empty.next()).toEqual({ done: true, value: undefined });
|
||||
const names = catalog.listDatabases();
|
||||
const pending = [names.next(), names.next(), names.next()];
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(names.pageToken()).toBeUndefined();
|
||||
const results = await Promise.all(pending);
|
||||
expect(
|
||||
results
|
||||
.filter((result) => !result.done)
|
||||
.map((result) => result.value)
|
||||
.sort(),
|
||||
).toEqual(["a", "b"]);
|
||||
expect(results.filter((result) => result.done)).toHaveLength(1);
|
||||
expect(requests).toHaveLength(2);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("validates endpoints and pagination", async () => {
|
||||
await expect(connectCatalog("/tmp/catalog")).rejects.toThrow();
|
||||
const catalog = await connectCatalog("http://127.0.0.1:1");
|
||||
for (const limit of [
|
||||
0,
|
||||
-1,
|
||||
1.5,
|
||||
2147483648,
|
||||
Number.NaN,
|
||||
Number.POSITIVE_INFINITY,
|
||||
]) {
|
||||
expect(() => catalog.listDatabases({ pageLimit: limit })).toThrow(
|
||||
"limit",
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it.each([
|
||||
"",
|
||||
".",
|
||||
"..",
|
||||
"/db",
|
||||
"db/",
|
||||
"a//b",
|
||||
"a/./b",
|
||||
"a/../b",
|
||||
"a b",
|
||||
"a:b",
|
||||
"a$b",
|
||||
"a%b",
|
||||
"a?b",
|
||||
"a#b",
|
||||
"a\\b",
|
||||
"café",
|
||||
"a\nb",
|
||||
])(
|
||||
"rejects invalid database name %j before sending requests",
|
||||
async (name) => {
|
||||
await withCatalog([], async (catalog, requests) => {
|
||||
for (const method of [
|
||||
"createDatabase",
|
||||
"connectDatabase",
|
||||
"dropDatabase",
|
||||
] as const) {
|
||||
await expect(catalog[method](name)).rejects.toThrow(
|
||||
"Invalid database name",
|
||||
);
|
||||
}
|
||||
expect(requests).toEqual([]);
|
||||
});
|
||||
},
|
||||
);
|
||||
});
|
||||
@@ -1,7 +1,8 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { readdirSync } from "fs";
|
||||
import { mkdirSync, readdirSync } from "fs";
|
||||
import { join } from "path";
|
||||
import { Field, Float64, Schema } from "apache-arrow";
|
||||
import * as tmp from "tmp";
|
||||
import {
|
||||
@@ -180,6 +181,32 @@ describe("given a connection", () => {
|
||||
expect(seen).toEqual(created);
|
||||
});
|
||||
|
||||
it("should list all names and use a default page size of 100", async () => {
|
||||
const names = Array.from(
|
||||
{ length: 130 },
|
||||
(_, i) => `t${i.toString().padStart(3, "0")}`,
|
||||
);
|
||||
for (const name of names) {
|
||||
mkdirSync(join(tmpDir.name, `${name}.lance`));
|
||||
}
|
||||
await expect(db.tableNames()).resolves.toEqual(names);
|
||||
await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]);
|
||||
const first = await db.listTables();
|
||||
expect(first.tables).toEqual(names.slice(0, 100));
|
||||
expect(first.pageToken).toBeDefined();
|
||||
const second = await db.listTables({ pageToken: first.pageToken });
|
||||
expect(second.tables).toEqual(names.slice(100));
|
||||
expect(second.pageToken).toBeUndefined();
|
||||
await expect(db.listTables({ limit: 200 })).resolves.toEqual({
|
||||
tables: names,
|
||||
pageToken: undefined,
|
||||
});
|
||||
await expect(db.listTables({ limit: 0 })).resolves.toEqual({
|
||||
tables: [],
|
||||
pageToken: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it("should list tables in a namespace", async () => {
|
||||
const db = await connect(tmpDir.name, {
|
||||
// biome-ignore lint/style/useNamingConvention: opaque backend property key, must match Rust
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
@rem SPDX-License-Identifier: Apache-2.0
|
||||
@rem SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
@exit /b 0
|
||||
@@ -16,18 +16,47 @@ describe("materialized views", () => {
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
db = await connect(tmpDir.name);
|
||||
await db.createTable(
|
||||
"people",
|
||||
[
|
||||
{ name: "ada", age: 36 },
|
||||
{ name: "kid", age: 7 },
|
||||
{ name: "grace", age: 85 },
|
||||
],
|
||||
{ storageOptions: { newTableEnableStableRowIds: "true" } },
|
||||
);
|
||||
await db.createTable("people", [
|
||||
{ name: "ada", age: 36 },
|
||||
{ name: "kid", age: 7 },
|
||||
{ name: "grace", age: 85 },
|
||||
]);
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("reads stored queries and legacy layouts", () => {
|
||||
const read = (stored: string) =>
|
||||
definitionFromMetadata(new Map([[DEFINITION_META_KEY, stored]]), "v");
|
||||
const query =
|
||||
"SELECT id, c.chunk FROM ns.docs, UNNEST(chunks) AS c WHERE id > 1";
|
||||
expect(read(`{"format":1,"query":${JSON.stringify(query)}}`).query).toBe(
|
||||
query,
|
||||
);
|
||||
|
||||
// The structured layout written before the format number reads as the
|
||||
// query it described, under either of its kind tags.
|
||||
expect(
|
||||
read(
|
||||
'{"kind":"namespaced_select","source_table":"people","source_namespace":["ns"],' +
|
||||
'"projections":[{"output":"name","expression":"`name`"},' +
|
||||
'{"output":"Shout","expression":"upper(name)"}],"filter":"age >= 18","limit":42}',
|
||||
).query,
|
||||
).toBe(
|
||||
"SELECT `name`, upper(name) AS `Shout` FROM ns.people WHERE age >= 18 LIMIT 42",
|
||||
);
|
||||
expect(read('{"kind":"select","source_table":"people"}').query).toBe(
|
||||
"SELECT * FROM people",
|
||||
);
|
||||
|
||||
// A newer writer's layout is reported, never guessed at.
|
||||
for (const newer of [
|
||||
`{"format":3,"query":${JSON.stringify(query)}}`,
|
||||
'{"kind":"select_v3","source_table":"people"}',
|
||||
]) {
|
||||
expect(() => read(newer)).toThrow(/cannot refresh/);
|
||||
}
|
||||
});
|
||||
|
||||
it("rejects a stored limit a number cannot carry", () => {
|
||||
const big = new Map([
|
||||
[
|
||||
@@ -38,36 +67,6 @@ describe("materialized views", () => {
|
||||
expect(() => definitionFromMetadata(big, "v")).toThrow(
|
||||
/too large to represent exactly/,
|
||||
);
|
||||
|
||||
const safe = new Map([
|
||||
[
|
||||
DEFINITION_META_KEY,
|
||||
'{"kind":"select","source_table":"people","limit":42}',
|
||||
],
|
||||
]);
|
||||
expect(definitionFromMetadata(safe, "v").limit).toBe(42);
|
||||
});
|
||||
|
||||
it("reads the namespaced select kind and refuses unknown kinds", () => {
|
||||
// "namespaced_select" is the namespaced form of "select": same shape, a
|
||||
// separate kind so readers that predate it refuse instead of resolving
|
||||
// the source at the root.
|
||||
const namespaced = new Map([
|
||||
[
|
||||
DEFINITION_META_KEY,
|
||||
'{"kind":"namespaced_select","source_table":"people","source_namespace":["ns"]}',
|
||||
],
|
||||
]);
|
||||
const definition = definitionFromMetadata(namespaced, "v");
|
||||
expect(definition.sourceTable).toBe("people");
|
||||
expect(definition.sourceNamespace).toEqual(["ns"]);
|
||||
|
||||
const unknown = new Map([
|
||||
[DEFINITION_META_KEY, '{"kind":"select_v3","source_table":"people"}'],
|
||||
]);
|
||||
expect(() => definitionFromMetadata(unknown, "v")).toThrow(
|
||||
/cannot refresh/,
|
||||
);
|
||||
});
|
||||
|
||||
it("creates, refreshes and queries a view", async () => {
|
||||
@@ -76,11 +75,7 @@ describe("materialized views", () => {
|
||||
where: "age >= 18",
|
||||
});
|
||||
expect(view.name).toBe("adults");
|
||||
expect(await view.table().countRows()).toBe(0);
|
||||
|
||||
const result = await view.refresh();
|
||||
expect(result.mode).toBe("rebuild");
|
||||
expect(Number(result.rowsWritten)).toBe(2);
|
||||
expect(await view.table().countRows()).toBe(2);
|
||||
|
||||
const rows = await view.table().query().toArray();
|
||||
expect(rows.map((r) => r.shout).sort()).toEqual(["ADA", "GRACE"]);
|
||||
@@ -92,37 +87,50 @@ describe("materialized views", () => {
|
||||
});
|
||||
const view = await db.openMaterializedView("adults");
|
||||
const definition = await view.definition();
|
||||
expect(definition.sourceTable).toBe("people");
|
||||
expect(definition.filter).toBe("age >= 18");
|
||||
expect(definition.projections).toEqual([
|
||||
["name", "`name`"],
|
||||
["age", "`age`"],
|
||||
]);
|
||||
expect(definition.inputs).toEqual(["age", "name"]);
|
||||
expect(definition.query).toBe(
|
||||
"SELECT name, age FROM people WHERE age >= 18",
|
||||
);
|
||||
});
|
||||
|
||||
it("refreshes incrementally after an append", async () => {
|
||||
const view = await db.createMaterializedView("copy", "people");
|
||||
it("fully refreshes after an append", async () => {
|
||||
const view = await db.createMaterializedView("copy", "people", {
|
||||
withNoData: true,
|
||||
});
|
||||
await view.refresh();
|
||||
|
||||
const people = await db.openTable("people");
|
||||
await people.add([{ name: "alan", age: 41 }]);
|
||||
const result = await view.refresh();
|
||||
expect(result.mode).toBe("incremental");
|
||||
expect(Number(result.rowsWritten)).toBe(1);
|
||||
expect(result.mode).toBe("rebuild");
|
||||
expect(Number(result.rowsWritten)).toBe(4);
|
||||
expect(await view.table().countRows()).toBe(4);
|
||||
|
||||
expect((await view.refresh()).mode).toBe("no_op");
|
||||
expect((await view.refresh()).mode).toBe("rebuild");
|
||||
});
|
||||
|
||||
it("lists views and rejects non-views", async () => {
|
||||
await db.createMaterializedView("adults", "people", {
|
||||
where: "age >= 18",
|
||||
});
|
||||
expect(await db.listMaterializedViews()).toEqual(["adults"]);
|
||||
expect(await collect(db.listMaterializedViews())).toEqual(["adults"]);
|
||||
await expect(db.openMaterializedView("people")).rejects.toThrow(
|
||||
"not a materialized view",
|
||||
);
|
||||
await expect(db.dropMaterializedView("people")).rejects.toThrow(
|
||||
"not a materialized view",
|
||||
);
|
||||
|
||||
await db.dropMaterializedView("adults");
|
||||
expect(await collect(db.listMaterializedViews())).toEqual([]);
|
||||
});
|
||||
|
||||
it("returns a job when dropping a view asynchronously", async () => {
|
||||
await db.createMaterializedView("adults", "people");
|
||||
|
||||
const job = await db.dropMaterializedViewAsync("adults");
|
||||
expect(job.id).toBeNull();
|
||||
await job.wait();
|
||||
expect(await collect(db.listMaterializedViews())).toEqual([]);
|
||||
});
|
||||
|
||||
it("rejects an invalid expression at create time", async () => {
|
||||
@@ -139,7 +147,7 @@ describe("materialized views", () => {
|
||||
db.createMaterializedView("bad", "people", { limit }),
|
||||
).rejects.toThrow("non-negative integer");
|
||||
}
|
||||
expect(await db.listMaterializedViews()).toEqual([]);
|
||||
expect(await collect(db.listMaterializedViews())).toEqual([]);
|
||||
|
||||
const view = await db.createMaterializedView("copy", "people");
|
||||
for (const sourceVersion of [-1, 1.5, Infinity, NaN]) {
|
||||
@@ -150,20 +158,24 @@ describe("materialized views", () => {
|
||||
});
|
||||
|
||||
it("quotes bare select names", async () => {
|
||||
await db.createTable("odd_names", [{ "order item": "widget" }], {
|
||||
storageOptions: { newTableEnableStableRowIds: "true" },
|
||||
});
|
||||
await db.createTable("odd_names", [{ "order item": "widget" }]);
|
||||
const view = await db.createMaterializedView("quoted", "odd_names", {
|
||||
select: ["order item"],
|
||||
withNoData: true,
|
||||
});
|
||||
const result = await view.refresh();
|
||||
expect(Number(result.rowsWritten)).toBe(1);
|
||||
});
|
||||
|
||||
it("requires stable row ids on the source", async () => {
|
||||
it("does not require stable row ids on the source", async () => {
|
||||
await db.createTable("plain", [{ x: 1 }]);
|
||||
await expect(db.createMaterializedView("v", "plain")).rejects.toThrow(
|
||||
"stable row ids",
|
||||
);
|
||||
const view = await db.createMaterializedView("v", "plain");
|
||||
expect(await view.table().countRows()).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
async function collect<T>(items: AsyncIterable<T>): Promise<T[]> {
|
||||
const result: T[] = [];
|
||||
for await (const item of items) result.push(item);
|
||||
return result;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as fs from "fs";
|
||||
import * as http from "http";
|
||||
import { execFileSync } from "node:child_process";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
import { OAuthConfig, OAuthFlowType, OAuthSession } from "../lancedb/oauth";
|
||||
|
||||
function tempCacheDir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "lancedb-oauth-cache-"));
|
||||
}
|
||||
|
||||
function deviceConfig(issuerUrl: string, cacheDir: string): OAuthConfig {
|
||||
return {
|
||||
issuerUrl,
|
||||
clientId: "client-id",
|
||||
scopes: ["openid"],
|
||||
flow: OAuthFlowType.DeviceCode,
|
||||
tokenCache: { cacheDir },
|
||||
};
|
||||
}
|
||||
|
||||
describe("OAuthSession", () => {
|
||||
beforeAll(() => {
|
||||
// Child processes inherit the real environment, just as Rust reads it.
|
||||
// Fail before login if the browser override only exists in Jest's sandbox.
|
||||
const browser = execFileSync(
|
||||
process.execPath,
|
||||
["-p", "process.env.LANCEDB_OAUTH_BROWSER ?? ''"],
|
||||
{ encoding: "utf8" },
|
||||
).trim();
|
||||
expect(browser).not.toBe("");
|
||||
expect(browser).toBe(process.env.LANCEDB_OAUTH_BROWSER);
|
||||
});
|
||||
|
||||
it("reports an absent session and logout is idempotent", async () => {
|
||||
const cacheDir = tempCacheDir();
|
||||
const session = new OAuthSession(
|
||||
deviceConfig("https://issuer.example.com", cacheDir),
|
||||
);
|
||||
|
||||
const status = await session.status();
|
||||
expect(status.refreshable).toBe(false);
|
||||
expect(status.issuerUrl).toBe("https://issuer.example.com");
|
||||
expect(status.clientId).toBe("client-id");
|
||||
expect(status.scopes).toEqual(["openid"]);
|
||||
expect(status.flow).toBe("device_code");
|
||||
expect(status.obtainedAt).toBeUndefined();
|
||||
|
||||
const logout = await session.logout();
|
||||
expect(logout.removed).toBe(false);
|
||||
});
|
||||
|
||||
it("requires token cache options", () => {
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://issuer.example.com",
|
||||
clientId: "client-id",
|
||||
scopes: ["openid"],
|
||||
flow: OAuthFlowType.DeviceCode,
|
||||
};
|
||||
expect(() => new OAuthSession(config)).toThrow(/token/);
|
||||
});
|
||||
|
||||
it("rejects azure managed identity persistence", () => {
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://login.microsoftonline.com/tenant/v2.0",
|
||||
clientId: "app-id",
|
||||
scopes: ["api://app/.default"],
|
||||
flow: OAuthFlowType.AzureManagedIdentity,
|
||||
tokenCache: { cacheDir: tempCacheDir() },
|
||||
};
|
||||
expect(() => new OAuthSession(config)).toThrow(/AzureManagedIdentity/);
|
||||
});
|
||||
|
||||
it.each([
|
||||
{},
|
||||
{
|
||||
resource: "https://api.example.com/a?x=1&y=two",
|
||||
audience: "audience + & / ü",
|
||||
},
|
||||
])(
|
||||
"logs in via device flow with target %j, caches, and logs out",
|
||||
async (target) => {
|
||||
const server = new MockIdp();
|
||||
await server.start();
|
||||
try {
|
||||
const cacheDir = tempCacheDir();
|
||||
const issuerUrl = server.issuerUrl();
|
||||
|
||||
const config = { ...deviceConfig(issuerUrl, cacheDir), ...target };
|
||||
const session = new OAuthSession(config);
|
||||
const status = await session.login();
|
||||
expect(status.refreshable).toBe(true);
|
||||
expect(status.resource).toBe(config.resource);
|
||||
expect(status.audience).toBe(config.audience);
|
||||
expect(status.obtainedAt).toBeGreaterThan(0);
|
||||
expect(server.state.deviceAuthorizations).toBe(1);
|
||||
|
||||
// An independent session (a fresh "process") sees the cached login.
|
||||
const other = new OAuthSession(config);
|
||||
const cached = await other.status();
|
||||
expect(cached.refreshable).toBe(true);
|
||||
|
||||
const logout = await other.logout();
|
||||
expect(logout.removed).toBe(true);
|
||||
const again = await session.logout();
|
||||
expect(again.removed).toBe(false);
|
||||
expect((await session.status()).refreshable).toBe(false);
|
||||
|
||||
// Only the initial login used the interactive device flow.
|
||||
expect(server.state.deviceAuthorizations).toBe(1);
|
||||
expect(server.state.refreshGrants).toBe(0);
|
||||
expect(server.requests).toHaveLength(2);
|
||||
for (const params of server.requests) {
|
||||
expect(params.getAll("resource")).toEqual(
|
||||
config.resource === undefined ? [] : [config.resource],
|
||||
);
|
||||
expect(params.getAll("audience")).toEqual(
|
||||
config.audience === undefined ? [] : [config.audience],
|
||||
);
|
||||
}
|
||||
} finally {
|
||||
server.close();
|
||||
}
|
||||
},
|
||||
15000,
|
||||
);
|
||||
});
|
||||
|
||||
/** Mock IdP with discovery, device authorization, and rotating refresh. */
|
||||
class MockIdp {
|
||||
readonly requests: URLSearchParams[] = [];
|
||||
readonly state = {
|
||||
deviceAuthorizations: 0,
|
||||
refreshGrants: 0,
|
||||
accessTokensIssued: 0,
|
||||
currentRefresh: null as string | null,
|
||||
};
|
||||
private server?: http.Server;
|
||||
private port = 0;
|
||||
|
||||
issuerUrl(): string {
|
||||
return `http://127.0.0.1:${this.port}`;
|
||||
}
|
||||
|
||||
async start(): Promise<void> {
|
||||
const server = http.createServer((req, res) => {
|
||||
const chunks: Buffer[] = [];
|
||||
req.on("data", (chunk) => chunks.push(chunk));
|
||||
req.on("end", () => {
|
||||
const body = Buffer.concat(chunks).toString();
|
||||
const params = new URLSearchParams(body);
|
||||
this.handle(req.url ?? "", params, res);
|
||||
});
|
||||
});
|
||||
await new Promise<void>((resolve) => {
|
||||
server.listen(0, "127.0.0.1", () => resolve());
|
||||
});
|
||||
const address = server.address();
|
||||
if (address && typeof address === "object") {
|
||||
this.port = address.port;
|
||||
}
|
||||
this.server = server;
|
||||
}
|
||||
|
||||
private handle(
|
||||
url: string,
|
||||
params: URLSearchParams,
|
||||
res: http.ServerResponse,
|
||||
): void {
|
||||
const respond = (status: number, payload: unknown): void => {
|
||||
const body = JSON.stringify(payload);
|
||||
res.writeHead(status, {
|
||||
"Content-Type": "application/json",
|
||||
"Content-Length": Buffer.byteLength(body),
|
||||
});
|
||||
res.end(body);
|
||||
};
|
||||
|
||||
if (url === "/.well-known/openid-configuration") {
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
token_endpoint: `${this.issuerUrl()}/token`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
device_authorization_endpoint: `${this.issuerUrl()}/device`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (url === "/device" || url === "/token") {
|
||||
this.requests.push(params);
|
||||
}
|
||||
|
||||
if (url === "/device") {
|
||||
this.state.deviceAuthorizations += 1;
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
device_code: "device-code",
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
user_code: "ABCD-EFGH",
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
verification_uri: `${this.issuerUrl()}/verify`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
expires_in: 60,
|
||||
interval: 1,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (url === "/token") {
|
||||
const grantType = params.get("grant_type") ?? "";
|
||||
if (grantType === "refresh_token") {
|
||||
this.state.refreshGrants += 1;
|
||||
if (params.get("refresh_token") !== this.state.currentRefresh) {
|
||||
respond(400, { error: "invalid_grant" });
|
||||
return;
|
||||
}
|
||||
} else if (!grantType.includes("device_code")) {
|
||||
respond(400, { error: "unsupported_grant_type" });
|
||||
return;
|
||||
}
|
||||
this.state.accessTokensIssued += 1;
|
||||
const number = this.state.accessTokensIssued;
|
||||
const refresh = `refresh-${number}`;
|
||||
this.state.currentRefresh = refresh;
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
access_token: `access-${number}`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
refresh_token: refresh,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
expires_in: 3600,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
respond(404, {});
|
||||
}
|
||||
|
||||
close(): void {
|
||||
this.server?.close();
|
||||
}
|
||||
}
|
||||
@@ -110,6 +110,129 @@ describe("Query outputSchema", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("Query empty results", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema(
|
||||
[
|
||||
new Field("id", new Int64(), false),
|
||||
new Field(
|
||||
"text",
|
||||
new Utf8(),
|
||||
true,
|
||||
new Map([["description", "searchable text"]]),
|
||||
),
|
||||
new Field(
|
||||
"vector",
|
||||
new FixedSizeList(2, new Field("item", new Float32())),
|
||||
true,
|
||||
),
|
||||
],
|
||||
new Map([["source", "empty-query-regression"]]),
|
||||
);
|
||||
table = await db.createTable(
|
||||
"test",
|
||||
makeArrowTable([{ id: 1n, text: "hello", vector: [1, 2] }], { schema }),
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{
|
||||
name: "filtered scan",
|
||||
query: (table: Table) => table.query().where("id < 0"),
|
||||
fields: ["id", "text", "vector"],
|
||||
},
|
||||
{
|
||||
name: "zero limit",
|
||||
query: (table: Table) => table.query().limit(0),
|
||||
fields: ["id", "text", "vector"],
|
||||
},
|
||||
{
|
||||
name: "column projection",
|
||||
query: (table: Table) =>
|
||||
table.query().where("id < 0").select(["text", "id"]),
|
||||
fields: ["text", "id"],
|
||||
},
|
||||
{
|
||||
name: "dynamic projection",
|
||||
query: (table: Table) =>
|
||||
table.query().where("id < 0").select({ doubled: "id * 2" }),
|
||||
fields: ["doubled"],
|
||||
},
|
||||
{
|
||||
name: "filtered vector search",
|
||||
query: (table: Table) =>
|
||||
table.search([0.5, 0.5]).where("id < 0").limit(3),
|
||||
fields: ["id", "text", "vector", "_distance"],
|
||||
},
|
||||
{
|
||||
name: "vector projection with row id",
|
||||
query: (table: Table) =>
|
||||
table
|
||||
.vectorSearch([0.5, 0.5])
|
||||
.where("id < 0")
|
||||
.select(["id"])
|
||||
.withRowId(),
|
||||
fields: ["id", "_distance", "_rowid"],
|
||||
},
|
||||
{
|
||||
name: "fast search without an index",
|
||||
query: (table: Table) => table.search([0.5, 0.5]).fastSearch(),
|
||||
fields: ["id", "text", "vector", "_distance"],
|
||||
},
|
||||
{
|
||||
name: "take with no offsets",
|
||||
query: (table: Table) => table.takeOffsets([]).select(["text"]),
|
||||
fields: ["text"],
|
||||
},
|
||||
])("preserves the output schema for $name", async ({ query, fields }) => {
|
||||
const result = await query(table).toArrow({ maxBatchLength: 1 });
|
||||
|
||||
expect(result.numRows).toBe(0);
|
||||
expect(result.schema.fields.map((field) => field.name)).toEqual(fields);
|
||||
for (const name of fields) {
|
||||
expect(result.getChild(name)).not.toBeNull();
|
||||
expect(result.getChild(name)?.length).toBe(0);
|
||||
}
|
||||
});
|
||||
|
||||
it("preserves types, nullability, and metadata", async () => {
|
||||
const result = await table.query().where("id < 0").toArrow();
|
||||
const schema = await table.schema();
|
||||
|
||||
expect(result.schema).toEqual(schema);
|
||||
expect(result.getChild("id")?.type.toString()).toBe("Int64");
|
||||
expect(result.schema.fields[0].nullable).toBe(false);
|
||||
expect(result.schema.fields[1].metadata.get("description")).toBe(
|
||||
"searchable text",
|
||||
);
|
||||
expect(result.schema.metadata.get("source")).toBe("empty-query-regression");
|
||||
});
|
||||
|
||||
it("preserves the score column for full-text search with no matches", async () => {
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
const result = await table
|
||||
.search("missing", "fts")
|
||||
.select(["id"])
|
||||
.toArrow();
|
||||
|
||||
expect(result.numRows).toBe(0);
|
||||
expect(result.schema.fields.map((field) => field.name)).toEqual([
|
||||
"id",
|
||||
"_score",
|
||||
]);
|
||||
expect(result.getChild("_score")?.type.toString()).toBe("Float32");
|
||||
});
|
||||
});
|
||||
|
||||
describe("Search pagination", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
@@ -404,3 +527,106 @@ describe("Query orderBy", () => {
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("IVF search effort", () => {
|
||||
let directory: tmp.DirResult;
|
||||
let table: Table;
|
||||
const vector = Array.from({ length: 16 }, (_, i) => (i === 0 ? 1 : 0));
|
||||
|
||||
beforeEach(async () => {
|
||||
directory = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(directory.name);
|
||||
const rows = Array.from({ length: 512 }, (_, id) => ({
|
||||
id,
|
||||
vector: Array.from(
|
||||
{ length: 16 },
|
||||
(_, j) =>
|
||||
(Math.floor(id / 32) === j ? 1 : 0) + Math.sin(id * 17 + j) * 0.001,
|
||||
),
|
||||
}));
|
||||
table = await db.createTable("effort", rows.slice(0, 256));
|
||||
await table.add(rows.slice(256));
|
||||
await table.createIndex("vector", {
|
||||
config: Index.ivfFlat({ numPartitions: 16 }),
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
directory.removeCallback();
|
||||
});
|
||||
|
||||
it.each([
|
||||
[null, 6],
|
||||
[0, 1],
|
||||
[0.25, 3],
|
||||
[0.5, 6],
|
||||
[0.75, 10],
|
||||
[1, 16],
|
||||
])(
|
||||
"preserves effort %s through execution and probe precedence",
|
||||
async (effort, expected) => {
|
||||
const partitions = async (query: ReturnType<Table["vectorSearch"]>) => {
|
||||
const plan = await query.limit(8).analyzePlan();
|
||||
const matches = [...plan.matchAll(/partitions_searched=(\d+)/g)];
|
||||
expect(matches.length).toBeGreaterThan(0);
|
||||
return matches.reduce((total, match) => total + Number(match[1]), 0);
|
||||
};
|
||||
expect(
|
||||
await partitions(table.vectorSearch(vector).searchEffort(effort)),
|
||||
).toBe(expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).nprobes(1).searchEffort(effort),
|
||||
),
|
||||
).toBe(effort === null ? 1 : expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).searchEffort(effort).nprobes(1),
|
||||
),
|
||||
).toBe(effort === null ? 1 : expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table
|
||||
.vectorSearch(vector)
|
||||
.searchEffort(effort)
|
||||
.nprobes(1)
|
||||
.searchEffort(null),
|
||||
),
|
||||
).toBe(1);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).searchEffort(effort).maximumNprobes(2),
|
||||
),
|
||||
).toBeLessThanOrEqual(2);
|
||||
const actual = await table
|
||||
.vectorSearch(vector)
|
||||
.searchEffort(effort)
|
||||
.limit(8)
|
||||
.toArray();
|
||||
const exact = await table
|
||||
.vectorSearch(vector)
|
||||
.bypassVectorIndex()
|
||||
.limit(8)
|
||||
.toArray();
|
||||
const ids = new Set(exact.map((row) => row.id));
|
||||
expect(
|
||||
actual.filter((row) => ids.has(row.id)).length / 8,
|
||||
).toBeGreaterThanOrEqual(0.5);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
-0.1,
|
||||
1.1,
|
||||
Number.NaN,
|
||||
Number.POSITIVE_INFINITY,
|
||||
Number.NEGATIVE_INFINITY,
|
||||
])("rejects invalid effort %s before execution", (effort) => {
|
||||
expect(() => table.vectorSearch(vector).searchEffort(effort)).toThrow(
|
||||
"search_effort must be finite and in [0, 1]",
|
||||
);
|
||||
expect(() =>
|
||||
table.vectorSearch(vector).fullTextSearch("dog").searchEffort(effort),
|
||||
).toThrow("search_effort must be finite and in [0, 1]");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -3,13 +3,10 @@
|
||||
|
||||
import * as apiArrow from "apache-arrow";
|
||||
|
||||
import * as arrow15 from "apache-arrow-15";
|
||||
import * as arrow16 from "apache-arrow-16";
|
||||
import * as arrow17 from "apache-arrow-17";
|
||||
import * as arrow18 from "apache-arrow-18";
|
||||
|
||||
import * as tmp from "tmp";
|
||||
|
||||
import { arrowVersions } from "./arrow_versions";
|
||||
|
||||
import { connect } from "../lancedb";
|
||||
import {
|
||||
EmbeddingFunction,
|
||||
@@ -18,7 +15,7 @@ import {
|
||||
} from "../lancedb/embedding";
|
||||
import { getRegistry, register } from "../lancedb/embedding/registry";
|
||||
|
||||
describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => {
|
||||
describe.each(arrowVersions)("LanceSchema", (arrow) => {
|
||||
test("should preserve input order", async () => {
|
||||
const schema = LanceSchema({
|
||||
id: new arrow.Int32(),
|
||||
@@ -29,7 +26,7 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => {
|
||||
});
|
||||
});
|
||||
|
||||
describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
|
||||
describe.each(arrowVersions)("Registry", (arrow) => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
@@ -240,7 +237,7 @@ describe("Registry.setVar", () => {
|
||||
return data.map(() => [1, 2, 3]);
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow18.Float32() as apiArrow.Float;
|
||||
return new apiArrow.Float32();
|
||||
}
|
||||
protected getSensitiveKeys() {
|
||||
return ["secretKey"];
|
||||
|
||||
+606
-10
@@ -5,18 +5,30 @@ 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 {
|
||||
Table as ArrowTable,
|
||||
Field,
|
||||
Int64,
|
||||
RecordBatch,
|
||||
Schema,
|
||||
tableToIPC,
|
||||
} from "../lancedb/arrow";
|
||||
import {
|
||||
HeaderProvider,
|
||||
OAuthHeaderProvider,
|
||||
StaticHeaderProvider,
|
||||
} from "../lancedb/header";
|
||||
import { Index } from "../lancedb/indices";
|
||||
import { LocalTable } from "../lancedb/table";
|
||||
|
||||
// Test-only header providers
|
||||
class CustomProvider extends HeaderProvider {
|
||||
@@ -82,21 +94,314 @@ async function withMockDatabase(
|
||||
}
|
||||
|
||||
describe("remote connection", () => {
|
||||
it("refuses materialized views before issuing any request", async () => {
|
||||
const paths: string[] = [];
|
||||
it("reports unsupported operations with backend-neutral errors", async () => {
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
paths.push(req.url ?? "");
|
||||
expect(req.url).toBe("/v1/table/test/describe/");
|
||||
res.writeHead(200, { "Content-Type": "application/json" }).end(
|
||||
JSON.stringify({
|
||||
version: 1,
|
||||
schema: {
|
||||
fields: [
|
||||
{ name: "id", type: { type: "int64" }, nullable: false },
|
||||
{
|
||||
name: "vector",
|
||||
type: {
|
||||
type: "fixed_size_list",
|
||||
fields: [
|
||||
{ name: "item", type: { type: "float" }, nullable: true },
|
||||
],
|
||||
length: 2,
|
||||
},
|
||||
nullable: false,
|
||||
},
|
||||
],
|
||||
},
|
||||
}),
|
||||
);
|
||||
},
|
||||
async (db) => {
|
||||
const table = (await db.openTable("test")) as LocalTable;
|
||||
const cases: [() => Promise<unknown>, string][] = [
|
||||
[() => table.usesV2ManifestPaths(), "uses_v2_manifest_paths"],
|
||||
[() => table.migrateManifestPathsV2(), "migrate_manifest_paths_v2"],
|
||||
[() => table.optimize(), "optimize"],
|
||||
[
|
||||
() => table.createIndex("vector", { config: Index.hnswPq() }),
|
||||
"IVF_HNSW_PQ",
|
||||
],
|
||||
];
|
||||
for (const [run, operation] of cases) {
|
||||
let message = `LanceDBError: not supported: ${operation} is not supported for remote tables.`;
|
||||
if (operation === "IVF_HNSW_PQ") {
|
||||
message += " Please use IVF_HNSW_SQ instead.";
|
||||
}
|
||||
await expect(run()).rejects.toThrow(new Error(message));
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects nprobes(0) before sending a query", async () => {
|
||||
const requests: string[] = [];
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
requests.push(req.url ?? "");
|
||||
if (req.url === "/v1/table/test/describe/") {
|
||||
res
|
||||
.writeHead(200, { "Content-Type": "application/json" })
|
||||
.end(JSON.stringify({ version: 1, schema: { fields: [] } }));
|
||||
} else {
|
||||
res.writeHead(404).end();
|
||||
}
|
||||
},
|
||||
async (db) => {
|
||||
const table = await db.openTable("test");
|
||||
expect(() => table.vectorSearch([0, 0]).nprobes(0)).toThrow(
|
||||
"Invalid input, nprobes must be greater than 0",
|
||||
);
|
||||
expect(() =>
|
||||
table.query().nearestTo([0, 0]).fullTextSearch("dog").nprobes(0),
|
||||
).toThrow("Invalid input, nprobes must be greater than 0");
|
||||
await expect(
|
||||
table
|
||||
.vectorSearch(Promise.resolve([0, 0]))
|
||||
.nprobes(0)
|
||||
.toArrow(),
|
||||
).rejects.toThrow("Invalid input, nprobes must be greater than 0");
|
||||
},
|
||||
);
|
||||
expect(requests).toEqual(["/v1/table/test/describe/"]);
|
||||
});
|
||||
|
||||
it.each([false, true])(
|
||||
"preserves an empty query's schema with an empty batch: %s",
|
||||
async (withEmptyBatch) => {
|
||||
const schema = new Schema(
|
||||
[new Field("doubled", new Int64(), false)],
|
||||
new Map([["source", "query-output"]]),
|
||||
);
|
||||
const result = new ArrowTable(
|
||||
schema,
|
||||
withEmptyBatch ? [new RecordBatch(schema, undefined)] : [],
|
||||
);
|
||||
const response = Buffer.from(tableToIPC(result, "stream"));
|
||||
let queryRequests = 0;
|
||||
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
if (req.url?.endsWith("/describe/")) {
|
||||
res.writeHead(200, { "Content-Type": "application/json" }).end(
|
||||
JSON.stringify({
|
||||
name: "items",
|
||||
version: 1,
|
||||
schema: {
|
||||
fields: [
|
||||
{ name: "id", type: { type: "int64" }, nullable: false },
|
||||
],
|
||||
},
|
||||
}),
|
||||
);
|
||||
} else if (req.url?.endsWith("/query/")) {
|
||||
queryRequests++;
|
||||
req.resume();
|
||||
req.on("end", () => {
|
||||
res
|
||||
.writeHead(200, {
|
||||
"Content-Type": "application/vnd.apache.arrow.stream",
|
||||
})
|
||||
.end(response);
|
||||
});
|
||||
} else {
|
||||
res.writeHead(404).end();
|
||||
}
|
||||
},
|
||||
async (db) => {
|
||||
const table = await db.openTable("items");
|
||||
const result = await table
|
||||
.query()
|
||||
.where("id < 0")
|
||||
.select({ doubled: "id * 2" })
|
||||
.toArrow();
|
||||
|
||||
expect(result.numRows).toBe(0);
|
||||
expect(result.schema).toEqual(schema);
|
||||
expect(result.getChild("doubled")?.length).toBe(0);
|
||||
},
|
||||
);
|
||||
|
||||
expect(queryRequests).toBe(1);
|
||||
},
|
||||
);
|
||||
|
||||
it("rejects the external blob opt-in without blocking regular adds", async () => {
|
||||
let insertRequests = 0;
|
||||
const describeRequests: string[] = [];
|
||||
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
const requestPath = req.url ?? "";
|
||||
if (requestPath.endsWith("/describe/")) {
|
||||
describeRequests.push(requestPath);
|
||||
res.writeHead(200, { "Content-Type": "application/json" }).end(
|
||||
JSON.stringify({
|
||||
name: "items",
|
||||
version: 1,
|
||||
schema: {
|
||||
fields: [
|
||||
{
|
||||
name: "id",
|
||||
type: { type: "int64" },
|
||||
nullable: true,
|
||||
},
|
||||
],
|
||||
},
|
||||
}),
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (requestPath.endsWith("/insert/")) {
|
||||
insertRequests++;
|
||||
req.resume();
|
||||
req.on("end", () => {
|
||||
res
|
||||
.writeHead(200, { "Content-Type": "application/json" })
|
||||
.end(JSON.stringify({ version: insertRequests + 1 }));
|
||||
});
|
||||
return;
|
||||
}
|
||||
res.writeHead(404).end();
|
||||
},
|
||||
async (db) => {
|
||||
await expect(db.openMaterializedView("secret_table")).rejects.toThrow(
|
||||
/only on local databases/,
|
||||
const table = await db.openTable("items");
|
||||
|
||||
await expect(
|
||||
table.add([{ id: 1n }], {
|
||||
allowExternalBlobOutsideBases: true,
|
||||
}),
|
||||
).rejects.toThrow("only supported on local tables");
|
||||
expect(insertRequests).toBe(0);
|
||||
|
||||
await expect(table.add([{ id: 2n }])).resolves.toMatchObject({
|
||||
version: 2,
|
||||
});
|
||||
await expect(
|
||||
table.add([{ id: 3n }], {
|
||||
allowExternalBlobOutsideBases: false,
|
||||
}),
|
||||
).resolves.toMatchObject({ version: 3 });
|
||||
},
|
||||
);
|
||||
|
||||
expect(describeRequests.length).toBeGreaterThan(0);
|
||||
expect(insertRequests).toBe(2);
|
||||
});
|
||||
|
||||
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 collect(db.listMaterializedViews())).toEqual([
|
||||
"daily_sales",
|
||||
]);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("creates a view and decodes the schema it resolved to", async () => {
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
expect(req.method).toBe("POST");
|
||||
expect(req.url).toBe("/v1/view/analytics$adults/create");
|
||||
res.writeHead(200, { "content-type": "application/json" }).end(
|
||||
JSON.stringify({
|
||||
name: "adults",
|
||||
namespace: ["analytics"],
|
||||
query: "SELECT name FROM people",
|
||||
// biome-ignore lint/style/useNamingConvention: the wire field is snake_case
|
||||
default_database: "db",
|
||||
// biome-ignore lint/style/useNamingConvention: the wire field is snake_case
|
||||
default_namespace: ["analytics"],
|
||||
schema: {
|
||||
fields: [
|
||||
{ name: "name", nullable: true, type: { type: "utf8" } },
|
||||
],
|
||||
},
|
||||
}),
|
||||
);
|
||||
await expect(db.listMaterializedViews()).rejects.toThrow(
|
||||
/only on local databases/,
|
||||
);
|
||||
expect(paths).toEqual([]);
|
||||
},
|
||||
async (db) => {
|
||||
const view = await db.createView("adults", "SELECT name FROM people", [
|
||||
"analytics",
|
||||
]);
|
||||
expect(view.name).toBe("adults");
|
||||
expect(view.namespacePath).toEqual(["analytics"]);
|
||||
expect(view.query).toBe("SELECT name FROM people");
|
||||
expect(view.defaultDatabase).toBe("db");
|
||||
expect(view.defaultNamespacePath).toEqual(["analytics"]);
|
||||
expect(view.schema.fields.map((f) => f.name)).toEqual(["name"]);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("lists and drops views through their own routes", async () => {
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
expect(req.url).toBe("/v1/namespace/$/view/list");
|
||||
res
|
||||
.writeHead(200, { "content-type": "application/json" })
|
||||
.end(JSON.stringify({ views: ["adults"] }));
|
||||
},
|
||||
async (db) => {
|
||||
expect(await collect(db.listViews())).toEqual(["adults"]);
|
||||
},
|
||||
);
|
||||
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
expect(req.method).toBe("POST");
|
||||
expect(req.url).toBe("/v1/view/adults/drop");
|
||||
res.writeHead(200, { "content-type": "application/json" }).end("{}");
|
||||
},
|
||||
async (db) => {
|
||||
await db.dropView("adults");
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("reports the cleanup job when a view drop is accepted", async () => {
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
expect(req.method).toBe("POST");
|
||||
expect(req.url).toBe("/v1/view/adults/drop");
|
||||
res
|
||||
.writeHead(202, { "content-type": "application/json" })
|
||||
.end('{"job_id": "j1-do-abc"}');
|
||||
},
|
||||
async (db) => {
|
||||
const job = await db.dropViewAsync("adults");
|
||||
expect(job.id).toBe("j1-do-abc");
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("reports a finished job when a view drop had nothing to delete", async () => {
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
expect(req.url).toBe("/v1/view/adults/drop");
|
||||
res.writeHead(200, { "content-type": "application/json" }).end("{}");
|
||||
},
|
||||
async (db) => {
|
||||
// A 200 means the name was not bound, so there is no cleanup to wait on.
|
||||
const job = await db.dropViewAsync("adults");
|
||||
expect(job.id).toBeNull();
|
||||
await job.wait();
|
||||
},
|
||||
);
|
||||
});
|
||||
@@ -171,6 +476,15 @@ describe("remote connection", () => {
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects a zero blob request concurrency limit", async () => {
|
||||
await expect(
|
||||
connect("db://dev", {
|
||||
apiKey: "fake",
|
||||
clientConfig: { blobRequestConcurrency: 0 },
|
||||
}),
|
||||
).rejects.toThrow("blob_request_concurrency");
|
||||
});
|
||||
|
||||
it("shows the full error messages on retry errors", async () => {
|
||||
await withMockDatabase(
|
||||
(_req, res) => {
|
||||
@@ -196,6 +510,66 @@ describe("remote connection", () => {
|
||||
);
|
||||
});
|
||||
|
||||
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) => {
|
||||
@@ -442,6 +816,40 @@ describe("remote connection", () => {
|
||||
]);
|
||||
});
|
||||
|
||||
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 = {
|
||||
@@ -998,6 +1406,24 @@ describe("remote connection jobs surface", () => {
|
||||
res
|
||||
.writeHead(200, { "Content-Type": "application/json" })
|
||||
.end('{"job_id": "job-1"}');
|
||||
} else if (req.url === "/v1/jobs/pause") {
|
||||
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", "paused": true}');
|
||||
} else if (req.url === "/v1/jobs/resume") {
|
||||
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", "resumed": false, "still_pausing": true}',
|
||||
);
|
||||
} else if (req.url === "/v1/jobs/query_events") {
|
||||
queryEventsPayloads.push(payload);
|
||||
res
|
||||
@@ -1011,7 +1437,7 @@ describe("remote connection jobs surface", () => {
|
||||
});
|
||||
},
|
||||
async (db) => {
|
||||
const jobs = await db.listJobs();
|
||||
const jobs = await collect(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");
|
||||
@@ -1019,6 +1445,9 @@ describe("remote connection jobs surface", () => {
|
||||
expect(await db.cancelJob("job-1")).toBe(true);
|
||||
expect(await db.cancelJob("missing")).toBe(false);
|
||||
|
||||
expect(await db.pauseJob("job-1")).toEqual("pausing");
|
||||
expect(await db.resumeJob("job-1")).toEqual("still_pausing");
|
||||
|
||||
// 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");
|
||||
@@ -1080,4 +1509,171 @@ describe("remote connection jobs surface", () => {
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("leaves numDeletedRows absent when the server omits it", async () => {
|
||||
// A server that predates num_deleted_rows says nothing about deleted rows
|
||||
const statsReply =
|
||||
'{"total_bytes":1,"num_rows":3,"num_indices":0,"fragment_stats":' +
|
||||
'{"num_fragments":1,"num_small_fragments":0,' +
|
||||
'"lengths":{"min":3,"max":3,"mean":3,"p25":3,"p50":3,"p75":3,"p99":3}}}';
|
||||
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
const path = req.url ?? "";
|
||||
const body = path.endsWith("/describe/")
|
||||
? JSON.stringify({ name: "t", version: 1, schema: { fields: [] } })
|
||||
: statsReply;
|
||||
res.writeHead(200, { "Content-Type": "application/json" }).end(body);
|
||||
},
|
||||
async (db) => {
|
||||
const stats = await (await db.openTable("t")).stats();
|
||||
expect("numDeletedRows" in stats).toBe(false);
|
||||
expect(stats.numDeletedRows).toBeUndefined();
|
||||
},
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
async function collect<T>(items: AsyncIterable<T>): Promise<T[]> {
|
||||
const result: T[] = [];
|
||||
for await (const item of items) result.push(item);
|
||||
return result;
|
||||
}
|
||||
|
||||
describe("lazy resource listings", () => {
|
||||
for (const kind of ["views", "materializedViews", "jobs"] as const) {
|
||||
const listing = (
|
||||
db: Connection,
|
||||
options: { pageToken?: string; pageLimit?: number } = {},
|
||||
) => {
|
||||
if (kind === "jobs") return db.listJobs(options);
|
||||
if (kind === "views") return db.listViews(["team"], options);
|
||||
return db.listMaterializedViews(options);
|
||||
};
|
||||
const page = (names: string[], token?: string) => {
|
||||
const items =
|
||||
kind === "jobs"
|
||||
? names.map((name) => ({
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
job_id: name,
|
||||
table: "t",
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
job_type: "create_index",
|
||||
state: "in_progress",
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
created_at_millis: 1,
|
||||
}))
|
||||
: names;
|
||||
return {
|
||||
[kind === "jobs" ? "jobs" : "views"]: items,
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
page_token: token,
|
||||
};
|
||||
};
|
||||
|
||||
it(`${kind} fetches lazily, skips empty pages, and resumes`, async () => {
|
||||
const requests: { url: string; body: Record<string, unknown> }[] = [];
|
||||
const responses = [
|
||||
page(["a", "b"], "empty"),
|
||||
page([], "last"),
|
||||
page(["c", "d"], ""),
|
||||
];
|
||||
await withMockDatabase(
|
||||
async (req, res) => {
|
||||
const chunks: Buffer[] = [];
|
||||
for await (const chunk of req) chunks.push(Buffer.from(chunk));
|
||||
const body = Buffer.concat(chunks).toString();
|
||||
requests.push({
|
||||
url: req.url ?? "",
|
||||
body: body ? JSON.parse(body) : {},
|
||||
});
|
||||
res
|
||||
.writeHead(200, { "content-type": "application/json" })
|
||||
.end(JSON.stringify(responses.shift()));
|
||||
},
|
||||
async (db) => {
|
||||
const items = listing(db, { pageToken: "start/token", pageLimit: 2 });
|
||||
expect(requests).toHaveLength(0);
|
||||
expect(items.pageToken()).toBe("start/token");
|
||||
expect(items.numPageResults()).toBe(0);
|
||||
expect((await items.next()).done).toBe(false);
|
||||
expect(items.numPageResults()).toBe(1);
|
||||
expect(items.pageToken()).toBe("empty");
|
||||
await items.next();
|
||||
expect(requests).toHaveLength(1);
|
||||
const resumed = listing(db, {
|
||||
pageToken: items.pageToken(),
|
||||
pageLimit: 2,
|
||||
});
|
||||
await resumed.next();
|
||||
expect(resumed.numPageResults()).toBe(1);
|
||||
expect(resumed.pageToken()).toBeUndefined();
|
||||
expect(requests).toHaveLength(3);
|
||||
await resumed.next();
|
||||
expect((await resumed.next()).done).toBe(true);
|
||||
expect((await resumed.next()).done).toBe(true);
|
||||
for (const [i, request] of requests.entries()) {
|
||||
const token = ["start/token", "empty", "last"][i];
|
||||
const url = new URL(request.url, "http://localhost");
|
||||
if (kind === "jobs") {
|
||||
expect(url.pathname).toBe("/v1/jobs/list");
|
||||
expect(request.body).toEqual({
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
page_token: token,
|
||||
limit: 2,
|
||||
});
|
||||
} else {
|
||||
expect(url.pathname).toBe(
|
||||
kind === "views"
|
||||
? "/v1/namespace/team/view/list"
|
||||
: "/v1/namespace/$/materialized_view/list",
|
||||
);
|
||||
expect(url.searchParams.get("page_token")).toBe(token);
|
||||
expect(url.searchParams.get("limit")).toBe("2");
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it(`${kind} retains failed tokens and serializes concurrent advances`, async () => {
|
||||
let requests = 0;
|
||||
await withMockDatabase(
|
||||
(_req, res) => {
|
||||
requests += 1;
|
||||
res
|
||||
.writeHead(requests === 1 ? 401 : 200, {
|
||||
"content-type": "application/json",
|
||||
})
|
||||
.end(JSON.stringify(requests === 1 ? {} : page(["a", "b"])));
|
||||
},
|
||||
async (db) => {
|
||||
const failed = listing(db, { pageToken: "retry" });
|
||||
await expect(failed.next()).rejects.toThrow();
|
||||
expect(failed.pageToken()).toBe("retry");
|
||||
expect((await failed.next()).done).toBe(true);
|
||||
expect(requests).toBe(1);
|
||||
const resumed = listing(db, { pageToken: failed.pageToken() });
|
||||
const results = await Promise.all([
|
||||
resumed.next(),
|
||||
resumed.next(),
|
||||
resumed.next(),
|
||||
]);
|
||||
expect(results.filter((result) => result.done)).toHaveLength(1);
|
||||
expect(requests).toBe(2);
|
||||
for (const limit of [
|
||||
0,
|
||||
-1,
|
||||
1.5,
|
||||
Number.NaN,
|
||||
Number.POSITIVE_INFINITY,
|
||||
2147483648,
|
||||
]) {
|
||||
expect(() => listing(db, { pageLimit: limit })).toThrow("limit");
|
||||
}
|
||||
expect(requests).toBe(2);
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
@@ -2,7 +2,13 @@
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as arrow from "../lancedb/arrow";
|
||||
import { sanitizeField, sanitizeMap, sanitizeType } from "../lancedb/sanitize";
|
||||
import {
|
||||
sanitizeField,
|
||||
sanitizeMap,
|
||||
sanitizeSchema,
|
||||
sanitizeType,
|
||||
} from "../lancedb/sanitize";
|
||||
import { arrow21 } from "./arrow_versions";
|
||||
|
||||
describe("sanitize", function () {
|
||||
describe("sanitizeType function", function () {
|
||||
@@ -59,6 +65,8 @@ describe("sanitize", function () {
|
||||
["durationmicrosecond", new arrow.DurationMicrosecond()],
|
||||
["durationmillisecond", new arrow.DurationMillisecond()],
|
||||
["durationsecond", new arrow.DurationSecond()],
|
||||
["largebinary", new arrow.LargeBinary()],
|
||||
["largeutf8", new arrow.LargeUtf8()],
|
||||
] as const;
|
||||
|
||||
it.each(allTypeNameTestCases)(
|
||||
@@ -89,6 +97,110 @@ describe("sanitize", function () {
|
||||
"Unrecognized type name in schema: invalid_type",
|
||||
);
|
||||
});
|
||||
|
||||
// These types only exist in Arrow 20 (IntervalMonthDayNano) and 21 (the
|
||||
// rest), so the objects come from a second copy of the newest release. The
|
||||
// library itself may be running on an older Arrow (see the CI matrix), in
|
||||
// which case sanitizing them has to fail with a clear message instead.
|
||||
const hostArrow = arrow as unknown as Record<string, unknown>;
|
||||
const hostSupportsNewerTypes = typeof hostArrow.Utf8View === "function";
|
||||
const whenSupported = hostSupportsNewerTypes ? describe : describe.skip;
|
||||
const whenUnsupported = hostSupportsNewerTypes ? describe.skip : describe;
|
||||
|
||||
const newerTypeTestCases = [
|
||||
["Utf8View", () => new arrow21.Utf8View()],
|
||||
["BinaryView", () => new arrow21.BinaryView()],
|
||||
["IntervalMonthDayNano", () => new arrow21.IntervalMonthDayNano()],
|
||||
[
|
||||
"LargeList",
|
||||
() =>
|
||||
new arrow21.LargeList(new arrow21.Field("item", new arrow21.Utf8())),
|
||||
],
|
||||
] as const;
|
||||
|
||||
whenSupported("types added in Arrow 20 and 21", function () {
|
||||
it.each(newerTypeTestCases)(
|
||||
"should sanitize a %s type from another copy of Arrow",
|
||||
function (name, makeForeign) {
|
||||
const foreign = makeForeign();
|
||||
const result = sanitizeType(foreign);
|
||||
// Interval subclasses share `Type.Interval`, so the sanitized type
|
||||
// is the generic `Interval` carrying the same unit, as for the
|
||||
// other units.
|
||||
const expected = name === "IntervalMonthDayNano" ? "Interval" : name;
|
||||
expect(result).toBeInstanceOf(hostArrow[expected] as () => unknown);
|
||||
expect(result.typeId).toBe(foreign.typeId);
|
||||
expect(result.toString()).toBe(foreign.toString());
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
["binaryview", "BinaryView"],
|
||||
["utf8view", "Utf8View"],
|
||||
["intervalmonthdaynano", "Interval"],
|
||||
])('should map type name "%s" to %s', function (name, expected) {
|
||||
expect(sanitizeType(name)).toBeInstanceOf(
|
||||
hostArrow[expected] as () => unknown,
|
||||
);
|
||||
});
|
||||
|
||||
it("should sanitize the child field of a LargeList", function () {
|
||||
const foreign = new arrow21.LargeList(
|
||||
new arrow21.Field("item", new arrow21.Int32(), false),
|
||||
);
|
||||
const result = sanitizeType(foreign) as arrow.List;
|
||||
expect(result.children[0]).toBeInstanceOf(arrow.Field);
|
||||
expect(result.children[0].name).toBe("item");
|
||||
expect(result.children[0].nullable).toBe(false);
|
||||
const child = result.children[0].type;
|
||||
expect(arrow.DataType.isInt(child)).toBe(true);
|
||||
expect((child as arrow.Int).bitWidth).toBe(32);
|
||||
expect((child as arrow.Int).isSigned).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
whenUnsupported("types newer than the installed Arrow", function () {
|
||||
it.each(
|
||||
newerTypeTestCases.filter(([name]) => name !== "IntervalMonthDayNano"),
|
||||
)(
|
||||
"should explain that %s needs a newer apache-arrow",
|
||||
function (name, makeForeign) {
|
||||
expect(() => sanitizeType(makeForeign())).toThrow(
|
||||
`The Arrow type ${name} requires apache-arrow`,
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
it("should sanitize IntervalMonthDayNano to a generic Interval", () => {
|
||||
// The subclass shares `Type.Interval`, which every supported Arrow
|
||||
// release has; only the unit is new.
|
||||
const result = sanitizeType(new arrow21.IntervalMonthDayNano());
|
||||
expect(result).toBeInstanceOf(arrow.Interval);
|
||||
expect((result as arrow.Interval).unit).toBe(
|
||||
arrow21.IntervalUnit.MONTH_DAY_NANO,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it("should rebuild a schema from a second copy of the same Arrow version", function () {
|
||||
// Arrow 21 makes `instanceof Schema` succeed across copies of the library
|
||||
// but not `instanceof <DataType>`; the schema must still be rebuilt.
|
||||
const foreign = new arrow21.Schema([
|
||||
new arrow21.Field("x", new arrow21.Float32(), true),
|
||||
]);
|
||||
// Only Arrow 21+ overrides `Symbol.hasInstance` this way.
|
||||
const crossCopyInstanceof = Object.getOwnPropertySymbols(
|
||||
arrow.Schema,
|
||||
).includes(Symbol.hasInstance);
|
||||
expect(foreign instanceof arrow.Schema).toBe(crossCopyInstanceof);
|
||||
const result = sanitizeSchema(foreign);
|
||||
expect(result).not.toBe(foreign);
|
||||
const type = result.fields[0].type;
|
||||
expect(type).toBeInstanceOf(arrow.DataType);
|
||||
expect(arrow.DataType.isFloat(type)).toBe(true);
|
||||
expect((type as arrow.Float).precision).toBe(arrow.Precision.SINGLE);
|
||||
expect(() => arrow.makeBuilder({ type })).not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe("sanitizeField function", function () {
|
||||
|
||||
+928
-15
File diff suppressed because it is too large
Load Diff
@@ -14,6 +14,9 @@ test.each([
|
||||
[new ArrayBuffer(8), "X'0000000000000000'"],
|
||||
[Buffer.from("hello"), "X'68656c6c6f'"],
|
||||
["Hello 'world'", "'Hello ''world'''"],
|
||||
[NaN, "CAST('NaN' AS DOUBLE)"],
|
||||
[Infinity, "CAST('Infinity' AS DOUBLE)"],
|
||||
[-Infinity, "CAST('-Infinity' AS DOUBLE)"],
|
||||
])("toSQL(%p) === %p", (value, expected) => {
|
||||
expect(toSQL(value)).toBe(expected);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user