317 lines
11 KiB
TypeScript
317 lines
11 KiB
TypeScript
import { expect, test, vi } from "vitest";
|
|
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import type { CatalogDatabaseClient, CatalogPostgresAccess } from "../src/catalog/postgres-access.js";
|
|
import { ConcreteSensitivityValueSource } from "../src/catalog/sensitivity-value-source.js";
|
|
import type { CatalogColumn, CatalogTable, WorkspaceDatabase } from "../src/catalog/types.js";
|
|
import type { WorkspaceSecretStore } from "../src/workspaces/secret-store.js";
|
|
import { CATALOG_SECRET_IDS } from "../src/catalog/secrets.js";
|
|
|
|
const database = {
|
|
id: "11111111-1111-4111-8111-111111111111",
|
|
workspaceId: "psd-clinical",
|
|
engine: "postgres",
|
|
databaseName: "warehouse",
|
|
schema: 'clinical"data',
|
|
version: 1,
|
|
createdAt: "2026-09-02T08:00:00Z",
|
|
updatedAt: "2026-09-02T08:00:00Z",
|
|
connectionStatus: "reachable",
|
|
binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
|
|
} satisfies WorkspaceDatabase;
|
|
|
|
const table = {
|
|
id: "22222222-2222-4222-8222-222222222222",
|
|
databaseId: database.id,
|
|
name: 'patient"facts',
|
|
sourceComment: null,
|
|
description: null,
|
|
generatedDescription: null,
|
|
lastSyncedDatabaseVersion: 1,
|
|
lastSyncedAt: "2026-09-02T08:00:00Z",
|
|
version: 1,
|
|
createdAt: "2026-09-02T08:00:00Z",
|
|
updatedAt: "2026-09-02T08:00:00Z",
|
|
} satisfies CatalogTable;
|
|
|
|
function column(id: string, name: string): CatalogColumn {
|
|
return {
|
|
id,
|
|
tableId: table.id,
|
|
name,
|
|
ordinalPosition: 1,
|
|
dataType: "text",
|
|
isNullable: true,
|
|
defaultExpression: null,
|
|
primaryKeyPosition: null,
|
|
isPrimaryKey: false,
|
|
isForeignKey: false,
|
|
foreignKeyCount: 0,
|
|
sourceComment: null,
|
|
description: null,
|
|
generatedDescription: null,
|
|
sensitive: false,
|
|
lastSyncedDatabaseVersion: 1,
|
|
lastSyncedAt: "2026-09-02T08:00:00Z",
|
|
version: 1,
|
|
createdAt: "2026-09-02T08:00:00Z",
|
|
updatedAt: "2026-09-02T08:00:00Z",
|
|
};
|
|
}
|
|
|
|
test("switches from a bounded full scan to a read-only PostgreSQL sample", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const contact = column("44444444-4444-4444-8444-444444444444", 'contact"value');
|
|
const fullRows = Array.from({ length: 200 }, () => ({
|
|
__value_0: "ordinary",
|
|
__length_0: "8",
|
|
__value_1: null,
|
|
__length_1: null,
|
|
}));
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.includes("TABLESAMPLE")) {
|
|
return { rows: [{ __value_0: "sample", __length_0: 6, __value_1: "x", __length_1: 1 }] };
|
|
}
|
|
if (sql.startsWith("FETCH FORWARD")) return { rows: fullRows };
|
|
return { rows: [] };
|
|
});
|
|
const end = vi.fn(async () => undefined);
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
|
|
};
|
|
let clockCalls = 0;
|
|
const values = new ConcreteSensitivityValueSource(access, undefined, {
|
|
now: () => clockCalls++ < 3 ? 1_000 : 6_100,
|
|
});
|
|
const consumed: unknown[] = [];
|
|
|
|
const coverage = await values.scanTable({
|
|
database,
|
|
table,
|
|
columns: [note, contact],
|
|
fullScanBudgetMs: 5_000,
|
|
deadline: 61_000,
|
|
}, (batch) => consumed.push(...batch), new AbortController().signal);
|
|
|
|
expect(coverage).toEqual({ kind: "sampled", observedRows: 201 });
|
|
expect(consumed).toContainEqual({ columnId: note.id, value: "ordinary", characterLength: 8 });
|
|
expect(consumed).toContainEqual({ columnId: contact.id, value: null, characterLength: null });
|
|
expect(consumed).toContainEqual({ columnId: contact.id, value: "x", characterLength: 1 });
|
|
expect(query.mock.calls[0]).toEqual(["BEGIN TRANSACTION READ ONLY", []]);
|
|
expect(query.mock.calls.some(([sql]) => (
|
|
String(sql).startsWith("DECLARE sensitivity_full_scan_cursor NO SCROLL CURSOR FOR SELECT")
|
|
))).toBe(true);
|
|
expect(query.mock.calls.some(([sql]) => String(sql) === (
|
|
"FETCH FORWARD 200 FROM sensitivity_full_scan_cursor"
|
|
))).toBe(true);
|
|
expect(query.mock.calls.some(([sql]) => String(sql).includes(" OFFSET "))).toBe(false);
|
|
expect(query.mock.calls.some(([sql]) => (
|
|
String(sql).includes('FROM "clinical""data"."patient""facts" TABLESAMPLE SYSTEM')
|
|
))).toBe(true);
|
|
expect(query.mock.calls.at(-1)).toEqual(["ROLLBACK", []]);
|
|
expect(end).toHaveBeenCalledOnce();
|
|
});
|
|
|
|
test("reports complete coverage when the final full-scan page is short", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const query = vi.fn(async (sql: string) => sql.startsWith("FETCH FORWARD")
|
|
? { rows: [{ __value_0: "ordinary", __length_0: 8 }] }
|
|
: { rows: [] });
|
|
const end = vi.fn(async () => undefined);
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
|
|
};
|
|
const values = new ConcreteSensitivityValueSource(access);
|
|
const consume = vi.fn();
|
|
|
|
const coverage = await values.scanTable({
|
|
database,
|
|
table,
|
|
columns: [note],
|
|
fullScanBudgetMs: 5_000,
|
|
deadline: Date.now() + 60_000,
|
|
}, consume, new AbortController().signal);
|
|
|
|
expect(coverage).toEqual({ kind: "complete", observedRows: 1 });
|
|
expect(query.mock.calls.filter(([sql]) => (
|
|
String(sql) === "FETCH FORWARD 200 FROM sensitivity_full_scan_cursor"
|
|
))).toHaveLength(1);
|
|
expect(consume).toHaveBeenCalledWith([
|
|
{ columnId: note.id, value: "ordinary", characterLength: 8 },
|
|
]);
|
|
});
|
|
|
|
test("falls back to sampling when PostgreSQL cancels the bounded full scan", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
let fullScanAttempts = 0;
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.includes("TABLESAMPLE")) {
|
|
return { rows: [{ __value_0: "sample", __length_0: 6 }] };
|
|
}
|
|
if (sql.startsWith("FETCH FORWARD")) {
|
|
fullScanAttempts += 1;
|
|
throw Object.assign(new Error("statement timeout"), { code: "57014" });
|
|
}
|
|
return { rows: [] };
|
|
});
|
|
const end = vi.fn(async () => undefined);
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
|
|
};
|
|
const values = new ConcreteSensitivityValueSource(access);
|
|
const consume = vi.fn();
|
|
|
|
const coverage = await values.scanTable({
|
|
database,
|
|
table,
|
|
columns: [note],
|
|
fullScanBudgetMs: 5_000,
|
|
deadline: Date.now() + 60_000,
|
|
}, consume, new AbortController().signal);
|
|
|
|
expect(fullScanAttempts).toBe(1);
|
|
expect(coverage).toEqual({ kind: "sampled", observedRows: 1 });
|
|
expect(query.mock.calls.map(([sql]) => String(sql))).toEqual(expect.arrayContaining([
|
|
"SAVEPOINT sensitivity_full_scan",
|
|
"ROLLBACK TO SAVEPOINT sensitivity_full_scan",
|
|
]));
|
|
expect(consume).toHaveBeenCalledWith([
|
|
{ columnId: note.id, value: "sample", characterLength: 6 },
|
|
]);
|
|
});
|
|
|
|
test("scans a REST run_query binding without using PostgreSQL-wire access", async () => {
|
|
const root = mkdtempSync(join(tmpdir(), "tht-sensitivity-rest-"));
|
|
const credentialFile = join(root, "api-key");
|
|
writeFileSync(credentialFile, "test-api-key\n", { mode: 0o600 });
|
|
const release = vi.fn();
|
|
const secretStore = {
|
|
materialize: vi.fn(() => ({
|
|
files: new Map([[CATALOG_SECRET_IDS.apiKey, credentialFile]]),
|
|
release,
|
|
})),
|
|
} as unknown as WorkspaceSecretStore;
|
|
const fetchMock = vi.fn(async () => new Response(JSON.stringify([
|
|
{ __value_0: "mario.rossi@example.it", __length_0: 23 },
|
|
]), { status: 200, headers: { "content-type": "application/json" } }));
|
|
vi.stubGlobal("fetch", fetchMock);
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => { throw new Error("PostgreSQL access must not be used"); }),
|
|
};
|
|
const values = new ConcreteSensitivityValueSource(access, secretStore);
|
|
const restDatabase: WorkspaceDatabase = {
|
|
...database,
|
|
binding: {
|
|
transport: "rest_api",
|
|
baseUrl: "https://dwh.example.test/root/",
|
|
restPath: "/health",
|
|
restAuth: "x-api-key",
|
|
},
|
|
};
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const consume = vi.fn();
|
|
|
|
try {
|
|
await expect(values.scanTable({
|
|
database: restDatabase,
|
|
table,
|
|
columns: [note],
|
|
fullScanBudgetMs: 5_000,
|
|
deadline: Date.now() + 60_000,
|
|
}, consume, new AbortController().signal)).resolves.toEqual({
|
|
kind: "complete",
|
|
observedRows: 1,
|
|
});
|
|
expect(access.connect).not.toHaveBeenCalled();
|
|
expect(fetchMock).toHaveBeenCalledWith(
|
|
"https://dwh.example.test/root/rpc/run_query",
|
|
expect.objectContaining({
|
|
method: "POST",
|
|
headers: { "content-type": "application/json", "x-api-key": "test-api-key" },
|
|
}),
|
|
);
|
|
const body = JSON.parse(String(fetchMock.mock.calls[0]![1]!.body));
|
|
expect(body.query_text).toContain('FROM "clinical""data"."patient""facts" LIMIT 200 OFFSET 0');
|
|
expect(consume).toHaveBeenCalledWith([
|
|
{ columnId: note.id, value: "mario.rossi@example.it", characterLength: 23 },
|
|
]);
|
|
expect(release).toHaveBeenCalledOnce();
|
|
} finally {
|
|
vi.unstubAllGlobals();
|
|
rmSync(root, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test("keeps multi-request REST scans conservative without a source transaction", async () => {
|
|
const root = mkdtempSync(join(tmpdir(), "tht-sensitivity-rest-pages-"));
|
|
const credentialFile = join(root, "api-key");
|
|
writeFileSync(credentialFile, "test-api-key\n", { mode: 0o600 });
|
|
const secretStore = {
|
|
materialize: vi.fn(() => ({
|
|
files: new Map([[CATALOG_SECRET_IDS.apiKey, credentialFile]]),
|
|
release: vi.fn(),
|
|
})),
|
|
} as unknown as WorkspaceSecretStore;
|
|
const fetchMock = vi.fn()
|
|
.mockResolvedValueOnce(new Response(JSON.stringify([
|
|
{ __value_0: "ordinary", __length_0: 8 },
|
|
]), { status: 200 }))
|
|
.mockResolvedValueOnce(new Response(JSON.stringify([]), { status: 200 }));
|
|
vi.stubGlobal("fetch", fetchMock);
|
|
const values = new ConcreteSensitivityValueSource({
|
|
connect: vi.fn(async () => { throw new Error("PostgreSQL access must not be used"); }),
|
|
}, secretStore, { batchRows: 1 });
|
|
const restDatabase: WorkspaceDatabase = {
|
|
...database,
|
|
binding: {
|
|
transport: "rest_api",
|
|
baseUrl: "https://dwh.example.test/root",
|
|
restPath: "/health",
|
|
restAuth: "x-api-key",
|
|
},
|
|
};
|
|
|
|
try {
|
|
await expect(values.scanTable({
|
|
database: restDatabase,
|
|
table,
|
|
columns: [column("33333333-3333-4333-8333-333333333333", "note")],
|
|
fullScanBudgetMs: 5_000,
|
|
deadline: Date.now() + 60_000,
|
|
}, vi.fn(), new AbortController().signal)).resolves.toEqual({
|
|
kind: "sampled",
|
|
observedRows: 1,
|
|
});
|
|
expect(fetchMock).toHaveBeenCalledTimes(2);
|
|
} finally {
|
|
vi.unstubAllGlobals();
|
|
rmSync(root, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test("does not start a PostgreSQL transaction when connecting consumed the run deadline", 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 now = vi.fn()
|
|
.mockReturnValueOnce(1_000)
|
|
.mockReturnValue(61_000);
|
|
const values = new ConcreteSensitivityValueSource(access, undefined, { now });
|
|
|
|
await expect(values.scanTable({
|
|
database,
|
|
table,
|
|
columns: [column("33333333-3333-4333-8333-333333333333", "note")],
|
|
fullScanBudgetMs: 5_000,
|
|
deadline: 60_000,
|
|
}, vi.fn(), new AbortController().signal)).resolves.toEqual({
|
|
kind: "sampled",
|
|
observedRows: 0,
|
|
});
|
|
expect(query).not.toHaveBeenCalled();
|
|
expect(end).toHaveBeenCalledOnce();
|
|
});
|