306 lines
11 KiB
TypeScript
306 lines
11 KiB
TypeScript
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import { expect, test, vi } from "vitest";
|
|
import type { CatalogDatabaseClient, CatalogPostgresAccess } from "../src/catalog/postgres-access.js";
|
|
import { CATALOG_SECRET_IDS } from "../src/catalog/secrets.js";
|
|
import { ConcreteSensitivityValueSource } from "../src/catalog/sensitivity-value-source.js";
|
|
import {
|
|
CatalogConnectorError,
|
|
type CatalogColumn,
|
|
type CatalogTable,
|
|
type WorkspaceDatabase,
|
|
} from "../src/catalog/types.js";
|
|
import type { WorkspaceSecretStore } from "../src/workspaces/secret-store.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",
|
|
};
|
|
}
|
|
|
|
function request(columns: readonly CatalogColumn[], overrides: Record<string, unknown> = {}) {
|
|
return {
|
|
database,
|
|
table,
|
|
columns,
|
|
valuesPerColumn: 300,
|
|
sampleOffset: 0,
|
|
sampleSeed: 37,
|
|
queryTimeoutMs: 5_000,
|
|
fullScanThreshold: 1_000,
|
|
...overrides,
|
|
};
|
|
}
|
|
|
|
test("uses bounded read-only PostgreSQL sampling for tables above 1,000 rows", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const contact = column("44444444-4444-4444-8444-444444444444", 'contact"value');
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.startsWith("SELECT 1 AS __present")) {
|
|
return { rows: Array.from({ length: 1_001 }, () => ({ __present: 1 })) };
|
|
}
|
|
if (sql.startsWith("WITH sampled")) {
|
|
return { rows: [
|
|
{ __column_index: 0, __value: "ordinary", __length: "8" },
|
|
{ __column_index: 1, __value: "mario.rossi@example.it", __length: 23 },
|
|
] };
|
|
}
|
|
return { rows: [] };
|
|
});
|
|
const end = vi.fn(async () => undefined);
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
|
|
};
|
|
const consume = vi.fn();
|
|
|
|
await expect(new ConcreteSensitivityValueSource(access).scanTable(
|
|
request([note, contact]),
|
|
consume,
|
|
new AbortController().signal,
|
|
)).resolves.toEqual({ kind: "sampled", observedValues: 2 });
|
|
|
|
expect(query.mock.calls[0]).toEqual(["BEGIN TRANSACTION READ ONLY", []]);
|
|
expect(query).toHaveBeenCalledWith("SELECT set_config('statement_timeout', $1, true)", ["5000ms"]);
|
|
const sampleSql = query.mock.calls.map(([sql]) => String(sql)).find((sql) => sql.startsWith("WITH sampled"));
|
|
expect(sampleSql).toContain('FROM "clinical""data"."patient""facts" TABLESAMPLE SYSTEM (30)');
|
|
expect(sampleSql).toContain("REPEATABLE (37)");
|
|
expect(sampleSql).toContain("LIMIT 3000 OFFSET 0");
|
|
expect(sampleSql).toContain("CROSS JOIN LATERAL");
|
|
expect(sampleSql).toContain("WHERE __rank <= 300");
|
|
expect(consume).toHaveBeenCalledWith([
|
|
{ columnId: note.id, value: "ordinary", characterLength: 8 },
|
|
{ columnId: contact.id, value: "mario.rossi@example.it", characterLength: 23 },
|
|
]);
|
|
expect(query.mock.calls.at(-1)).toEqual(["ROLLBACK", []]);
|
|
expect(end).toHaveBeenCalledOnce();
|
|
});
|
|
|
|
test("fully scans a table when the 1,001-row probe proves it is small", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.startsWith("SELECT 1 AS __present")) return { rows: [{ __present: 1 }] };
|
|
if (sql.startsWith("WITH sampled")) {
|
|
return { rows: [{ __column_index: 0, __value: "ordinary", __length: 8 }] };
|
|
}
|
|
return { rows: [] };
|
|
});
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
|
|
};
|
|
const consume = vi.fn();
|
|
|
|
await expect(new ConcreteSensitivityValueSource(access).scanTable(
|
|
request([note]),
|
|
consume,
|
|
new AbortController().signal,
|
|
)).resolves.toEqual({ kind: "complete", observedValues: 1 });
|
|
|
|
const valueSql = query.mock.calls.map(([sql]) => String(sql)).find((sql) => sql.startsWith("WITH sampled"));
|
|
expect(valueSql).not.toContain("TABLESAMPLE");
|
|
expect(valueSql).toContain("WHERE __rank <= 1000");
|
|
expect(consume).toHaveBeenCalledWith([
|
|
{ columnId: note.id, value: "ordinary", characterLength: 8 },
|
|
]);
|
|
});
|
|
|
|
test("falls back to sampling when the small-table probe reaches its query timeout", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.startsWith("SELECT 1 AS __present")) {
|
|
throw Object.assign(new Error("statement timeout"), { code: "57014" });
|
|
}
|
|
if (sql.startsWith("WITH sampled")) {
|
|
return { rows: [{ __column_index: 0, __value: "sample", __length: 6 }] };
|
|
}
|
|
return { rows: [] };
|
|
});
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
|
|
};
|
|
const consume = vi.fn();
|
|
|
|
await expect(new ConcreteSensitivityValueSource(access).scanTable(
|
|
request([note]),
|
|
consume,
|
|
new AbortController().signal,
|
|
)).resolves.toEqual({ kind: "sampled", observedValues: 1 });
|
|
expect(query.mock.calls.map(([sql]) => String(sql))).toContain(
|
|
"ROLLBACK TO SAVEPOINT sensitivity_scan_1",
|
|
);
|
|
});
|
|
|
|
test("limits each source query to at most 25 columns", async () => {
|
|
const columns = Array.from({ length: 26 }, (_, index) => column(
|
|
`00000000-0000-4000-8000-${(index + 1).toString().padStart(12, "0")}`,
|
|
`attribute_${index + 1}`,
|
|
));
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.startsWith("SELECT 1 AS __present")) {
|
|
return { rows: Array.from({ length: 1_001 }, () => ({ __present: 1 })) };
|
|
}
|
|
if (sql.startsWith("WITH sampled")) {
|
|
return { rows: [{ __column_index: 0, __value: "ordinary", __length: 8 }] };
|
|
}
|
|
return { rows: [] };
|
|
});
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
|
|
};
|
|
|
|
await new ConcreteSensitivityValueSource(access).scanTable(
|
|
request(columns),
|
|
vi.fn(),
|
|
new AbortController().signal,
|
|
);
|
|
|
|
expect(query.mock.calls.filter(([sql]) => String(sql).startsWith("WITH sampled"))).toHaveLength(2);
|
|
});
|
|
|
|
test("scans a REST run_query binding without 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([
|
|
{ __column_index: 0, __value: "mario.rossi@example.it", __length: 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(request([note], {
|
|
database: restDatabase,
|
|
fullScanThreshold: undefined,
|
|
}), consume, new AbortController().signal)).resolves.toEqual({
|
|
kind: "sampled",
|
|
observedValues: 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" TABLESAMPLE SYSTEM (30)');
|
|
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("falls back to a sequential bounded sample when randomized sampling times out", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.startsWith("WITH sampled") && sql.includes("TABLESAMPLE")) {
|
|
throw Object.assign(new Error("raw source detail"), { code: "57014" });
|
|
}
|
|
if (sql.startsWith("WITH sampled")) {
|
|
return { rows: [{ __column_index: 0, __value: "ordinary", __length: 8 }] };
|
|
}
|
|
return { rows: [] };
|
|
});
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
|
|
};
|
|
|
|
await expect(new ConcreteSensitivityValueSource(access).scanTable(
|
|
request([note], { fullScanThreshold: undefined }),
|
|
vi.fn(),
|
|
new AbortController().signal,
|
|
)).resolves.toEqual({ kind: "sampled", observedValues: 1 });
|
|
expect(query.mock.calls.filter(([sql]) => String(sql).startsWith("WITH sampled"))).toHaveLength(2);
|
|
});
|
|
|
|
test("fails explicitly when both randomized and sequential sample queries time out", async () => {
|
|
const note = column("33333333-3333-4333-8333-333333333333", "note");
|
|
const query = vi.fn(async (sql: string) => {
|
|
if (sql.startsWith("WITH sampled")) {
|
|
throw Object.assign(new Error("raw source detail"), { code: "57014" });
|
|
}
|
|
return { rows: [] };
|
|
});
|
|
const access: CatalogPostgresAccess = {
|
|
connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) as CatalogDatabaseClient),
|
|
};
|
|
|
|
await expect(new ConcreteSensitivityValueSource(access).scanTable(
|
|
request([note], { fullScanThreshold: undefined }),
|
|
vi.fn(),
|
|
new AbortController().signal,
|
|
)).rejects.toEqual(new CatalogConnectorError("Sensitivity sample query timed out"));
|
|
});
|