feat: protect sensitive catalog samples

This commit is contained in:
Codex
2026-08-30 12:14:23 +02:00
parent 6278ee9d81
commit 0736983bc5
28 changed files with 1162 additions and 128 deletions
@@ -12,6 +12,7 @@ import type {
DescriptionSourceSampleValue,
DescriptionTargetSourceSample,
} from "./description-source-sampler.js";
import { syntheticSampleValue } from "./synthetic-sample-value.js";
import {
DescriptionGenerationRunActiveError,
type CatalogColumn,
@@ -306,35 +307,55 @@ function sourceSampleFor(
): PromptSourceSample | undefined {
const targetId = target.kind === "column" ? target.column.id : target.table.id;
const sample = samples.find((candidate) => candidate.targetId === targetId);
if (!sample) return undefined;
const relevantColumns = new Set(
target.kind === "column" ? [target.column.name] : target.columns.map((column) => column.name),
);
const rows = sample.rows.slice(0, budget.rows).map((row) => {
const columns = target.kind === "column" ? [target.column] : target.columns;
const sourceRows = sample?.rows ?? [];
const hasSensitiveColumns = columns.some((column) => column.sensitive);
const hasNonSensitiveColumns = columns.some((column) => !column.sensitive);
const realRowCount = hasNonSensitiveColumns
? Math.min(sourceRows.length, budget.rows)
: 0;
const rowCount = hasSensitiveColumns ? MAX_SAMPLE_ROWS_PER_REQUEST : realRowCount;
const rows = Array.from({ length: rowCount }, (_, rowIndex) => {
const row = rowIndex < realRowCount ? sourceRows[rowIndex] : undefined;
const sourceFields = new Map((row?.fields ?? []).map((field) => [field.name, field.value]));
const fields: Array<{ name: string; value: PromptSampleValue }> = [];
const seen = new Set<string>();
for (const field of row.fields) {
if (!relevantColumns.has(field.name) || seen.has(field.name)) continue;
seen.add(field.name);
for (const column of columns) {
const value = column.sensitive
? syntheticSampleValue(column, rowIndex + 1)
: sourceFields.get(column.name);
if (value === undefined) continue;
fields.push({
name: boundedJsonText(field.name, MAX_IDENTIFIER_JSON_BYTES),
value: promptSampleValue(field.value),
name: boundedJsonText(column.name, MAX_IDENTIFIER_JSON_BYTES),
value: promptSampleValue(value),
});
if (fields.length === MAX_SAMPLE_FIELDS_PER_ROW) break;
}
return { fields };
});
budget.rows -= rows.length;
budget.rows -= realRowCount;
const representativeValues: PromptSourceSample["representativeValues"] = [];
const seenColumns = new Set<string>();
let remainingRepresentativeValues = budget.representativeValues;
for (const examples of sample.representativeValues) {
if (remainingRepresentativeValues === 0) break;
if (!relevantColumns.has(examples.column) || seenColumns.has(examples.column)) continue;
seenColumns.add(examples.column);
let remainingRealRepresentativeValues = budget.representativeValues;
let remainingSyntheticRepresentativeValues = MAX_REPRESENTATIVE_VALUES_PER_REQUEST;
for (const column of columns) {
if (seenColumns.has(column.name)) continue;
const remainingRepresentativeValues = column.sensitive
? remainingSyntheticRepresentativeValues
: remainingRealRepresentativeValues;
if (remainingRepresentativeValues === 0) continue;
const sourceExamples = sample?.representativeValues.find(
(examples) => examples.column === column.name,
);
const exampleValues = column.sensitive
? Array.from(
{ length: Math.min(Math.max(rowCount, 1), remainingRepresentativeValues) },
(_, index) => syntheticSampleValue(column, index + 1),
)
: sourceExamples?.values ?? [];
seenColumns.add(column.name);
const values: Array<Exclude<PromptSampleValue, null>> = [];
const seenValues = new Set<string>();
for (const value of examples.values) {
for (const value of exampleValues) {
const normalized = promptSampleValue(value);
if (normalized === null) continue;
const key = JSON.stringify([typeof normalized, normalized]);
@@ -345,11 +366,15 @@ function sourceSampleFor(
}
if (values.length > 0) {
representativeValues.push({
column: boundedJsonText(examples.column, MAX_IDENTIFIER_JSON_BYTES),
column: boundedJsonText(column.name, MAX_IDENTIFIER_JSON_BYTES),
values,
});
remainingRepresentativeValues -= values.length;
budget.representativeValues -= values.length;
if (column.sensitive) {
remainingSyntheticRepresentativeValues -= values.length;
} else {
remainingRealRepresentativeValues -= values.length;
budget.representativeValues -= values.length;
}
}
if (representativeValues.length === MAX_SAMPLE_COLUMNS) break;
}
@@ -933,8 +958,9 @@ export class DescriptionGenerationWorker {
targetId: target.kind === "column" ? target.column.id : target.table.id,
tableName: target.table.name,
columnNames: target.kind === "column"
? [target.column.name]
: target.columns.map((column) => column.name),
? target.column.sensitive ? [] : [target.column.name]
: target.columns.filter((column) => !column.sensitive)
.map((column) => column.name),
})),
signal,
);
+10 -1
View File
@@ -202,10 +202,18 @@ export class MemoryCatalogRepository implements CatalogRepository {
expectedVersion: number,
description: string | null,
generatedDescription: string | null,
sensitive?: boolean,
): Promise<CatalogColumn | undefined> {
const current = await this.getColumn(databaseId, tableId, columnId);
if (!current || current.version !== expectedVersion) return undefined;
const updated = { ...current, description, generatedDescription, version: current.version + 1, updatedAt: new Date().toISOString() };
const updated = {
...current,
description,
generatedDescription,
sensitive: sensitive ?? current.sensitive,
version: current.version + 1,
updatedAt: new Date().toISOString(),
};
this.columns.set(columnId, updated);
return structuredClone(updated);
}
@@ -621,6 +629,7 @@ export class MemoryCatalogRepository implements CatalogRepository {
sourceComment: observed.sourceComment,
description: null,
generatedDescription: null,
sensitive: false,
lastSyncedDatabaseVersion: expectedDatabaseVersion,
lastSyncedAt: now,
version: 1,
+2
View File
@@ -8,6 +8,7 @@ import * as catalogTablesMigration from "./migrations/002_catalog_tables.js";
import * as catalogSchemaSyncMigration from "./migrations/003_catalog_schema_sync.js";
import * as catalogRuntimeSequencePrivilegesMigration from "./migrations/004_catalog_runtime_sequence_privileges.js";
import * as descriptionGenerationRunsMigration from "./migrations/005_description_generation_runs.js";
import * as sensitiveDataFlagMigration from "./migrations/006_sensitive_data_flag.js";
const connectionString = process.env.THT_CATALOG_MIGRATOR_DATABASE_URL;
const host = process.env.THT_CATALOG_DB_HOST;
@@ -38,6 +39,7 @@ const provider: MigrationProvider = {
"003_catalog_schema_sync": catalogSchemaSyncMigration,
"004_catalog_runtime_sequence_privileges": catalogRuntimeSequencePrivilegesMigration,
"005_description_generation_runs": descriptionGenerationRunsMigration,
"006_sensitive_data_flag": sensitiveDataFlagMigration,
};
},
};
@@ -0,0 +1,12 @@
import type { Kysely } from "kysely";
import type { CatalogDatabase } from "../repository.js";
export async function up(db: Kysely<CatalogDatabase>): Promise<void> {
await db.schema.alterTable("catalog_columns")
.addColumn("sensitive", "boolean", (column) => column.notNull().defaultTo(false))
.execute();
}
export async function down(db: Kysely<CatalogDatabase>): Promise<void> {
await db.schema.alterTable("catalog_columns").dropColumn("sensitive").execute();
}
+4
View File
@@ -108,6 +108,7 @@ interface CatalogColumnTable {
sourceComment: string | null;
description: string | null;
generatedDescription: string | null;
sensitive: Generated<boolean>;
lastSyncedDatabaseVersion: number | null;
lastSyncedAt: Timestamp | null;
version: Generated<number>;
@@ -290,6 +291,7 @@ function serializeColumn(row: Selectable<CatalogColumnTable>, foreignKeyCount =
sourceComment: row.sourceComment,
description: row.description,
generatedDescription: row.generatedDescription,
sensitive: row.sensitive,
lastSyncedDatabaseVersion: row.lastSyncedDatabaseVersion,
lastSyncedAt: row.lastSyncedAt === null ? null : new Date(row.lastSyncedAt).toISOString(),
version: row.version,
@@ -561,6 +563,7 @@ export class KyselyCatalogRepository implements CatalogRepository {
expectedVersion: number,
description: string | null,
generatedDescription: string | null,
sensitive?: boolean,
): Promise<CatalogColumn | undefined> {
const belongs = await this.db.selectFrom("catalogTables").select("id")
.where("id", "=", tableId).where("databaseId", "=", databaseId).executeTakeFirst();
@@ -568,6 +571,7 @@ export class KyselyCatalogRepository implements CatalogRepository {
const row = await this.db.updateTable("catalogColumns").set({
description,
generatedDescription,
...(sensitive === undefined ? {} : { sensitive }),
version: sql`version + 1`,
updatedAt: sql`now()`,
}).where("id", "=", columnId).where("tableId", "=", tableId)
@@ -0,0 +1,103 @@
import { z } from "zod";
import type { MetadataGenerationModels } from "./metadata-generation-models.js";
import type { ModelCompleter } from "./model-completer.js";
import type { CatalogRepository } from "./types.js";
const MAX_COLUMNS = 10_000;
const responseSchema = z.object({
suggestions: z.array(z.object({
columnId: z.uuid(),
sensitive: z.boolean(),
}).strict()).max(MAX_COLUMNS),
}).strict();
export interface SensitiveDataSuggestion {
columnId: string;
sensitive: boolean;
}
export class SensitiveDataSuggestionTargetNotFoundError extends Error {
constructor() {
super("database not found");
this.name = "SensitiveDataSuggestionTargetNotFoundError";
}
}
export class SensitiveDataSuggestionInvalidResponseError extends Error {
constructor() {
super("sensitive-data suggestion response is invalid");
this.name = "SensitiveDataSuggestionInvalidResponseError";
}
}
export class SensitiveDataSuggester {
constructor(
private readonly repository: CatalogRepository,
private readonly models: MetadataGenerationModels,
private readonly completer: ModelCompleter,
) {}
async suggest(
databaseId: string,
modelId: string,
signal: AbortSignal,
): Promise<readonly SensitiveDataSuggestion[]> {
const database = await this.repository.get(databaseId);
if (!database) throw new SensitiveDataSuggestionTargetNotFoundError();
const tables = await this.repository.listTables(databaseId);
const columns = (await Promise.all(tables.map(async (table) => ({
table,
columns: await this.repository.listColumns(databaseId, table.id),
})))).flatMap(({ table, columns: tableColumns }) => tableColumns.map((column) => ({
columnId: column.id,
table: table.name,
column: column.name,
dataType: column.dataType,
nullable: column.isNullable,
primaryKey: column.isPrimaryKey,
foreignKey: column.isForeignKey,
})));
if (columns.length === 0) return [];
if (columns.length > MAX_COLUMNS) throw new SensitiveDataSuggestionInvalidResponseError();
const content = await this.completer.complete({
model: this.models.resolve(modelId),
signal,
messages: [
{
role: "system",
content: [
"Classify whether each database column is likely to contain sensitive source values.",
"Use only the supplied structural metadata. Return strict JSON with this exact shape:",
'{"suggestions":[{"columnId":"uuid","sensitive":true}]}',
"Return every supplied column exactly once. Do not add explanations or markdown.",
].join("\n"),
},
{
role: "user",
content: JSON.stringify({
database: database.databaseName,
schema: database.schema,
columns,
}),
},
],
});
try {
const parsed = responseSchema.parse(JSON.parse(content));
const expected = new Set(columns.map((column) => column.columnId));
const received = new Set(parsed.suggestions.map((suggestion) => suggestion.columnId));
if (received.size !== parsed.suggestions.length
|| received.size !== expected.size
|| [...received].some((columnId) => !expected.has(columnId))) {
throw new SensitiveDataSuggestionInvalidResponseError();
}
return parsed.suggestions;
} catch (error) {
if (error instanceof SensitiveDataSuggestionInvalidResponseError) throw error;
throw new SensitiveDataSuggestionInvalidResponseError();
}
}
}
@@ -0,0 +1,67 @@
import type { DescriptionSourceSampleValue } from "./description-source-sampler.js";
const FIRST_NAMES = ["marta", "luca", "elena", "paolo", "giulia"] as const;
const LAST_NAMES = ["rossi", "bianchi", "conti", "romano", "ferrari"] as const;
const NUMERIC_TYPE = /(int|numeric|decimal|real|double|float|money)/;
const BOOLEAN_TYPE = /(bool)/;
function nameAt(index: number): string {
const offset = Math.max(0, index - 1);
return `${FIRST_NAMES[offset % FIRST_NAMES.length]} ${LAST_NAMES[offset % LAST_NAMES.length]}`;
}
export function syntheticSampleValue(
column: { name: string; dataType: string },
index: number,
): DescriptionSourceSampleValue {
const ordinal = Math.max(1, index);
const name = column.name.toLocaleLowerCase("en-US");
const type = column.dataType.toLocaleLowerCase("en-US");
const person = nameAt(ordinal).split(" ");
const safeName = name.replace(/[^a-z0-9]+/g, "_").replace(/^_+|_+$/g, "") || "value";
if (type.endsWith("[]") || type.startsWith("_") || /\barray\b/.test(type)) {
if (NUMERIC_TYPE.test(type)) {
return `{${1000 + ordinal},${1001 + ordinal}}`;
}
if (BOOLEAN_TYPE.test(type)) return `{${ordinal % 2 === 1},${ordinal % 2 !== 1}}`;
return `{${safeName}_${String(ordinal).padStart(3, "0")},${safeName}_${String(ordinal + 1).padStart(3, "0")}}`;
}
if (/^jsonb?$/.test(type)) {
return JSON.stringify({ example: `${safeName}_${String(ordinal).padStart(3, "0")}`, sequence: ordinal });
}
if (/(uuid|uniqueidentifier)/.test(type)) {
return `00000000-0000-4000-8000-${String(ordinal).padStart(12, "0")}`;
}
if (/\bcidr\b/.test(type)) return "192.0.2.0/24";
if (/\binet\b/.test(type)) return `192.0.2.${((ordinal - 1) % 254) + 1}`;
if (/(timestamp|datetime)/.test(type)) {
return `2024-01-${String(Math.min(ordinal, 28)).padStart(2, "0")}T10:30:00.000Z`;
}
if (/\bdate\b/.test(type)) {
return `198${ordinal % 10}-01-${String(Math.min(ordinal, 28)).padStart(2, "0")}`;
}
if (/\btime\b/.test(type)) {
return `10:30:${String(ordinal % 60).padStart(2, "0")}`;
}
if (/\binterval\b/.test(type)) {
return `${ordinal} days ${String(ordinal % 24).padStart(2, "0")}:00:00`;
}
if (/\bbytea\b/.test(type)) return `\\x${ordinal.toString(16).padStart(8, "0")}`;
if (BOOLEAN_TYPE.test(type)) return ordinal % 2 === 1;
if (NUMERIC_TYPE.test(type)) return 1000 + ordinal;
if (/e[-_]?mail/.test(name)) {
return `${person[0]}.${person[1]}@example.com`;
}
if (/(phone|mobile|cell|telefono|telefono_mobile|tel_)/.test(name)) {
return `+39 02 5550 ${String(1000 + ordinal).padStart(4, "0")}`;
}
if (/(first_?name|given_?name|nome)/.test(name)) return person[0]!;
if (/(last_?name|family_?name|surname|cognome)/.test(name)) return person[1]!;
if (/(full_?name|patient_?name|person_?name)/.test(name)) return nameAt(ordinal);
if (/(birth|dob|data_nascita)/.test(name)) {
return `198${ordinal % 10}-01-${String(Math.min(ordinal, 28)).padStart(2, "0")}`;
}
return `${safeName}_${String(ordinal).padStart(3, "0")}`;
}
+2
View File
@@ -88,6 +88,7 @@ export interface CatalogColumn {
sourceComment: string | null;
description: string | null;
generatedDescription: string | null;
sensitive: boolean;
lastSyncedDatabaseVersion: number | null;
lastSyncedAt: string | null;
version: number;
@@ -349,6 +350,7 @@ export interface CatalogRepository {
expectedVersion: number,
description: string | null,
generatedDescription: string | null,
sensitive?: boolean,
): Promise<CatalogColumn | undefined>;
consolidateGeneratedDescriptions(
databaseId: string,