dbx/tests/sqlCompletion.test.ts

import { strict as assert } from "node:assert";
import test from "node:test";
import {
  buildSqlCompletionItems,
  type SqlCompletionColumn,
  type SqlCompletionTable,
} from "../src/lib/sqlCompletion.ts";

const tables: SqlCompletionTable[] = [
  { name: "users", schema: "public", type: "table" },
  { name: "user_profiles", schema: "public", type: "table" },
  { name: "orders", schema: "public", type: "table" },
];

const columnsByTable = new Map<string, SqlCompletionColumn[]>([
  [
    "public.users",
    [
      { name: "id", table: "users", schema: "public", dataType: "bigint" },
      { name: "name", table: "users", schema: "public", dataType: "varchar" },
      { name: "email", table: "users", schema: "public", dataType: "varchar" },
    ],
  ],
  [
    "public.orders",
    [
      { name: "id", table: "orders", schema: "public", dataType: "bigint" },
      { name: "user_id", table: "orders", schema: "public", dataType: "bigint" },
      { name: "status", table: "orders", schema: "public", dataType: "varchar" },
    ],
  ],
]);

test("suggests SQL keywords for generic keyword input", () => {
  const items = buildSqlCompletionItems("sel", 3, {
    tables,
    columnsByTable,
  });

  assert.equal(items[0]?.label, "SELECT");
  assert.equal(items[0]?.type, "keyword");
});

test("suggests matching table names after FROM", () => {
  const sql = "select * from us";
  const items = buildSqlCompletionItems(sql, sql.length, {
    tables,
    columnsByTable,
  });

  assert.deepEqual(
    items.slice(0, 2).map((item) => item.label),
    ["users", "user_profiles"],
  );
});

test("suggests columns for an explicit alias qualifier", () => {
  const sql = "select u. from public.users u";
  const cursor = "select u.".length;
  const items = buildSqlCompletionItems(sql, cursor, {
    tables,
    columnsByTable,
  });

  const columnItems = items.filter((item) => item.type === "column");
  assert.deepEqual(
    columnItems.map((item) => item.label),
    ["id", "name", "email"],
  );
});

test("suggests columns from referenced tables in select list", () => {
  const sql = "select na from public.users u join public.orders o on u.id = o.user_id";
  const cursor = "select na".length;
  const items = buildSqlCompletionItems(sql, cursor, {
    tables,
    columnsByTable,
  });

  assert.equal(items[0]?.label, "name");
  assert.equal(items[0]?.type, "column");
});

test("suggests tables after LEFT JOIN", () => {
  const sql = "select * from users left join us";
  const items = buildSqlCompletionItems(sql, sql.length, {
    tables,
    columnsByTable,
  });

  assert.ok(items.some((item) => item.label === "users" && item.type === "table"));
  assert.ok(items.some((item) => item.label === "user_profiles" && item.type === "table"));
});

test("suggests tables after comma in FROM clause", () => {
  const sql = "select * from users, or";
  const items = buildSqlCompletionItems(sql, sql.length, {
    tables,
    columnsByTable,
  });

  assert.ok(items.some((item) => item.label === "orders" && item.type === "table"));
});

test("suggests keywords when typing without context", () => {
  const sql = "us";
  const items = buildSqlCompletionItems(sql, sql.length, {
    tables,
    columnsByTable,
  });

  assert.ok(items.some((item) => item.type === "keyword" && item.label === "USING"));
});

test("always includes keywords alongside table suggestions", () => {
  const sql = "select * from us";
  const items = buildSqlCompletionItems(sql, sql.length, {
    tables,
    columnsByTable,
  });

  assert.ok(items.some((item) => item.type === "table"));
  assert.ok(items.some((item) => item.type === "keyword" && item.label === "USING"));
});

test("limits table suggestions for large schemas after filtering by prefix", () => {
  const largeTables: SqlCompletionTable[] = Array.from({ length: 500 }, (_, index) => ({
    name: `erp_invoice_${String(index).padStart(4, "0")}`,
    schema: "dbo",
    type: "table",
  }));

  const sql = "select * from erp_invoice_";
  const items = buildSqlCompletionItems(sql, sql.length, {
    tables: largeTables,
    columnsByTable,
  });

  const tableItems = items.filter((item) => item.type === "table");
  assert.equal(tableItems.length, 200);
  assert.equal(tableItems[0]?.label, "erp_invoice_0000");
  assert.equal(tableItems.at(-1)?.label, "erp_invoice_0199");
});