Fix ignored insecure option and add unit and integration tests

.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

README.md

@@ -29,6 +29,10 @@ const client = await Client.connect({
 });
 ```
 
+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:
 
 ```javascript

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,36 +29,37 @@ async function getServer(): Promise<string> {
   }
 }
 
+// The test server (sqlflite) does not use TLS.
+const CREDENTIALS = { username: "lancedb", password: "password", insecure: true };
+
 describe("initially we", () => {
   let server: string;
   beforeEach(async () => {
     server = await getServer();
   });
 
   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<string, unknown>;
+
+    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<never>((_, 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);
 });

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<Stream<FlightData, FlightData>> {
+  const call = new FakeCall();
+  const stream = new Stream<FlightData, FlightData>(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<FlightData, FlightData>(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" },
+    ]);
+  });
+});