feat: sample sensitive columns progressively
This commit is contained in:
Codex
@@ -8,82 +8,117 @@ import type {
|
||||
SensitivityValueObservation,
|
||||
SensitivityValueSource,
|
||||
} from "./sensitivity-classifier.js";
|
||||
import { CatalogConnectorError } from "./types.js";
|
||||
import { CatalogConnectorError, type CatalogColumn } from "./types.js";
|
||||
|
||||
const MAX_VALUE_CHARACTERS = 501;
|
||||
const DEFAULT_BATCH_ROWS = 200;
|
||||
const DEFAULT_SAMPLE_ROWS = 200;
|
||||
const MAX_COLUMNS_PER_QUERY = 25;
|
||||
const SAMPLE_OVERSCAN_FACTOR = 10;
|
||||
|
||||
function quoteIdentifier(identifier: string): string {
|
||||
return `"${identifier.replaceAll('"', '""')}"`;
|
||||
}
|
||||
|
||||
function projections(request: SensitivityScanRequest): string {
|
||||
return request.columns.flatMap((column, index) => {
|
||||
function chunks<T>(items: readonly T[], size: number): T[][] {
|
||||
const result: T[][] = [];
|
||||
for (let offset = 0; offset < items.length; offset += size) {
|
||||
result.push(items.slice(offset, offset + size));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
function tableReference(request: SensitivityScanRequest): string {
|
||||
return `${quoteIdentifier(request.database.schema)}.${quoteIdentifier(request.table.name)}`;
|
||||
}
|
||||
|
||||
function samplePercentage(valuesPerColumn: number): number {
|
||||
if (valuesPerColumn <= 300) return 30;
|
||||
if (valuesPerColumn <= 700) return 70;
|
||||
return 100;
|
||||
}
|
||||
|
||||
function flatValueQuery(
|
||||
request: SensitivityScanRequest,
|
||||
columns: readonly CatalogColumn[],
|
||||
options: { complete: boolean; randomized: boolean },
|
||||
): string {
|
||||
const projections = columns.map((column) => quoteIdentifier(column.name)).join(", ");
|
||||
const perColumnLimit = options.complete
|
||||
? request.fullScanThreshold ?? request.valuesPerColumn
|
||||
: request.valuesPerColumn;
|
||||
const rowLimit = Math.max(perColumnLimit, perColumnLimit * SAMPLE_OVERSCAN_FACTOR);
|
||||
const sample = options.complete
|
||||
? `SELECT ${projections} FROM ${tableReference(request)}`
|
||||
: [
|
||||
`SELECT ${projections} FROM ${tableReference(request)}`,
|
||||
...(options.randomized
|
||||
? [`TABLESAMPLE SYSTEM (${samplePercentage(request.valuesPerColumn)}) REPEATABLE (${request.sampleSeed})`]
|
||||
: []),
|
||||
`LIMIT ${rowLimit} OFFSET ${request.sampleOffset}`,
|
||||
].join(" ");
|
||||
const values = columns.map((column, index) => {
|
||||
const identifier = quoteIdentifier(column.name);
|
||||
return [
|
||||
`LEFT((${identifier})::text, ${MAX_VALUE_CHARACTERS}) AS "__value_${index}"`,
|
||||
`CASE WHEN ${identifier} IS NULL THEN NULL ELSE char_length((${identifier})::text) END AS "__length_${index}"`,
|
||||
];
|
||||
`(${index}, LEFT((sampled.${identifier})::text, ${MAX_VALUE_CHARACTERS}),`,
|
||||
`CASE WHEN sampled.${identifier} IS NULL THEN NULL`,
|
||||
`ELSE char_length((sampled.${identifier})::text) END)`,
|
||||
].join(" ");
|
||||
}).join(", ");
|
||||
return [
|
||||
`WITH sampled AS MATERIALIZED (${sample}),`,
|
||||
"ranked AS (",
|
||||
"SELECT value.__column_index, value.__value, value.__length,",
|
||||
"row_number() OVER (PARTITION BY value.__column_index) AS __rank",
|
||||
"FROM sampled",
|
||||
`CROSS JOIN LATERAL (VALUES ${values}) AS value(__column_index, __value, __length)`,
|
||||
"WHERE value.__value IS NOT NULL",
|
||||
")",
|
||||
"SELECT __column_index, __value, __length FROM ranked",
|
||||
`WHERE __rank <= ${perColumnLimit}`,
|
||||
].join(" ");
|
||||
}
|
||||
|
||||
function observations(
|
||||
request: SensitivityScanRequest,
|
||||
columns: readonly CatalogColumn[],
|
||||
rows: readonly Record<string, unknown>[],
|
||||
): SensitivityValueObservation[] {
|
||||
return rows.flatMap((row) => request.columns.map((column, index) => {
|
||||
const sourceValue = row[`__value_${index}`];
|
||||
const sourceLength = row[`__length_${index}`];
|
||||
const value = sourceValue === null || sourceValue === undefined ? null : String(sourceValue);
|
||||
const parsedLength = sourceLength === null || sourceLength === undefined
|
||||
return rows.flatMap((row) => {
|
||||
const index = Number(row.__column_index);
|
||||
const column = Number.isSafeInteger(index) && index >= 0 ? columns[index] : undefined;
|
||||
if (!column || row.__value === null || row.__value === undefined) return [];
|
||||
const value = String(row.__value);
|
||||
const parsedLength = row.__length === null || row.__length === undefined
|
||||
? null
|
||||
: Number(sourceLength);
|
||||
return {
|
||||
: Number(row.__length);
|
||||
return [{
|
||||
columnId: column.id,
|
||||
value,
|
||||
characterLength: parsedLength !== null && Number.isSafeInteger(parsedLength) && parsedLength >= 0
|
||||
? parsedLength
|
||||
: value?.length ?? null,
|
||||
};
|
||||
}));
|
||||
: value.length,
|
||||
}];
|
||||
});
|
||||
}
|
||||
|
||||
function cancelled(error: unknown): boolean {
|
||||
return Boolean(error && typeof error === "object" && "code" in error && error.code === "57014");
|
||||
}
|
||||
|
||||
interface SensitivityValueSourceOptions {
|
||||
now?: () => number;
|
||||
batchRows?: number;
|
||||
sampleRows?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* PostgreSQL value adapter. It owns bounded read mechanics and emits normalized values, never a
|
||||
* sensitivity decision.
|
||||
* Database-specific sampling adapter. Policy stays in SensitivityClassifier; this module only
|
||||
* produces bounded, normalized non-null observations without persisting or logging values.
|
||||
*/
|
||||
export class ConcreteSensitivityValueSource implements SensitivityValueSource {
|
||||
private readonly now: () => number;
|
||||
private readonly batchRows: number;
|
||||
private readonly sampleRows: number;
|
||||
|
||||
constructor(
|
||||
private readonly access: CatalogPostgresAccess,
|
||||
private readonly secretStore?: Pick<WorkspaceSecretStore, "materialize">,
|
||||
options: SensitivityValueSourceOptions = {},
|
||||
) {
|
||||
this.now = options.now ?? Date.now;
|
||||
this.batchRows = options.batchRows ?? DEFAULT_BATCH_ROWS;
|
||||
this.sampleRows = options.sampleRows ?? DEFAULT_SAMPLE_ROWS;
|
||||
}
|
||||
) {}
|
||||
|
||||
async scanTable(
|
||||
request: SensitivityScanRequest,
|
||||
consume: (batch: readonly SensitivityValueObservation[]) => void | Promise<void>,
|
||||
signal: AbortSignal,
|
||||
): Promise<SensitivityScanCoverage> {
|
||||
if (request.columns.length === 0) return { kind: "unavailable", observedRows: 0 };
|
||||
if (request.columns.length === 0) return { kind: "complete", observedValues: 0 };
|
||||
if (request.database.binding.transport === "rest_api") {
|
||||
return await this.scanRest(request, consume, signal);
|
||||
}
|
||||
@@ -97,69 +132,63 @@ export class ConcreteSensitivityValueSource implements SensitivityValueSource {
|
||||
): Promise<SensitivityScanCoverage> {
|
||||
const client = await this.access.connect(request.database, signal);
|
||||
let transactionOpen = false;
|
||||
const startedAt = this.now();
|
||||
const fullDeadline = Math.min(request.deadline, startedAt + request.fullScanBudgetMs);
|
||||
let observedRows = 0;
|
||||
let cursorOpen = false;
|
||||
let savepointSequence = 0;
|
||||
let observedValues = 0;
|
||||
try {
|
||||
if (signal.aborted || this.now() >= request.deadline) {
|
||||
return { kind: "sampled", observedRows: 0 };
|
||||
}
|
||||
signal.throwIfAborted();
|
||||
await client.query("BEGIN TRANSACTION READ ONLY", []);
|
||||
transactionOpen = true;
|
||||
await client.query("SELECT set_config('statement_timeout', $1, true)", [
|
||||
`${Math.max(1, Math.floor(fullDeadline - startedAt))}ms`,
|
||||
`${Math.max(1, Math.floor(request.queryTimeoutMs))}ms`,
|
||||
]);
|
||||
await client.query("SAVEPOINT sensitivity_full_scan", []);
|
||||
const cursor = [
|
||||
"DECLARE sensitivity_full_scan_cursor NO SCROLL CURSOR FOR",
|
||||
`SELECT ${projections(request)}`,
|
||||
`FROM ${quoteIdentifier(request.database.schema)}.${quoteIdentifier(request.table.name)}`,
|
||||
].join(" ");
|
||||
await client.query(cursor, []);
|
||||
cursorOpen = true;
|
||||
while (!signal.aborted && this.now() < fullDeadline) {
|
||||
let rows: Array<Record<string, unknown>>;
|
||||
const boundedQuery = async (sql: string): Promise<Array<Record<string, unknown>> | undefined> => {
|
||||
signal.throwIfAborted();
|
||||
savepointSequence += 1;
|
||||
const savepoint = `sensitivity_scan_${savepointSequence}`;
|
||||
await client.query(`SAVEPOINT ${savepoint}`, []);
|
||||
try {
|
||||
await client.query("SELECT set_config('statement_timeout', $1, true)", [
|
||||
`${Math.max(1, Math.floor(fullDeadline - this.now()))}ms`,
|
||||
]);
|
||||
rows = (await client.query(
|
||||
`FETCH FORWARD ${this.batchRows} FROM sensitivity_full_scan_cursor`,
|
||||
[],
|
||||
)).rows;
|
||||
return (await client.query(sql, [])).rows;
|
||||
} catch (error) {
|
||||
if (!cancelled(error)) throw error;
|
||||
await client.query("ROLLBACK TO SAVEPOINT sensitivity_full_scan", []);
|
||||
cursorOpen = false;
|
||||
break;
|
||||
}
|
||||
if (rows.length > 0) {
|
||||
observedRows += rows.length;
|
||||
await consume(observations(request, rows));
|
||||
}
|
||||
if (rows.length < this.batchRows) {
|
||||
return { kind: "complete", observedRows };
|
||||
await client.query(`ROLLBACK TO SAVEPOINT ${savepoint}`, []);
|
||||
return undefined;
|
||||
} finally {
|
||||
await client.query(`RELEASE SAVEPOINT ${savepoint}`, []).catch(() => undefined);
|
||||
}
|
||||
};
|
||||
|
||||
let complete = false;
|
||||
if (request.fullScanThreshold !== undefined) {
|
||||
const probe = await boundedQuery(
|
||||
`SELECT 1 AS __present FROM ${tableReference(request)} LIMIT ${request.fullScanThreshold + 1}`,
|
||||
);
|
||||
complete = probe !== undefined && probe.length <= request.fullScanThreshold;
|
||||
}
|
||||
if (signal.aborted || this.now() >= request.deadline) {
|
||||
return { kind: "sampled", observedRows };
|
||||
for (const columnChunk of chunks(request.columns, MAX_COLUMNS_PER_QUERY)) {
|
||||
signal.throwIfAborted();
|
||||
let rows = await boundedQuery(flatValueQuery(request, columnChunk, {
|
||||
complete,
|
||||
randomized: !complete,
|
||||
}));
|
||||
if (rows === undefined && complete) {
|
||||
complete = false;
|
||||
rows = await boundedQuery(flatValueQuery(request, columnChunk, {
|
||||
complete: false,
|
||||
randomized: true,
|
||||
}));
|
||||
}
|
||||
if (!complete && (rows === undefined || rows.length === 0)) {
|
||||
rows = await boundedQuery(flatValueQuery(request, columnChunk, {
|
||||
complete: false,
|
||||
randomized: false,
|
||||
}));
|
||||
}
|
||||
if (rows === undefined) throw new CatalogConnectorError("Sensitivity sample query timed out");
|
||||
const batch = observations(columnChunk, rows);
|
||||
observedValues += batch.length;
|
||||
if (batch.length > 0) await consume(batch);
|
||||
}
|
||||
if (cursorOpen) await client.query("CLOSE sensitivity_full_scan_cursor", []);
|
||||
await client.query("RELEASE SAVEPOINT sensitivity_full_scan", []);
|
||||
await client.query("SELECT set_config('statement_timeout', $1, true)", [
|
||||
`${Math.max(1, Math.floor(request.deadline - this.now()))}ms`,
|
||||
]);
|
||||
const sampleSql = [
|
||||
`SELECT ${projections(request)}`,
|
||||
`FROM ${quoteIdentifier(request.database.schema)}.${quoteIdentifier(request.table.name)}`,
|
||||
"TABLESAMPLE SYSTEM (1) REPEATABLE (37)",
|
||||
"LIMIT $1",
|
||||
].join(" ");
|
||||
const sampledRows = (await client.query(sampleSql, [this.sampleRows])).rows;
|
||||
observedRows += sampledRows.length;
|
||||
if (sampledRows.length > 0) await consume(observations(request, sampledRows));
|
||||
return { kind: "sampled", observedRows };
|
||||
return { kind: complete ? "complete" : "sampled", observedValues };
|
||||
} catch (error) {
|
||||
if (error instanceof CatalogConnectorError) throw error;
|
||||
throw new CatalogConnectorError("Sensitivity source scan failed");
|
||||
@@ -180,9 +209,7 @@ export class ConcreteSensitivityValueSource implements SensitivityValueSource {
|
||||
request.database.workspaceId,
|
||||
auth === "none" ? [] : [CATALOG_SECRET_IDS.apiKey],
|
||||
);
|
||||
const startedAt = this.now();
|
||||
const fullDeadline = Math.min(request.deadline, startedAt + request.fullScanBudgetMs);
|
||||
let observedRows = 0;
|
||||
let observedValues = 0;
|
||||
try {
|
||||
const headers: Record<string, string> = { "content-type": "application/json" };
|
||||
if (auth !== "none") {
|
||||
@@ -194,61 +221,66 @@ export class ConcreteSensitivityValueSource implements SensitivityValueSource {
|
||||
}
|
||||
const baseUrl = request.database.binding.baseUrl?.replace(/\/+$/u, "");
|
||||
if (!baseUrl) throw new CatalogConnectorError("Database binding is incomplete");
|
||||
const runQuery = async (sql: string, deadline: number): Promise<Array<Record<string, unknown>>> => {
|
||||
const response = await fetch(`${baseUrl}/rpc/run_query`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify({ query_text: sql }),
|
||||
signal: AbortSignal.any([
|
||||
signal,
|
||||
AbortSignal.timeout(Math.max(1, Math.floor(deadline - this.now()))),
|
||||
]),
|
||||
});
|
||||
if (!response.ok) throw new CatalogConnectorError("REST sensitivity source scan failed");
|
||||
const body: unknown = await response.json();
|
||||
if (!Array.isArray(body)
|
||||
|| body.some((row) => !row || typeof row !== "object" || Array.isArray(row))) {
|
||||
throw new CatalogConnectorError("REST sensitivity source response is invalid");
|
||||
const runQuery = async (sql: string): Promise<Array<Record<string, unknown>> | undefined> => {
|
||||
const timeout = AbortSignal.timeout(Math.max(1, Math.floor(request.queryTimeoutMs)));
|
||||
try {
|
||||
const response = await fetch(`${baseUrl}/rpc/run_query`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify({ query_text: sql }),
|
||||
signal: AbortSignal.any([signal, timeout]),
|
||||
});
|
||||
if (!response.ok) throw new CatalogConnectorError("REST sensitivity source scan failed");
|
||||
const body: unknown = await response.json();
|
||||
if (!Array.isArray(body)
|
||||
|| body.some((row) => !row || typeof row !== "object" || Array.isArray(row))) {
|
||||
throw new CatalogConnectorError("REST sensitivity source response is invalid");
|
||||
}
|
||||
return body as Array<Record<string, unknown>>;
|
||||
} catch (error) {
|
||||
if (signal.aborted) throw error;
|
||||
if (timeout.aborted) return undefined;
|
||||
throw error;
|
||||
}
|
||||
return body as Array<Record<string, unknown>>;
|
||||
};
|
||||
|
||||
let offset = 0;
|
||||
const baseSelect = [
|
||||
`SELECT ${projections(request)}`,
|
||||
`FROM ${quoteIdentifier(request.database.schema)}.${quoteIdentifier(request.table.name)}`,
|
||||
].join(" ");
|
||||
while (!signal.aborted) {
|
||||
let rows: Array<Record<string, unknown>>;
|
||||
try {
|
||||
rows = await runQuery(
|
||||
`${baseSelect} LIMIT ${this.batchRows} OFFSET ${offset}`,
|
||||
fullDeadline,
|
||||
);
|
||||
} catch (error) {
|
||||
if (signal.aborted || this.now() < fullDeadline) throw error;
|
||||
break;
|
||||
}
|
||||
observedRows += rows.length;
|
||||
if (rows.length > 0) await consume(observations(request, rows));
|
||||
if (rows.length < this.batchRows) {
|
||||
return { kind: offset === 0 ? "complete" : "sampled", observedRows };
|
||||
}
|
||||
offset += rows.length;
|
||||
if (this.now() >= fullDeadline) break;
|
||||
let complete = false;
|
||||
if (request.fullScanThreshold !== undefined) {
|
||||
const probe = await runQuery(
|
||||
`SELECT 1 AS __present FROM ${tableReference(request)} LIMIT ${request.fullScanThreshold + 1}`,
|
||||
);
|
||||
complete = probe !== undefined && probe.length <= request.fullScanThreshold;
|
||||
}
|
||||
if (signal.aborted || this.now() >= request.deadline) {
|
||||
return { kind: "sampled", observedRows };
|
||||
let requestCount = request.fullScanThreshold === undefined ? 0 : 1;
|
||||
for (const columnChunk of chunks(request.columns, MAX_COLUMNS_PER_QUERY)) {
|
||||
signal.throwIfAborted();
|
||||
let rows = await runQuery(flatValueQuery(request, columnChunk, {
|
||||
complete,
|
||||
randomized: !complete,
|
||||
}));
|
||||
requestCount += 1;
|
||||
if (rows === undefined && complete) {
|
||||
complete = false;
|
||||
rows = await runQuery(flatValueQuery(request, columnChunk, {
|
||||
complete: false,
|
||||
randomized: true,
|
||||
}));
|
||||
requestCount += 1;
|
||||
}
|
||||
if (!complete && (rows === undefined || rows.length === 0)) {
|
||||
rows = await runQuery(flatValueQuery(request, columnChunk, {
|
||||
complete: false,
|
||||
randomized: false,
|
||||
}));
|
||||
requestCount += 1;
|
||||
}
|
||||
if (rows === undefined) throw new CatalogConnectorError("REST sensitivity sample query timed out");
|
||||
const batch = observations(columnChunk, rows);
|
||||
observedValues += batch.length;
|
||||
if (batch.length > 0) await consume(batch);
|
||||
}
|
||||
const sampleSql = [
|
||||
baseSelect,
|
||||
"TABLESAMPLE SYSTEM (1) REPEATABLE (37)",
|
||||
`LIMIT ${this.sampleRows}`,
|
||||
].join(" ");
|
||||
const sampledRows = await runQuery(sampleSql, request.deadline);
|
||||
observedRows += sampledRows.length;
|
||||
if (sampledRows.length > 0) await consume(observations(request, sampledRows));
|
||||
return { kind: "sampled", observedRows };
|
||||
// Multiple HTTP requests cannot share a source snapshot, so only one-request reads are complete.
|
||||
return { kind: complete && requestCount === 1 ? "complete" : "sampled", observedValues };
|
||||
} catch (error) {
|
||||
if (error instanceof CatalogConnectorError) throw error;
|
||||
throw new CatalogConnectorError("REST sensitivity source scan failed");
|
||||
|
||||
Reference in New Issue
Block a user