ThothII/backend/test/catalog-description-source-sampler.test.ts

import { expect, test, vi } from "vitest";
import {
  PostgresDescriptionSourceSampler,
  type DescriptionSourceSamplingTarget,
} from "../src/catalog/description-source-sampler.js";
import type {
  CatalogDatabaseClient,
  CatalogPostgresAccess,
} from "../src/catalog/postgres-access.js";
import type { WorkspaceDatabase } from "../src/catalog/types.js";

const database: WorkspaceDatabase = {
  id: "11111111-1111-4111-8111-111111111111",
  workspaceId: "psd-clinical",
  engine: "postgres",
  databaseName: "warehouse",
  schema: 'clinical"data',
  version: 1,
  createdAt: "2026-08-28T08:00:00Z",
  updatedAt: "2026-08-28T08:00:00Z",
  connectionStatus: "reachable",
  binding: {
    transport: "postgres_direct",
    host: "db.internal",
    port: 5432,
    username: "reader",
  },
};

const target: DescriptionSourceSamplingTarget = {
  targetId: "22222222-2222-4222-8222-222222222222",
  tableName: 'patient"facts',
  columnNames: ['status"code', "ward"],
};

test("samples at most five source rows and five distinct non-null examples in a read-only transaction", async () => {
  const query = vi.fn(async (sql: string) => {
    if (!sql.startsWith("SELECT")) return { rows: [] };
    return {
      rows: [
        { 'status"code': "active", ward: null },
        { 'status"code': "pending", ward: "A" },
        { 'status"code': "closed", ward: "A" },
        { 'status"code': "transferred", ward: "B" },
        { 'status"code': "unknown", ward: "C" },
        { 'status"code': "must-not-be-sampled", ward: "D" },
      ],
    };
  });
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  const sampler = new PostgresDescriptionSourceSampler(access);
  const controller = new AbortController();

  const samples = await sampler.sample(database, [target], controller.signal);

  expect(samples).toEqual([{
    targetId: target.targetId,
    tableName: target.tableName,
    rows: [
      { fields: [{ name: 'status"code', value: "active" }, { name: "ward", value: null }] },
      { fields: [{ name: 'status"code', value: "pending" }, { name: "ward", value: "A" }] },
      { fields: [{ name: 'status"code', value: "closed" }, { name: "ward", value: "A" }] },
      { fields: [{ name: 'status"code', value: "transferred" }, { name: "ward", value: "B" }] },
      { fields: [{ name: 'status"code', value: "unknown" }, { name: "ward", value: "C" }] },
    ],
    representativeValues: [
      {
        column: 'status"code',
        values: ["active", "pending", "closed", "transferred"],
      },
      { column: "ward", values: ["A"] },
    ],
  }]);
  expect(access.connect).toHaveBeenCalledWith(database, controller.signal);
  expect(samples[0]!.representativeValues.flatMap((entry) => entry.values)).toHaveLength(5);
  expect(query.mock.calls).toEqual([
    ["BEGIN TRANSACTION READ ONLY", []],
    [
      'SELECT LEFT(("status""code")::text, $1) AS "status""code", LEFT(("ward")::text, $1) AS "ward" FROM "clinical""data"."patient""facts" LIMIT $2',
      [256, 5],
    ],
    ["ROLLBACK", []],
  ]);
  expect(query.mock.calls.map(([sql]) => String(sql).split(" ")[0])).toEqual([
    "BEGIN",
    "SELECT",
    "ROLLBACK",
  ]);
  expect(end).toHaveBeenCalledOnce();
});

test("does not issue a SELECT when a protected target has no source columns", async () => {
  const query = vi.fn(async () => ({ rows: [] }));
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  const sampler = new PostgresDescriptionSourceSampler(access);

  const samples = await sampler.sample(database, [{
    targetId: target.targetId,
    tableName: target.tableName,
    columnNames: [],
  }], new AbortController().signal);

  expect(samples).toEqual([{
    targetId: target.targetId,
    tableName: target.tableName,
    rows: [],
    representativeValues: [],
  }]);
  expect(query.mock.calls).toEqual([
    ["BEGIN TRANSACTION READ ONLY", []],
    ["ROLLBACK", []],
  ]);
  expect(query.mock.calls.some(([sql]) => String(sql).startsWith("SELECT"))).toBe(false);
  expect(end).toHaveBeenCalledOnce();
});

test("rolls back and closes the source connection when sampling fails", async () => {
  const query = vi.fn(async (sql: string) => {
    if (sql.startsWith("SELECT")) throw new Error("distinctive-source-secret");
    return { rows: [] };
  });
  const end = vi.fn(async () => undefined);
  const access: CatalogPostgresAccess = {
    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
  };
  const sampler = new PostgresDescriptionSourceSampler(access);
  const controller = new AbortController();

  await expect(sampler.sample(database, [target], controller.signal)).rejects.toThrow();

  expect(query).toHaveBeenCalledWith("ROLLBACK", []);
  expect(end).toHaveBeenCalledOnce();
});