feat: sample sensitive columns progressively
This commit is contained in:
Codex
@@ -1,12 +1,17 @@
|
||||
import { expect, test, vi } from "vitest";
|
||||
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 { 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";
|
||||
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",
|
||||
@@ -60,128 +65,145 @@ function column(id: string, name: string): CatalogColumn {
|
||||
};
|
||||
}
|
||||
|
||||
test("switches from a bounded full scan to a read-only PostgreSQL sample", async () => {
|
||||
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 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("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 },
|
||||
] };
|
||||
}
|
||||
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 consume = vi.fn();
|
||||
|
||||
const coverage = await values.scanTable({
|
||||
database,
|
||||
table,
|
||||
columns: [note, contact],
|
||||
fullScanBudgetMs: 5_000,
|
||||
deadline: 61_000,
|
||||
}, (batch) => consumed.push(...batch), new AbortController().signal);
|
||||
await expect(new ConcreteSensitivityValueSource(access).scanTable(
|
||||
request([note, contact]),
|
||||
consume,
|
||||
new AbortController().signal,
|
||||
)).resolves.toEqual({ kind: "sampled", observedValues: 2 });
|
||||
|
||||
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).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("reports complete coverage when the final full-scan page is short", async () => {
|
||||
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) => sql.startsWith("FETCH FORWARD")
|
||||
? { rows: [{ __value_0: "ordinary", __length_0: 8 }] }
|
||||
: { rows: [] });
|
||||
const end = vi.fn(async () => undefined);
|
||||
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 }) as CatalogDatabaseClient),
|
||||
connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) 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);
|
||||
await expect(new ConcreteSensitivityValueSource(access).scanTable(
|
||||
request([note]),
|
||||
consume,
|
||||
new AbortController().signal,
|
||||
)).resolves.toEqual({ kind: "complete", observedValues: 1 });
|
||||
|
||||
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);
|
||||
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 PostgreSQL cancels the bounded full scan", async () => {
|
||||
test("falls back to sampling when the small-table probe reaches its query timeout", 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;
|
||||
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 end = vi.fn(async () => undefined);
|
||||
const access: CatalogPostgresAccess = {
|
||||
connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
|
||||
connect: vi.fn(async () => ({ query, end: vi.fn(async () => undefined) }) 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 },
|
||||
]);
|
||||
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("scans a REST run_query binding without using PostgreSQL-wire access", async () => {
|
||||
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 });
|
||||
@@ -193,7 +215,7 @@ test("scans a REST run_query binding without using PostgreSQL-wire access", asyn
|
||||
})),
|
||||
} as unknown as WorkspaceSecretStore;
|
||||
const fetchMock = vi.fn(async () => new Response(JSON.stringify([
|
||||
{ __value_0: "mario.rossi@example.it", __length_0: 23 },
|
||||
{ __column_index: 0, __value: "mario.rossi@example.it", __length: 23 },
|
||||
]), { status: 200, headers: { "content-type": "application/json" } }));
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
const access: CatalogPostgresAccess = {
|
||||
@@ -213,15 +235,12 @@ test("scans a REST run_query binding without using PostgreSQL-wire access", asyn
|
||||
const consume = vi.fn();
|
||||
|
||||
try {
|
||||
await expect(values.scanTable({
|
||||
await expect(values.scanTable(request([note], {
|
||||
database: restDatabase,
|
||||
table,
|
||||
columns: [note],
|
||||
fullScanBudgetMs: 5_000,
|
||||
deadline: Date.now() + 60_000,
|
||||
}, consume, new AbortController().signal)).resolves.toEqual({
|
||||
kind: "complete",
|
||||
observedRows: 1,
|
||||
fullScanThreshold: undefined,
|
||||
}), consume, new AbortController().signal)).resolves.toEqual({
|
||||
kind: "sampled",
|
||||
observedValues: 1,
|
||||
});
|
||||
expect(access.connect).not.toHaveBeenCalled();
|
||||
expect(fetchMock).toHaveBeenCalledWith(
|
||||
@@ -232,7 +251,7 @@ test("scans a REST run_query binding without using PostgreSQL-wire access", asyn
|
||||
}),
|
||||
);
|
||||
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(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 },
|
||||
]);
|
||||
@@ -243,74 +262,44 @@ test("scans a REST run_query binding without using PostgreSQL-wire access", asyn
|
||||
}
|
||||
});
|
||||
|
||||
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,
|
||||
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: [] };
|
||||
});
|
||||
expect(query).not.toHaveBeenCalled();
|
||||
expect(end).toHaveBeenCalledOnce();
|
||||
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"));
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user