feat: classify sensitive columns locally

This commit is contained in:
Codex
2026-09-03 02:11:13 +02:00
parent 7b87e95427
commit f114d0065a
57 changed files with 4038 additions and 1149 deletions
+25
View File
@@ -12,14 +12,17 @@
"@types/pg": "^8.20.3",
"fastify": "^5.0.0",
"kysely": "^0.29.5",
"libphonenumber-js": "1.13.12",
"openid-client": "6.8.5",
"pg": "^8.22.0",
"validator": "13.15.35",
"yaml": "^2.9.0",
"zod": "^4.4.3"
},
"devDependencies": {
"@testcontainers/postgresql": "^12.1.0",
"@types/node": "24.13.3",
"@types/validator": "13.15.10",
"tsx": "^4.19.0",
"typescript": "^5.6.0",
"vitest": "^2.1.0"
@@ -1329,6 +1332,13 @@
"dev": true,
"license": "MIT"
},
"node_modules/@types/validator": {
"version": "13.15.10",
"resolved": "https://registry.npmjs.org/@types/validator/-/validator-13.15.10.tgz",
"integrity": "sha512-T8L6i7wCuyoK8A/ZeLYt1+q0ty3Zb9+qbSSvrIVitzT3YjZqkTZ40IbRsPanlB4h1QB3JVL1SYCdR6ngtFYcuA==",
"dev": true,
"license": "MIT"
},
"node_modules/@vitest/expect": {
"version": "2.1.9",
"resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-2.1.9.tgz",
@@ -2824,6 +2834,12 @@
"safe-buffer": "~5.1.0"
}
},
"node_modules/libphonenumber-js": {
"version": "1.13.12",
"resolved": "https://registry.npmjs.org/libphonenumber-js/-/libphonenumber-js-1.13.12.tgz",
"integrity": "sha512-uLVeV1c9OTk6qkdqnj+mpMD+ZdnZ0szVyWu58HwMmpwkHA1gCEkyjd3veZQXDnuw9KEwSRjcc9B1pS9XKIN1fA==",
"license": "MIT"
},
"node_modules/light-my-request": {
"version": "6.6.0",
"resolved": "https://registry.npmjs.org/light-my-request/-/light-my-request-6.6.0.tgz",
@@ -4121,6 +4137,15 @@
"dev": true,
"license": "MIT"
},
"node_modules/validator": {
"version": "13.15.35",
"resolved": "https://registry.npmjs.org/validator/-/validator-13.15.35.tgz",
"integrity": "sha512-TQ5pAGhd5whStmqWvYF4OjQROlmv9SMFVt37qoCBdqRffuuklWYQlCNnEs2ZaIBD1kZRNnikiZOS1eqgkar0iw==",
"license": "MIT",
"engines": {
"node": ">= 0.10"
}
},
"node_modules/vite": {
"version": "5.4.21",
"resolved": "https://registry.npmjs.org/vite/-/vite-5.4.21.tgz",
+4
View File
@@ -7,6 +7,7 @@
"prebuild": "node scripts/clean-dist.mjs",
"build": "tsc -p tsconfig.json",
"catalog:migrate": "node dist/catalog/migrate.js",
"sensitivity:shadow": "node dist/catalog/sensitivity-shadow.js",
"test": "vitest run",
"start": "node dist/server.js",
"test:schema-v4-verifier": "python3 -I -B scripts/test_revision_state_policy.py && node --test scripts/verify-workspace-descriptor-files.test.mjs scripts/revision-state-policy.test.mjs",
@@ -19,14 +20,17 @@
"@types/pg": "^8.20.3",
"fastify": "^5.0.0",
"kysely": "^0.29.5",
"libphonenumber-js": "1.13.12",
"openid-client": "6.8.5",
"pg": "^8.22.0",
"validator": "13.15.35",
"yaml": "^2.9.0",
"zod": "^4.4.3"
},
"devDependencies": {
"@testcontainers/postgresql": "^12.1.0",
"@types/node": "24.13.3",
"@types/validator": "13.15.10",
"tsx": "^4.19.0",
"typescript": "^5.6.0",
"vitest": "^2.1.0"
@@ -0,0 +1,8 @@
34448b82c17d60fec9b65b1f093c115ddbaadc04beb1b0140b6bfed2e012a930 ./.gitattributes
4d9344c58a2a2ea4bb4ff4f7c611a853cf413205fc10d0cace564eba06f73828 ./README.md
180f0a10d1d5ed5ce3318db0bcb0b1b7780d79a52f0a8fc3acbd27f74536d0e4 ./THOTHII_MODEL_REVISION
164f17362bcf9d114067d3465e7374bfdd79ce6b605acb745de5a49dabb9595c ./config.json
f27dd63cc43a248d2566f0b6ad7a115db353676ce0561dcbca45bac766464c1a ./encoder_config/config.json
0280f6f39f6012da50b6640bad438d9b7e763a1b0102094115d1b710c4dd79b6 ./model.safetensors
f6df10ec83bea993035b2dd7c39345a3d4fcf23421c2adb6cb4ffc1e6d1bc4b5 ./tokenizer.json
233beed1f1095cccfc7907cde31a8d90a0c6aa4fdfaf6493f8e55fd162e81ae6 ./tokenizer_config.json
@@ -0,0 +1,34 @@
# Optional offline CPU pack. Fully version-locked in its own venv; not part of the base image.
--extra-index-url https://download.pytorch.org/whl/cpu
accelerate==1.14.0
annotated-types==0.8.0
certifi==2026.7.22
charset-normalizer==3.5.1
filelock==3.32.5
fsspec==2026.7.0
gliner2[local]==2.0.0
hf-xet==1.6.0
huggingface-hub==0.36.2
idna==3.19
Jinja2==3.1.6
MarkupSafe==3.0.3
mpmath==1.3.0
networkx==3.6.1
numpy==2.5.2
packaging==26.3
peft==0.20.0
psutil==7.2.2
pydantic==2.13.5
pydantic-core==2.46.5
PyYAML==6.0.3
regex==2026.9.3
requests==2.34.2
safetensors==0.8.0
sympy==1.14.0
tokenizers==0.22.2
torch==2.14.0+cpu
tqdm==4.70.0
transformers==4.57.6
typing-extensions==4.16.0
typing-inspection==0.4.4
urllib3==2.7.0
+301
View File
@@ -0,0 +1,301 @@
"""Offline, CPU-only JSONL worker for optional sensitivity NER evidence."""
from __future__ import annotations
import argparse
import contextlib
import ctypes
import errno
import hashlib
import json
import os
import socket
import sys
import tempfile
from pathlib import Path
from typing import Any
PII_LABELS = [
"person",
"full_name",
"first_name",
"middle_name",
"last_name",
"date_of_birth",
"email",
"phone_number",
"address",
"street_address",
"city",
"state_or_region",
"postal_code",
"country",
"government_id",
"national_id_number",
"passport_number",
"drivers_license_number",
"license_number",
"tax_id",
"tax_number",
"bank_account",
"account_number",
"routing_number",
"iban",
"payment_card",
"card_number",
"card_expiry",
"card_cvv",
"username",
"ip_address",
"account_id",
"sensitive_account_id",
"password",
"secret",
"api_key",
"access_token",
"recovery_code",
"sensitive_date",
"document_date",
"expiration_date",
"transaction_date",
]
_MODEL_COMPAT_DIRECTORY: tempfile.TemporaryDirectory[str] | None = None
_EXPECTED_MODEL_REVISION = "c153999da5f4c509df4322b0c6a1baf3d2c284d7"
def _arguments() -> argparse.Namespace:
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument("--model", required=True)
parser.add_argument("--threads", type=int, default=2)
return parser.parse_args()
def _disable_network() -> None:
libc = ctypes.CDLL(None, use_errno=True)
libc.prctl.argtypes = [
ctypes.c_int,
ctypes.c_ulong,
ctypes.c_ulong,
ctypes.c_ulong,
ctypes.c_ulong,
]
libc.prctl.restype = ctypes.c_int
if libc.prctl(38, 1, 0, 0, 0) != 0: # PR_SET_NO_NEW_PRIVS
raise RuntimeError("cannot enable no-new-privileges for network isolation")
try:
seccomp = ctypes.CDLL("libseccomp.so.2", use_errno=True)
except OSError as error:
raise RuntimeError("libseccomp is required for network isolation") from error
seccomp.seccomp_init.argtypes = [ctypes.c_uint32]
seccomp.seccomp_init.restype = ctypes.c_void_p
seccomp.seccomp_syscall_resolve_name.argtypes = [ctypes.c_char_p]
seccomp.seccomp_syscall_resolve_name.restype = ctypes.c_int
seccomp.seccomp_rule_add.argtypes = [
ctypes.c_void_p,
ctypes.c_uint32,
ctypes.c_int,
ctypes.c_uint,
]
seccomp.seccomp_rule_add.restype = ctypes.c_int
seccomp.seccomp_load.argtypes = [ctypes.c_void_p]
seccomp.seccomp_load.restype = ctypes.c_int
seccomp.seccomp_release.argtypes = [ctypes.c_void_p]
seccomp.seccomp_release.restype = None
allow = 0x7FFF0000 # SCMP_ACT_ALLOW
deny = 0x00050000 | errno.EPERM # SCMP_ACT_ERRNO(EPERM)
filter_context = seccomp.seccomp_init(allow)
if not filter_context:
raise RuntimeError("cannot initialize network syscall filter")
try:
for syscall in (
"socket",
"connect",
"sendto",
"sendmsg",
"sendmmsg",
"bind",
"listen",
"accept",
"accept4",
):
syscall_number = seccomp.seccomp_syscall_resolve_name(syscall.encode("ascii"))
if syscall_number < 0:
raise RuntimeError(f"cannot resolve network syscall: {syscall}")
if seccomp.seccomp_rule_add(filter_context, deny, syscall_number, 0) != 0:
raise RuntimeError(f"cannot block network syscall: {syscall}")
if seccomp.seccomp_load(filter_context) != 0:
raise RuntimeError("cannot activate network syscall filter")
finally:
seccomp.seccomp_release(filter_context)
def blocked(*_args: Any, **_kwargs: Any) -> Any:
raise PermissionError(errno.EPERM, "network disabled")
socket.socket = blocked # type: ignore[assignment]
socket.create_connection = blocked # type: ignore[assignment]
def _verify_model(path: Path) -> None:
revision_path = path / "THOTHII_MODEL_REVISION"
try:
revision = revision_path.read_text(encoding="utf-8").strip()
except OSError as error:
raise RuntimeError("model revision marker is unavailable") from error
if revision != _EXPECTED_MODEL_REVISION:
raise RuntimeError("model revision is not approved")
manifest_path = Path(__file__).with_name("sensitivity-ner-model-sha256.txt")
try:
manifest = manifest_path.read_text(encoding="utf-8").splitlines()
except OSError as error:
raise RuntimeError("model checksum manifest is unavailable") from error
for line in manifest:
checksum, separator, relative_name = line.partition(" ")
if not separator or len(checksum) != 64 or not relative_name.startswith("./"):
raise RuntimeError("model checksum manifest is invalid")
relative_path = Path(relative_name[2:])
if relative_path.is_absolute() or ".." in relative_path.parts:
raise RuntimeError("model checksum path is invalid")
model_file = path / relative_path
if not model_file.is_file() or model_file.is_symlink():
raise RuntimeError("approved model file is unavailable")
digest = hashlib.sha256()
with model_file.open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
if digest.hexdigest() != checksum:
raise RuntimeError("approved model checksum does not match")
def _transformers4_model_path(path: Path) -> Path:
"""Adapt tokenizer metadata emitted by Transformers 5 without changing pinned weights.
GLiNER2 2.0.0 officially requires Transformers <5, while current Fastino checkpoints were
saved by Transformers 5.8.0. Transformers 4 calls the same list
``additional_special_tokens``; Transformers 5 renamed it to ``extra_special_tokens`` and
changed its type. Keep the downloaded model immutable and create a temporary symlink view
containing only the compatibility metadata needed by the supported GLiNER2 dependency set.
"""
tokenizer_path = path / "tokenizer_config.json"
try:
tokenizer = json.loads(tokenizer_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as error:
raise RuntimeError("invalid tokenizer configuration") from error
extra_tokens = tokenizer.get("extra_special_tokens")
if extra_tokens is None:
return path
if not isinstance(extra_tokens, list) or not all(isinstance(token, str) for token in extra_tokens):
raise RuntimeError("unsupported extra_special_tokens configuration")
if "additional_special_tokens" in tokenizer:
raise RuntimeError("ambiguous special-token configuration")
global _MODEL_COMPAT_DIRECTORY
_MODEL_COMPAT_DIRECTORY = tempfile.TemporaryDirectory(prefix="thothii-ner-model-")
compatible_path = Path(_MODEL_COMPAT_DIRECTORY.name)
for child in path.iterdir():
if child.name == tokenizer_path.name:
continue
(compatible_path / child.name).symlink_to(child, target_is_directory=child.is_dir())
tokenizer["additional_special_tokens"] = tokenizer.pop("extra_special_tokens")
(compatible_path / tokenizer_path.name).write_text(
json.dumps(tokenizer, ensure_ascii=False, indent=2) + "\n",
encoding="utf-8",
)
return compatible_path
def _load_model(model_path: str, threads: int) -> Any:
path = Path(model_path).resolve(strict=True)
if not path.is_dir():
raise RuntimeError("model path must be a local directory")
_verify_model(path)
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["HIP_VISIBLE_DEVICES"] = ""
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
import torch
from gliner2 import AutoExtractor
torch.set_num_threads(max(1, min(threads, 8)))
torch.set_num_interop_threads(1)
compatible_path = _transformers4_model_path(path)
with contextlib.redirect_stdout(sys.stderr):
model = AutoExtractor.from_pretrained(str(compatible_path), map_location="cpu")
_disable_network()
return model
def _request(value: Any) -> tuple[str, list[dict[str, str]]]:
if not isinstance(value, dict) or not isinstance(value.get("id"), str):
raise ValueError("invalid request")
candidates = value.get("candidates")
if not isinstance(candidates, list) or not 1 <= len(candidates) <= 128:
raise ValueError("invalid candidates")
parsed: list[dict[str, str]] = []
for candidate in candidates:
if not isinstance(candidate, dict):
raise ValueError("invalid candidate")
column_id = candidate.get("columnId")
text = candidate.get("text")
if not isinstance(column_id, str) or not isinstance(text, str) or not 1 <= len(text) <= 500:
raise ValueError("invalid candidate")
parsed.append({"columnId": column_id, "text": text})
return value["id"], parsed
def _detect(model: Any, candidates: list[dict[str, str]]) -> list[dict[str, Any]]:
evidence: list[dict[str, Any]] = []
for candidate in candidates:
result = model.extract_entities(
candidate["text"],
PII_LABELS,
threshold=0.5,
include_confidence=True,
)
entities = result.get("entities", {}) if isinstance(result, dict) else {}
best: tuple[str, float] | None = None
if isinstance(entities, dict):
for label, matches in entities.items():
if label not in PII_LABELS or not isinstance(matches, list):
continue
for match in matches:
if not isinstance(match, dict):
continue
confidence = match.get("confidence")
if not isinstance(confidence, (int, float)) or not 0 <= confidence <= 1:
continue
if best is None or confidence > best[1]:
best = (label, float(confidence))
if best is not None:
evidence.append(
{
"columnId": candidate["columnId"],
"label": best[0],
"confidence": best[1],
}
)
return evidence
def main() -> int:
args = _arguments()
model = _load_model(args.model, args.threads)
print(json.dumps({"ready": True}, separators=(",", ":")), flush=True)
for line in sys.stdin:
request_id = "invalid"
try:
request_id, candidates = _request(json.loads(line))
response = {"id": request_id, "ok": True, "evidence": _detect(model, candidates)}
except Exception:
response = {"id": request_id, "ok": False, "error": "detection_failed"}
print(json.dumps(response, separators=(",", ":")), flush=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+34 -8
View File
@@ -3,6 +3,7 @@ import cors from "@fastify/cors";
import cookie from "@fastify/cookie";
import rateLimit from "@fastify/rate-limit";
import { dirname, isAbsolute, join } from "node:path";
import { fileURLToPath } from "node:url";
import { tmpdir } from "node:os";
import type { AppConfig } from "./config.js";
import { ThtRunner } from "./tht/tht-runner.js";
@@ -57,8 +58,11 @@ import { metadataGenerationModelRoutes } from "./routes/metadata-generation-mode
import { catalogDescriptionConsolidationRoutes } from "./routes/catalog-description-consolidation.js";
import { PythonModelCompleter, type ModelCompleter } from "./catalog/model-completer.js";
import { DescriptionGenerationWorker } from "./catalog/description-generation-worker.js";
import { SensitiveDataSuggester } from "./catalog/sensitive-data-suggester.js";
import { SensitiveDataSuggestionRunner } from "./catalog/sensitive-data-suggestion-runner.js";
import { SensitivityAnalysisService } from "./catalog/sensitivity-analysis-service.js";
import { SensitivityAnalysisRunner } from "./catalog/sensitivity-analysis-runner.js";
import { SensitivityClassifier, type LocalNerDetector, type SensitivityValueSource } from "./catalog/sensitivity-classifier.js";
import { ConcreteSensitivityValueSource } from "./catalog/sensitivity-value-source.js";
import { PythonLocalNerDetector } from "./catalog/local-ner-detector.js";
import {
ConcreteDescriptionSourceSampler,
type DescriptionSourceSampler,
@@ -93,6 +97,8 @@ export interface BuildAppDeps {
runtimeModelCatalog?: RuntimeModelCatalog;
modelCompleter?: ModelCompleter;
descriptionSourceSampler?: DescriptionSourceSampler;
sensitivityValueSource?: SensitivityValueSource;
localNerDetector?: LocalNerDetector;
workspaceRuntimeSupport?: (workspace: WorkspaceDescriptor) => boolean;
maintenanceBarrier?: MaintenanceBarrier;
piManagement?: PiManagementService;
@@ -194,12 +200,24 @@ export function buildApp(config: AppConfig, deps?: BuildAppDeps): FastifyInstanc
catalogOperationCoordinator,
descriptionSourceSampler,
);
const sensitiveDataSuggester = new SensitiveDataSuggester(
const sensitivityValueSource = deps?.sensitivityValueSource
?? new ConcreteSensitivityValueSource(catalogPostgresAccess, workspaceSecretStore);
const configuredNerWorker = config.sensitivityNer?.workerScript
?? fileURLToPath(new URL("../python/sensitivity_ner_worker.py", import.meta.url));
const localNerDetector = deps?.localNerDetector ?? (config.sensitivityNer
? new PythonLocalNerDetector({
pythonExecutable: config.sensitivityNer.pythonExecutable,
workerScript: configuredNerWorker,
modelPath: config.sensitivityNer.modelPath,
cwd: dirname(configuredNerWorker),
threads: config.sensitivityNer.threads,
})
: undefined);
const sensitiveDataSuggester = new SensitivityAnalysisService(
catalogRepository,
metadataGenerationModels,
modelCompleter,
new SensitivityClassifier(sensitivityValueSource, localNerDetector),
);
const sensitiveDataSuggestionRunner = new SensitiveDataSuggestionRunner(
const sensitivityAnalysisRunner = new SensitivityAnalysisRunner(
catalogRepository,
sensitiveDataSuggester,
);
@@ -235,12 +253,20 @@ export function buildApp(config: AppConfig, deps?: BuildAppDeps): FastifyInstanc
);
app.addHook("onReady", async () => { await catalogSyncWorker.initialize(); });
app.addHook("onReady", async () => { await descriptionGenerationWorker.initialize(); });
app.addHook("onReady", async () => { await sensitiveDataSuggestionRunner.initialize(); });
app.addHook("onReady", async () => { await sensitivityAnalysisRunner.initialize(); });
if (localNerDetector?.warmup) {
app.addHook("onReady", async () => {
void localNerDetector.warmup?.().catch(() => undefined);
});
}
if (!deps?.catalogRepository && catalogRepository.close) {
app.addHook("onClose", async () => { await catalogRepository.close?.(); });
}
app.addHook("onClose", async () => { await catalogSyncWorker.stop(); });
app.addHook("onClose", async () => { await descriptionGenerationWorker.stop(); });
if (localNerDetector?.close) {
app.addHook("onClose", async () => { await localNerDetector.close?.(); });
}
const workspaceDiagnoser = deps?.workspaceDiagnoser
?? createProductionWorkspaceDiagnoser(config.workspaceDiagnosticTimeoutMs, undefined, {
internalQdrantUrl: config.internalQdrantUrl,
@@ -483,7 +509,7 @@ export function buildApp(config: AppConfig, deps?: BuildAppDeps): FastifyInstanc
catalogDescriptionGenerationRoutes(app, {
repository: catalogRepository,
worker: descriptionGenerationWorker,
sensitiveDataSuggestionRunner,
sensitivityAnalysisRunner,
});
settingsRoutes(app, { cfg: config, getSettings });
piManagementRoutes(app, { service: piManagement });
+254
View File
@@ -0,0 +1,254 @@
import { randomUUID } from "node:crypto";
import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
import { tmpdir } from "node:os";
import { z } from "zod";
import type {
LocalNerCandidate,
LocalNerDetector,
LocalNerEvidence,
} from "./sensitivity-classifier.js";
const MAX_LINE_BYTES = 64 * 1024;
const candidateSchema = z.object({
columnId: z.uuid(),
text: z.string().min(1).max(500),
}).strict();
const workerMessageSchema = z.union([
z.object({ ready: z.literal(true) }).strict(),
z.object({
id: z.uuid(),
ok: z.literal(true),
evidence: z.array(z.object({
columnId: z.uuid(),
label: z.string().min(1).max(80),
confidence: z.number().min(0).max(1),
}).strict()).max(1_000),
}).strict(),
z.object({ id: z.uuid(), ok: z.literal(false), error: z.string().min(1).max(80) }).strict(),
]);
export class LocalNerUnavailableError extends Error {
constructor() {
super("local NER is unavailable");
this.name = "LocalNerUnavailableError";
}
}
interface PendingRequest {
resolve: (value: readonly LocalNerEvidence[]) => void;
reject: (error: Error) => void;
timer: ReturnType<typeof setTimeout>;
signal: AbortSignal;
cancel: () => void;
}
/** Persistent JSONL adapter for the optional, CPU-only Python NER worker. */
export class PythonLocalNerDetector implements LocalNerDetector {
private child?: ChildProcessWithoutNullStreams;
private ready?: Promise<void>;
private readyResolve?: () => void;
private readyReject?: (error: Error) => void;
private workerReady = false;
private stdout = "";
private readonly pending = new Map<string, PendingRequest>();
constructor(private readonly options: {
pythonExecutable: string;
workerScript: string;
modelPath: string;
cwd: string;
threads?: number;
startupTimeoutMs?: number;
}) {}
async warmup(): Promise<void> {
await this.ensureStarted();
}
isReady(): boolean {
return this.workerReady
&& this.child !== undefined
&& this.child.exitCode === null
&& this.child.signalCode === null;
}
async detect(
candidates: readonly LocalNerCandidate[],
signal: AbortSignal,
deadline: number,
): Promise<readonly LocalNerEvidence[]> {
const parsed = z.array(candidateSchema).min(1).max(128).parse(candidates);
if (signal.aborted || deadline <= Date.now()) throw new LocalNerUnavailableError();
await this.ensureStartedWithin(signal, deadline);
if (!this.child || this.child.exitCode !== null || this.child.signalCode !== null) {
throw new LocalNerUnavailableError();
}
const id = randomUUID();
return await new Promise<readonly LocalNerEvidence[]>((resolve, reject) => {
const fail = () => {
this.finishPending(id);
reject(new LocalNerUnavailableError());
this.stopWorker();
};
const timer = setTimeout(fail, Math.max(1, Math.floor(deadline - Date.now())));
const cancel = fail;
const pending: PendingRequest = { resolve, reject, timer, signal, cancel };
this.pending.set(id, pending);
signal.addEventListener("abort", cancel, { once: true });
this.child!.stdin.write(`${JSON.stringify({ id, candidates: parsed })}\n`, (error) => {
if (error) fail();
});
});
}
async close(): Promise<void> {
const child = this.child;
if (!child || child.exitCode !== null || child.signalCode !== null) return;
await new Promise<void>((resolve) => {
child.once("close", () => resolve());
child.kill("SIGTERM");
setTimeout(() => {
if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL");
}, 250).unref();
});
}
private async ensureStarted(): Promise<void> {
if (this.ready) return await this.ready;
this.ready = new Promise<void>((resolve, reject) => {
this.readyResolve = resolve;
this.readyReject = reject;
});
const threads = String(this.options.threads ?? 2);
const inheritedRuntimeEnvironment = Object.fromEntries([
"PATH", "SystemRoot", "WINDIR", "PATHEXT", "TMPDIR", "TEMP", "TMP", "LANG", "LC_ALL",
].flatMap((name) => process.env[name] === undefined ? [] : [[name, process.env[name]!]]));
const child = spawn(this.options.pythonExecutable, [
"-I",
"-B",
this.options.workerScript,
"--model",
this.options.modelPath,
"--threads",
threads,
], {
cwd: this.options.cwd,
stdio: ["pipe", "pipe", "pipe"],
env: {
...inheritedRuntimeEnvironment,
HOME: process.env.HOME ?? tmpdir(),
CUDA_VISIBLE_DEVICES: "",
HIP_VISIBLE_DEVICES: "",
HF_HUB_OFFLINE: "1",
HF_HUB_DISABLE_TELEMETRY: "1",
TRANSFORMERS_OFFLINE: "1",
TOKENIZERS_PARALLELISM: "false",
PYTHONNOUSERSITE: "1",
OMP_NUM_THREADS: threads,
MKL_NUM_THREADS: threads,
OPENBLAS_NUM_THREADS: threads,
HTTP_PROXY: "",
HTTPS_PROXY: "",
ALL_PROXY: "",
NO_PROXY: "*",
},
});
this.child = child;
child.stdout.setEncoding("utf8");
child.stdout.on("data", (chunk: string) => this.receive(chunk));
child.stderr.resume();
child.once("error", () => this.failWorker());
child.once("close", () => this.failWorker());
const startupTimer = setTimeout(() => this.failWorker(), this.options.startupTimeoutMs ?? 120_000);
startupTimer.unref();
try {
await this.ready;
} finally {
clearTimeout(startupTimer);
}
}
private async ensureStartedWithin(signal: AbortSignal, deadline: number): Promise<void> {
const started = this.ensureStarted();
await new Promise<void>((resolve, reject) => {
let settled = false;
const finish = (error?: Error, stopWorker = false) => {
if (settled) return;
settled = true;
clearTimeout(timer);
signal.removeEventListener("abort", cancel);
if (stopWorker) this.failWorker();
if (error) reject(error);
else resolve();
};
const cancel = () => finish(new LocalNerUnavailableError(), true);
const timer = setTimeout(cancel, Math.max(1, Math.floor(deadline - Date.now())));
signal.addEventListener("abort", cancel, { once: true });
void started.then(
() => finish(),
() => finish(new LocalNerUnavailableError()),
);
});
}
private receive(chunk: string): void {
this.stdout += chunk;
if (Buffer.byteLength(this.stdout, "utf8") > MAX_LINE_BYTES) {
this.failWorker();
return;
}
let newline: number;
while ((newline = this.stdout.indexOf("\n")) >= 0) {
const line = this.stdout.slice(0, newline);
this.stdout = this.stdout.slice(newline + 1);
if (!line) continue;
try {
const message = workerMessageSchema.parse(JSON.parse(line));
if ("ready" in message) {
this.workerReady = true;
this.readyResolve?.();
this.readyResolve = undefined;
this.readyReject = undefined;
continue;
}
const pending = this.pending.get(message.id);
if (!pending) continue;
this.finishPending(message.id);
if (message.ok) pending.resolve(message.evidence);
else pending.reject(new LocalNerUnavailableError());
} catch {
this.failWorker();
return;
}
}
}
private finishPending(id: string): void {
const pending = this.pending.get(id);
if (!pending) return;
clearTimeout(pending.timer);
pending.signal.removeEventListener("abort", pending.cancel);
this.pending.delete(id);
}
private stopWorker(): void {
const child = this.child;
if (child && child.exitCode === null && child.signalCode === null) child.kill("SIGTERM");
}
private failWorker(): void {
const error = new LocalNerUnavailableError();
this.readyReject?.(error);
this.readyResolve = undefined;
this.readyReject = undefined;
for (const [id, pending] of this.pending) {
this.finishPending(id);
pending.reject(error);
}
this.stopWorker();
this.child = undefined;
this.ready = undefined;
this.workerReady = false;
this.stdout = "";
}
}
+50 -47
View File
@@ -35,10 +35,10 @@ import {
type DescriptionGenerationRun,
type DescriptionGenerationRunUpdate,
type DescriptionGenerationScope,
type SensitiveDataSuggestionEvent,
type SensitiveDataSuggestionRun,
type SensitiveDataSuggestionRunUpdate,
type SensitiveDataSuggestionScope,
type SensitivityAnalysisEvent,
type SensitivityAnalysisRun,
type SensitivityAnalysisRunUpdate,
type SensitivityAnalysisScope,
type TableSyncRepositoryResult,
type WorkspaceDatabase,
} from "./types.js";
@@ -56,8 +56,8 @@ export class MemoryCatalogRepository implements CatalogRepository {
private readonly logicalRelationships = new Map<string, CatalogLogicalRelationship>();
private readonly descriptionGenerationRuns = new Map<string, DescriptionGenerationRun>();
private readonly descriptionGenerationEvents = new Map<string, DescriptionGenerationEvent[]>();
private readonly sensitiveDataSuggestionRuns = new Map<string, SensitiveDataSuggestionRun>();
private readonly sensitiveDataSuggestionEvents = new Map<string, SensitiveDataSuggestionEvent[]>();
private readonly sensitivityAnalysisRuns = new Map<string, SensitivityAnalysisRun>();
private readonly sensitivityAnalysisEvents = new Map<string, SensitivityAnalysisEvent[]>();
private readonly syncRuns = new Map<string, CatalogSyncRun>();
private readonly syncEvents = new Map<string, CatalogSyncEvent[]>();
@@ -183,10 +183,10 @@ export class MemoryCatalogRepository implements CatalogRepository {
this.descriptionGenerationRuns.delete(runId);
this.descriptionGenerationEvents.delete(runId);
}
for (const [runId, run] of this.sensitiveDataSuggestionRuns) {
for (const [runId, run] of this.sensitivityAnalysisRuns) {
if (run.databaseId !== id) continue;
this.sensitiveDataSuggestionRuns.delete(runId);
this.sensitiveDataSuggestionEvents.delete(runId);
this.sensitivityAnalysisRuns.delete(runId);
this.sensitivityAnalysisEvents.delete(runId);
}
return this.records.delete(id);
}
@@ -433,93 +433,96 @@ export class MemoryCatalogRepository implements CatalogRepository {
.map((event) => structuredClone(event));
}
async createSensitiveDataSuggestionRun(
async createSensitivityAnalysisRun(
databaseId: string,
scope: SensitiveDataSuggestionScope,
modelId: string,
): Promise<SensitiveDataSuggestionRun> {
scope: SensitivityAnalysisScope,
origin: { engine: "llm"; modelId: string } | { engine: "local"; policyVersion: string },
): Promise<SensitivityAnalysisRun> {
const now = new Date().toISOString();
const run: SensitiveDataSuggestionRun = {
const run: SensitivityAnalysisRun = {
id: randomUUID(),
databaseId,
scope,
modelId,
engine: origin.engine,
modelId: origin.engine === "llm" ? origin.modelId : null,
policyVersion: origin.engine === "local" ? origin.policyVersion : null,
status: "running",
total: 0,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
suggestedNonSensitive: 0,
unknown: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
createdAt: now,
startedAt: now,
updatedAt: now,
finishedAt: null,
errorSummary: null,
};
this.sensitiveDataSuggestionRuns.set(run.id, run);
this.sensitivityAnalysisRuns.set(run.id, run);
return structuredClone(run);
}
async getSensitiveDataSuggestionRun(
async getSensitivityAnalysisRun(
runId: string,
): Promise<SensitiveDataSuggestionRun | undefined> {
const run = this.sensitiveDataSuggestionRuns.get(runId);
): Promise<SensitivityAnalysisRun | undefined> {
const run = this.sensitivityAnalysisRuns.get(runId);
return run ? structuredClone(run) : undefined;
}
async listSensitiveDataSuggestionRuns(limit = 50): Promise<SensitiveDataSuggestionRun[]> {
return [...this.sensitiveDataSuggestionRuns.values()]
async listSensitivityAnalysisRuns(limit = 50): Promise<SensitivityAnalysisRun[]> {
return [...this.sensitivityAnalysisRuns.values()]
.sort((a, b) => b.createdAt.localeCompare(a.createdAt) || b.id.localeCompare(a.id))
.slice(0, limit)
.map((run) => structuredClone(run));
}
async interruptActiveSensitiveDataSuggestionRuns(
async interruptActiveSensitivityAnalysisRuns(
errorSummary: string,
): Promise<SensitiveDataSuggestionRun[]> {
const interrupted: SensitiveDataSuggestionRun[] = [];
for (const run of this.sensitiveDataSuggestionRuns.values()) {
): Promise<SensitivityAnalysisRun[]> {
const interrupted: SensitivityAnalysisRun[] = [];
for (const run of this.sensitivityAnalysisRuns.values()) {
if (run.status !== "running") continue;
const now = new Date().toISOString();
const updated: SensitiveDataSuggestionRun = {
const updated: SensitivityAnalysisRun = {
...run,
status: "interrupted",
updatedAt: now,
finishedAt: now,
errorSummary,
};
this.sensitiveDataSuggestionRuns.set(run.id, updated);
this.sensitivityAnalysisRuns.set(run.id, updated);
interrupted.push(structuredClone(updated));
}
return interrupted;
}
async updateSensitiveDataSuggestionRun(
async updateSensitivityAnalysisRun(
runId: string,
update: SensitiveDataSuggestionRunUpdate,
): Promise<SensitiveDataSuggestionRun | undefined> {
const current = this.sensitiveDataSuggestionRuns.get(runId);
update: SensitivityAnalysisRunUpdate,
): Promise<SensitivityAnalysisRun | undefined> {
const current = this.sensitivityAnalysisRuns.get(runId);
if (!current) return undefined;
const updated = {
...current,
...structuredClone(update),
updatedAt: new Date().toISOString(),
};
this.sensitiveDataSuggestionRuns.set(runId, updated);
this.sensitivityAnalysisRuns.set(runId, updated);
return structuredClone(updated);
}
async appendSensitiveDataSuggestionEvent(
async appendSensitivityAnalysisEvent(
runId: string,
level: SensitiveDataSuggestionEvent["level"],
level: SensitivityAnalysisEvent["level"],
message: string,
): Promise<SensitiveDataSuggestionEvent> {
if (!this.sensitiveDataSuggestionRuns.has(runId)) {
throw new CatalogConflictError("Sensitive Data Suggestion Run does not exist");
): Promise<SensitivityAnalysisEvent> {
if (!this.sensitivityAnalysisRuns.has(runId)) {
throw new CatalogConflictError("Sensitivity Analysis Run does not exist");
}
const events = this.sensitiveDataSuggestionEvents.get(runId) ?? [];
const event: SensitiveDataSuggestionEvent = {
const events = this.sensitivityAnalysisEvents.get(runId) ?? [];
const event: SensitivityAnalysisEvent = {
runId,
sequence: events.length + 1,
level,
@@ -527,15 +530,15 @@ export class MemoryCatalogRepository implements CatalogRepository {
createdAt: new Date().toISOString(),
};
events.push(event);
this.sensitiveDataSuggestionEvents.set(runId, events);
this.sensitivityAnalysisEvents.set(runId, events);
return structuredClone(event);
}
async listSensitiveDataSuggestionEvents(
async listSensitivityAnalysisEvents(
runId: string,
afterSequence = 0,
): Promise<SensitiveDataSuggestionEvent[]> {
return (this.sensitiveDataSuggestionEvents.get(runId) ?? [])
): Promise<SensitivityAnalysisEvent[]> {
return (this.sensitivityAnalysisEvents.get(runId) ?? [])
.filter((event) => event.sequence > afterSequence)
.map((event) => structuredClone(event));
}
+4 -2
View File
@@ -9,10 +9,11 @@ import * as catalogSchemaSyncMigration from "./migrations/003_catalog_schema_syn
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";
import * as sensitiveDataSuggestionRunsMigration from "./migrations/007_sensitive_data_suggestion_runs.js";
import * as sensitivityAnalysisRunsMigration from "./migrations/007_sensitive_data_suggestion_runs.js";
import * as catalogLogicalRelationshipsMigration from "./migrations/008_catalog_logical_relationships.js";
import * as aiTokenUsageMigration from "./migrations/009_ai_token_usage.js";
import * as canonicalModelIdsMigration from "./migrations/010_canonical_model_ids.js";
import * as localSensitivityAnalysisMigration from "./migrations/011_local_sensitivity_analysis.js";
const connectionString = process.env.THT_CATALOG_MIGRATOR_DATABASE_URL;
const host = process.env.THT_CATALOG_DB_HOST;
@@ -44,10 +45,11 @@ const provider: MigrationProvider = {
"004_catalog_runtime_sequence_privileges": catalogRuntimeSequencePrivilegesMigration,
"005_description_generation_runs": descriptionGenerationRunsMigration,
"006_sensitive_data_flag": sensitiveDataFlagMigration,
"007_sensitive_data_suggestion_runs": sensitiveDataSuggestionRunsMigration,
"007_sensitive_data_suggestion_runs": sensitivityAnalysisRunsMigration,
"008_catalog_logical_relationships": catalogLogicalRelationshipsMigration,
"009_ai_token_usage": aiTokenUsageMigration,
"010_canonical_model_ids": canonicalModelIdsMigration,
"011_local_sensitivity_analysis": localSensitivityAnalysisMigration,
};
},
};
@@ -0,0 +1,42 @@
import { sql, type Kysely } from "kysely";
import type { CatalogDatabase } from "../repository.js";
export async function up(db: Kysely<CatalogDatabase>): Promise<void> {
await sql.raw(`alter table sensitive_data_suggestion_runs
alter column model_id drop not null,
add column engine text not null default 'llm',
add column policy_version text,
add column unknown integer not null default 0,
drop constraint sensitive_data_suggestion_runs_counters_check,
add constraint sensitive_data_suggestion_runs_counters_check
check (total >= 0
and suggested_sensitive >= 0
and suggested_non_sensitive >= 0
and unknown >= 0
and suggested_sensitive + suggested_non_sensitive + unknown <= total),
add constraint sensitive_data_suggestion_runs_engine_check
check (engine in ('llm', 'local')),
add constraint sensitive_data_suggestion_runs_origin_check
check ((engine = 'llm' and model_id is not null and policy_version is null)
or (engine = 'local' and model_id is null
and policy_version ~ '^[a-z][a-z0-9._-]{0,63}$'))`).execute(db);
}
export async function down(db: Kysely<CatalogDatabase>): Promise<void> {
await sql.raw(`alter table sensitive_data_suggestion_runs
drop constraint sensitive_data_suggestion_runs_origin_check,
drop constraint sensitive_data_suggestion_runs_engine_check,
drop constraint sensitive_data_suggestion_runs_counters_check`).execute(db);
await sql.raw(`update sensitive_data_suggestion_runs
set model_id = coalesce(model_id, 'local/sensitivity-v1')`).execute(db);
await sql.raw(`alter table sensitive_data_suggestion_runs
drop column unknown,
drop column policy_version,
drop column engine,
alter column model_id set not null,
add constraint sensitive_data_suggestion_runs_counters_check
check (total >= 0
and suggested_sensitive >= 0
and suggested_non_sensitive >= 0
and suggested_sensitive + suggested_non_sensitive <= total)`).execute(db);
}
+58 -51
View File
@@ -45,10 +45,10 @@ import {
type DescriptionGenerationScope,
type ObservedCatalogTable,
type ObservedSchemaSnapshot,
type SensitiveDataSuggestionEvent,
type SensitiveDataSuggestionRun,
type SensitiveDataSuggestionRunUpdate,
type SensitiveDataSuggestionScope,
type SensitivityAnalysisEvent,
type SensitivityAnalysisRun,
type SensitivityAnalysisRunUpdate,
type SensitivityAnalysisScope,
type TableSyncRepositoryResult,
type WorkspaceDatabase,
} from "./types.js";
@@ -189,15 +189,18 @@ interface DescriptionGenerationEventTable {
createdAt: Timestamp;
}
interface SensitiveDataSuggestionRunTable {
interface SensitivityAnalysisRunTable {
id: string;
databaseId: string;
scope: SensitiveDataSuggestionScope;
modelId: string;
status: SensitiveDataSuggestionRun["status"];
scope: SensitivityAnalysisScope;
engine: SensitivityAnalysisRun["engine"];
modelId: string | null;
policyVersion: string | null;
status: SensitivityAnalysisRun["status"];
total: number;
suggestedSensitive: number;
suggestedNonSensitive: number;
unknown: number;
inputTokens: number;
cacheReadTokens: number;
outputTokens: number;
@@ -208,10 +211,10 @@ interface SensitiveDataSuggestionRunTable {
errorSummary: string | null;
}
interface SensitiveDataSuggestionEventTable {
interface SensitivityAnalysisEventTable {
runId: string;
sequence: number;
level: SensitiveDataSuggestionEvent["level"];
level: SensitivityAnalysisEvent["level"];
message: string;
createdAt: Timestamp;
}
@@ -264,8 +267,9 @@ export interface CatalogDatabase {
catalogLogicalRelationships: CatalogLogicalRelationshipTable;
descriptionGenerationRuns: DescriptionGenerationRunTable;
descriptionGenerationEvents: DescriptionGenerationEventTable;
sensitiveDataSuggestionRuns: SensitiveDataSuggestionRunTable;
sensitiveDataSuggestionEvents: SensitiveDataSuggestionEventTable;
// Legacy physical table names retained for migration and storage compatibility.
sensitiveDataSuggestionRuns: SensitivityAnalysisRunTable;
sensitiveDataSuggestionEvents: SensitivityAnalysisEventTable;
catalogSyncRuns: CatalogSyncRunTable;
catalogSyncEvents: CatalogSyncEventTable;
}
@@ -402,9 +406,9 @@ function serializeDescriptionGenerationEvent(
return { ...row, createdAt: new Date(row.createdAt).toISOString() };
}
function serializeSensitiveDataSuggestionRun(
row: Selectable<SensitiveDataSuggestionRunTable>,
): SensitiveDataSuggestionRun {
function serializeSensitivityAnalysisRun(
row: Selectable<SensitivityAnalysisRunTable>,
): SensitivityAnalysisRun {
const stamp = (value: Date | string | null) => value === null ? null : new Date(value).toISOString();
return {
...row,
@@ -415,9 +419,9 @@ function serializeSensitiveDataSuggestionRun(
};
}
function serializeSensitiveDataSuggestionEvent(
row: Selectable<SensitiveDataSuggestionEventTable>,
): SensitiveDataSuggestionEvent {
function serializeSensitivityAnalysisEvent(
row: Selectable<SensitivityAnalysisEventTable>,
): SensitivityAnalysisEvent {
return { ...row, createdAt: new Date(row.createdAt).toISOString() };
}
@@ -938,52 +942,55 @@ export class KyselyCatalogRepository implements CatalogRepository {
return rows.map(serializeDescriptionGenerationEvent);
}
async createSensitiveDataSuggestionRun(
async createSensitivityAnalysisRun(
databaseId: string,
scope: SensitiveDataSuggestionScope,
modelId: string,
): Promise<SensitiveDataSuggestionRun> {
scope: SensitivityAnalysisScope,
origin: { engine: "llm"; modelId: string } | { engine: "local"; policyVersion: string },
): Promise<SensitivityAnalysisRun> {
const row = await this.db.insertInto("sensitiveDataSuggestionRuns").values({
id: randomUUID(),
databaseId,
scope,
modelId,
engine: origin.engine,
modelId: origin.engine === "llm" ? origin.modelId : null,
policyVersion: origin.engine === "local" ? origin.policyVersion : null,
status: "running",
total: 0,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
unknown: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
finishedAt: null,
errorSummary: null,
}).returningAll().executeTakeFirstOrThrow();
return serializeSensitiveDataSuggestionRun(row);
return serializeSensitivityAnalysisRun(row);
}
async getSensitiveDataSuggestionRun(
async getSensitivityAnalysisRun(
runId: string,
): Promise<SensitiveDataSuggestionRun | undefined> {
): Promise<SensitivityAnalysisRun | undefined> {
const row = await this.db.selectFrom("sensitiveDataSuggestionRuns")
.selectAll()
.where("id", "=", runId)
.executeTakeFirst();
return row ? serializeSensitiveDataSuggestionRun(row) : undefined;
return row ? serializeSensitivityAnalysisRun(row) : undefined;
}
async listSensitiveDataSuggestionRuns(limit = 50): Promise<SensitiveDataSuggestionRun[]> {
async listSensitivityAnalysisRuns(limit = 50): Promise<SensitivityAnalysisRun[]> {
const rows = await this.db.selectFrom("sensitiveDataSuggestionRuns")
.selectAll()
.orderBy("createdAt", "desc")
.orderBy("id", "desc")
.limit(limit)
.execute();
return rows.map(serializeSensitiveDataSuggestionRun);
return rows.map(serializeSensitivityAnalysisRun);
}
async interruptActiveSensitiveDataSuggestionRuns(
async interruptActiveSensitivityAnalysisRuns(
errorSummary: string,
): Promise<SensitiveDataSuggestionRun[]> {
): Promise<SensitivityAnalysisRun[]> {
const rows = await this.db.updateTable("sensitiveDataSuggestionRuns")
.set({
status: "interrupted",
@@ -994,34 +1001,34 @@ export class KyselyCatalogRepository implements CatalogRepository {
.where("status", "=", "running")
.returningAll()
.execute();
return rows.map(serializeSensitiveDataSuggestionRun);
return rows.map(serializeSensitivityAnalysisRun);
}
async updateSensitiveDataSuggestionRun(
async updateSensitivityAnalysisRun(
runId: string,
update: SensitiveDataSuggestionRunUpdate,
): Promise<SensitiveDataSuggestionRun | undefined> {
update: SensitivityAnalysisRunUpdate,
): Promise<SensitivityAnalysisRun | undefined> {
const values: any = { ...update, updatedAt: sql`now()` };
const row = await this.db.updateTable("sensitiveDataSuggestionRuns")
.set(values)
.where("id", "=", runId)
.returningAll()
.executeTakeFirst();
return row ? serializeSensitiveDataSuggestionRun(row) : undefined;
return row ? serializeSensitivityAnalysisRun(row) : undefined;
}
async appendSensitiveDataSuggestionEvent(
async appendSensitivityAnalysisEvent(
runId: string,
level: SensitiveDataSuggestionEvent["level"],
level: SensitivityAnalysisEvent["level"],
message: string,
): Promise<SensitiveDataSuggestionEvent> {
): Promise<SensitivityAnalysisEvent> {
return await this.db.transaction().execute(async (trx) => {
const run = await trx.selectFrom("sensitiveDataSuggestionRuns")
.select("id")
.where("id", "=", runId)
.forUpdate()
.executeTakeFirst();
if (!run) throw new CatalogConflictError("Sensitive Data Suggestion Run does not exist");
if (!run) throw new CatalogConflictError("Sensitivity Analysis Run does not exist");
const current = await trx.selectFrom("sensitiveDataSuggestionEvents")
.select(sql<number>`coalesce(max(sequence), 0)::int`.as("sequence"))
.where("runId", "=", runId)
@@ -1032,21 +1039,21 @@ export class KyselyCatalogRepository implements CatalogRepository {
level,
message,
}).returningAll().executeTakeFirstOrThrow();
return serializeSensitiveDataSuggestionEvent(row);
return serializeSensitivityAnalysisEvent(row);
});
}
async listSensitiveDataSuggestionEvents(
async listSensitivityAnalysisEvents(
runId: string,
afterSequence = 0,
): Promise<SensitiveDataSuggestionEvent[]> {
): Promise<SensitivityAnalysisEvent[]> {
const rows = await this.db.selectFrom("sensitiveDataSuggestionEvents")
.selectAll()
.where("runId", "=", runId)
.where("sequence", ">", afterSequence)
.orderBy("sequence")
.execute();
return rows.map(serializeSensitiveDataSuggestionEvent);
return rows.map(serializeSensitivityAnalysisEvent);
}
async listRelationships(databaseId: string): Promise<CatalogPhysicalRelationship[]> {
@@ -1792,13 +1799,13 @@ export class UnavailableCatalogRepository implements CatalogRepository {
async updateDescriptionGenerationRun(): Promise<DescriptionGenerationRun | undefined> { return this.fail(); }
async appendDescriptionGenerationEvent(): Promise<DescriptionGenerationEvent> { return this.fail(); }
async listDescriptionGenerationEvents(): Promise<DescriptionGenerationEvent[]> { return this.fail(); }
async createSensitiveDataSuggestionRun(): Promise<SensitiveDataSuggestionRun> { return this.fail(); }
async getSensitiveDataSuggestionRun(): Promise<SensitiveDataSuggestionRun | undefined> { return this.fail(); }
async listSensitiveDataSuggestionRuns(): Promise<SensitiveDataSuggestionRun[]> { return this.fail(); }
async interruptActiveSensitiveDataSuggestionRuns(): Promise<SensitiveDataSuggestionRun[]> { return this.fail(); }
async updateSensitiveDataSuggestionRun(): Promise<SensitiveDataSuggestionRun | undefined> { return this.fail(); }
async appendSensitiveDataSuggestionEvent(): Promise<SensitiveDataSuggestionEvent> { return this.fail(); }
async listSensitiveDataSuggestionEvents(): Promise<SensitiveDataSuggestionEvent[]> { return this.fail(); }
async createSensitivityAnalysisRun(): Promise<SensitivityAnalysisRun> { return this.fail(); }
async getSensitivityAnalysisRun(): Promise<SensitivityAnalysisRun | undefined> { return this.fail(); }
async listSensitivityAnalysisRuns(): Promise<SensitivityAnalysisRun[]> { return this.fail(); }
async interruptActiveSensitivityAnalysisRuns(): Promise<SensitivityAnalysisRun[]> { return this.fail(); }
async updateSensitivityAnalysisRun(): Promise<SensitivityAnalysisRun | undefined> { return this.fail(); }
async appendSensitivityAnalysisEvent(): Promise<SensitivityAnalysisEvent> { return this.fail(); }
async listSensitivityAnalysisEvents(): Promise<SensitivityAnalysisEvent[]> { return this.fail(); }
async listRelationships(): Promise<CatalogPhysicalRelationship[]> { return this.fail(); }
async listLogicalRelationships(): Promise<CatalogLogicalRelationship[]> { return this.fail(); }
async getLogicalRelationshipContext(): Promise<CatalogLogicalRelationshipContext | undefined> { return this.fail(); }
@@ -1,254 +0,0 @@
import { z } from "zod";
import type { MetadataGenerationModels } from "./metadata-generation-models.js";
import type { ModelCompleter, ModelCompletionMessage, ModelCompletionResult, ModelCompletionUsage } from "./model-completer.js";
import type {
CatalogColumn,
CatalogRepository,
CatalogTable,
SensitiveDataSuggestionScope,
} from "./types.js";
export type { SensitiveDataSuggestionScope } from "./types.js";
// The helper accepts at most 64 KiB per message. Keep the same safety margin used by
// Description Generation so UTF-8 structural metadata never reaches that hard limit.
const MAX_USER_MESSAGE_BYTES = 60 * 1024;
// Preserve ThothAI's proven completion granularity: small batches keep generation time and
// structured-output accuracy predictable even when the helper byte limit would allow much more.
const MAX_COLUMNS_PER_BATCH = 10;
const responseSchema = z.object({
suggestions: z.array(z.object({
columnId: z.uuid(),
sensitive: z.boolean(),
}).strict()),
}).strict();
interface StructuralColumn {
columnId: string;
tableId: string;
table: string;
column: string;
dataType: string;
nullable: boolean;
primaryKey: boolean;
foreignKey: boolean;
version: number;
currentSensitive: boolean;
}
export interface SensitiveDataSuggestion {
columnId: string;
tableId: string;
tableName: string;
columnName: string;
version: number;
currentSensitive: boolean;
sensitive: boolean;
}
export class SensitiveDataSuggestionTargetNotFoundError extends Error {
constructor(readonly target: "database" | "table" | "column") {
super(`${target} not found`);
this.name = "SensitiveDataSuggestionTargetNotFoundError";
}
}
export class SensitiveDataSuggestionDuplicateTargetIdsError extends Error {
constructor() {
super("sensitive-data suggestion target IDs must be unique");
this.name = "SensitiveDataSuggestionDuplicateTargetIdsError";
}
}
export class SensitiveDataSuggestionNoEligibleColumnsError extends Error {
constructor(readonly scope: SensitiveDataSuggestionScope) {
super("selected scope has no catalog columns");
this.name = "SensitiveDataSuggestionNoEligibleColumnsError";
}
}
export class SensitiveDataSuggestionPayloadTooLargeError extends Error {
constructor() {
super("sensitive-data suggestion structural metadata is too large");
this.name = "SensitiveDataSuggestionPayloadTooLargeError";
}
}
export class SensitiveDataSuggestionInvalidResponseError extends Error {
constructor() {
super("sensitive-data suggestion response is invalid");
this.name = "SensitiveDataSuggestionInvalidResponseError";
}
}
function userContent(
database: { databaseName: string; schema: string },
columns: readonly StructuralColumn[],
): string {
return JSON.stringify({
database: database.databaseName,
schema: database.schema,
columns: columns.map((column) => ({
columnId: column.columnId,
table: column.table,
column: column.column,
dataType: column.dataType,
nullable: column.nullable,
primaryKey: column.primaryKey,
foreignKey: column.foreignKey,
})),
});
}
function batchesFor(
database: { databaseName: string; schema: string },
columns: readonly StructuralColumn[],
): StructuralColumn[][] {
const batches: StructuralColumn[][] = [];
let current: StructuralColumn[] = [];
for (const column of columns) {
if (current.length === MAX_COLUMNS_PER_BATCH) {
batches.push(current);
current = [];
}
const candidate = [...current, column];
if (Buffer.byteLength(userContent(database, candidate), "utf8") <= MAX_USER_MESSAGE_BYTES) {
current = candidate;
continue;
}
if (current.length === 0) throw new SensitiveDataSuggestionPayloadTooLargeError();
batches.push(current);
current = [column];
if (Buffer.byteLength(userContent(database, current), "utf8") > MAX_USER_MESSAGE_BYTES) {
throw new SensitiveDataSuggestionPayloadTooLargeError();
}
}
if (current.length > 0) batches.push(current);
return batches;
}
function structuralColumn(table: CatalogTable, column: CatalogColumn): StructuralColumn {
return {
columnId: column.id,
tableId: table.id,
table: table.name,
column: column.name,
dataType: column.dataType,
nullable: column.isNullable,
primaryKey: column.isPrimaryKey,
foreignKey: column.isForeignKey,
version: column.version,
currentSensitive: column.sensitive,
};
}
const systemMessage: ModelCompletionMessage = {
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"),
};
export class SensitiveDataSuggester {
constructor(
private readonly repository: CatalogRepository,
private readonly models: MetadataGenerationModels,
private readonly completer: ModelCompleter,
) {}
private async selectColumns(
databaseId: string,
scope: SensitiveDataSuggestionScope,
targetIds: readonly string[],
): Promise<StructuralColumn[]> {
if (new Set(targetIds).size !== targetIds.length) {
throw new SensitiveDataSuggestionDuplicateTargetIdsError();
}
const tables = await this.repository.listTables(databaseId);
const tableIds = new Set(targetIds);
const selectedTables = scope === "selected_tables"
? tables.filter((table) => tableIds.has(table.id))
: tables;
if (scope === "selected_tables" && selectedTables.length !== targetIds.length) {
throw new SensitiveDataSuggestionTargetNotFoundError("table");
}
const columns = (await Promise.all(selectedTables.map(async (table) => (
(await this.repository.listColumns(databaseId, table.id)).map((column) => (
structuralColumn(table, column)
))
)))).flat();
const columnIds = new Set(targetIds);
const selectedColumns = scope === "selected_columns"
? columns.filter((column) => columnIds.has(column.columnId))
: columns;
if (scope === "selected_columns" && selectedColumns.length !== targetIds.length) {
throw new SensitiveDataSuggestionTargetNotFoundError("column");
}
if (selectedColumns.length === 0) {
throw new SensitiveDataSuggestionNoEligibleColumnsError(scope);
}
return selectedColumns;
}
async suggest(
databaseId: string,
modelId: string,
scope: SensitiveDataSuggestionScope,
targetIds: readonly string[],
signal: AbortSignal,
onPrepared?: (total: number) => void | Promise<void>,
onProgress?: (processed: number, suggestions: readonly SensitiveDataSuggestion[]) => void | Promise<void>,
onUsage?: (usage: ModelCompletionUsage) => void | Promise<void>,
): Promise<readonly SensitiveDataSuggestion[]> {
const database = await this.repository.get(databaseId);
if (!database) throw new SensitiveDataSuggestionTargetNotFoundError("database");
const columns = await this.selectColumns(databaseId, scope, targetIds);
await onPrepared?.(columns.length);
const model = this.models.resolve(modelId);
const suggestions: SensitiveDataSuggestion[] = [];
for (const batch of batchesFor(database, columns)) {
let received: Map<string, { columnId: string; sensitive: boolean }> | undefined;
for (let attempt = 0; attempt < 2 && !received; attempt += 1) {
const completion = await this.completer.complete({
model,
signal,
messages: [systemMessage, { role: "user", content: userContent(database, batch) }],
});
const result: ModelCompletionResult = typeof completion === "string"
? { content: completion, usage: { input: 0, cacheRead: 0, output: 0 } }
: completion;
await onUsage?.(result.usage);
const content = result.content;
try {
const parsed = responseSchema.parse(JSON.parse(content));
const expected = new Set(batch.map((column) => column.columnId));
const candidate = new Map(parsed.suggestions.map((suggestion) => [suggestion.columnId, suggestion]));
if (candidate.size !== parsed.suggestions.length
|| candidate.size !== expected.size
|| [...candidate.keys()].some((columnId) => !expected.has(columnId))) {
throw new SensitiveDataSuggestionInvalidResponseError();
}
received = candidate;
} catch {
if (attempt === 1) throw new SensitiveDataSuggestionInvalidResponseError();
}
}
suggestions.push(...batch.map((column) => ({
columnId: column.columnId,
tableId: column.tableId,
tableName: column.table,
columnName: column.column,
version: column.version,
currentSensitive: column.currentSensitive,
sensitive: received!.get(column.columnId)!.sensitive,
})));
await onProgress?.(suggestions.length, suggestions.slice(-batch.length));
}
return suggestions;
}
}
@@ -1,136 +0,0 @@
import type {
SensitiveDataSuggestion,
} from "./sensitive-data-suggester.js";
import {
SensitiveDataSuggester,
SensitiveDataSuggestionTargetNotFoundError,
} from "./sensitive-data-suggester.js";
import type {
CatalogRepository,
SensitiveDataSuggestionRun,
SensitiveDataSuggestionScope,
} from "./types.js";
import type { ModelCompletionUsage } from "./model-completer.js";
const interruptedMessage = "Sensitive-field suggestion generation was interrupted by backend restart.";
const failedMessage = "Sensitive-field suggestion generation failed.";
export interface SensitiveDataSuggestionRunResult {
suggestions: readonly SensitiveDataSuggestion[];
run: SensitiveDataSuggestionRun;
}
export class SensitiveDataSuggestionRunner {
constructor(
private readonly repository: CatalogRepository,
private readonly suggester: SensitiveDataSuggester,
) {}
async initialize(): Promise<void> {
if (!(await this.repository.available())) return;
const interrupted = await this.repository.interruptActiveSensitiveDataSuggestionRuns(
interruptedMessage,
);
for (const run of interrupted) {
await this.repository.appendSensitiveDataSuggestionEvent(
run.id,
"warning",
interruptedMessage,
);
}
}
async run(
databaseId: string,
modelId: string,
scope: SensitiveDataSuggestionScope,
targetIds: readonly string[],
signal: AbortSignal,
): Promise<SensitiveDataSuggestionRunResult> {
if (!(await this.repository.get(databaseId))) {
throw new SensitiveDataSuggestionTargetNotFoundError("database");
}
const started = await this.repository.createSensitiveDataSuggestionRun(
databaseId,
scope,
modelId,
);
try {
await this.repository.appendSensitiveDataSuggestionEvent(
started.id,
"info",
"Sensitive-field suggestion generation started.",
);
const suggestions = await this.suggester.suggest(
databaseId,
modelId,
scope,
targetIds,
signal,
async (total) => {
const prepared = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
total,
});
if (!prepared) throw new Error("Sensitive Data Suggestion Run disappeared");
},
async (processed, batch) => {
const suggestedSensitive = batch.filter((suggestion) => suggestion.sensitive).length;
const suggestedNonSensitive = batch.length - suggestedSensitive;
const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
const progress = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
suggestedSensitive: current.suggestedSensitive + suggestedSensitive,
suggestedNonSensitive: current.suggestedNonSensitive + suggestedNonSensitive,
});
if (!progress) throw new Error("Sensitive Data Suggestion Run disappeared");
await this.repository.appendSensitiveDataSuggestionEvent(
started.id,
"info",
`Classified ${processed} of ${progress.total} columns.`,
);
},
async (usage: ModelCompletionUsage) => {
const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
await this.repository.updateSensitiveDataSuggestionRun(started.id, {
inputTokens: current.inputTokens + usage.input,
cacheReadTokens: current.cacheReadTokens + usage.cacheRead,
outputTokens: current.outputTokens + usage.output,
});
},
);
const suggestedSensitive = suggestions.filter((suggestion) => suggestion.sensitive).length;
const suggestedNonSensitive = suggestions.length - suggestedSensitive;
await this.repository.appendSensitiveDataSuggestionEvent(
started.id,
"info",
`Sensitive-field suggestion generation completed for ${suggestions.length} column${
suggestions.length === 1 ? "" : "s"
}.`,
);
const completed = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
status: "completed",
total: suggestions.length,
suggestedSensitive,
suggestedNonSensitive,
finishedAt: new Date().toISOString(),
errorSummary: null,
});
if (!completed) throw new Error("Sensitive Data Suggestion Run disappeared");
return { suggestions, run: completed };
} catch (error) {
await this.repository.updateSensitiveDataSuggestionRun(started.id, {
status: "failed",
finishedAt: new Date().toISOString(),
errorSummary: failedMessage,
}).catch(() => undefined);
await this.repository.appendSensitiveDataSuggestionEvent(
started.id,
"error",
failedMessage,
).catch(() => undefined);
throw error;
}
}
}
@@ -0,0 +1,172 @@
import type {
SensitivityReviewItem,
} from "./sensitivity-analysis-service.js";
import {
SENSITIVITY_POLICY_VERSION,
SensitivityAnalysisInterruptedError,
SensitivityAnalysisService,
SensitivityAnalysisTargetNotFoundError,
} from "./sensitivity-analysis-service.js";
import type {
CatalogRepository,
SensitivityAnalysisRun,
SensitivityAnalysisScope,
} from "./types.js";
const interruptedMessage = "Local sensitivity analysis was interrupted by backend restart.";
const deadlineMessage = "Local sensitivity analysis reached its time limit.";
const failedMessage = "Local sensitivity analysis failed.";
function ensureActive(signal: AbortSignal): void {
if (signal.aborted) throw new SensitivityAnalysisInterruptedError();
}
export interface SensitivityAnalysisRunResult {
suggestions: readonly SensitivityReviewItem[];
run: SensitivityAnalysisRun;
}
export class SensitivityAnalysisRunner {
constructor(
private readonly repository: CatalogRepository,
private readonly analysis: SensitivityAnalysisService,
) {}
async initialize(): Promise<void> {
if (!(await this.repository.available())) return;
const interrupted = await this.repository.interruptActiveSensitivityAnalysisRuns(
interruptedMessage,
);
for (const run of interrupted) {
await this.repository.appendSensitivityAnalysisEvent(
run.id,
"warning",
interruptedMessage,
);
}
}
async run(
databaseId: string,
scope: SensitivityAnalysisScope,
targetIds: readonly string[],
signal: AbortSignal,
): Promise<SensitivityAnalysisRunResult> {
ensureActive(signal);
const database = await this.repository.get(databaseId);
ensureActive(signal);
if (!database) {
throw new SensitivityAnalysisTargetNotFoundError("database");
}
ensureActive(signal);
const started = await this.repository.createSensitivityAnalysisRun(
databaseId,
scope,
{ engine: "local", policyVersion: SENSITIVITY_POLICY_VERSION },
);
let preparedTotal = 0;
let processedSensitive = 0;
let processedNonSensitive = 0;
try {
ensureActive(signal);
await this.repository.appendSensitivityAnalysisEvent(
started.id,
"info",
"Local sensitivity analysis started.",
);
ensureActive(signal);
const suggestions = await this.analysis.analyze(
databaseId,
scope,
targetIds,
signal,
async (total) => {
ensureActive(signal);
preparedTotal = total;
const prepared = await this.repository.updateSensitivityAnalysisRun(started.id, {
total,
});
ensureActive(signal);
if (!prepared) throw new Error("Sensitivity Analysis Run disappeared");
},
async (processed, batch) => {
ensureActive(signal);
const suggestedSensitive = batch.filter(
(suggestion) => suggestion.assessment === "sensitive",
).length;
const suggestedNonSensitive = batch.filter(
(suggestion) => suggestion.assessment === "non_sensitive",
).length;
const unknown = batch.filter((suggestion) => suggestion.assessment === "unknown").length;
const current = await this.repository.getSensitivityAnalysisRun(started.id);
ensureActive(signal);
if (!current) throw new Error("Sensitivity Analysis Run disappeared");
const progress = await this.repository.updateSensitivityAnalysisRun(started.id, {
suggestedSensitive: current.suggestedSensitive + suggestedSensitive,
suggestedNonSensitive: current.suggestedNonSensitive + suggestedNonSensitive,
unknown: current.unknown + unknown,
});
if (!progress) throw new Error("Sensitivity Analysis Run disappeared");
processedSensitive += suggestedSensitive;
processedNonSensitive += suggestedNonSensitive;
ensureActive(signal);
await this.repository.appendSensitivityAnalysisEvent(
started.id,
"info",
`Assessed ${processed} of ${progress.total} columns locally.`,
);
ensureActive(signal);
},
);
ensureActive(signal);
const suggestedSensitive = suggestions.filter(
(suggestion) => suggestion.assessment === "sensitive",
).length;
const suggestedNonSensitive = suggestions.filter(
(suggestion) => suggestion.assessment === "non_sensitive",
).length;
const unknown = suggestions.filter((suggestion) => suggestion.assessment === "unknown").length;
await this.repository.appendSensitivityAnalysisEvent(
started.id,
"info",
`Local sensitivity analysis completed for ${suggestions.length} column${
suggestions.length === 1 ? "" : "s"
}.`,
);
ensureActive(signal);
const completed = await this.repository.updateSensitivityAnalysisRun(started.id, {
status: "completed",
total: suggestions.length,
suggestedSensitive,
suggestedNonSensitive,
unknown,
finishedAt: new Date().toISOString(),
errorSummary: null,
});
ensureActive(signal);
if (!completed) throw new Error("Sensitivity Analysis Run disappeared");
return { suggestions, run: completed };
} catch (error) {
const interrupted = signal.aborted || error instanceof SensitivityAnalysisInterruptedError;
const message = interrupted ? deadlineMessage : failedMessage;
await this.repository.updateSensitivityAnalysisRun(started.id, {
status: interrupted ? "interrupted" : "failed",
...(interrupted ? {
total: preparedTotal,
suggestedSensitive: processedSensitive,
suggestedNonSensitive: processedNonSensitive,
unknown: Math.max(0, preparedTotal - processedSensitive - processedNonSensitive),
} : {}),
finishedAt: new Date().toISOString(),
errorSummary: message,
}).catch(() => undefined);
await this.repository.appendSensitivityAnalysisEvent(
started.id,
interrupted ? "warning" : "error",
message,
).catch(() => undefined);
throw error;
}
}
}
@@ -0,0 +1,175 @@
import type {
SensitivityClassifier,
SensitivityColumnAssessment,
SensitivityEvidence,
SensitivityNerBudget,
} from "./sensitivity-classifier.js";
import type {
CatalogColumn,
CatalogRepository,
CatalogTable,
SensitivityAnalysisScope,
} from "./types.js";
export type { SensitivityAnalysisScope } from "./types.js";
export const SENSITIVITY_POLICY_VERSION = "sensitivity-v1";
interface SelectedColumn {
table: CatalogTable;
column: CatalogColumn;
}
export interface SensitivityReviewItem {
columnId: string;
tableId: string;
tableName: string;
columnName: string;
version: number;
currentSensitive: boolean;
sensitive: boolean;
assessment: SensitivityColumnAssessment["assessment"];
evidence: readonly SensitivityEvidence[];
observedValues: number;
}
export class SensitivityAnalysisTargetNotFoundError extends Error {
constructor(readonly target: "database" | "table" | "column") {
super(`${target} not found`);
this.name = "SensitivityAnalysisTargetNotFoundError";
}
}
export class SensitivityAnalysisDuplicateTargetIdsError extends Error {
constructor() {
super("sensitivity analysis target IDs must be unique");
this.name = "SensitivityAnalysisDuplicateTargetIdsError";
}
}
export class SensitivityAnalysisInterruptedError extends Error {
constructor() {
super("sensitivity analysis deadline exceeded");
this.name = "SensitivityAnalysisInterruptedError";
}
}
function ensureActive(signal: AbortSignal): void {
if (signal.aborted) throw new SensitivityAnalysisInterruptedError();
}
export class SensitivityAnalysisNoEligibleColumnsError extends Error {
constructor(readonly scope: SensitivityAnalysisScope) {
super("selected scope has no catalog columns");
this.name = "SensitivityAnalysisNoEligibleColumnsError";
}
}
/** Selection and table orchestration around the single SensitivityClassifier decision module. */
export class SensitivityAnalysisService {
constructor(
private readonly repository: CatalogRepository,
private readonly classifier: SensitivityClassifier,
private readonly options: { runBudgetMs?: number; nerBudgetMs?: number; now?: () => number } = {},
) {}
private async selectColumns(
databaseId: string,
scope: SensitivityAnalysisScope,
targetIds: readonly string[],
signal: AbortSignal,
): Promise<readonly SelectedColumn[]> {
ensureActive(signal);
if (new Set(targetIds).size !== targetIds.length) {
throw new SensitivityAnalysisDuplicateTargetIdsError();
}
const tables = await this.repository.listTables(databaseId);
ensureActive(signal);
const tableIds = new Set(targetIds);
const selectedTables = scope === "selected_tables"
? tables.filter((table) => tableIds.has(table.id))
: tables;
if (scope === "selected_tables" && selectedTables.length !== targetIds.length) {
throw new SensitivityAnalysisTargetNotFoundError("table");
}
const columns = (await Promise.all(selectedTables.map(async (table) => (
(await this.repository.listColumns(databaseId, table.id)).map((column) => ({ table, column }))
)))).flat();
ensureActive(signal);
const columnIds = new Set(targetIds);
const selectedColumns = scope === "selected_columns"
? columns.filter(({ column }) => columnIds.has(column.id))
: columns;
if (scope === "selected_columns" && selectedColumns.length !== targetIds.length) {
throw new SensitivityAnalysisTargetNotFoundError("column");
}
if (selectedColumns.length === 0) {
throw new SensitivityAnalysisNoEligibleColumnsError(scope);
}
return selectedColumns;
}
async analyze(
databaseId: string,
scope: SensitivityAnalysisScope,
targetIds: readonly string[],
signal: AbortSignal,
onPrepared?: (total: number) => void | Promise<void>,
onProgress?: (processed: number, suggestions: readonly SensitivityReviewItem[]) => void | Promise<void>,
): Promise<readonly SensitivityReviewItem[]> {
const now = this.options.now ?? Date.now;
const deadline = now() + (this.options.runBudgetMs ?? 60_000);
const configuredNerBudget = this.options.nerBudgetMs ?? 10_000;
const nerBudget: SensitivityNerBudget = {
remainingMs: Number.isFinite(configuredNerBudget) && configuredNerBudget >= 0
? configuredNerBudget
: 10_000,
};
ensureActive(signal);
const database = await this.repository.get(databaseId);
ensureActive(signal);
if (!database) throw new SensitivityAnalysisTargetNotFoundError("database");
const selected = await this.selectColumns(databaseId, scope, targetIds, signal);
await onPrepared?.(selected.length);
ensureActive(signal);
const byTable = new Map<string, SelectedColumn[]>();
for (const item of selected) {
const items = byTable.get(item.table.id) ?? [];
items.push(item);
byTable.set(item.table.id, items);
}
const suggestions: SensitivityReviewItem[] = [];
for (const items of byTable.values()) {
ensureActive(signal);
const first = items[0]!;
const assessments = await this.classifier.assessTable({
database,
table: first.table,
columns: items.map(({ column }) => column),
}, signal, deadline, nerBudget);
ensureActive(signal);
const assessmentById = new Map(assessments.map((assessment) => [
assessment.columnId,
assessment,
]));
const batch = items.map(({ table, column }) => {
const assessment = assessmentById.get(column.id)!;
return {
columnId: column.id,
tableId: table.id,
tableName: table.name,
columnName: column.name,
version: column.version,
currentSensitive: column.sensitive,
sensitive: assessment.proposedSensitive,
assessment: assessment.assessment,
evidence: assessment.evidence,
observedValues: assessment.observedValues,
};
});
suggestions.push(...batch);
await onProgress?.(suggestions.length, batch);
ensureActive(signal);
}
return suggestions;
}
}
@@ -0,0 +1,441 @@
import { CatalogConnectorError, type CatalogColumn, type CatalogTable, type WorkspaceDatabase } from "./types.js";
import { findPhoneNumbersInText } from "libphonenumber-js/max";
import validator from "validator";
export type SensitivityAssessment = "sensitive" | "non_sensitive" | "unknown";
export interface SensitivityEvidence {
kind: "metadata" | "content" | "length" | "ner" | "coverage";
ruleId: string;
label?: string;
confidence?: number;
}
export interface SensitivityValueObservation {
columnId: string;
value: string | null;
characterLength: number | null;
}
export interface SensitivityScanCoverage {
kind: "complete" | "sampled" | "unavailable";
observedRows: number;
}
export interface SensitivityTableScan {
batches: readonly (readonly SensitivityValueObservation[])[];
coverage: SensitivityScanCoverage;
}
export interface SensitivityScanRequest {
database: WorkspaceDatabase;
table: CatalogTable;
columns: readonly CatalogColumn[];
fullScanBudgetMs: number;
deadline: number;
}
export interface SensitivityValueSource {
scanTable(
request: SensitivityScanRequest,
consume: (batch: readonly SensitivityValueObservation[]) => void | Promise<void>,
signal: AbortSignal,
): Promise<SensitivityScanCoverage>;
}
export interface LocalNerCandidate {
columnId: string;
text: string;
}
export interface LocalNerEvidence {
columnId: string;
label: string;
confidence: number;
}
export interface SensitivityNerBudget {
remainingMs: number;
}
/** Optional local detector. It returns evidence only; it never decides a column assessment. */
export interface LocalNerDetector {
warmup?(): Promise<void>;
isReady?(): boolean;
detect(
candidates: readonly LocalNerCandidate[],
signal: AbortSignal,
deadline: number,
): Promise<readonly LocalNerEvidence[]>;
close?(): Promise<void>;
}
export interface SensitivityColumnAssessment {
columnId: string;
assessment: SensitivityAssessment;
proposedSensitive: boolean;
evidence: readonly SensitivityEvidence[];
observedValues: number;
}
export interface SensitivityTableTarget {
database: WorkspaceDatabase;
table: CatalogTable;
columns: readonly CatalogColumn[];
}
const EMAIL = /(?<![\p{L}\p{N}._%+-])[\p{L}\p{N}._%+-]+@[\p{L}\p{N}.-]+\.[\p{L}]{2,63}(?![\p{L}\p{N}._%+-])/giu;
const DIRECT_IDENTIFIER_NAMES = new Set([
"address", "birth_date", "codice_fiscale", "date_of_birth", "dob", "email", "e_mail",
"bic", "first_name", "fiscal_code", "full_name", "iban", "indirizzo", "last_name", "mobile",
"nome", "passport", "phone", "surname", "swift", "swift_code", "tax_id", "telefono",
]);
const CREDENTIAL_NAME = /(?:^|_)(?:api_key|credential|password|passwd|private_key|pwd|secret|token)(?:_|$)/u;
const HEALTH_NAME = /(?:^|_)(?:anamnesi|clinical|diagnos(?:i|is)|health|medical|patient|patologia|therapy|terapia)(?:_|$)/u;
const CLINICAL_TERM = /(?:^|[^\p{L}])(?:allergi[ae]|anamnesi|carcinoma|chemioterapia|diabete|diagnos[ei]|epatite|farmac[io]|gravidanza|hiv|metastasi|neoplasia|patologia|radioterapia|referto|terapia|tumore)(?:$|[^\p{L}])/iu;
const UNSUPPORTED_BINARY_TYPE = /(?:^|\s)(?:binary|blob|bytea|image|varbinary)(?:\s|$|\()/iu;
const MAX_NER_CANDIDATES_PER_REQUEST = 128;
function normalizedName(value: string): string {
return value.normalize("NFKD")
.replace(/[\u0300-\u036f]/g, "")
.replace(/([a-z0-9])([A-Z])/g, "$1_$2")
.toLocaleLowerCase("en-US")
.replace(/[^a-z0-9]+/g, "_")
.replace(/^_+|_+$/g, "");
}
function boundedCount(value: number | undefined, fallback: number, maximum: number): number {
return value === undefined || !Number.isSafeInteger(value)
? fallback
: Math.max(1, Math.min(value, maximum));
}
function metadataEvidence(column: CatalogColumn): SensitivityEvidence | undefined {
const ruleId = sensitiveNameRule(column.name);
return ruleId ? { kind: "metadata", ruleId } : undefined;
}
function sensitiveNameRule(value: string): string | undefined {
const name = normalizedName(value);
if (DIRECT_IDENTIFIER_NAMES.has(name)) {
return "metadata.direct_identifier";
}
if (CREDENTIAL_NAME.test(name)) {
return "metadata.credential";
}
if (HEALTH_NAME.test(name)) {
return "metadata.health";
}
return undefined;
}
const ITALIAN_FISCAL_CODE = /(?<![A-Z0-9])[A-Z]{6}[0-9LMNPQRSTUV]{2}[ABCDEHLMPRST][0-9LMNPQRSTUV]{2}[A-Z][0-9LMNPQRSTUV]{3}[A-Z](?![A-Z0-9])/giu;
const FISCAL_ODD: Record<string, number> = {
"0": 1, "1": 0, "2": 5, "3": 7, "4": 9, "5": 13, "6": 15, "7": 17, "8": 19, "9": 21,
A: 1, B: 0, C: 5, D: 7, E: 9, F: 13, G: 15, H: 17, I: 19, J: 21,
K: 2, L: 4, M: 18, N: 20, O: 11, P: 3, Q: 6, R: 8, S: 12, T: 14,
U: 16, V: 10, W: 22, X: 25, Y: 24, Z: 23,
};
function validItalianFiscalCode(candidate: string): boolean {
const value = candidate.toUpperCase();
if (value.length !== 16) return false;
let sum = 0;
for (let index = 0; index < 15; index += 1) {
const character = value[index]!;
if (index % 2 === 0) sum += FISCAL_ODD[character] ?? -1000;
else sum += /\d/u.test(character) ? Number(character) : character.charCodeAt(0) - 65;
}
return String.fromCharCode(65 + (sum % 26)) === value[15];
}
function validIban(candidate: string): boolean {
const value = candidate.replace(/\s+/gu, "").toUpperCase();
if (!/^[A-Z]{2}\d{2}[A-Z0-9]{11,30}$/u.test(value)) return false;
const rearranged = value.slice(4) + value.slice(0, 4);
let remainder = 0;
for (const character of rearranged) {
const digits = /\d/u.test(character) ? character : String(character.charCodeAt(0) - 55);
for (const digit of digits) remainder = (remainder * 10 + Number(digit)) % 97;
}
return remainder === 1;
}
function validPaymentCard(candidate: string): boolean {
const digits = candidate.replace(/[ -]/gu, "");
if (!/^\d{13,19}$/u.test(digits) || /^(\d)\1+$/u.test(digits)) return false;
let sum = 0;
let double = false;
for (let index = digits.length - 1; index >= 0; index -= 1) {
let digit = Number(digits[index]);
if (double) {
digit *= 2;
if (digit > 9) digit -= 9;
}
sum += digit;
double = !double;
}
return sum % 10 === 0;
}
function jsonHasSensitiveKey(value: string): boolean {
const trimmed = value.trim();
if (!(trimmed.startsWith("{") || trimmed.startsWith("["))) return false;
try {
const pending: Array<{ value: unknown; depth: number }> = [{ value: JSON.parse(trimmed), depth: 0 }];
let visited = 0;
while (pending.length > 0 && visited < 1_000) {
const item = pending.pop()!;
visited += 1;
if (item.depth > 8 || item.value === null || typeof item.value !== "object") continue;
if (Array.isArray(item.value)) {
for (const child of item.value) pending.push({ value: child, depth: item.depth + 1 });
continue;
}
for (const [key, child] of Object.entries(item.value)) {
if (sensitiveNameRule(key)) return true;
pending.push({ value: child, depth: item.depth + 1 });
}
}
} catch {
return false;
}
return false;
}
function contentEvidence(value: string): SensitivityEvidence | undefined {
if (/-----BEGIN (?:[A-Z0-9]+ )?PRIVATE KEY-----/u.test(value)) {
return { kind: "content", ruleId: "credential.private_key" };
}
if (/(?:^|[^A-Z0-9])AKIA[A-Z0-9]{16}(?![A-Z0-9])/u.test(value)
|| /(?:^|[^A-Za-z0-9_])gh[pousr]_[A-Za-z0-9_]{30,}(?![A-Za-z0-9_])/u.test(value)
|| /(?:^|[^A-Za-z0-9_-])eyJ[A-Za-z0-9_-]{5,}\.[A-Za-z0-9_-]{5,}\.[A-Za-z0-9_-]{5,}(?![A-Za-z0-9_-])/u.test(value)) {
return { kind: "content", ruleId: "credential.access_key" };
}
if (/(?:^|[^\p{L}\p{N}_])(?:api[_ -]?key|access[_ -]?token|password|passwd|pwd|secret)\s*[:=]\s*[^\s,;]{4,}/iu.test(value)) {
return { kind: "content", ruleId: "credential.key_value" };
}
if (CLINICAL_TERM.test(value)) return { kind: "content", ruleId: "health.clinical_term" };
for (const match of value.matchAll(EMAIL)) {
if (validator.isEmail(match[0])) return { kind: "content", ruleId: "pii.email" };
}
for (const match of value.matchAll(ITALIAN_FISCAL_CODE)) {
if (validItalianFiscalCode(match[0])) {
return { kind: "content", ruleId: "pii.italian_fiscal_code" };
}
}
for (const match of value.matchAll(/\b(?:passaporto|passport)(?:\s+(?:numero|number|n\.?))?\s*[:#-]?\s*([A-Z0-9]{9})\b/giu)) {
if (validator.isPassportNumber(match[1]!, "IT")) {
return { kind: "content", ruleId: "pii.passport_number" };
}
}
for (const match of value.matchAll(/\bC[A-Z]\d{5}[A-Z]{2}\b/giu)) {
if (validator.isIdentityCard(match[0], "IT")) {
return { kind: "content", ruleId: "pii.identity_card" };
}
}
if (/\b(?:patente(?:\s+di\s+guida)?|driving\s+licen[cs]e)(?:\s+(?:numero|number|n\.?))?\s*[:#-]?\s*[A-Z0-9]{8,12}\b/iu.test(value)) {
return { kind: "content", ruleId: "pii.drivers_license_number" };
}
for (const match of value.matchAll(/(?<![A-Z0-9])[A-Z]{2}\d{2}(?:\s?[A-Z0-9]){11,30}(?![A-Z0-9])/giu)) {
if (validIban(match[0])) return { kind: "content", ruleId: "financial.iban" };
}
for (const match of value.matchAll(/(?<!\d)(?:\d[ -]?){13,19}(?!\d)/gu)) {
if (validPaymentCard(match[0])) {
return { kind: "content", ruleId: "financial.payment_card" };
}
}
for (const match of value.matchAll(/(?<![A-Z0-9])[A-Z]{6}[A-Z0-9]{2}(?:[A-Z0-9]{3})?(?![A-Z0-9])/giu)) {
const before = value.slice(Math.max(0, (match.index ?? 0) - 24), match.index ?? 0);
if (/\b(?:bic|swift)\s*[:=-]?\s*$/iu.test(before) && validator.isBIC(match[0])) {
return { kind: "content", ruleId: "financial.bic" };
}
}
for (const match of value.matchAll(/(?<!\d)(?:IT[ .-]?)?\d{11}(?!\d)/giu)) {
const candidate = match[0].replace(/[ .-]/gu, "");
if (validator.isVAT(candidate.replace(/^IT/iu, ""), "IT")) {
return { kind: "content", ruleId: "pii.italian_vat" };
}
}
for (const match of value.matchAll(/(?<![A-F0-9])(?:[A-F0-9]{2}[:-]){5}[A-F0-9]{2}(?![A-F0-9])/giu)) {
if (validator.isMACAddress(match[0])) {
return { kind: "content", ruleId: "network.mac_address" };
}
}
for (const match of value.matchAll(/(?<![A-F0-9:.])[A-F0-9:.]{3,45}(?![A-F0-9:.])/giu)) {
if (validator.isIP(match[0])) return { kind: "content", ruleId: "network.ip_address" };
}
for (const match of value.matchAll(/(?<![A-F0-9-])[0-9A-F]{8}-[0-9A-F]{4}-[1-8][0-9A-F]{3}-[89AB][0-9A-F]{3}-[0-9A-F]{12}(?![A-F0-9-])/giu)) {
if (validator.isUUID(match[0])) return { kind: "content", ruleId: "pii.uuid" };
}
for (const match of value.matchAll(/\b(?:https?|ftp):\/\/[^\s<>"']+/giu)) {
const candidate = match[0].replace(/[.,;:!?\])}]+$/u, "");
if (validator.isURL(candidate, { require_protocol: true })) {
return { kind: "content", ruleId: "network.url" };
}
}
if (findPhoneNumbersInText(value, "IT").some((match) => match.number.isValid())) {
return { kind: "content", ruleId: "pii.phone_number" };
}
if (jsonHasSensitiveKey(value)) {
return { kind: "content", ruleId: "pii.json_sensitive_key" };
}
return undefined;
}
/** Sole decision module for local column-level sensitivity assessments. */
export class SensitivityClassifier {
constructor(
private readonly values: SensitivityValueSource,
private readonly detector?: LocalNerDetector,
private readonly options: {
fullScanBudgetMs?: number;
runBudgetMs?: number;
nerConfidenceThreshold?: number;
maxNerValuesPerColumn?: number;
maxNerCandidatesPerTable?: number;
now?: () => number;
} = {},
) {}
async assessTable(
target: SensitivityTableTarget,
signal: AbortSignal,
runDeadline?: number,
nerBudget?: SensitivityNerBudget,
): Promise<readonly SensitivityColumnAssessment[]> {
const now = this.options.now ?? Date.now;
const deadline = runDeadline ?? now() + (this.options.runBudgetMs ?? 60_000);
const evidence = new Map(target.columns.map((column) => {
const match = metadataEvidence(column);
return [column.id, match ? [match] : [] as SensitivityEvidence[]];
}));
const observed = new Map(target.columns.map((column) => [column.id, 0]));
const nerCandidates = new Map(target.columns.map((column) => [column.id, [] as string[]]));
const maxNerValuesPerColumn = boundedCount(this.options.maxNerValuesPerColumn, 8, 8);
const unsupported = new Set(target.columns
.filter((column) => UNSUPPORTED_BINARY_TYPE.test(column.dataType))
.map((column) => column.id));
const scannableColumns = target.columns.filter((column) => (
!unsupported.has(column.id) && evidence.get(column.id)!.length === 0
));
let coverage: SensitivityScanCoverage = { kind: "unavailable", observedRows: 0 };
if (scannableColumns.length > 0 && now() < deadline) {
try {
coverage = await this.values.scanTable({
...target,
columns: scannableColumns,
fullScanBudgetMs: this.options.fullScanBudgetMs ?? 5_000,
deadline,
}, (batch) => {
for (const item of batch) {
if (!evidence.has(item.columnId) || item.value === null) continue;
observed.set(item.columnId, (observed.get(item.columnId) ?? 0) + 1);
const matches = evidence.get(item.columnId)!;
if (matches.length === 0 && (item.characterLength ?? item.value.length) > 500) {
matches.push({ kind: "length", ruleId: "text.over_500_characters" });
} else if (matches.length === 0) {
const match = contentEvidence(item.value);
if (match) matches.push(match);
else {
const candidates = nerCandidates.get(item.columnId)!;
if (candidates.length < maxNerValuesPerColumn && !candidates.includes(item.value)) {
candidates.push(item.value);
}
}
}
}
}, signal);
} catch (error) {
if (!(error instanceof CatalogConnectorError)) throw error;
}
}
if (this.detector && (this.detector.isReady?.() ?? true) && !signal.aborted
&& now() < deadline && (nerBudget?.remainingMs ?? 1) > 0) {
const candidates: LocalNerCandidate[] = [];
const maxCandidates = boundedCount(this.options.maxNerCandidatesPerTable, 2, 1_024);
candidateSelection: for (let valueIndex = 0; valueIndex < maxNerValuesPerColumn; valueIndex += 1) {
for (const column of target.columns) {
if (evidence.get(column.id)!.length > 0) continue;
const text = nerCandidates.get(column.id)![valueIndex];
if (text === undefined) continue;
candidates.push({ columnId: column.id, text });
if (candidates.length >= maxCandidates) break candidateSelection;
}
}
if (candidates.length > 0) {
const threshold = this.options.nerConfidenceThreshold ?? 0.8;
const nerStartedAt = now();
const allowedNerMs = nerBudget
? Math.max(0, nerBudget.remainingMs)
: Math.max(0, deadline - nerStartedAt);
const nerDeadline = Math.min(deadline, nerStartedAt + allowedNerMs);
try {
for (let offset = 0; offset < candidates.length; offset += MAX_NER_CANDIDATES_PER_REQUEST) {
if (signal.aborted || now() >= nerDeadline) break;
try {
const detected = await this.detector.detect(
candidates.slice(offset, offset + MAX_NER_CANDIDATES_PER_REQUEST),
signal,
nerDeadline,
);
for (const item of detected) {
const matches = evidence.get(item.columnId);
if (!matches || matches.length > 0 || !Number.isFinite(item.confidence)
|| item.confidence < threshold || item.confidence > 1) continue;
const label = normalizedName(item.label).slice(0, 80);
if (!label) continue;
matches.push({
kind: "ner",
ruleId: "ner.entity",
label,
confidence: item.confidence,
});
}
} catch {
// NER is optional: deterministic findings and scan coverage remain authoritative.
break;
}
}
} finally {
if (nerBudget) {
const elapsedMs = Math.max(1, now() - nerStartedAt);
nerBudget.remainingMs = Math.max(0, nerBudget.remainingMs - elapsedMs);
}
}
}
}
return target.columns.map((column) => {
const matches = evidence.get(column.id)!;
const count = observed.get(column.id) ?? 0;
const assessment: SensitivityAssessment = matches.length > 0
? "sensitive"
: unsupported.has(column.id) || count === 0 || coverage.kind !== "complete"
? "unknown"
: "non_sensitive";
return {
columnId: column.id,
assessment,
proposedSensitive: assessment === "unknown" ? column.sensitive : assessment === "sensitive",
evidence: matches.length > 0
? matches
: assessment === "unknown"
? [{
kind: "coverage",
ruleId: unsupported.has(column.id)
? "coverage.unsupported_type"
: coverage.kind === "unavailable"
? "coverage.unavailable"
: count === 0
? "coverage.no_values"
: "coverage.incomplete",
}]
: [],
observedValues: count,
};
});
}
}
+94
View File
@@ -0,0 +1,94 @@
import { dirname } from "node:path";
import { fileURLToPath } from "node:url";
import { loadConfig } from "../config.js";
import { WorkspaceSecretStore } from "../workspaces/secret-store.js";
import { PythonLocalNerDetector } from "./local-ner-detector.js";
import { ConcreteCatalogPostgresAccess } from "./postgres-access.js";
import { createCatalogRepository } from "./repository.js";
import { SensitivityAnalysisService } from "./sensitivity-analysis-service.js";
import { SensitivityClassifier } from "./sensitivity-classifier.js";
import { ConcreteSensitivityValueSource } from "./sensitivity-value-source.js";
const WORKSPACE_ID = /^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$/u;
async function main(): Promise<void> {
const workspaceId = process.argv[2];
if (!workspaceId || !WORKSPACE_ID.test(workspaceId)) {
process.stderr.write("Usage: sensitivity-shadow <workspace-id>\n");
process.exitCode = 2;
return;
}
let detector: PythonLocalNerDetector | undefined;
let stage = "configuration";
try {
const config = loadConfig(process.env);
stage = "catalog";
const repository = createCatalogRepository(config.catalogDatabase);
if (!(await repository.available())) throw new Error("catalog unavailable");
const database = await repository.getByWorkspace(workspaceId);
if (!database) throw new Error("database unavailable");
stage = "source";
const secretStore = new WorkspaceSecretStore({
root: config.workspaceSecretStoreRoot,
runtimeRoot: config.workspaceSecretRuntimeRoot,
installationId: config.workspaceRegistry.installationId,
});
const access = new ConcreteCatalogPostgresAccess(secretStore, {
connectTimeoutMs: config.workspaceDiagnosticTimeoutMs,
});
const source = new ConcreteSensitivityValueSource(access, secretStore);
if (config.sensitivityNer) {
const workerScript = config.sensitivityNer.workerScript
?? fileURLToPath(new URL("../../python/sensitivity_ner_worker.py", import.meta.url));
detector = new PythonLocalNerDetector({
pythonExecutable: config.sensitivityNer.pythonExecutable,
workerScript,
modelPath: config.sensitivityNer.modelPath,
cwd: dirname(workerScript),
threads: config.sensitivityNer.threads,
});
try {
await detector.warmup();
} catch {
await detector.close();
detector = undefined;
}
}
const startedAt = Date.now();
stage = "analysis";
const suggestions = await new SensitivityAnalysisService(
repository,
new SensitivityClassifier(source, detector),
).analyze(database.id, "all", [], AbortSignal.timeout(65_000));
const assessments = { sensitive: 0, nonSensitive: 0, unknown: 0 };
const rules = new Map<string, number>();
for (const suggestion of suggestions) {
if (suggestion.assessment === "sensitive") assessments.sensitive += 1;
else if (suggestion.assessment === "non_sensitive") assessments.nonSensitive += 1;
else assessments.unknown += 1;
for (const evidence of suggestion.evidence) {
rules.set(evidence.ruleId, (rules.get(evidence.ruleId) ?? 0) + 1);
}
}
process.stdout.write(`${JSON.stringify({
ok: true,
policyVersion: "sensitivity-v1",
nerEnabled: detector !== undefined,
total: suggestions.length,
assessments,
rules: Object.fromEntries([...rules].sort(([left], [right]) => left.localeCompare(right))),
elapsedMs: Date.now() - startedAt,
})}\n`);
} catch {
process.stdout.write(`${JSON.stringify({
ok: false,
code: `sensitivity_shadow_${stage}_failed`,
})}\n`);
process.exitCode = 1;
} finally {
await detector?.close();
}
}
await main();
@@ -0,0 +1,259 @@
import { readFile } from "node:fs/promises";
import type { WorkspaceSecretStore } from "../workspaces/secret-store.js";
import { CATALOG_SECRET_IDS } from "./secrets.js";
import type { CatalogPostgresAccess } from "./postgres-access.js";
import type {
SensitivityScanCoverage,
SensitivityScanRequest,
SensitivityValueObservation,
SensitivityValueSource,
} from "./sensitivity-classifier.js";
import { CatalogConnectorError } from "./types.js";
const MAX_VALUE_CHARACTERS = 501;
const DEFAULT_BATCH_ROWS = 200;
const DEFAULT_SAMPLE_ROWS = 200;
function quoteIdentifier(identifier: string): string {
return `"${identifier.replaceAll('"', '""')}"`;
}
function projections(request: SensitivityScanRequest): string {
return request.columns.flatMap((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}"`,
];
}).join(", ");
}
function observations(
request: SensitivityScanRequest,
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
? null
: Number(sourceLength);
return {
columnId: column.id,
value,
characterLength: parsedLength !== null && Number.isSafeInteger(parsedLength) && parsedLength >= 0
? parsedLength
: value?.length ?? null,
};
}));
}
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.
*/
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.database.binding.transport === "rest_api") {
return await this.scanRest(request, consume, signal);
}
return await this.scanPostgres(request, consume, signal);
}
private async scanPostgres(
request: SensitivityScanRequest,
consume: (batch: readonly SensitivityValueObservation[]) => void | Promise<void>,
signal: AbortSignal,
): 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;
try {
if (signal.aborted || this.now() >= request.deadline) {
return { kind: "sampled", observedRows: 0 };
}
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`,
]);
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>>;
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;
} 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 };
}
}
if (signal.aborted || this.now() >= request.deadline) {
return { kind: "sampled", observedRows };
}
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 };
} catch (error) {
if (error instanceof CatalogConnectorError) throw error;
throw new CatalogConnectorError("Sensitivity source scan failed");
} finally {
if (transactionOpen) await client.query("ROLLBACK", []).catch(() => undefined);
await client.end().catch(() => undefined);
}
}
private async scanRest(
request: SensitivityScanRequest,
consume: (batch: readonly SensitivityValueObservation[]) => void | Promise<void>,
signal: AbortSignal,
): Promise<SensitivityScanCoverage> {
if (!this.secretStore) throw new CatalogConnectorError("REST sensitivity scanning is not configured");
const auth = request.database.binding.restAuth ?? "bearer";
const materialized = this.secretStore.materialize(
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;
try {
const headers: Record<string, string> = { "content-type": "application/json" };
if (auth !== "none") {
const credentialFile = materialized.files.get(CATALOG_SECRET_IDS.apiKey);
if (!credentialFile) throw new CatalogConnectorError("REST API key is not configured");
const credential = (await readFile(credentialFile, "utf8")).trim();
if (auth === "bearer") headers.authorization = `Bearer ${credential}`;
else headers["x-api-key"] = credential;
}
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");
}
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;
}
if (signal.aborted || this.now() >= request.deadline) {
return { kind: "sampled", observedRows };
}
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 };
} catch (error) {
if (error instanceof CatalogConnectorError) throw error;
throw new CatalogConnectorError("REST sensitivity source scan failed");
} finally {
materialized.release();
}
}
}
+29 -25
View File
@@ -266,18 +266,21 @@ export interface DescriptionGenerationEvent {
createdAt: string;
}
export type SensitiveDataSuggestionScope = "all" | "selected_tables" | "selected_columns";
export type SensitiveDataSuggestionStatus = "running" | "completed" | "failed" | "interrupted";
export type SensitivityAnalysisScope = "all" | "selected_tables" | "selected_columns";
export type SensitivityAnalysisStatus = "running" | "completed" | "failed" | "interrupted";
export interface SensitiveDataSuggestionRun {
export interface SensitivityAnalysisRun {
id: string;
databaseId: string;
scope: SensitiveDataSuggestionScope;
modelId: string;
status: SensitiveDataSuggestionStatus;
scope: SensitivityAnalysisScope;
engine: "llm" | "local";
modelId: string | null;
policyVersion: string | null;
status: SensitivityAnalysisStatus;
total: number;
suggestedSensitive: number;
suggestedNonSensitive: number;
unknown: number;
inputTokens: number;
cacheReadTokens: number;
outputTokens: number;
@@ -288,11 +291,12 @@ export interface SensitiveDataSuggestionRun {
errorSummary: string | null;
}
export interface SensitiveDataSuggestionRunUpdate {
status?: SensitiveDataSuggestionStatus;
export interface SensitivityAnalysisRunUpdate {
status?: SensitivityAnalysisStatus;
total?: number;
suggestedSensitive?: number;
suggestedNonSensitive?: number;
unknown?: number;
finishedAt?: string | null;
errorSummary?: string | null;
inputTokens?: number;
@@ -300,7 +304,7 @@ export interface SensitiveDataSuggestionRunUpdate {
outputTokens?: number;
}
export interface SensitiveDataSuggestionEvent {
export interface SensitivityAnalysisEvent {
runId: string;
sequence: number;
level: "info" | "warning" | "error";
@@ -491,29 +495,29 @@ export interface CatalogRepository {
runId: string,
afterSequence?: number,
): Promise<DescriptionGenerationEvent[]>;
createSensitiveDataSuggestionRun(
createSensitivityAnalysisRun(
databaseId: string,
scope: SensitiveDataSuggestionScope,
modelId: string,
): Promise<SensitiveDataSuggestionRun>;
getSensitiveDataSuggestionRun(runId: string): Promise<SensitiveDataSuggestionRun | undefined>;
listSensitiveDataSuggestionRuns(limit?: number): Promise<SensitiveDataSuggestionRun[]>;
interruptActiveSensitiveDataSuggestionRuns(
scope: SensitivityAnalysisScope,
origin: { engine: "llm"; modelId: string } | { engine: "local"; policyVersion: string },
): Promise<SensitivityAnalysisRun>;
getSensitivityAnalysisRun(runId: string): Promise<SensitivityAnalysisRun | undefined>;
listSensitivityAnalysisRuns(limit?: number): Promise<SensitivityAnalysisRun[]>;
interruptActiveSensitivityAnalysisRuns(
errorSummary: string,
): Promise<SensitiveDataSuggestionRun[]>;
updateSensitiveDataSuggestionRun(
): Promise<SensitivityAnalysisRun[]>;
updateSensitivityAnalysisRun(
runId: string,
update: SensitiveDataSuggestionRunUpdate,
): Promise<SensitiveDataSuggestionRun | undefined>;
appendSensitiveDataSuggestionEvent(
update: SensitivityAnalysisRunUpdate,
): Promise<SensitivityAnalysisRun | undefined>;
appendSensitivityAnalysisEvent(
runId: string,
level: SensitiveDataSuggestionEvent["level"],
level: SensitivityAnalysisEvent["level"],
message: string,
): Promise<SensitiveDataSuggestionEvent>;
listSensitiveDataSuggestionEvents(
): Promise<SensitivityAnalysisEvent>;
listSensitivityAnalysisEvents(
runId: string,
afterSequence?: number,
): Promise<SensitiveDataSuggestionEvent[]>;
): Promise<SensitivityAnalysisEvent[]>;
listRelationships(databaseId: string): Promise<CatalogPhysicalRelationship[]>;
listLogicalRelationships(databaseId: string): Promise<CatalogLogicalRelationship[]>;
getLogicalRelationshipContext(databaseId: string): Promise<CatalogLogicalRelationshipContext | undefined>;
+36
View File
@@ -33,6 +33,12 @@ export interface AppConfig {
secretsFile?: string;
installationConfigFile?: string;
modelCatalogFile?: string;
sensitivityNer?: {
pythonExecutable: string;
modelPath: string;
workerScript?: string;
threads: number;
};
piAuthFile?: string;
secretFiles: Readonly<Record<string, string | undefined>>;
modelApiKeyFile?: string;
@@ -354,6 +360,35 @@ export function loadConfig(
|| modelCatalogFile.includes("\0")
|| !path.isAbsolute(modelCatalogFile)
)) throw new Error("runtime model catalog configuration is invalid");
const sensitivityNerModelPath = env.THT_SENSITIVITY_NER_MODEL_PATH;
const sensitivityNerPython = env.THT_SENSITIVITY_NER_PYTHON;
const sensitivityNerWorker = env.THT_SENSITIVITY_NER_WORKER;
for (const [value, label] of [
[sensitivityNerModelPath, "model path"],
[sensitivityNerPython, "Python executable"],
[sensitivityNerWorker, "worker path"],
] as const) {
if (value !== undefined && (
value.length === 0 || value.trim() !== value || value.includes("\0") || !path.isAbsolute(value)
)) throw new Error(`sensitivity NER ${label} configuration is invalid`);
}
if (sensitivityNerModelPath === undefined && (
sensitivityNerPython !== undefined
|| sensitivityNerWorker !== undefined
|| env.THT_SENSITIVITY_NER_THREADS !== undefined
)) throw new Error("sensitivity NER settings require a model path");
const sensitivityNerThreads = Number(env.THT_SENSITIVITY_NER_THREADS ?? 2);
if (!Number.isSafeInteger(sensitivityNerThreads) || sensitivityNerThreads < 1 || sensitivityNerThreads > 8) {
throw new Error("sensitivity NER thread configuration is invalid");
}
const sensitivityNer = sensitivityNerModelPath === undefined
? undefined
: {
modelPath: sensitivityNerModelPath,
pythonExecutable: sensitivityNerPython ?? "/opt/sensitivity-ner/bin/python",
...(sensitivityNerWorker ? { workerScript: sensitivityNerWorker } : {}),
threads: sensitivityNerThreads,
};
const piAuthFile = env.THT_PI_AUTH_FILE;
if (piAuthFile !== undefined && (
piAuthFile.trim() !== piAuthFile || piAuthFile.length === 0 || piAuthFile.includes("\0")
@@ -442,6 +477,7 @@ export function loadConfig(
secretsFile,
installationConfigFile,
modelCatalogFile,
sensitivityNer,
piAuthFile,
secretFiles,
modelApiKeyFile,
@@ -11,38 +11,35 @@ import {
type DescriptionGenerationWorker,
} from "../catalog/description-generation-worker.js";
import { MetadataGenerationModelUnavailableError } from "../catalog/metadata-generation-models.js";
import { ModelCompletionProviderError } from "../catalog/model-completer.js";
import {
SensitiveDataSuggestionDuplicateTargetIdsError,
SensitiveDataSuggestionInvalidResponseError,
SensitiveDataSuggestionNoEligibleColumnsError,
SensitiveDataSuggestionPayloadTooLargeError,
SensitiveDataSuggestionTargetNotFoundError,
} from "../catalog/sensitive-data-suggester.js";
import type { SensitiveDataSuggestionRunner } from "../catalog/sensitive-data-suggestion-runner.js";
SensitivityAnalysisDuplicateTargetIdsError,
SensitivityAnalysisInterruptedError,
SensitivityAnalysisNoEligibleColumnsError,
SensitivityAnalysisTargetNotFoundError,
} from "../catalog/sensitivity-analysis-service.js";
import type { SensitivityAnalysisRunner } from "../catalog/sensitivity-analysis-runner.js";
import {
CatalogOperationInProgressError,
CatalogConnectorError,
CatalogUnavailableError,
DescriptionGenerationRunActiveError,
type CatalogRepository,
type DescriptionGenerationEvent,
type DescriptionGenerationRun,
type SensitiveDataSuggestionEvent,
type SensitiveDataSuggestionRun,
type SensitivityAnalysisEvent,
type SensitivityAnalysisRun,
} from "../catalog/types.js";
const idSchema = z.uuid();
const modelIdSchema = z.string().regex(/^[a-z][a-z0-9._-]{0,63}\/[A-Za-z0-9][A-Za-z0-9._:-]{0,255}$/);
const selectedTargetIdsSchema = z.array(idSchema).min(1);
const suggestionSchema = z.discriminatedUnion("scope", [
z.object({ modelId: modelIdSchema, scope: z.literal("all") }).strict(),
z.object({ scope: z.literal("all") }).strict(),
z.object({
modelId: modelIdSchema,
scope: z.literal("selected_tables"),
targetIds: selectedTargetIdsSchema,
}).strict(),
z.object({
modelId: modelIdSchema,
scope: z.literal("selected_columns"),
targetIds: selectedTargetIdsSchema,
}).strict(),
@@ -112,7 +109,7 @@ function publicRun(run: DescriptionGenerationRun) {
};
}
function publicSensitiveDataSuggestionEvent(event: SensitiveDataSuggestionEvent) {
function publicSensitivityAnalysisEvent(event: SensitivityAnalysisEvent) {
return {
runId: event.runId,
sequence: event.sequence,
@@ -122,16 +119,19 @@ function publicSensitiveDataSuggestionEvent(event: SensitiveDataSuggestionEvent)
};
}
function publicSensitiveDataSuggestionRun(run: SensitiveDataSuggestionRun) {
function publicSensitivityAnalysisRun(run: SensitivityAnalysisRun) {
return {
id: run.id,
databaseId: run.databaseId,
scope: run.scope,
engine: run.engine,
modelId: run.modelId,
policyVersion: run.policyVersion,
status: run.status,
total: run.total,
suggestedSensitive: run.suggestedSensitive,
suggestedNonSensitive: run.suggestedNonSensitive,
unknown: run.unknown,
inputTokens: run.inputTokens,
cacheReadTokens: run.cacheReadTokens,
outputTokens: run.outputTokens,
@@ -232,16 +232,10 @@ function safeSuggestionError(reply: FastifyReply, error: unknown) {
if (error instanceof CatalogUnavailableError) {
return reply.code(503).send({
code: "catalog_unavailable",
message: "The database catalog is unavailable, so no sensitive-field suggestions were prepared.",
message: "The database catalog is unavailable, so no sensitivity assessments were prepared.",
});
}
if (error instanceof MetadataGenerationModelUnavailableError) {
return reply.code(409).send({
code: "metadata_generation_model_unavailable",
message: "The selected metadata-generation model is unavailable.",
});
}
if (error instanceof SensitiveDataSuggestionTargetNotFoundError) {
if (error instanceof SensitivityAnalysisTargetNotFoundError) {
const code = error.target === "database"
? "database_not_found"
: error.target === "table"
@@ -254,45 +248,65 @@ function safeSuggestionError(reply: FastifyReply, error: unknown) {
: "One or more selected Catalog Columns were not found in this database.";
return reply.code(404).send({ code, message });
}
if (error instanceof SensitiveDataSuggestionDuplicateTargetIdsError) {
if (error instanceof SensitivityAnalysisDuplicateTargetIdsError) {
return reply.code(400).send({
code: "sensitive_data_suggestion_target_ids_duplicate",
message: "Each selected table or column must appear only once.",
});
}
if (error instanceof SensitiveDataSuggestionNoEligibleColumnsError) {
if (error instanceof SensitivityAnalysisNoEligibleColumnsError) {
return reply.code(409).send({
code: "sensitive_data_suggestion_no_columns",
message: "The selected scope contains no Catalog Columns to classify.",
message: "The selected scope contains no Catalog Columns to assess.",
});
}
if (error instanceof SensitiveDataSuggestionPayloadTooLargeError) {
return reply.code(413).send({
code: "sensitive_data_suggestion_payload_too_large",
message: "The selected structural metadata cannot be divided into safe LLM requests.",
if (error instanceof SensitivityAnalysisInterruptedError) {
return reply.code(504).send({
code: "sensitivity_analysis_timeout",
message: "Sensitivity analysis reached its time limit. No assessments were applied.",
});
}
if (error instanceof SensitiveDataSuggestionInvalidResponseError) {
if (error instanceof CatalogConnectorError) {
return reply.code(502).send({
code: "sensitive_data_suggestion_invalid_response",
message: "The LLM returned an incomplete or invalid classification. No suggestions were applied.",
});
}
if (error instanceof ModelCompletionProviderError) {
return reply.code(502).send({
code: "sensitive_data_suggestion_provider_unavailable",
message: "The selected LLM service could not complete the request. No suggestions were applied.",
code: "sensitivity_source_unavailable",
message: "The source values could not be inspected safely. No assessments were applied.",
});
}
if (error instanceof z.ZodError) {
return reply.code(400).send({
code: "sensitive_data_suggestion_request_invalid",
message: "Choose a database, one or more tables, or one or more columns to classify.",
message: "Choose a database, one or more tables, or one or more columns to assess.",
});
}
return reply.code(500).send({
code: "sensitive_data_suggestion_failed",
message: "Sensitive-field suggestions failed before review. No changes were applied.",
message: "Local sensitivity analysis failed before review. No changes were applied.",
});
}
function untilAborted<T>(operation: Promise<T>, signal: AbortSignal): Promise<T> {
if (signal.aborted) {
void operation.catch(() => undefined);
return Promise.reject(new SensitivityAnalysisInterruptedError());
}
return new Promise<T>((resolve, reject) => {
const abort = () => reject(new SensitivityAnalysisInterruptedError());
signal.addEventListener("abort", abort, { once: true });
if (signal.aborted) {
void operation.catch(() => undefined);
abort();
return;
}
operation.then(
(value) => {
signal.removeEventListener("abort", abort);
resolve(value);
},
(error: unknown) => {
signal.removeEventListener("abort", abort);
reject(error);
},
);
});
}
@@ -300,18 +314,18 @@ function safeSuggestionHistoryError(reply: FastifyReply, error: unknown) {
if (error instanceof CatalogUnavailableError) {
return reply.code(503).send({
code: "catalog_unavailable",
message: "Sensitive Data Suggestion history is unavailable because the database catalog is unavailable.",
message: "Sensitivity Analysis history is unavailable because the database catalog is unavailable.",
});
}
if (error instanceof z.ZodError) {
return reply.code(400).send({
code: "sensitive_data_suggestion_history_request_invalid",
message: "Sensitive Data Suggestion history parameters are invalid.",
message: "Sensitivity Analysis history parameters are invalid.",
});
}
return reply.code(500).send({
code: "sensitive_data_suggestion_history_failed",
message: "Sensitive Data Suggestion history could not be loaded.",
message: "Sensitivity Analysis history could not be loaded.",
});
}
@@ -320,7 +334,7 @@ export function catalogDescriptionGenerationRoutes(
deps: {
repository: CatalogRepository;
worker: DescriptionGenerationWorker;
sensitiveDataSuggestionRunner: SensitiveDataSuggestionRunner;
sensitivityAnalysisRunner: SensitivityAnalysisRunner;
},
): void {
app.post("/catalog/databases/:databaseId/sensitive-data-suggestions", async (request, reply) => {
@@ -328,16 +342,16 @@ export function catalogDescriptionGenerationRoutes(
try {
const databaseId = idSchema.parse((request.params as { databaseId?: unknown }).databaseId);
const input = suggestionSchema.parse(request.body);
const result = await deps.sensitiveDataSuggestionRunner.run(
const signal = AbortSignal.timeout(60_000);
const result = await untilAborted(deps.sensitivityAnalysisRunner.run(
databaseId,
input.modelId,
input.scope,
"targetIds" in input ? input.targetIds : [],
new AbortController().signal,
);
signal,
), signal);
return {
suggestions: result.suggestions,
run: publicSensitiveDataSuggestionRun(result.run),
run: publicSensitivityAnalysisRun(result.run),
};
} catch (error) {
return safeSuggestionError(reply, error);
@@ -348,8 +362,8 @@ export function catalogDescriptionGenerationRoutes(
if (!manage(request, reply)) return reply;
try {
const { limit } = historyQuerySchema.parse(request.query);
return (await deps.repository.listSensitiveDataSuggestionRuns(limit))
.map(publicSensitiveDataSuggestionRun);
return (await deps.repository.listSensitivityAnalysisRuns(limit))
.map(publicSensitivityAnalysisRun);
} catch (error) {
return safeSuggestionHistoryError(reply, error);
}
@@ -359,12 +373,12 @@ export function catalogDescriptionGenerationRoutes(
if (!manage(request, reply)) return reply;
try {
const runId = idSchema.parse((request.params as { runId?: unknown }).runId);
const run = await deps.repository.getSensitiveDataSuggestionRun(runId);
const run = await deps.repository.getSensitivityAnalysisRun(runId);
if (!run) return reply.code(404).send({
code: "sensitive_data_suggestion_run_not_found",
message: "Sensitive Data Suggestion Run was not found.",
message: "Sensitivity Analysis Run was not found.",
});
return publicSensitiveDataSuggestionRun(run);
return publicSensitivityAnalysisRun(run);
} catch (error) {
return safeSuggestionHistoryError(reply, error);
}
@@ -375,14 +389,14 @@ export function catalogDescriptionGenerationRoutes(
try {
const runId = idSchema.parse((request.params as { runId?: unknown }).runId);
const { after } = eventQuerySchema.parse(request.query);
if (!(await deps.repository.getSensitiveDataSuggestionRun(runId))) {
if (!(await deps.repository.getSensitivityAnalysisRun(runId))) {
return reply.code(404).send({
code: "sensitive_data_suggestion_run_not_found",
message: "Sensitive Data Suggestion Run was not found.",
message: "Sensitivity Analysis Run was not found.",
});
}
return (await deps.repository.listSensitiveDataSuggestionEvents(runId, after))
.map(publicSensitiveDataSuggestionEvent);
return (await deps.repository.listSensitivityAnalysisEvents(runId, after))
.map(publicSensitivityAnalysisEvent);
} catch (error) {
return safeSuggestionHistoryError(reply, error);
}
@@ -14,6 +14,7 @@ import {
type ModelCompletionRequest,
} from "../src/catalog/model-completer.js";
import { CatalogOperationCoordinator } from "../src/catalog/operation-coordinator.js";
import type { SensitivityValueSource } from "../src/catalog/sensitivity-classifier.js";
import type {
CatalogDatabaseClient,
CatalogPostgresAccess,
@@ -69,6 +70,16 @@ async function setup(
sample: vi.fn(async () => []),
},
catalogPostgresAccess?: CatalogPostgresAccess,
sensitivityValueSource: SensitivityValueSource = {
scanTable: vi.fn(async (request, consume) => {
await consume(request.columns.map((column) => ({
columnId: column.id,
value: "ordinary",
characterLength: 8,
})));
return { kind: "complete", observedRows: 1 };
}),
},
) {
const repository = new MemoryCatalogRepository();
const database = await repository.create({
@@ -115,10 +126,11 @@ async function setup(
catalogOperationCoordinator: operations,
metadataGenerationModels: models(),
modelCompleter,
sensitivityValueSource,
...(descriptionSourceSampler ? { descriptionSourceSampler } : {}),
...(catalogPostgresAccess ? { catalogPostgresAccess } : {}),
});
return { app, repository, database, table, column, operations };
return { app, repository, database, table, column, operations, sensitivityValueSource };
}
async function waitForTerminalRun(app: ReturnType<typeof buildApp>, runId: string) {
@@ -136,22 +148,15 @@ async function waitForTerminalRun(app: ReturnType<typeof buildApp>, runId: strin
throw new Error(`Description Generation Run ${runId} did not finish`);
}
test("suggests sensitive flags from structural metadata without persisting them", async () => {
const modelCompleter = {
complete: vi.fn(async () => JSON.stringify({
suggestions: [{ columnId: expect.any(String), sensitive: true }],
})),
};
test("assesses sensitive flags locally without persisting them or calling an LLM", async () => {
const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") };
const { app, repository, database, table, column } = await setup(modelCompleter);
modelCompleter.complete.mockResolvedValueOnce(JSON.stringify({
suggestions: [{ columnId: column.id, sensitive: true }],
}));
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
payload: { scope: "all" },
});
expect(response.statusCode).toBe(200);
@@ -160,11 +165,14 @@ test("suggests sensitive flags from structural metadata without persisting them"
run: {
databaseId: database.id,
scope: "all",
modelId: configuredModel.id,
engine: "local",
modelId: null,
policyVersion: "sensitivity-v1",
status: "completed",
total: 1,
suggestedSensitive: 1,
suggestedNonSensitive: 0,
unknown: 0,
errorSummary: null,
},
suggestions: [{
@@ -175,6 +183,8 @@ test("suggests sensitive flags from structural metadata without persisting them"
version: column.version,
currentSensitive: false,
sensitive: true,
assessment: "sensitive",
evidence: [{ kind: "metadata", ruleId: "metadata.direct_identifier" }],
}],
});
expect(await repository.getColumn(database.id, column.tableId, column.id))
@@ -207,49 +217,56 @@ test("suggests sensitive flags from structural metadata without persisting them"
runId: responseBody.run.id,
sequence: 1,
level: "info",
message: "Sensitive-field suggestion generation started.",
message: "Local sensitivity analysis started.",
},
{
runId: responseBody.run.id,
sequence: 2,
level: "info",
message: "Classified 1 of 1 columns.",
message: "Assessed 1 of 1 columns locally.",
},
{
runId: responseBody.run.id,
sequence: 3,
level: "info",
message: "Sensitive-field suggestion generation completed for 1 column.",
message: "Local sensitivity analysis completed for 1 column.",
},
]);
const request = modelCompleter.complete.mock.calls[0]![0] as ModelCompletionRequest;
const prompt = request.messages.map((message) => message.content).join("\n");
expect(prompt).toContain("patients");
expect(prompt).toContain("birth_date");
expect(prompt).toContain("date");
expect(prompt).not.toContain("Patient date of birth");
expect(prompt).not.toContain("test-provider-secret");
expect(modelCompleter.complete).not.toHaveBeenCalled();
} finally {
await app.close();
}
});
test("limits sensitive-data suggestions to the selected tables or columns", async () => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async (request) => {
const payload = JSON.parse(request.messages.find((message) => message.role === "user")!.content) as {
columns: Array<{ columnId: string; column: string }>;
};
return JSON.stringify({
suggestions: payload.columns.map((column) => ({
columnId: column.columnId,
sensitive: column.column.includes("name") || column.column.includes("note"),
})),
});
}),
};
const { app, repository, database } = await setup(modelCompleter);
test("stops sensitivity analysis at the HTTP deadline without creating a review", async () => {
const controller = new AbortController();
controller.abort();
const timeout = vi.spyOn(AbortSignal, "timeout").mockReturnValue(controller.signal);
const { app, repository, database } = await setup({ complete: vi.fn(async () => "unused") });
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { scope: "all" },
});
expect(response.statusCode).toBe(504);
expect(response.json()).toEqual({
code: "sensitivity_analysis_timeout",
message: "Sensitivity analysis reached its time limit. No assessments were applied.",
});
expect(await repository.listSensitivityAnalysisRuns()).toEqual([]);
} finally {
timeout.mockRestore();
await app.close();
}
});
test("limits sensitivity analysis to the selected tables or columns", async () => {
const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") };
const { app, repository, database, sensitivityValueSource } = await setup(modelCompleter);
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
@@ -281,7 +298,6 @@ test("limits sensitive-data suggestions to the selected tables or columns", asyn
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_tables",
targetIds: [visits.id, patients.id],
},
@@ -301,7 +317,6 @@ test("limits sensitive-data suggestions to the selected tables or columns", asyn
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_columns",
targetIds: [clinicalNote.id, status.id],
},
@@ -313,44 +328,41 @@ test("limits sensitive-data suggestions to the selected tables or columns", asyn
expect.objectContaining({ tableId: visits.id, columnId: clinicalNote.id, sensitive: true }),
]));
const prompts = vi.mocked(modelCompleter.complete).mock.calls.map(([request]) => (
JSON.parse(request.messages.find((message) => message.role === "user")!.content) as {
columns: Array<{ columnId: string }>;
}
));
expect(prompts[0]!.columns.map((column) => column.columnId).sort()).toEqual(
[...patientColumns, ...visitColumns].map((column) => column.id).sort(),
);
expect(prompts[0]!.columns.map((column) => column.columnId)).not.toContain(billingColumns[0]!.id);
expect(prompts[1]!.columns.map((column) => column.columnId).sort()).toEqual(
[status.id, clinicalNote.id].sort(),
const scannedColumnIds = vi.mocked(sensitivityValueSource.scanTable).mock.calls.flatMap(
([request]) => request.columns.map((column) => column.id),
);
expect(scannedColumnIds).toEqual([status.id, status.id]);
expect(scannedColumnIds).not.toContain(patientColumns.find(
(column) => column.name === "patient_name",
)!.id);
expect(scannedColumnIds).not.toContain(clinicalNote.id);
expect(scannedColumnIds).not.toContain(billingColumns[0]!.id);
expect(modelCompleter.complete).not.toHaveBeenCalled();
} finally {
await app.close();
}
});
test("explains invalid sensitive-data suggestion selections without calling the model", async () => {
test("explains invalid sensitivity-analysis selections without reading source values", async () => {
const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") };
const { app, database, table } = await setup(modelCompleter);
const { app, database, table, sensitivityValueSource } = await setup(modelCompleter);
try {
const empty = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "selected_tables", targetIds: [] },
payload: { scope: "selected_tables", targetIds: [] },
});
expect(empty.statusCode).toBe(400);
expect(empty.json()).toEqual({
code: "sensitive_data_suggestion_request_invalid",
message: "Choose a database, one or more tables, or one or more columns to classify.",
message: "Choose a database, one or more tables, or one or more columns to assess.",
});
const duplicate = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_tables",
targetIds: [table.id, table.id],
},
@@ -365,7 +377,6 @@ test("explains invalid sensitive-data suggestion selections without calling the
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_tables",
targetIds: ["00000000-0000-4000-8000-000000000001"],
},
@@ -380,7 +391,6 @@ test("explains invalid sensitive-data suggestion selections without calling the
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_columns",
targetIds: ["00000000-0000-4000-8000-000000000002"],
},
@@ -391,205 +401,7 @@ test("explains invalid sensitive-data suggestion selections without calling the
message: "One or more selected Catalog Columns were not found in this database.",
});
expect(modelCompleter.complete).not.toHaveBeenCalled();
} finally {
await app.close();
}
});
test("batches sensitive-data suggestions for schemas larger than one helper message", async () => {
const maxHelperMessageBytes = 64 * 1024;
const seenColumnIds: string[] = [];
const modelCompleter: ModelCompleter = {
complete: vi.fn(async (request) => {
const userMessage = request.messages.find((message) => message.role === "user")!;
expect(Buffer.byteLength(userMessage.content, "utf8")).toBeLessThanOrEqual(maxHelperMessageBytes);
const payload = JSON.parse(userMessage.content) as {
columns: Array<{ columnId: string; column: string }>;
};
expect(payload.columns.length).toBeLessThanOrEqual(10);
seenColumnIds.push(...payload.columns.map((column) => column.columnId));
return JSON.stringify({
suggestions: payload.columns.map((column) => ({
columnId: column.columnId,
sensitive: column.column.endsWith("_private"),
})),
});
}),
};
const { app, repository, database } = await setup(modelCompleter);
const columnCount = 900;
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables: [{ name: "wide_table", sourceComment: null }],
columns: Array.from({ length: columnCount }, (_, index) => ({
tableName: "wide_table",
name: `field_${index.toString().padStart(4, "0")}${index % 10 === 0 ? "_private" : ""}`,
ordinalPosition: index + 1,
dataType: "character varying(255)",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
})),
relationships: [],
});
const wideTable = (await repository.listTables(database.id)).find((table) => table.name === "wide_table")!;
const expectedColumnIds = (await repository.listColumns(database.id, wideTable.id)).map((column) => column.id);
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(200);
const suggestions = response.json().suggestions as Array<{
columnName: string;
currentSensitive: boolean;
sensitive: boolean;
}>;
expect(suggestions).toHaveLength(columnCount);
expect(suggestions).toEqual(expect.arrayContaining([
expect.objectContaining({ columnName: "field_0000_private", currentSensitive: false, sensitive: true }),
expect.objectContaining({ columnName: "field_0001", currentSensitive: false, sensitive: false }),
]));
expect(vi.mocked(modelCompleter.complete).mock.calls.length).toBeGreaterThan(1);
expect(seenColumnIds.slice().sort()).toEqual(expectedColumnIds.slice().sort());
expect(new Set(seenColumnIds).size).toBe(columnCount);
} finally {
await app.close();
}
});
test("retries one invalid sensitive-data classification before returning the review draft", async () => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async () => "unused"),
};
const { app, database, column } = await setup(modelCompleter);
vi.mocked(modelCompleter.complete)
.mockResolvedValueOnce("not-json")
.mockResolvedValueOnce(JSON.stringify({
suggestions: [{ columnId: column.id, sensitive: true }],
}));
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(200);
expect(response.json().suggestions).toEqual([
expect.objectContaining({ columnId: column.id, sensitive: true }),
]);
expect(modelCompleter.complete).toHaveBeenCalledTimes(2);
} finally {
await app.close();
}
});
test.each(["malformed", "incomplete", "duplicate"] as const)(
"fails safely when sensitive-data suggestions are %s",
async (kind) => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async () => "unused"),
};
const { app, repository, database, column } = await setup(modelCompleter);
const rawResponse = kind === "malformed"
? "RAW_PROVIDER_RESPONSE_DO_NOT_EXPOSE_{"
: kind === "incomplete"
? JSON.stringify({ suggestions: [] })
: JSON.stringify({
suggestions: [
{ columnId: column.id, sensitive: true },
{ columnId: column.id, sensitive: true },
],
});
vi.mocked(modelCompleter.complete).mockResolvedValueOnce(rawResponse);
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(502);
expect(response.json()).toEqual({
code: "sensitive_data_suggestion_invalid_response",
message: "The LLM returned an incomplete or invalid classification. No suggestions were applied.",
});
expect(response.body).not.toContain(rawResponse);
expect(await repository.getColumn(database.id, column.tableId, column.id))
.toMatchObject({ sensitive: false });
} finally {
await app.close();
}
},
);
test("explains a sensitive-data suggestion provider failure without exposing provider details", async () => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async () => {
throw new ModelCompletionProviderError();
}),
};
const { app, repository, database, column } = await setup(modelCompleter);
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(502);
expect(response.json()).toEqual({
code: "sensitive_data_suggestion_provider_unavailable",
message: "The selected LLM service could not complete the request. No suggestions were applied.",
});
expect(response.body).not.toContain("model completion failed");
expect(await repository.getColumn(database.id, column.tableId, column.id))
.toMatchObject({ sensitive: false });
const history = await app.inject({
method: "GET",
url: "/catalog/sensitive-data-suggestion-runs",
});
expect(history.statusCode).toBe(200);
const [failedRun] = history.json();
expect(failedRun).toMatchObject({
databaseId: database.id,
status: "failed",
total: 1,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
errorSummary: "Sensitive-field suggestion generation failed.",
});
const events = await app.inject({
method: "GET",
url: `/catalog/sensitive-data-suggestion-runs/${failedRun.id}/events-list`,
});
expect(events.statusCode).toBe(200);
expect(events.json()).toMatchObject([
{
runId: failedRun.id,
sequence: 1,
level: "info",
message: "Sensitive-field suggestion generation started.",
},
{
runId: failedRun.id,
sequence: 2,
level: "error",
message: "Sensitive-field suggestion generation failed.",
},
]);
expect(events.body).not.toContain("model completion failed");
expect(sensitivityValueSource.scanTable).not.toHaveBeenCalled();
} finally {
await app.close();
}
@@ -15,6 +15,7 @@ import { up as upSensitiveDataFlag } from "../src/catalog/migrations/006_sensiti
import { up as upSensitiveSuggestionRuns } from "../src/catalog/migrations/007_sensitive_data_suggestion_runs.js";
import { up as upAiTokenUsage } from "../src/catalog/migrations/009_ai_token_usage.js";
import { up as upCanonicalModelIds } from "../src/catalog/migrations/010_canonical_model_ids.js";
import { up as upLocalSensitivityAnalysis } from "../src/catalog/migrations/011_local_sensitivity_analysis.js";
import { KyselyCatalogRepository, type CatalogDatabase } from "../src/catalog/repository.js";
import { loadConfig } from "../src/config.js";
import type { WorkspaceRegistry } from "../src/workspaces/registry.js";
@@ -52,6 +53,7 @@ test.skipIf(!dockerAvailable)("Fastify persists Description Generation success a
await upSensitiveSuggestionRuns(db);
await upAiTokenUsage(db);
await upCanonicalModelIds(db);
await upLocalSensitivityAnalysis(db);
const repository = new KyselyCatalogRepository(db);
const database = await repository.create({
workspaceId: "psd-clinical",
@@ -0,0 +1,131 @@
import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, expect, test, vi } from "vitest";
import { PythonLocalNerDetector } from "../src/catalog/local-ner-detector.js";
const roots: string[] = [];
afterEach(() => {
vi.unstubAllEnvs();
for (const root of roots.splice(0)) rmSync(root, { recursive: true, force: true });
});
test("keeps a CPU-only local worker warm and returns sanitized evidence", async () => {
vi.stubEnv("THT_MODEL_API_KEY", "must-not-reach-worker");
const root = mkdtempSync(join(tmpdir(), "thothii-local-ner-"));
roots.push(root);
const helper = join(root, "fake_ner_worker.py");
writeFileSync(helper, `
import json
import os
import pathlib
import sys
root = pathlib.Path.cwd()
root.joinpath("runtime.json").write_text(json.dumps({
"argv": sys.argv,
"cuda": os.environ.get("CUDA_VISIBLE_DEVICES"),
"hip": os.environ.get("HIP_VISIBLE_DEVICES"),
"offline": os.environ.get("HF_HUB_OFFLINE"),
"inherited_secret": os.environ.get("THT_MODEL_API_KEY"),
"pid": os.getpid(),
}), encoding="utf-8")
print(json.dumps({"ready": True}), flush=True)
for line in sys.stdin:
request = json.loads(line)
root.joinpath("request.json").write_text(json.dumps(request), encoding="utf-8")
print(json.dumps({
"id": request["id"],
"ok": True,
"evidence": [{
"columnId": request["candidates"][0]["columnId"],
"label": "person",
"confidence": 0.93,
}],
}), flush=True)
`, "utf8");
const detector = new PythonLocalNerDetector({
pythonExecutable: "python3",
workerScript: helper,
modelPath: join(root, "pinned-model"),
cwd: root,
threads: 2,
startupTimeoutMs: 5_000,
});
const candidate = {
columnId: "33333333-3333-4333-8333-333333333333",
text: "Dimesso Mario Rossi",
};
try {
expect(detector.isReady()).toBe(false);
await detector.warmup();
expect(detector.isReady()).toBe(true);
expect(existsSync(join(root, "request.json"))).toBe(false);
await expect(detector.detect(
[candidate],
new AbortController().signal,
Date.now() + 5_000,
)).resolves.toEqual([{
columnId: candidate.columnId,
label: "person",
confidence: 0.93,
}]);
const firstRuntime = JSON.parse(readFileSync(join(root, "runtime.json"), "utf8"));
expect(firstRuntime).toMatchObject({
cuda: "",
hip: "",
offline: "1",
inherited_secret: null,
});
expect(JSON.stringify(firstRuntime.argv)).not.toContain(candidate.text);
expect(JSON.parse(readFileSync(join(root, "request.json"), "utf8")).candidates).toEqual([candidate]);
await detector.detect([candidate], new AbortController().signal, Date.now() + 5_000);
const secondRuntime = JSON.parse(readFileSync(join(root, "runtime.json"), "utf8"));
expect(secondRuntime.pid).toBe(firstRuntime.pid);
} finally {
await detector.close();
}
});
test("bounds worker startup by the caller deadline", async () => {
const root = mkdtempSync(join(tmpdir(), "thothii-local-ner-deadline-"));
roots.push(root);
const helper = join(root, "slow_ner_worker.py");
writeFileSync(helper, `
import json
import sys
import time
time.sleep(2)
print(json.dumps({"ready": True}), flush=True)
for line in sys.stdin:
request = json.loads(line)
print(json.dumps({"id": request["id"], "ok": True, "evidence": []}), flush=True)
`, "utf8");
const detector = new PythonLocalNerDetector({
pythonExecutable: "python3",
workerScript: helper,
modelPath: join(root, "pinned-model"),
cwd: root,
startupTimeoutMs: 5_000,
});
const startedAt = Date.now();
try {
await expect(detector.detect(
[{
columnId: "33333333-3333-4333-8333-333333333333",
text: "Dimesso Mario Rossi",
}],
new AbortController().signal,
startedAt + 50,
)).rejects.toThrow("local NER is unavailable");
expect(Date.now() - startedAt).toBeLessThan(1_000);
} finally {
await detector.close();
}
});
@@ -16,6 +16,7 @@ import { up as upSensitiveSuggestionRuns } from "../src/catalog/migrations/007_s
import { up as upLogicalRelationships } from "../src/catalog/migrations/008_catalog_logical_relationships.js";
import { up as upAiTokenUsage } from "../src/catalog/migrations/009_ai_token_usage.js";
import { up as upCanonicalModelIds } from "../src/catalog/migrations/010_canonical_model_ids.js";
import { up as upLocalSensitivityAnalysis } from "../src/catalog/migrations/011_local_sensitivity_analysis.js";
const dockerAvailable = spawnSync("docker", ["info"], { stdio: "ignore" }).status === 0;
@@ -47,12 +48,23 @@ test.skipIf(!dockerAvailable)("PostgreSQL migration enforces one database per wo
modelId: "openai-mini", language: "en", status: "completed", total: 1,
processed: 1, generated: 1,
}).execute();
const historicalSuggestionRunId = randomUUID();
await db.insertInto("sensitiveDataSuggestionRuns").values({
id: randomUUID(), databaseId: historicalDatabaseId, scope: "all",
id: historicalSuggestionRunId, databaseId: historicalDatabaseId, scope: "all",
modelId: "openai-mini", status: "completed", total: 1,
suggestedSensitive: 1,
}).execute();
await upCanonicalModelIds(db);
await upLocalSensitivityAnalysis(db);
await expect(db.selectFrom("sensitiveDataSuggestionRuns")
.select(["engine", "modelId", "policyVersion", "unknown"])
.where("id", "=", historicalSuggestionRunId)
.executeTakeFirstOrThrow()).resolves.toMatchObject({
engine: "llm",
modelId: "openai-mini",
policyVersion: null,
unknown: 0,
});
await expect(db.insertInto("descriptionGenerationRuns").values({
id: randomUUID(), databaseId: historicalDatabaseId, scope: "all",
modelId: "openai/gpt-5-mini", language: "en", status: "completed", total: 1,
@@ -418,6 +430,7 @@ test.skipIf(!dockerAvailable)("PostgreSQL repository persists description and se
await upSensitiveSuggestionRuns(db);
await upAiTokenUsage(db);
await upCanonicalModelIds(db);
await upLocalSensitivityAnalysis(db);
const repository = new KyselyCatalogRepository(db);
const firstDatabase = await repository.create({
workspaceId: "generation-one",
@@ -586,10 +599,10 @@ test.skipIf(!dockerAvailable)("PostgreSQL repository persists description and se
]);
expect(await repository.getActiveDescriptionGenerationRun()).toBeUndefined();
const suggestionRun = await repository.createSensitiveDataSuggestionRun(
const suggestionRun = await repository.createSensitivityAnalysisRun(
firstDatabase.id,
"selected_columns",
"openai/gpt-4.1-mini",
{ engine: "local", policyVersion: "sensitivity-v1" },
);
expect(suggestionRun).toMatchObject({
databaseId: firstDatabase.id,
@@ -597,43 +610,49 @@ test.skipIf(!dockerAvailable)("PostgreSQL repository persists description and se
total: 0,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
unknown: 0,
engine: "local",
modelId: null,
policyVersion: "sensitivity-v1",
startedAt: expect.any(String),
});
await repository.appendSensitiveDataSuggestionEvent(
await repository.appendSensitivityAnalysisEvent(
suggestionRun.id,
"info",
"Sensitive-field suggestion generation started.",
);
await repository.appendSensitiveDataSuggestionEvent(
await repository.appendSensitivityAnalysisEvent(
suggestionRun.id,
"info",
"Sensitive-field suggestion generation completed for 2 columns.",
);
expect(await repository.updateSensitiveDataSuggestionRun(suggestionRun.id, {
expect(await repository.updateSensitivityAnalysisRun(suggestionRun.id, {
status: "completed",
total: 2,
suggestedSensitive: 1,
suggestedNonSensitive: 1,
suggestedNonSensitive: 0,
unknown: 1,
finishedAt: new Date().toISOString(),
})).toMatchObject({
status: "completed",
total: 2,
suggestedSensitive: 1,
suggestedNonSensitive: 1,
suggestedNonSensitive: 0,
unknown: 1,
});
expect(await repository.listSensitiveDataSuggestionEvents(suggestionRun.id, 1)).toEqual([
expect(await repository.listSensitivityAnalysisEvents(suggestionRun.id, 1)).toEqual([
expect.objectContaining({ sequence: 2, level: "info" }),
]);
expect((await repository.listSensitiveDataSuggestionRuns(1))[0]).toMatchObject({
expect((await repository.listSensitivityAnalysisRuns(1))[0]).toMatchObject({
id: suggestionRun.id,
});
const interruptedSuggestionRun = await repository.createSensitiveDataSuggestionRun(
const interruptedSuggestionRun = await repository.createSensitivityAnalysisRun(
secondDatabase.id,
"all",
"openai/gpt-4.1-mini",
{ engine: "local", policyVersion: "sensitivity-v1" },
);
expect(await repository.interruptActiveSensitiveDataSuggestionRuns(
expect(await repository.interruptActiveSensitivityAnalysisRuns(
"Sensitive-field suggestion generation was interrupted by backend restart.",
)).toEqual([
expect.objectContaining({
@@ -0,0 +1,99 @@
import { expect, test, vi } from "vitest";
import {
SensitivityAnalysisInterruptedError,
SensitivityAnalysisService,
} from "../src/catalog/sensitivity-analysis-service.js";
import { SensitivityAnalysisRunner } from "../src/catalog/sensitivity-analysis-runner.js";
import type { SensitivityClassifier } from "../src/catalog/sensitivity-classifier.js";
import type {
CatalogRepository,
SensitivityAnalysisRun,
WorkspaceDatabase,
} from "../src/catalog/types.js";
const database = {
id: "11111111-1111-4111-8111-111111111111",
workspaceId: "psd-clinical",
engine: "postgres",
databaseName: "warehouse",
schema: "public",
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 running: SensitivityAnalysisRun = {
id: "22222222-2222-4222-8222-222222222222",
databaseId: database.id,
scope: "all",
engine: "local",
modelId: null,
policyVersion: "sensitivity-v1",
status: "running",
total: 0,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
unknown: 0,
inputTokens: 0,
cacheReadTokens: 0,
outputTokens: 0,
createdAt: "2026-09-02T08:00:00Z",
startedAt: "2026-09-02T08:00:00Z",
updatedAt: "2026-09-02T08:00:00Z",
finishedAt: null,
errorSummary: null,
};
test("stops catalog selection when the request expires during a catalog read", async () => {
const controller = new AbortController();
const listTables = vi.fn();
const repository = {
get: vi.fn(async () => {
controller.abort();
return database;
}),
listTables,
} as unknown as CatalogRepository;
const classifier = { assessTable: vi.fn() } as unknown as SensitivityClassifier;
const analysis = new SensitivityAnalysisService(repository, classifier);
await expect(analysis.analyze(
database.id,
"all",
[],
controller.signal,
)).rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
expect(listTables).not.toHaveBeenCalled();
expect(classifier.assessTable).not.toHaveBeenCalled();
});
test("marks a created run interrupted if the request deadline expires during persistence", async () => {
const controller = new AbortController();
const update = vi.fn(async (_runId: string, changes: Partial<SensitivityAnalysisRun>) => ({
...running,
...changes,
}));
const repository = {
get: vi.fn(async () => database),
createSensitivityAnalysisRun: vi.fn(async () => {
controller.abort();
return running;
}),
updateSensitivityAnalysisRun: update,
appendSensitivityAnalysisEvent: vi.fn(async () => undefined),
} as unknown as CatalogRepository;
const analysis = { analyze: vi.fn() } as unknown as SensitivityAnalysisService;
const runner = new SensitivityAnalysisRunner(repository, analysis);
await expect(runner.run(database.id, "all", [], controller.signal))
.rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
expect(analysis.analyze).not.toHaveBeenCalled();
expect(update).toHaveBeenCalledWith(running.id, expect.objectContaining({
status: "interrupted",
total: 0,
unknown: 0,
errorSummary: "Local sensitivity analysis reached its time limit.",
}));
});
@@ -0,0 +1,450 @@
import { expect, test, vi } from "vitest";
import {
SensitivityClassifier,
type LocalNerDetector,
type SensitivityNerBudget,
type SensitivityTableScan,
type SensitivityValueSource,
} from "../src/catalog/sensitivity-classifier.js";
import type { CatalogColumn, CatalogTable, WorkspaceDatabase } from "../src/catalog/types.js";
import { CatalogConnectorError } from "../src/catalog/types.js";
const database = {
id: "11111111-1111-4111-8111-111111111111",
workspaceId: "psd-clinical",
engine: "postgres",
databaseName: "warehouse",
schema: "public",
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: "observations",
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(overrides: Partial<CatalogColumn> = {}): CatalogColumn {
return {
id: "33333333-3333-4333-8333-333333333333",
tableId: table.id,
name: "note",
ordinalPosition: 1,
dataType: "character varying",
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",
...overrides,
};
}
function source(scan: SensitivityTableScan): SensitivityValueSource {
return { scanTable: vi.fn(async (_request, consume) => {
for (const batch of scan.batches) await consume(batch);
return scan.coverage;
}) };
}
test("one email hidden in a generically named column makes the whole column sensitive", async () => {
const target = column();
const values = source({
batches: [[
{ columnId: target.id, value: "nessun contatto", characterLength: 16 },
{ columnId: target.id, value: "mario.rossi@example.it", characterLength: 23 },
]],
coverage: { kind: "complete", observedRows: 2 },
});
const classifier = new SensitivityClassifier(values);
const [assessment] = await classifier.assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({
columnId: target.id,
assessment: "sensitive",
proposedSensitive: true,
evidence: [{ kind: "content", ruleId: "pii.email" }],
});
});
test("one text value longer than 500 characters makes the whole column sensitive", async () => {
const target = column({ name: "comment" });
const values = source({
batches: [[{ columnId: target.id, value: "x".repeat(501), characterLength: 743 }]],
coverage: { kind: "sampled", observedRows: 1 },
});
const [assessment] = await new SensitivityClassifier(values).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({
assessment: "sensitive",
proposedSensitive: true,
evidence: [{ kind: "length", ruleId: "text.over_500_characters" }],
});
});
test("complete coverage permits non-sensitive while empty columns remain unknown", async () => {
const benign = column({ id: "44444444-4444-4444-8444-444444444444", name: "status" });
const empty = column({ id: "55555555-5555-4555-8555-555555555555", name: "optional_note" });
const humanProtected = column({
id: "66666666-6666-4666-8666-666666666666",
name: "category",
sensitive: true,
});
const values = source({
batches: [[
{ columnId: benign.id, value: "active", characterLength: 6 },
{ columnId: empty.id, value: null, characterLength: null },
{ columnId: humanProtected.id, value: "administrative", characterLength: 14 },
]],
coverage: { kind: "complete", observedRows: 1 },
});
const assessments = await new SensitivityClassifier(values).assessTable(
{ database, table, columns: [benign, empty, humanProtected] },
new AbortController().signal,
);
expect(assessments).toEqual([
expect.objectContaining({ columnId: benign.id, assessment: "non_sensitive", proposedSensitive: false }),
expect.objectContaining({
columnId: empty.id,
assessment: "unknown",
proposedSensitive: false,
evidence: [{ kind: "coverage", ruleId: "coverage.no_values" }],
}),
expect.objectContaining({
columnId: humanProtected.id,
assessment: "non_sensitive",
proposedSensitive: false,
}),
]);
});
test("sampled coverage without a match is unknown and preserves the current human flag", async () => {
const target = column({ sensitive: true });
const values = source({
batches: [[{ columnId: target.id, value: "ordinary", characterLength: 8 }]],
coverage: { kind: "sampled", observedRows: 1 },
});
const [assessment] = await new SensitivityClassifier(values).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({
assessment: "unknown",
proposedSensitive: true,
evidence: [{ kind: "coverage", ruleId: "coverage.incomplete" }],
});
});
test("an unavailable source produces sanitized unknown evidence without losing metadata findings", async () => {
const unresolved = column();
const metadataMatch = column({
id: "44444444-4444-4444-8444-444444444444",
name: "codice_fiscale",
});
const values: SensitivityValueSource = {
scanTable: vi.fn(async () => {
throw new CatalogConnectorError("upstream detail must not escape");
}),
};
const assessments = await new SensitivityClassifier(values).assessTable(
{ database, table, columns: [unresolved, metadataMatch] },
new AbortController().signal,
);
expect(assessments).toEqual([
expect.objectContaining({
columnId: unresolved.id,
assessment: "unknown",
evidence: [{ kind: "coverage", ruleId: "coverage.unavailable" }],
}),
expect.objectContaining({
columnId: metadataMatch.id,
assessment: "sensitive",
evidence: [{ kind: "metadata", ruleId: "metadata.direct_identifier" }],
}),
]);
});
test("strong Italian PII metadata is sensitive even when the source column is empty", async () => {
const target = column({ name: "codice_fiscale" });
const values = source({
batches: [],
coverage: { kind: "complete", observedRows: 0 },
});
const [assessment] = await new SensitivityClassifier(values).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({
assessment: "sensitive",
proposedSensitive: true,
evidence: [{ kind: "metadata", ruleId: "metadata.direct_identifier" }],
});
expect(values.scanTable).not.toHaveBeenCalled();
});
test.each([
["RSSMRA85T10A562S", "pii.italian_fiscal_code"],
["IT60 X054 2811 1010 0000 0123 456", "financial.iban"],
["4111 1111 1111 1111", "financial.payment_card"],
["SWIFT DEUTDEFF500", "financial.bic"],
["Partita IVA 00743110157", "pii.italian_vat"],
["Passaporto YA1234567", "pii.passport_number"],
["Carta d'identità CA12345AA", "pii.identity_card"],
["Patente di guida U11234567A", "pii.drivers_license_number"],
["Chiamare +39 347 123 4567", "pii.phone_number"],
["Client 192.168.1.5", "network.ip_address"],
["Device 00:1B:44:11:3A:B7", "network.mac_address"],
["https://example.org/profiles/mario", "network.url"],
["550e8400-e29b-41d4-a716-446655440000", "pii.uuid"],
["AWS key AKIAIOSFODNN7EXAMPLE", "credential.access_key"],
["Diagnosi: carcinoma mammario con metastasi ossee", "health.clinical_term"],
["-----BEGIN PRIVATE KEY----- secret -----END PRIVATE KEY-----", "credential.private_key"],
['{"profile":{"email":"not yet supplied"}}', "pii.json_sensitive_key"],
] as const)("recognizes validated sensitive content without relying on the column name: %s", async (
value,
ruleId,
) => {
const target = column();
const values = source({
batches: [[{ columnId: target.id, value, characterLength: value.length }]],
coverage: { kind: "complete", observedRows: 1 },
});
const [assessment] = await new SensitivityClassifier(values).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({
assessment: "sensitive",
evidence: [{ kind: "content", ruleId }],
});
});
test("does not make a malformed email decisive", async () => {
const target = column();
const [assessment] = await new SensitivityClassifier(source({
batches: [[{
columnId: target.id,
value: "contatto a@b..com non valido",
characterLength: 28,
}]],
coverage: { kind: "complete", observedRows: 1 },
})).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({ assessment: "non_sensitive", evidence: [] });
});
test("finds a valid email after a malformed candidate in the same value", async () => {
const target = column();
const [assessment] = await new SensitivityClassifier(source({
batches: [[{
columnId: target.id,
value: "contatto a@b..com; indirizzo valido mario.rossi@example.it",
characterLength: 58,
}]],
coverage: { kind: "complete", observedRows: 1 },
})).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({
assessment: "sensitive",
evidence: [{ kind: "content", ruleId: "pii.email" }],
});
});
test("optional local NER evidence can make otherwise ambiguous Italian text sensitive", async () => {
const target = column();
const values = source({
batches: [[{ columnId: target.id, value: "Dimesso Mario Rossi", characterLength: 19 }]],
coverage: { kind: "sampled", observedRows: 1 },
});
const detector: LocalNerDetector = {
detect: vi.fn(async () => [{ columnId: target.id, label: "person_name", confidence: 0.91 }]),
};
const [assessment] = await new SensitivityClassifier(values, detector).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(detector.detect).toHaveBeenCalledWith(
[{ columnId: target.id, text: "Dimesso Mario Rossi" }],
expect.any(AbortSignal),
expect.any(Number),
);
expect(assessment).toMatchObject({
assessment: "sensitive",
evidence: [{ kind: "ner", ruleId: "ner.entity", label: "person_name", confidence: 0.91 }],
});
});
test("does not wait for an optional NER worker that is still warming", async () => {
const target = column();
const detector: LocalNerDetector = {
isReady: () => false,
detect: vi.fn(async () => [{ columnId: target.id, label: "person", confidence: 0.99 }]),
};
const [assessment] = await new SensitivityClassifier(source({
batches: [[{ columnId: target.id, value: "Dimesso Mario Rossi", characterLength: 19 }]],
coverage: { kind: "sampled", observedRows: 1 },
}), detector).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(detector.detect).not.toHaveBeenCalled();
expect(assessment).toMatchObject({ assessment: "unknown" });
});
test("bounds each optional NER request when an installation raises the per-table work limit", async () => {
const columns = Array.from({ length: 17 }, (_, index) => column({
id: `00000000-0000-4000-8000-${(index + 1).toString(16).padStart(12, "0")}`,
name: `attribute_${index + 1}`,
ordinalPosition: index + 1,
}));
const observations = columns.flatMap((item, columnIndex) => Array.from(
{ length: 8 },
(_, valueIndex) => ({
columnId: item.id,
value: `ordinary-${columnIndex}-${valueIndex}`,
characterLength: 13,
}),
));
const detector: LocalNerDetector = { detect: vi.fn(async () => []) };
await new SensitivityClassifier(source({
batches: [observations],
coverage: { kind: "complete", observedRows: 8 },
}), detector, { maxNerCandidatesPerTable: 136 }).assessTable(
{ database, table, columns },
new AbortController().signal,
);
expect(detector.detect).toHaveBeenCalledTimes(2);
expect(vi.mocked(detector.detect).mock.calls.map(([candidates]) => candidates.length)).toEqual([
128,
8,
]);
});
test("limits default NER work to two candidates spread across a wide table", async () => {
const columns = Array.from({ length: 10 }, (_, index) => column({
id: `10000000-0000-4000-8000-${(index + 1).toString(16).padStart(12, "0")}`,
name: `attribute_${index + 1}`,
ordinalPosition: index + 1,
}));
const detector: LocalNerDetector = { detect: vi.fn(async () => []) };
await new SensitivityClassifier(source({
batches: [columns.flatMap((item, columnIndex) => [0, 1].map((valueIndex) => ({
columnId: item.id,
value: `ordinary-${columnIndex}-${valueIndex}`,
characterLength: 13,
})))],
coverage: { kind: "complete", observedRows: 2 },
}), detector).assessTable(
{ database, table, columns },
new AbortController().signal,
);
expect(detector.detect).toHaveBeenCalledOnce();
const submitted = vi.mocked(detector.detect).mock.calls[0]![0];
expect(submitted).toHaveLength(2);
expect(new Set(submitted.map((candidate) => candidate.columnId)).size).toBe(2);
});
test("shares a bounded NER time allowance across tables in one analysis run", async () => {
const target = column();
const values = source({
batches: [[{ columnId: target.id, value: "Dimesso Mario Rossi", characterLength: 19 }]],
coverage: { kind: "sampled", observedRows: 1 },
});
const detector: LocalNerDetector = {
detect: vi.fn(async () => {
await new Promise((resolve) => setTimeout(resolve, 20));
return [];
}),
};
const classifier = new SensitivityClassifier(values, detector);
const nerBudget: SensitivityNerBudget = { remainingMs: 1 };
await classifier.assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
Date.now() + 1_000,
nerBudget,
);
await classifier.assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
Date.now() + 1_000,
nerBudget,
);
expect(detector.detect).toHaveBeenCalledOnce();
expect(nerBudget.remainingMs).toBe(0);
});
test("uninterpretable binary content remains unknown after complete coverage", async () => {
const target = column({ dataType: "bytea" });
const values = source({
batches: [[{ columnId: target.id, value: "\\xdeadbeef", characterLength: 10 }]],
coverage: { kind: "complete", observedRows: 1 },
});
const [assessment] = await new SensitivityClassifier(values).assessTable(
{ database, table, columns: [target] },
new AbortController().signal,
);
expect(assessment).toMatchObject({
assessment: "unknown",
proposedSensitive: false,
evidence: [{ kind: "coverage", ruleId: "coverage.unsupported_type" }],
});
});
@@ -0,0 +1,316 @@
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();
});
+29
View File
@@ -81,6 +81,35 @@ test("loadConfig keeps local development defaults", () => {
expect(loadConfig({}).dataRoot).toBeUndefined();
});
test("loadConfig keeps local NER disabled unless an absolute model path is configured", () => {
expect(loadConfig({}).sensitivityNer).toBeUndefined();
expect(loadConfig({
THT_SENSITIVITY_NER_MODEL_PATH: "/models/gliner2-pii",
}).sensitivityNer?.pythonExecutable).toBe("/opt/sensitivity-ner/bin/python");
expect(loadConfig({
THT_SENSITIVITY_NER_MODEL_PATH: "/models/gliner2-pii",
THT_SENSITIVITY_NER_PYTHON: "/opt/sensitivity-ner/bin/python",
THT_SENSITIVITY_NER_WORKER: "/app/backend/python/sensitivity_ner_worker.py",
THT_SENSITIVITY_NER_THREADS: "3",
}).sensitivityNer).toEqual({
modelPath: "/models/gliner2-pii",
pythonExecutable: "/opt/sensitivity-ner/bin/python",
workerScript: "/app/backend/python/sensitivity_ner_worker.py",
threads: 3,
});
});
test("loadConfig rejects ambiguous or unsafe local NER configuration", () => {
expect(() => loadConfig({ THT_SENSITIVITY_NER_MODEL_PATH: "fastino/model" }))
.toThrow("sensitivity NER model path configuration is invalid");
expect(() => loadConfig({
THT_SENSITIVITY_NER_MODEL_PATH: "/models/gliner2-pii",
THT_SENSITIVITY_NER_THREADS: "0",
})).toThrow("sensitivity NER thread configuration is invalid");
expect(() => loadConfig({ THT_SENSITIVITY_NER_PYTHON: "/opt/ner/bin/python" }))
.toThrow("sensitivity NER settings require a model path");
});
test("loadConfig allows none and mock only outside production when auth.yaml is absent", () => {
const originalNodeEnvironment = process.env.NODE_ENV;
delete process.env.NODE_ENV;