From fed65024563386fa8e4fdafbf8e719374f79e1b5 Mon Sep 17 00:00:00 2001 From: Dan Tasse Date: Thu, 1 Oct 2026 10:51:57 -0500 Subject: [PATCH 1/2] test: add lots of tests --- README.md | 2 + integration_test/client.test.ts | 149 ++++++++++++++++++++++++-- src/arrow_util.test.ts | 183 ++++++++++++++++++++++++++++++++ src/async_util.test.ts | 98 ++++++++++++++++- src/client.test.ts | 93 ++++++++++++++++ src/client.ts | 3 + src/flight.test.ts | 31 ++++++ src/flight.ts | 11 +- src/flightsql.test.ts | 174 ++++++++++++++++++++++++++++++ src/flightsql.ts | 5 +- src/grpc_util.test.ts | 139 ++++++++++++++++++++++++ 11 files changed, 868 insertions(+), 20 deletions(-) create mode 100644 src/arrow_util.test.ts create mode 100644 src/client.test.ts create mode 100644 src/flight.test.ts create mode 100644 src/flightsql.test.ts create mode 100644 src/grpc_util.test.ts diff --git a/README.md b/README.md index ed4890b..526ad83 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,8 @@ const client = await Client.connect({ }); ``` +Connections use TLS by default. If your server does not use TLS, pass `insecure: true` in the connection options. + Once you have connected to your database, you can run queries: ```javascript diff --git a/integration_test/client.test.ts b/integration_test/client.test.ts index dbf1cd8..08042bb 100644 --- a/integration_test/client.test.ts +++ b/integration_test/client.test.ts @@ -1,4 +1,5 @@ import { connect } from "node:net"; +import { execFileSync } from "node:child_process"; import { describe, expect, test, beforeEach } from "@jest/globals"; import { Client } from "@lancedb/arrow-flight-sql-client"; @@ -28,6 +29,9 @@ async function getServer(): Promise { } } +// The test server (sqlflite) does not use TLS. +const CREDENTIALS = { username: "lancedb", password: "password", insecure: true }; + describe("initially we", () => { let server: string; beforeEach(async () => { @@ -35,29 +39,27 @@ describe("initially we", () => { }); test("can connect to a client", async () => { - const client = await Client.connect({ - host: server, - username: "lancedb", - password: "password", - }); + const client = await Client.connect({ host: server, ...CREDENTIALS }); expect(client).toBeDefined(); }); + + test("are rejected with bad credentials", async () => { + await expect(Client.connect({ host: server, ...CREDENTIALS, password: "wrong" })).rejects.toThrow( + /UNAUTHENTICATED/, + ); + }); }); describe("with a client we can", () => { let client: Client; beforeEach(async () => { - client = await Client.connect({ - host: await getServer(), - username: "lancedb", - password: "password", - }); + client = await Client.connect({ host: await getServer(), ...CREDENTIALS }); }); test("issue a simple SQL query", async () => { interface ExpectedRecord { - l_orderkey: number; + l_orderkey: bigint; l_returnflag: string; } @@ -77,4 +79,129 @@ describe("with a client we can", () => { const batch = result[0]; expect(batch.numRows).toBe(10); }); + + test("get an empty result set", async () => { + const queryResult = await client.query("SELECT * FROM lineitem WHERE 1 = 0"); + const batches = await queryResult.collectToArrow(); + for (const batch of batches) { + expect(batch.numRows).toBe(0); + expect(batch.schema.fields.map((f: { name: string }) => f.name)).toContain("l_orderkey"); + } + + const rows = await (await client.query("SELECT * FROM lineitem WHERE 1 = 0")).collectToObjects(); + expect(rows).toEqual([]); + }); + + test("read every batch of a multi-batch result", async () => { + const [{ n }] = (await (await client.query("SELECT COUNT(*) AS n FROM lineitem")).collectToObjects()) as { + n: bigint; + }[]; + const expectedRows = Number(n); + expect(expectedRows).toBeGreaterThan(1000); + + const batches = await (await client.query("SELECT * FROM lineitem")).collectToArrow(); + expect(batches.length).toBeGreaterThan(1); + expect(batches.reduce((total: number, batch: { numRows: number }) => total + batch.numRows, 0)).toBe(expectedRows); + + const rows = await (await client.query("SELECT * FROM lineitem")).collectToObjects(); + expect(rows.length).toBe(expectedRows); + }); + + test("map SQL types to JavaScript values", async () => { + // sqlflite is backed by DuckDB, so this uses DuckDB's literal and cast syntax. + const sql = `SELECT + 1::INTEGER AS i, + 42::BIGINT AS big, + 1.5::DOUBLE AS d, + 'abc' AS s, + 'héllo ✓' AS uni, + TRUE AS b, + NULL AS nul, + DATE '2024-01-02' AS dt, + TIMESTAMP '2024-01-02 03:04:05' AS ts, + CAST(1.25 AS DECIMAL(10,2)) AS dec`; + const rows = await (await client.query(sql)).collectToObjects(); + expect(rows).toHaveLength(1); + const row = rows[0] as Record; + + expect(row.i).toBe(1); + expect(row.big).toBe(BigInt(42)); + expect(row.d).toBe(1.5); + expect(row.s).toBe("abc"); + expect(row.uni).toBe("héllo ✓"); + expect(row.b).toBe(true); + expect(row.nul).toBeNull(); + // Dates and timestamps come back as epoch milliseconds. + expect(row.dt).toBe(Date.UTC(2024, 0, 2)); + expect(row.ts).toBe(Date.UTC(2024, 0, 2, 3, 4, 5)); + // Decimals are returned as their scaled integer digits. + expect(String(row.dec)).toBe("125"); + }); + + test("get a server error for invalid SQL", async () => { + await expect(client.query("SELECT * FROM no_such_table")).rejects.toThrow(/no_such_table/); + }); + + test("stop reading a stream early and keep using the client", async () => { + const queryResult = await client.query("SELECT * FROM lineitem"); + let batchesSeen = 0; + for await (const batch of queryResult.toArrowStream()) { + expect(batch.numRows).toBeGreaterThan(0); + batchesSeen += 1; + break; + } + expect(batchesSeen).toBe(1); + + const rows = await (await client.query("SELECT 1 AS one")).collectToObjects(); + expect(rows).toHaveLength(1); + }); + + test("issue concurrent queries without hanging", async () => { + // sqlflite cannot execute statements concurrently on its single DuckDB connection, so some of + // these may fail server-side. What we verify here is the client's behaviour: every query settles, + // successes return the right rows, failures are real errors, and the client stays usable afterward. + const timeout = new Promise((_, reject) => + setTimeout(() => reject(new Error("concurrent queries did not settle")), 20_000).unref(), + ); + const queries = [5, 10, 15, 20].map(async (limit) => { + const rows = await (await client.query(`SELECT * FROM lineitem LIMIT ${limit}`)).collectToObjects(); + return { limit, rows: rows.length }; + }); + const outcomes = await Promise.race([Promise.allSettled(queries), timeout]); + + for (const outcome of outcomes) { + if (outcome.status === "fulfilled") { + expect(outcome.value.rows).toBe(outcome.value.limit); + } else { + expect(outcome.reason).toBeInstanceOf(Error); + } + } + + const rows = await (await client.query("SELECT 1 AS one")).collectToObjects(); + expect(rows).toHaveLength(1); + }); +}); + +describe("the published package", () => { + test("can be imported from an ES module", async () => { + const script = ` + import { Client } from "@lancedb/arrow-flight-sql-client"; + const client = await Client.connect({ + host: process.env.TEST_SERVER, + username: "lancedb", + password: "password", + insecure: true, + }); + const rows = await (await client.query("SELECT 1 AS one")).collectToObjects(); + console.log("esm-ok", rows.length, rows[0].one); + process.exit(0); + `; + const stdout = execFileSync(process.execPath, ["--input-type=module", "-e", script], { + cwd: __dirname, + env: { ...process.env, TEST_SERVER: await getServer() }, + encoding: "utf8", + timeout: 30_000, + }); + expect(stdout).toContain("esm-ok 1 1"); + }, 40_000); }); diff --git a/src/arrow_util.test.ts b/src/arrow_util.test.ts new file mode 100644 index 0000000..9b510af --- /dev/null +++ b/src/arrow_util.test.ts @@ -0,0 +1,183 @@ +import { describe, expect, jest, test } from "@jest/globals"; +import { EventEmitter } from "node:events"; + +import { Metadata as GrpcMetadata } from "@grpc/grpc-js"; +import { MessageReader, Table, Utf8, makeVector, tableFromArrays, tableToIPC, vectorFromArray } from "apache-arrow"; + +import { RecordBatchStream } from "./arrow_util"; +import { arrow } from "./generated/flight"; +import { Stream } from "./grpc_util"; + +type FlightData = arrow.flight.protocol.IFlightData; + +class FakeCall extends EventEmitter { + cancel = jest.fn(); +} + +// The private framing helper, reached through element access so the test can drive it directly. +const createIpcMessage = (RecordBatchStream as unknown as { createIpcMessage: (data: FlightData) => Uint8Array }) + .createIpcMessage; + +/** + * Split an Arrow IPC stream into (header, body) pairs, which is how Flight transmits it: + * each IPC message becomes one FlightData with the flatbuffer header in `dataHeader` + * and the buffers in `dataBody`. + */ +function toFlightData(ipcStream: Uint8Array): FlightData[] { + const reader = new MessageReader(ipcStream); + const frames: FlightData[] = []; + let offset = 0; + for (;;) { + const message = reader.readMessage(); + if (!message) { + break; + } + const headerLength = new DataView(ipcStream.buffer, ipcStream.byteOffset + offset + 4, 4).getInt32(0, true); + const dataHeader = ipcStream.subarray(offset + 8, offset + 8 + headerLength); + const dataBody = reader.readMessageBody(message.bodyLength); + frames.push({ dataHeader, dataBody }); + offset += 8 + headerLength + message.bodyLength; + } + return frames; +} + +async function streamFromFrames(frames: FlightData[]): Promise> { + const call = new FakeCall(); + const stream = new Stream(call); + for (const frame of frames) { + call.emit("data", frame); + } + call.emit("end"); + return stream; +} + +async function collectRows(batches: RecordBatchStream): Promise<{ rows: unknown[]; batchCount: number }> { + const rows: unknown[] = []; + let batchCount = 0; + for await (const batch of batches) { + batchCount += 1; + rows.push(...batch.toArray().map((row) => row.toJSON())); + } + return { rows, batchCount }; +} + +describe("createIpcMessage", () => { + test("frames the header and body with a continuation marker and 8-byte padding", () => { + const dataHeader = Uint8Array.from([1, 2, 3, 4, 5]); + const dataBody = Uint8Array.from([9, 8, 7]); + + const message = createIpcMessage({ dataHeader, dataBody }); + + // 8 byte prefix + 5 byte header + 3 bytes of padding + 3 byte body + expect(message.length).toBe(8 + 5 + 3 + 3); + expect(Array.from(message.subarray(0, 4))).toEqual([0xff, 0xff, 0xff, 0xff]); + expect(new DataView(message.buffer, 4, 4).getInt32(0, true)).toBe(5); + expect(Array.from(message.subarray(8, 13))).toEqual([1, 2, 3, 4, 5]); + expect(Array.from(message.subarray(13, 16))).toEqual([0, 0, 0]); + expect(Array.from(message.subarray(16))).toEqual([9, 8, 7]); + }); + + test("adds no padding when the header is already 8-byte aligned", () => { + const dataHeader = new Uint8Array(16).fill(1); + const dataBody = Uint8Array.from([42]); + const message = createIpcMessage({ dataHeader, dataBody }); + expect(message.length).toBe(8 + 16 + 1); + expect(message[24]).toBe(42); + }); + + test("handles an empty body, as sent for schema messages", () => { + const dataHeader = Uint8Array.from([1, 2, 3]); + const message = createIpcMessage({ dataHeader, dataBody: new Uint8Array(0) }); + expect(message.length).toBe(8 + 3 + 5); + }); + + test("rejects messages with a missing header or body", () => { + expect(() => createIpcMessage({ dataBody: new Uint8Array(1) })).toThrow("Data header or body missing"); + expect(() => createIpcMessage({ dataHeader: new Uint8Array(1) })).toThrow("Data header or body missing"); + }); +}); + +describe("RecordBatchStream", () => { + // Plain Utf8 rather than tableFromArrays, which dictionary-encodes string columns. + const table = new Table({ + id: makeVector(Int32Array.from([1, 2, 3])), + name: vectorFromArray(["a", "b", "c"], new Utf8()), + }); + + test("reassembles Flight data frames into record batches", async () => { + const frames = toFlightData(tableToIPC(table, "stream")); + // schema message + one record batch + expect(frames).toHaveLength(2); + + const batches = await RecordBatchStream.create(await streamFromFrames(frames), table.schema); + expect(batches.schema).toBe(table.schema); + + const { rows, batchCount } = await collectRows(batches); + expect(batchCount).toBe(1); + expect(rows).toEqual([ + { id: 1, name: "a" }, + { id: 2, name: "b" }, + { id: 3, name: "c" }, + ]); + }); + + test("preserves batch boundaries across multiple record batches", async () => { + const second = new Table({ + id: makeVector(Int32Array.from([4, 5])), + name: vectorFromArray(["d", "e"], new Utf8()), + }); + const combined = table.concat(second); + const frames = toFlightData(tableToIPC(combined, "stream")); + // schema message + two record batches + expect(frames).toHaveLength(3); + + const batches = await RecordBatchStream.create(await streamFromFrames(frames), combined.schema); + const { rows, batchCount } = await collectRows(batches); + expect(batchCount).toBe(2); + expect(rows.map((row) => (row as { id: number }).id)).toEqual([1, 2, 3, 4, 5]); + }); + + test("ignores metadata-only envelopes mixed into the stream", async () => { + const frames = toFlightData(tableToIPC(table, "stream")); + const call = new FakeCall(); + const stream = new Stream(call); + call.emit("metadata", new GrpcMetadata()); + call.emit("data", frames[0]); + call.emit("status", { code: 0, metadata: new GrpcMetadata() }); + call.emit("data", frames[1]); + call.emit("end"); + + const batches = await RecordBatchStream.create(stream, table.schema); + const { rows } = await collectRows(batches); + expect(rows).toHaveLength(3); + }); + + test("yields a single empty batch for a schema-only stream", async () => { + // This is how the Arrow reader represents an empty result: one zero-row batch carrying the schema. + const frames = toFlightData(tableToIPC(table, "stream")).slice(0, 1); + const batches = await RecordBatchStream.create(await streamFromFrames(frames), table.schema); + let numRows = -1; + let batchCount = 0; + for await (const batch of batches) { + batchCount += 1; + numRows = batch.numRows; + expect(batch.schema.fields.map((f) => f.name)).toEqual(["id", "name"]); + } + expect(batchCount).toBe(1); + expect(numRows).toBe(0); + }); + + test("decodes dictionary-encoded columns, which arrive as an extra dictionary batch", async () => { + const encoded = tableFromArrays({ id: Int32Array.from([1, 2]), name: ["x", "y"] }); + const frames = toFlightData(tableToIPC(encoded, "stream")); + // schema message + dictionary batch + record batch + expect(frames).toHaveLength(3); + + const batches = await RecordBatchStream.create(await streamFromFrames(frames), encoded.schema); + const { rows } = await collectRows(batches); + expect(rows).toEqual([ + { id: 1, name: "x" }, + { id: 2, name: "y" }, + ]); + }); +}); diff --git a/src/async_util.test.ts b/src/async_util.test.ts index 14e1c46..3612e8c 100644 --- a/src/async_util.test.ts +++ b/src/async_util.test.ts @@ -1,6 +1,6 @@ -import { beforeEach, describe, expect, test } from "@jest/globals"; +import { beforeEach, describe, expect, jest, test } from "@jest/globals"; -import { lastValueFrom, firstValueFrom } from "./async_util"; +import { SimpleChannel, lastValueFrom, firstValueFrom } from "./async_util"; describe("given an async iterable", () => { class CustomIterable { @@ -57,3 +57,97 @@ describe("given an async iterable", () => { expect(customIterable.returned).toBe(true); }); }); + +describe("SimpleChannel", () => { + test("delivers values that were pushed before anyone asked", async () => { + const channel = new SimpleChannel(null); + channel.push(1); + channel.push(2); + expect(await channel.next()).toEqual({ done: false, value: 1 }); + expect(await channel.next()).toEqual({ done: false, value: 2 }); + }); + + test("delivers values to a consumer that is already waiting", async () => { + const channel = new SimpleChannel(null); + const pending = channel.next(); + channel.push(5); + expect(await pending).toEqual({ done: false, value: 5 }); + }); + + test("delivers errors as rejections, in order with values", async () => { + const channel = new SimpleChannel(null); + channel.push(1); + channel.push_err(new Error("bad")); + channel.push(2); + expect(await channel.next()).toEqual({ done: false, value: 1 }); + await expect(channel.next()).rejects.toThrow("bad"); + expect(await channel.next()).toEqual({ done: false, value: 2 }); + }); + + test("rejects a waiting consumer when an error arrives", async () => { + const channel = new SimpleChannel(null); + const pending = channel.next(); + channel.push_err(new Error("late")); + await expect(pending).rejects.toThrow("late"); + }); + + test("close resolves waiting consumers with done", async () => { + const channel = new SimpleChannel(null); + const pending = channel.next(); + channel.close(); + expect(await pending).toEqual({ done: true, value: undefined }); + expect(await channel.next()).toEqual({ done: true, value: undefined }); + }); + + test("values queued before close are still drained", async () => { + const channel = new SimpleChannel(null); + channel.push(1); + channel.close(); + expect(await channel.next()).toEqual({ done: false, value: 1 }); + expect(await channel.next()).toEqual({ done: true, value: undefined }); + }); + + test("ignores pushes after close", async () => { + const channel = new SimpleChannel(null); + channel.close(); + channel.push(1); + channel.push_err(new Error("ignored")); + expect(await channel.next()).toEqual({ done: true, value: undefined }); + }); + + test("return invokes the cancel callback and ends iteration", async () => { + const onCancel = jest.fn(); + const channel = new SimpleChannel(onCancel); + channel.push(1); + expect(await channel.return()).toEqual({ done: true, value: undefined }); + expect(onCancel).toHaveBeenCalledTimes(1); + }); + + test("throw also invokes the cancel callback", async () => { + const onCancel = jest.fn(); + const channel = new SimpleChannel(onCancel); + expect(await channel.throw()).toEqual({ done: true, value: undefined }); + expect(onCancel).toHaveBeenCalledTimes(1); + }); + + test("breaking out of for-await cancels the channel", async () => { + const onCancel = jest.fn(); + const channel = new SimpleChannel(onCancel); + channel.push(1); + channel.push(2); + for await (const value of channel) { + expect(value).toBe(1); + break; + } + expect(onCancel).toHaveBeenCalledTimes(1); + }); + + test("exhausting the channel normally does not cancel", async () => { + const onCancel = jest.fn(); + const channel = new SimpleChannel(onCancel); + channel.push(1); + channel.close(); + expect(await lastValueFrom(channel)).toBe(1); + expect(onCancel).not.toHaveBeenCalled(); + }); +}); diff --git a/src/client.test.ts b/src/client.test.ts new file mode 100644 index 0000000..c72756e --- /dev/null +++ b/src/client.test.ts @@ -0,0 +1,93 @@ +import { beforeEach, describe, expect, jest, test } from "@jest/globals"; + +import { RecordBatch, tableFromArrays } from "apache-arrow"; + +import { RecordBatchStream } from "./arrow_util"; +import { Client, QueryResult } from "./client"; +import { FlightSqlClient } from "./flightsql"; + +jest.mock("./flightsql", () => ({ + FlightSqlClient: { connect: jest.fn() }, +})); + +const mockedConnect = jest.mocked(FlightSqlClient.connect); + +function batchOf(values: number[]): RecordBatch { + return tableFromArrays({ a: Int32Array.from(values) }).batches[0]; +} + +function streamOf(...batches: RecordBatch[]): RecordBatchStream { + return { + schema: batches[0]?.schema, + async *[Symbol.asyncIterator]() { + for (const batch of batches) { + yield batch; + } + }, + } as unknown as RecordBatchStream; +} + +function plain(rows: unknown[]): unknown[] { + return rows.map((row) => (row as { toJSON(): unknown }).toJSON()); +} + +beforeEach(() => { + mockedConnect.mockReset(); +}); + +describe("Client.connect", () => { + // Regression test: the insecure option used to be dropped on the way to the Flight client. + test("passes every option, including insecure, through to the Flight SQL client", async () => { + const sql = {} as FlightSqlClient; + mockedConnect.mockResolvedValue(sql); + + const client = await Client.connect({ + host: "db.example:443", + username: "alice", + password: "s3cret", + defaultDatabase: "analytics", + insecure: true, + }); + + expect(client).toBeInstanceOf(Client); + expect(mockedConnect).toHaveBeenCalledWith("db.example:443", "alice", "s3cret", "analytics", true); + }); +}); + +describe("QueryResult", () => { + test("returns empty results when there are no batches", async () => { + const result = new QueryResult(streamOf()); + expect(await result.collectToArrow()).toEqual([]); + expect(await result.collectToObjects()).toEqual([]); + }); + + test("collects a single batch", async () => { + const result = new QueryResult(streamOf(batchOf([1, 2]))); + const batches = await result.collectToArrow(); + expect(batches).toHaveLength(1); + expect(batches[0].numRows).toBe(2); + + const rows = await new QueryResult(streamOf(batchOf([1, 2]))).collectToObjects(); + expect(plain(rows)).toEqual([{ a: 1 }, { a: 2 }]); + }); + + test("concatenates several batches in order", async () => { + const result = new QueryResult(streamOf(batchOf([1, 2]), batchOf([3]), batchOf([4, 5, 6]))); + const rows = await result.collectToObjects(); + expect(plain(rows)).toEqual([{ a: 1 }, { a: 2 }, { a: 3 }, { a: 4 }, { a: 5 }, { a: 6 }]); + }); + + test("skips empty batches when collecting objects", async () => { + const result = new QueryResult(streamOf(batchOf([]), batchOf([7]))); + expect(plain(await result.collectToObjects())).toEqual([{ a: 7 }]); + }); + + test("toArrowStream can be consumed incrementally", async () => { + const result = new QueryResult(streamOf(batchOf([1]), batchOf([2]))); + const seen: number[] = []; + for await (const batch of result.toArrowStream()) { + seen.push(batch.numRows); + } + expect(seen).toEqual([1, 1]); + }); +}); diff --git a/src/client.ts b/src/client.ts index afa4eca..1644255 100644 --- a/src/client.ts +++ b/src/client.ts @@ -26,6 +26,8 @@ export interface ClientOptions { defaultDatabase?: string; /** * Is the server using TLS? If not, this must be set to true. + * + * Defaults to false, meaning the connection is made with TLS. */ insecure?: boolean; } @@ -127,6 +129,7 @@ export class Client { options.username, options.password, options.defaultDatabase, + options.insecure, ); return new Client(sql); } diff --git a/src/flight.test.ts b/src/flight.test.ts new file mode 100644 index 0000000..b31e3ce --- /dev/null +++ b/src/flight.test.ts @@ -0,0 +1,31 @@ +import { afterEach, describe, expect, jest, test } from "@jest/globals"; + +import { credentials } from "@grpc/grpc-js"; + +import { FlightClient } from "./flight"; + +afterEach(() => { + jest.restoreAllMocks(); +}); + +describe("FlightClient", () => { + test("uses TLS channel credentials by default", () => { + const createSsl = jest.spyOn(credentials, "createSsl"); + const createInsecure = jest.spyOn(credentials, "createInsecure"); + + new FlightClient("db.example:443"); + + expect(createSsl).toHaveBeenCalledTimes(1); + expect(createInsecure).not.toHaveBeenCalled(); + }); + + test("uses insecure channel credentials when asked", () => { + const createSsl = jest.spyOn(credentials, "createSsl"); + const createInsecure = jest.spyOn(credentials, "createInsecure"); + + new FlightClient("localhost:31337", true); + + expect(createInsecure).toHaveBeenCalledTimes(1); + expect(createSsl).not.toHaveBeenCalled(); + }); +}); diff --git a/src/flight.ts b/src/flight.ts index 3dc8f8b..3fbca81 100644 --- a/src/flight.ts +++ b/src/flight.ts @@ -117,8 +117,9 @@ export class FlightClient { * No actual messages are sent yet. * * @param host The hostname / port of the server, separated by a colon + * @param insecure If true, connect without TLS. Defaults to false (TLS is used). */ - public constructor(host: string) { + public constructor(host: string, insecure: boolean = false) { const packageDefinition = loadSync("Flight.proto", { longs: String, enums: String, @@ -132,11 +133,8 @@ export class FlightClient { const flight = arrow.flight as GrpcObject; const flight_protocol = flight.protocol as GrpcObject; const flight_service = flight_protocol.FlightService as ServiceClientConstructor; - this.client = new flight_service( - host, - credentials.createInsecure(), - {}, - ) as unknown as arrow.flight.protocol.FlightService; + const channelCredentials = insecure ? credentials.createInsecure() : credentials.createSsl(); + this.client = new flight_service(host, channelCredentials, {}) as unknown as arrow.flight.protocol.FlightService; } // Merges any call-specific metadata with the client's default metadata @@ -192,6 +190,7 @@ export class FlightClient { (err: unknown, info: arrow.flight.protocol.IFlightInfo) => { if (err) { reject(err); + return; } resolve(new FlightInfo(info)); }, diff --git a/src/flightsql.test.ts b/src/flightsql.test.ts new file mode 100644 index 0000000..43e0e94 --- /dev/null +++ b/src/flightsql.test.ts @@ -0,0 +1,174 @@ +import { beforeEach, describe, expect, jest, test } from "@jest/globals"; + +import { Metadata as GrpcMetadata } from "@grpc/grpc-js"; +import { Field, Int32, Schema } from "apache-arrow"; + +import { RecordBatchStream } from "./arrow_util"; +import { FlightClient, FlightInfo } from "./flight"; +import { FlightSqlClient } from "./flightsql"; +import { google } from "./generated/any"; +import { arrow } from "./generated/flight"; +import { arrow as fsql } from "./generated/flightsql"; +import { Envelope, Metadata } from "./grpc_util"; + +jest.mock("./flight", () => ({ + FlightClient: jest.fn(), +})); + +const MockedFlightClient = jest.mocked(FlightClient); + +type HandshakeResponse = arrow.flight.protocol.IHandshakeResponse; + +async function* envelopes(...items: Envelope[]): AsyncGenerator> { + for (const item of items) { + yield item; + } +} + +function metadataWith(values: Record): Metadata { + const inner = new GrpcMetadata(); + for (const [key, value] of Object.entries(values)) { + inner.set(key, value); + } + return new Metadata(inner); +} + +function makeFakeFlight(handshakeResponses: Envelope[]) { + const handshakeCall = { + send: jest.fn(), + cancel: jest.fn(), + responses: envelopes(...handshakeResponses), + }; + const flight = { + handshake: jest.fn<(metadata: Record) => typeof handshakeCall>(() => handshakeCall), + set_default_metadata: jest.fn(), + getFlightInfo: jest.fn<(descriptor: arrow.flight.protocol.IFlightDescriptor) => Promise>(), + doGet: jest.fn<(ticket: arrow.flight.protocol.ITicket, schema: Schema) => Promise>(), + }; + (MockedFlightClient as unknown as jest.Mock).mockImplementation(() => flight); + return { flight, handshakeCall }; +} + +const tokenPayload = (token: string): Envelope => ({ + data: { payload: new TextEncoder().encode(token) }, +}); + +beforeEach(() => { + MockedFlightClient.mockReset(); +}); + +describe("FlightSqlClient.connect", () => { + test("sends Basic auth on the handshake and keeps the payload as a Bearer token", async () => { + const { flight, handshakeCall } = makeFakeFlight([tokenPayload("tok123")]); + + await FlightSqlClient.connect("db.example:443", "alice", "s3cret"); + + expect(flight.handshake).toHaveBeenCalledWith({ authorization: "Basic " + btoa("alice:s3cret") }); + expect(handshakeCall.send).toHaveBeenCalledWith({ payload: new Uint8Array(), protocolVersion: 0 }); + expect(flight.set_default_metadata).toHaveBeenCalledWith({ authorization: "Bearer tok123" }); + expect(handshakeCall.cancel).toHaveBeenCalledTimes(1); + }); + + test("falls back to the authorization header when the server answers with metadata", async () => { + const { flight } = makeFakeFlight([{ metadata: metadataWith({ authorization: "Bearer from-metadata" }) }]); + + await FlightSqlClient.connect("db.example:443", "alice", "s3cret"); + + expect(flight.set_default_metadata).toHaveBeenCalledWith({ authorization: "Bearer from-metadata" }); + }); + + test("rejects when metadata arrives without an authorization header", async () => { + const { flight, handshakeCall } = makeFakeFlight([{ metadata: metadataWith({ "x-other": "1" }) }]); + + await expect(FlightSqlClient.connect("db.example:443", "alice", "s3cret")).rejects.toThrow( + "no authorization header present", + ); + expect(flight.set_default_metadata).not.toHaveBeenCalled(); + expect(handshakeCall.cancel).toHaveBeenCalledTimes(1); + }); + + test("rejects when the handshake completes without any response", async () => { + const { handshakeCall } = makeFakeFlight([]); + + await expect(FlightSqlClient.connect("db.example:443", "alice", "s3cret")).rejects.toThrow( + "no metadata or data received", + ); + expect(handshakeCall.cancel).toHaveBeenCalledTimes(1); + }); + + test("includes the default database alongside the token", async () => { + const { flight } = makeFakeFlight([tokenPayload("tok")]); + + await FlightSqlClient.connect("db.example:443", "alice", "s3cret", "analytics"); + + expect(flight.set_default_metadata).toHaveBeenCalledWith({ database: "analytics", authorization: "Bearer tok" }); + }); + + test("uses TLS by default and plaintext only when insecure is set", async () => { + makeFakeFlight([tokenPayload("tok")]); + await FlightSqlClient.connect("db.example:443", "alice", "s3cret"); + expect(MockedFlightClient).toHaveBeenLastCalledWith("db.example:443", false); + + makeFakeFlight([tokenPayload("tok")]); + await FlightSqlClient.connect("localhost:31337", "alice", "s3cret", undefined, true); + expect(MockedFlightClient).toHaveBeenLastCalledWith("localhost:31337", true); + }); +}); + +describe("FlightSqlClient.statementQuery", () => { + const schema = new Schema([new Field("a", new Int32())]); + const ticket = { ticket: Uint8Array.from([1, 2, 3]) }; + + async function connectedClient(flightInfo: Partial) { + const { flight } = makeFakeFlight([tokenPayload("tok")]); + flight.getFlightInfo.mockResolvedValue(flightInfo as FlightInfo); + const client = await FlightSqlClient.connect("db.example:443", "alice", "s3cret"); + return { client, flight }; + } + + test("wraps the SQL in a CommandStatementQuery inside an Any and fetches the ticket", async () => { + const { client, flight } = await connectedClient({ decodedSchema: schema, endpoint: [{ ticket }] }); + const stream = {} as RecordBatchStream; + flight.doGet.mockResolvedValue(stream); + + const result = await client.statementQuery({ query: "SELECT 1" }); + + expect(result).toBe(stream); + expect(flight.getFlightInfo).toHaveBeenCalledTimes(1); + const descriptor = flight.getFlightInfo.mock.calls[0][0]; + expect(descriptor.type).toBe(arrow.flight.protocol.FlightDescriptor.DescriptorType.CMD); + const any = google.protobuf.Any.decode(descriptor.cmd!); + expect(any.typeUrl).toMatch(/arrow\.flight\.protocol\.sql\.CommandStatementQuery$/); + const command = fsql.flight.protocol.sql.CommandStatementQuery.decode(any.value); + expect(command.query).toBe("SELECT 1"); + expect(flight.doGet).toHaveBeenCalledWith(ticket, schema); + }); + + test("rejects when the server returns no schema", async () => { + const { client } = await connectedClient({ decodedSchema: null, endpoint: [{ ticket }] }); + await expect(client.statementQuery({ query: "SELECT 1" })).rejects.toThrow("No schema provided by server"); + }); + + test("rejects when the server returns no endpoints", async () => { + const { client } = await connectedClient({ decodedSchema: schema }); + await expect(client.statementQuery({ query: "SELECT 1" })).rejects.toThrow("No endpoint provided by server"); + }); + + test("rejects when the server returns more than one endpoint", async () => { + const { client } = await connectedClient({ decodedSchema: schema, endpoint: [{ ticket }, { ticket }] }); + await expect(client.statementQuery({ query: "SELECT 1" })).rejects.toThrow("Expected exactly one endpoint"); + }); + + test("rejects endpoints that point at a remote location", async () => { + const { client } = await connectedClient({ + decodedSchema: schema, + endpoint: [{ ticket, location: [{ uri: "grpc://elsewhere:31337" }] }], + }); + await expect(client.statementQuery({ query: "SELECT 1" })).rejects.toThrow("Cannot handle remote location"); + }); + + test("rejects when the endpoint has no ticket", async () => { + const { client } = await connectedClient({ decodedSchema: schema, endpoint: [{}] }); + await expect(client.statementQuery({ query: "SELECT 1" })).rejects.toThrow("No ticket provided by server"); + }); +}); diff --git a/src/flightsql.ts b/src/flightsql.ts index fcbe1be..2a40620 100644 --- a/src/flightsql.ts +++ b/src/flightsql.ts @@ -92,6 +92,8 @@ export class FlightSqlClient { * @param host The hostname / port of the server, separated by a colon * @param username The username to use for the handshake * @param password The password to use for the handshake + * @param defaultDatabase The default database to use, if the server supports it + * @param insecure If true, connect without TLS. Defaults to false (TLS is used). * @returns A client that can be used to execute queries */ public static async connect( @@ -99,8 +101,9 @@ export class FlightSqlClient { username: string, password: string, defaultDatabase?: string, + insecure?: boolean, ): Promise { - const sql = new FlightClient(host); + const sql = new FlightClient(host, insecure ?? false); const client = new FlightSqlClient(sql); await client.login(username, password, defaultDatabase); return client; diff --git a/src/grpc_util.test.ts b/src/grpc_util.test.ts new file mode 100644 index 0000000..b09cf69 --- /dev/null +++ b/src/grpc_util.test.ts @@ -0,0 +1,139 @@ +import { afterEach, describe, expect, jest, test } from "@jest/globals"; +import { EventEmitter } from "node:events"; + +import { Metadata as GrpcMetadata } from "@grpc/grpc-js"; + +import { firstValueFrom } from "./async_util"; +import { Envelope, Metadata, Stream } from "./grpc_util"; + +// Mimics the parts of a grpc-js call object that Stream and Bidirectional rely on. +class FakeCall extends EventEmitter { + cancel = jest.fn(); + write = jest.fn(); +} + +async function collect(iterable: AsyncIterable): Promise { + const items: T[] = []; + for await (const item of iterable) { + items.push(item); + } + return items; +} + +afterEach(() => { + jest.restoreAllMocks(); +}); + +describe("Metadata", () => { + test("get returns every value for a key", () => { + const inner = new GrpcMetadata(); + inner.add("k", "a"); + inner.add("k", "b"); + expect(new Metadata(inner).get("k")).toEqual(["a", "b"]); + }); + + test("get_first_instance returns the first value or null", () => { + const inner = new GrpcMetadata(); + inner.add("k", "a"); + inner.add("k", "b"); + const metadata = new Metadata(inner); + expect(metadata.get_first_instance("k")).toBe("a"); + expect(metadata.get_first_instance("missing")).toBeNull(); + }); + + test("get_first_string decodes binary values and passes strings through", () => { + const inner = new GrpcMetadata(); + inner.set("token-bin", Buffer.from("secret", "utf8")); + inner.set("plain", "value"); + const metadata = new Metadata(inner); + expect(metadata.get_first_string("token-bin")).toBe("secret"); + expect(metadata.get_first_string("plain")).toBe("value"); + expect(metadata.get_first_string("missing")).toBeNull(); + }); +}); + +describe("Stream", () => { + test("yields data envelopes in order and finishes on end", async () => { + const call = new FakeCall(); + const stream = new Stream(call); + call.emit("data", 1); + call.emit("data", 2); + call.emit("end"); + expect(await collect(stream.responses)).toEqual([{ data: 1 }, { data: 2 }]); + }); + + test("applies the message decoder to each message", async () => { + const call = new FakeCall(); + const stream = new Stream(call, (n) => `item-${n}`); + call.emit("data", 1); + call.emit("data", 2); + call.emit("end"); + expect(await collect(stream.responses)).toEqual([{ data: "item-1" }, { data: "item-2" }]); + }); + + test("surfaces initial metadata and trailing metadata from an OK status", async () => { + const call = new FakeCall(); + const stream = new Stream(call); + const initial = new GrpcMetadata(); + initial.set("authorization", "Bearer abc"); + const trailing = new GrpcMetadata(); + trailing.set("x-trailer", "done"); + + call.emit("metadata", initial); + call.emit("data", 1); + call.emit("status", { code: 0, metadata: trailing }); + call.emit("end"); + + const envelopes = await collect(stream.responses); + expect(envelopes).toHaveLength(3); + expect(envelopes[0].metadata).toBeInstanceOf(Metadata); + expect(envelopes[0].metadata?.get_first_string("authorization")).toBe("Bearer abc"); + expect(envelopes[1]).toEqual({ data: 1 }); + expect(envelopes[2].metadata?.get_first_string("x-trailer")).toBe("done"); + }); + + test("propagates a call error to the consumer and does not treat the status as metadata", async () => { + const call = new FakeCall(); + const stream = new Stream(call); + const iterator = stream.responses[Symbol.asyncIterator](); + + // grpc-js emits "error" and then "status" for a failed call, and "end" once the readable drains. + call.emit("data", 1); + call.emit("error", new Error("boom")); + call.emit("status", { code: 2, metadata: new GrpcMetadata() }); + call.emit("end"); + + expect(await iterator.next()).toEqual({ done: false, value: { data: 1 } }); + await expect(iterator.next()).rejects.toThrow("boom"); + expect(await iterator.next()).toEqual({ done: true, value: undefined }); + }); + + test("an error with no preceding data rejects the first read", async () => { + const call = new FakeCall(); + const stream = new Stream(call); + call.emit("error", new Error("unavailable")); + call.emit("end"); + await expect(collect(stream.responses)).rejects.toThrow("unavailable"); + }); + + test("stopping iteration early cancels the underlying call", async () => { + const call = new FakeCall(); + const stream = new Stream(call); + call.emit("data", 1); + call.emit("data", 2); + + const first = await firstValueFrom(stream.responses); + expect(first).toEqual({ data: 1 }); + expect(call.cancel).toHaveBeenCalledTimes(1); + }); + + test("data pushed after the consumer has gone away is dropped", async () => { + const call = new FakeCall(); + const stream = new Stream(call); + call.emit("data", 1); + await firstValueFrom(stream.responses); + call.emit("data", 2); + const leftover: Envelope[] = await collect(stream.responses); + expect(leftover).toEqual([]); + }); +}); From 0320b4df9c4a78753114859a3eae2270047a615a Mon Sep 17 00:00:00 2001 From: Dan Tasse Date: Mon, 5 Oct 2026 14:50:53 -0400 Subject: [PATCH 2/2] make default insecure=true, add tsc and npm test to CI --- .github/workflows/test.yml | 4 +++- README.md | 4 +++- src/client.test.ts | 27 +++++++++++++++++++++++++-- src/client.ts | 15 +++++++++++++-- src/flight.test.ts | 16 ++++++++-------- src/flight.ts | 5 +++-- src/flightsql.test.ts | 10 +++++----- src/flightsql.ts | 5 +++-- tsconfig.json | 4 +++- 9 files changed, 66 insertions(+), 24 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e960507..22261f0 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -12,7 +12,7 @@ concurrency: jobs: lint: - name: Lint & Formatting + name: Lint, Formatting & Unit Tests runs-on: ubuntu-latest defaults: run: @@ -27,6 +27,8 @@ jobs: - run: npm ci - run: npx eslint - run: npx prettier --check src + - run: npx tsc --noEmit + - run: npm test test: name: Integration tests runs-on: ubuntu-latest diff --git a/README.md b/README.md index 526ad83..e24d6dc 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,9 @@ const client = await Client.connect({ }); ``` -Connections use TLS by default. If your server does not use TLS, pass `insecure: true` in the connection options. +Connections do not use TLS by default. If your server uses TLS, pass `insecure: false` in the connection options. +The default will change to TLS in a future release, so the client logs a warning when `insecure` is not set. +Set it explicitly (`insecure: true` for plaintext servers) to keep your behavior stable and silence the warning. Once you have connected to your database, you can run queries: diff --git a/src/client.test.ts b/src/client.test.ts index c72756e..61d0529 100644 --- a/src/client.test.ts +++ b/src/client.test.ts @@ -46,11 +46,34 @@ describe("Client.connect", () => { username: "alice", password: "s3cret", defaultDatabase: "analytics", - insecure: true, + insecure: false, }); expect(client).toBeInstanceOf(Client); - expect(mockedConnect).toHaveBeenCalledWith("db.example:443", "alice", "s3cret", "analytics", true); + expect(mockedConnect).toHaveBeenCalledWith("db.example:443", "alice", "s3cret", "analytics", false); + }); + + test("warns about the upcoming default change when insecure is not set", async () => { + mockedConnect.mockResolvedValue({} as FlightSqlClient); + const warn = jest.spyOn(console, "warn").mockImplementation(() => {}); + + await Client.connect({ host: "localhost:31337", username: "alice", password: "s3cret" }); + + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0][0]).toMatch(/insecure/); + expect(mockedConnect).toHaveBeenCalledWith("localhost:31337", "alice", "s3cret", undefined, undefined); + warn.mockRestore(); + }); + + test("does not warn when insecure is set explicitly", async () => { + mockedConnect.mockResolvedValue({} as FlightSqlClient); + const warn = jest.spyOn(console, "warn").mockImplementation(() => {}); + + await Client.connect({ host: "localhost:31337", username: "alice", password: "s3cret", insecure: true }); + await Client.connect({ host: "db.example:443", username: "alice", password: "s3cret", insecure: false }); + + expect(warn).not.toHaveBeenCalled(); + warn.mockRestore(); }); }); diff --git a/src/client.ts b/src/client.ts index 1644255..b6c58a4 100644 --- a/src/client.ts +++ b/src/client.ts @@ -25,9 +25,13 @@ export interface ClientOptions { */ defaultDatabase?: string; /** - * Is the server using TLS? If not, this must be set to true. + * Is the server using TLS? If so, this must be set to false. * - * Defaults to false, meaning the connection is made with TLS. + * Defaults to true, meaning the connection is made without TLS. This matches the behavior of + * earlier releases, which always connected without TLS. + * + * The default will change to false (TLS) in a future release. Set this option explicitly to + * avoid a warning and to keep your connection behavior stable across that change. */ insecure?: boolean; } @@ -124,6 +128,13 @@ export class Client { * @returns A client that can be used to execute queries */ public static async connect(options: ClientOptions): Promise { + if (options.insecure === undefined) { + console.warn( + "@lancedb/arrow-flight-sql-client: the `insecure` option was not set, so the connection will be made without TLS. " + + "This default will change to TLS (`insecure: false`) in a future release. " + + "Set `insecure` explicitly to silence this warning.", + ); + } const sql = await FlightSqlClient.connect( options.host, options.username, diff --git a/src/flight.test.ts b/src/flight.test.ts index b31e3ce..873faa9 100644 --- a/src/flight.test.ts +++ b/src/flight.test.ts @@ -9,23 +9,23 @@ afterEach(() => { }); describe("FlightClient", () => { - test("uses TLS channel credentials by default", () => { + test("uses insecure channel credentials by default", () => { const createSsl = jest.spyOn(credentials, "createSsl"); const createInsecure = jest.spyOn(credentials, "createInsecure"); - new FlightClient("db.example:443"); + new FlightClient("localhost:31337"); - expect(createSsl).toHaveBeenCalledTimes(1); - expect(createInsecure).not.toHaveBeenCalled(); + expect(createInsecure).toHaveBeenCalledTimes(1); + expect(createSsl).not.toHaveBeenCalled(); }); - test("uses insecure channel credentials when asked", () => { + test("uses TLS channel credentials when insecure is false", () => { const createSsl = jest.spyOn(credentials, "createSsl"); const createInsecure = jest.spyOn(credentials, "createInsecure"); - new FlightClient("localhost:31337", true); + new FlightClient("db.example:443", false); - expect(createInsecure).toHaveBeenCalledTimes(1); - expect(createSsl).not.toHaveBeenCalled(); + expect(createSsl).toHaveBeenCalledTimes(1); + expect(createInsecure).not.toHaveBeenCalled(); }); }); diff --git a/src/flight.ts b/src/flight.ts index 3fbca81..554c2d8 100644 --- a/src/flight.ts +++ b/src/flight.ts @@ -117,9 +117,10 @@ export class FlightClient { * No actual messages are sent yet. * * @param host The hostname / port of the server, separated by a colon - * @param insecure If true, connect without TLS. Defaults to false (TLS is used). + * @param insecure If true, connect without TLS. Defaults to true (no TLS), matching the behavior of + * earlier releases. Set to false to connect with TLS. */ - public constructor(host: string, insecure: boolean = false) { + public constructor(host: string, insecure: boolean = true) { const packageDefinition = loadSync("Flight.proto", { longs: String, enums: String, diff --git a/src/flightsql.test.ts b/src/flightsql.test.ts index 43e0e94..2067719 100644 --- a/src/flightsql.test.ts +++ b/src/flightsql.test.ts @@ -104,14 +104,14 @@ describe("FlightSqlClient.connect", () => { expect(flight.set_default_metadata).toHaveBeenCalledWith({ database: "analytics", authorization: "Bearer tok" }); }); - test("uses TLS by default and plaintext only when insecure is set", async () => { + test("uses plaintext by default and TLS only when insecure is false", async () => { makeFakeFlight([tokenPayload("tok")]); - await FlightSqlClient.connect("db.example:443", "alice", "s3cret"); - expect(MockedFlightClient).toHaveBeenLastCalledWith("db.example:443", false); + await FlightSqlClient.connect("localhost:31337", "alice", "s3cret"); + expect(MockedFlightClient).toHaveBeenLastCalledWith("localhost:31337", true); makeFakeFlight([tokenPayload("tok")]); - await FlightSqlClient.connect("localhost:31337", "alice", "s3cret", undefined, true); - expect(MockedFlightClient).toHaveBeenLastCalledWith("localhost:31337", true); + await FlightSqlClient.connect("db.example:443", "alice", "s3cret", undefined, false); + expect(MockedFlightClient).toHaveBeenLastCalledWith("db.example:443", false); }); }); diff --git a/src/flightsql.ts b/src/flightsql.ts index 2a40620..12dc0a7 100644 --- a/src/flightsql.ts +++ b/src/flightsql.ts @@ -93,7 +93,8 @@ export class FlightSqlClient { * @param username The username to use for the handshake * @param password The password to use for the handshake * @param defaultDatabase The default database to use, if the server supports it - * @param insecure If true, connect without TLS. Defaults to false (TLS is used). + * @param insecure If true, connect without TLS. Defaults to true (no TLS), matching the behavior of + * earlier releases. Set to false to connect with TLS. * @returns A client that can be used to execute queries */ public static async connect( @@ -103,7 +104,7 @@ export class FlightSqlClient { defaultDatabase?: string, insecure?: boolean, ): Promise { - const sql = new FlightClient(host, insecure ?? false); + const sql = new FlightClient(host, insecure ?? true); const client = new FlightSqlClient(sql); await client.login(username, password, defaultDatabase); return client; diff --git a/tsconfig.json b/tsconfig.json index 665a4f7..6b7c7e4 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -107,5 +107,7 @@ /* Completeness */ // "skipDefaultLibCheck": true, /* Skip type checking .d.ts files that are included with TypeScript. */ "skipLibCheck": true /* Skip type checking all .d.ts files. */ - } + }, + /* integration_test/ imports the published package and has its own package.json, so it is not part of this program. */ + "exclude": ["node_modules", "dist", "integration_test"] }