fix(snowflake): unblock multi-schema ingest and relationship discovery (#204)

* feat(setup): drop redundant Snowflake schema prompt; fall back to free-text on listSchemas failure

Snowflake setup previously asked for a single schema as free text, then
ran a multiselect against the discovered schemas — two schema questions
back-to-back, with the first being only a session bootstrap. The SDK's
`schema` is optional, so the bootstrap step is unnecessary.

- Remove the free-text Snowflake schema prompt; only pass `schema` to
  snowflake-sdk when one is configured.
- When `listSchemas()` fails (e.g. role lacks SHOW SCHEMAS), prompt the
  user for a comma-separated list, persist it as `schema_names`, and use
  it as both the table-list filter and the multiselect default. Applies
  to every driver with a scope-discovery spec, not just Snowflake.
- Update docs to lead with `schema_names`; keep `schema_name` as a
  documented single-schema shorthand.

* fix(snowflake): keep introspecting when primary-key discovery is denied

The PK query joins INFORMATION_SCHEMA.TABLE_CONSTRAINTS and
INFORMATION_SCHEMA.KEY_COLUMN_USAGE, which require grants the
connection role may not have. Previously a 'SQL compilation error:
Object ANALYTICS.INFORMATION_SCHEMA.KEY_COLUMN_USAGE does not exist
or not authorized' aborted the entire introspect — schemas, columns,
and row counts were all discarded over a missing nice-to-have.

Wrap the constraint query in try/catch, log a one-line warning per
schema, and return an empty PK map. Columns end up with
primaryKey=false; relationship inference still has FK and profiling
to fall back on.

* fix(scan): unblock relationship discovery on Snowflake

Two adjacent bugs prevented the scan's relationship pipeline from producing
any joins on a Snowflake warehouse:

- relationship-profiling.ts fell through to a default `GROUP_CONCAT` branch
  for unknown drivers. Snowflake has no GROUP_CONCAT, so every per-table
  profile query failed with "Unknown function GROUP_CONCAT". Add an explicit
  Snowflake branch that uses LISTAGG with a literal '\x1f' delimiter
  (Snowflake requires the delimiter to be a constant, so CHR(31) is rejected).
- description-generation.ts destructured `connector.sampleTable` and
  `connector.sampleColumn` into bare locals, losing the `this` binding when
  the class-method connectors (Snowflake, Postgres, MySQL) were invoked.
  Every sample call threw "Cannot read properties of undefined (reading
  'assertConnection')" and degraded LLM descriptions to metadata-only
  prompts. Call the methods through the connector instead.

Without these, even after the primary-key probe is allowed to fail softly,
the scan ends up with 0 validated relationships and an empty `joins:` block
in every shard YAML.

* test(scan): cover table-ref helpers

* feat(scan): plumb tableScope through live-database introspection port

* feat(scan): apply tableScope during metadata fetch

* feat(scan): enforce table scope at fetch boundary

* feat(scan): pool Snowflake sessions and batch enrichment for faster ingest (#206)

* feat(cli): add RSA key-pair auth option to Snowflake setup wizard

Extends the interactive Snowflake setup flow with an authentication-method
prompt (password vs RSA/JWT key-pair). The RSA branch collects a private-key
path (env/file/absolute) and an optional passphrase; the resulting connection
config records `authMethod: 'rsa'` with `privateKey` and `passphrase` instead
of `password`.

* feat(scan): pool Snowflake sessions

* fix(scan): reuse structural snapshots and cleanup connectors

* feat(scan): parallelize relationship profiling

* feat(scan): batch table description generation

* docs: document Snowflake ingest concurrency knobs

* fix(scan): close Snowflake ingest perf verification gaps

* fix(scan): keep batched description failure bounded

* feat(scan): dispatch query-history probes by connection driver

Extract historic-sql dialect resolution into a shared helper so the
status-project readiness check and the local ingest factory agree on
which connections enable query history and which probe to run. The
status command now picks the postgres/snowflake/bigquery probe based on
the connection's driver instead of always reporting against postgres,
which previously caused snowflake connections with queryHistory.enabled
to surface a misleading "driver is snowflake" failure.

Also drops a noisy console.warn from Snowflake primary-key discovery —
INFORMATION_SCHEMA.KEY_COLUMN_USAGE is commonly ungranted for read-only
roles and the FK + profiling paths handle the empty PK map already.

* fix(llm): allow StructuredOutput tool and raise maxTurns for generateObject

The Claude Code agent SDK announces an internal pseudo-tool named
StructuredOutput in the system/init message whenever outputFormat is set
to { type: 'json_schema' }. The runtime's isolation check built its
allowedToolIds set only from MCP tool ids and treated StructuredOutput
as an unexpected host-injected tool, so every generateObject call threw
"Claude Code runtime isolation failed: tools=StructuredOutput ..." and
the table-descriptions and relationship-LLM-proposal enrichment stages
recorded null output across the board.

Whitelist StructuredOutput specifically in generateObject's
allowedToolIds — the check also enforces missing_tools symmetry, so
generateText and runAgentLoop, which do not see StructuredOutput, must
not require it.

generateObject also ran with maxTurns: 1, which the model intermittently
breached when it emitted thinking text before the structured response.
Raised to 5 to give the schema-bound call enough headroom without
allowing unbounded loops. The existing tests now exercise the path with
an init message that announces StructuredOutput so the regression cannot
slip back in.

* chore(scripts): add ktx-reset.sh project-cleanup helper

Convenience script for repeatable ingest testing: takes a project
directory and prunes everything except ktx.yaml and .ktx/secrets/, so
the next ktx setup or ktx ingest run starts from a known-clean state.
This commit is contained in:
Andrey Avtomonov 2026-05-23 10:41:30 +02:00 committed by GitHub
parent b0dd13ce7c
commit 394a985d2a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
72 changed files with 3508 additions and 655 deletions

View file

@ -5,6 +5,10 @@ import type { KtxConfigIssue, KtxProjectConfig, KtxProjectConnectionConfig, KtxP
import type { KtxLocalProject } from './context/project/project.js';
import { ktxLocalStateDbPath } from './context/project/local-state-db.js';
import type { PostgresPgssProbeResult } from './context/ingest/adapters/historic-sql/types.js';
import {
isQueryHistoryEnabled,
queryHistoryDialectForConnection,
} from './context/ingest/adapters/historic-sql/connection-dialect.js';
import {
formatClaudeCodePromptCachingFix,
formatClaudeCodePromptCachingWarning,
@ -47,7 +51,8 @@ interface ConnectionStatus extends ProjectStatusLine {
interface QueryHistoryStatus extends ProjectStatusLine {
connection: string;
dialect: 'postgres';
driver: string;
dialect: string;
}
interface PipelineStatus {
@ -396,45 +401,44 @@ function buildConnectionStatus(
}
}
interface PostgresQueryHistoryProbeInput {
interface QueryHistoryProbeInput {
projectDir: string;
connectionId: string;
connection: KtxProjectConnectionConfig;
env: NodeJS.ProcessEnv;
}
type PostgresQueryHistoryProbe = (
input: PostgresQueryHistoryProbeInput,
) => Promise<PostgresPgssProbeResult>;
function recordValue(value: unknown): Record<string, unknown> | null {
return value && typeof value === 'object' && !Array.isArray(value) ? (value as Record<string, unknown>) : null;
interface GenericProbeResult {
warnings: string[];
info?: string[];
}
function queryHistoryRecord(connection: KtxProjectConnectionConfig): Record<string, unknown> | null {
const context = recordValue(connection.context);
return recordValue(context?.queryHistory);
}
type PostgresQueryHistoryProbe = (input: QueryHistoryProbeInput) => Promise<PostgresPgssProbeResult>;
type SnowflakeQueryHistoryProbe = (input: QueryHistoryProbeInput) => Promise<GenericProbeResult>;
type BigQueryQueryHistoryProbe = (input: QueryHistoryProbeInput) => Promise<GenericProbeResult>;
function legacyHistoricSqlRecord(connection: KtxProjectConnectionConfig): Record<string, unknown> | null {
return recordValue(connection.historicSql);
}
function isEnabledPostgresQueryHistory(connection: KtxProjectConnectionConfig): boolean {
const queryHistory = queryHistoryRecord(connection);
if (queryHistory) {
return queryHistory.enabled === true;
function failureDetail(error: unknown): string {
if (error instanceof Error && error.message.trim().length > 0) {
return error.message.trim().split('\n')[0] ?? error.message.trim();
}
const legacy = legacyHistoricSqlRecord(connection);
return legacy?.enabled === true && legacy.dialect === 'postgres';
return String(error);
}
function isPostgresDriver(connection: KtxProjectConnectionConfig): boolean {
const driver = String(connection.driver ?? '').toLowerCase();
return driver === 'postgres' || driver === 'postgresql';
function postgresReadinessDetail(result: PostgresPgssProbeResult): string {
const warningText = result.warnings.length > 0 ? ` with warnings: ${result.warnings.join('; ')}` : '';
const info = result.info ?? [];
const infoText = info.length > 0 ? `; info: ${info.join('; ')}` : '';
return `pg_stat_statements ready (${result.pgServerVersion})${warningText}${infoText}`;
}
function queryHistoryFailureFix(error: unknown, connectionId: string, projectDir: string): string {
function genericReadinessDetail(label: string, result: GenericProbeResult): string {
const warningText = result.warnings.length > 0 ? ` with warnings: ${result.warnings.join('; ')}` : '';
const info = result.info ?? [];
const infoText = info.length > 0 ? `; info: ${info.join('; ')}` : '';
return `${label} ready${warningText}${infoText}`;
}
function probeFailureFix(error: unknown, dialect: string, connectionId: string, projectDir: string): string {
if (error instanceof Error && error.name === 'HistoricSqlExtensionMissingError' && 'remediation' in error) {
return String(error.remediation);
}
@ -444,25 +448,11 @@ function queryHistoryFailureFix(error: unknown, connectionId: string, projectDir
if (error instanceof Error && error.name === 'HistoricSqlVersionUnsupportedError') {
return 'Use PostgreSQL 14 or newer, or disable query history for this connection';
}
return `Fix connections.${connectionId} Postgres settings, then rerun \`ktx status --project-dir ${projectDir}\``;
}
function failureDetail(error: unknown): string {
if (error instanceof Error && error.message.trim().length > 0) {
return error.message.trim().split('\n')[0] ?? error.message.trim();
}
return String(error);
}
function readinessDetail(result: PostgresPgssProbeResult): string {
const warningText = result.warnings.length > 0 ? ` with warnings: ${result.warnings.join('; ')}` : '';
const info = result.info ?? [];
const infoText = info.length > 0 ? `; info: ${info.join('; ')}` : '';
return `pg_stat_statements ready (${result.pgServerVersion})${warningText}${infoText}`;
return `Fix connections.${connectionId} ${dialect} settings, then rerun \`ktx status --project-dir ${projectDir}\``;
}
async function defaultPostgresQueryHistoryProbe(
input: PostgresQueryHistoryProbeInput,
input: QueryHistoryProbeInput,
): Promise<PostgresPgssProbeResult> {
const [{ PostgresPgssReader }, { KtxPostgresHistoricSqlQueryClient }, { isKtxPostgresConnectionConfig }] =
await Promise.all([
@ -488,63 +478,225 @@ async function defaultPostgresQueryHistoryProbe(
}
}
async function defaultSnowflakeQueryHistoryProbe(
input: QueryHistoryProbeInput,
): Promise<GenericProbeResult> {
const [{ SnowflakeHistoricSqlQueryHistoryReader }, { KtxSnowflakeHistoricSqlQueryClient }, { isKtxSnowflakeConnectionConfig }] =
await Promise.all([
import('./context/ingest/adapters/historic-sql/snowflake-query-history-reader.js'),
import('./connectors/snowflake/historic-sql-query-client.js'),
import('./connectors/snowflake/connector.js'),
]);
const inputDriver = input.connection.driver ?? 'unknown';
if (!isKtxSnowflakeConnectionConfig(input.connection)) {
throw new Error(`Native Snowflake connector cannot run driver "${inputDriver}"`);
}
const client = new KtxSnowflakeHistoricSqlQueryClient({
connectionId: input.connectionId,
connection: input.connection,
projectDir: input.projectDir,
env: input.env,
});
try {
return await new SnowflakeHistoricSqlQueryHistoryReader().probe(client);
} finally {
await client.cleanup();
}
}
async function defaultBigQueryQueryHistoryProbe(
input: QueryHistoryProbeInput,
): Promise<GenericProbeResult> {
const [
{ BigQueryHistoricSqlQueryHistoryReader },
{ KtxBigQueryScanConnector, isKtxBigQueryConnectionConfig },
{ resolveKtxConfigReference },
] = await Promise.all([
import('./context/ingest/adapters/historic-sql/bigquery-query-history-reader.js'),
import('./connectors/bigquery/connector.js'),
import('./context/core/config-reference.js'),
]);
const inputDriver = input.connection.driver ?? 'unknown';
if (!isKtxBigQueryConnectionConfig(input.connection)) {
throw new Error(`Native BigQuery connector cannot run driver "${inputDriver}"`);
}
const rawCredentials = typeof input.connection.credentials_json === 'string' ? input.connection.credentials_json : '';
const resolvedCredentials = resolveKtxConfigReference(rawCredentials, input.env);
if (!resolvedCredentials) {
throw new Error(`Query history BigQuery connection ${input.connectionId} requires credentials_json`);
}
const parsed = JSON.parse(resolvedCredentials) as { project_id?: unknown };
if (typeof parsed.project_id !== 'string' || parsed.project_id.trim().length === 0) {
throw new Error(`Query history BigQuery connection ${input.connectionId} requires credentials_json.project_id`);
}
const region =
typeof input.connection.location === 'string' && input.connection.location.trim().length > 0
? input.connection.location.trim()
: 'us';
const connector = new KtxBigQueryScanConnector({
connectionId: input.connectionId,
connection: input.connection,
});
try {
return await new BigQueryHistoricSqlQueryHistoryReader({
projectId: parsed.project_id,
region,
}).probe({
async executeQuery(sql: string) {
const result = await connector.executeReadOnly({ connectionId: input.connectionId, sql }, {} as never);
return {
headers: result.headers,
rows: result.rows,
totalRows: result.totalRows,
};
},
});
} finally {
await connector.cleanup();
}
}
interface DispatchedProbe {
label: string;
spinnerLabel: string;
fastSkipDetail: string;
run: () => Promise<{ status: ProjectStatusLevel; detail: string; fix?: string }>;
}
function postgresProbeDispatch(
input: QueryHistoryProbeInput,
probe: PostgresQueryHistoryProbe,
): DispatchedProbe {
return {
label: 'postgres',
spinnerLabel: `Probing pg_stat_statements on ${input.connectionId}`,
fastSkipDetail: 'pg_stat_statements probe skipped (--fast)',
run: async () => {
const result = await probe(input);
return {
status: result.warnings.length > 0 ? 'warn' : 'ok',
detail: postgresReadinessDetail(result),
...(result.warnings.length > 0
? {
fix: `Update the Postgres parameter group or config, then rerun \`ktx status --project-dir ${input.projectDir}\``,
}
: {}),
};
},
};
}
function snowflakeProbeDispatch(
input: QueryHistoryProbeInput,
probe: SnowflakeQueryHistoryProbe,
): DispatchedProbe {
return {
label: 'snowflake',
spinnerLabel: `Probing SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY on ${input.connectionId}`,
fastSkipDetail: 'SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY probe skipped (--fast)',
run: async () => {
const result = await probe(input);
return {
status: result.warnings.length > 0 ? 'warn' : 'ok',
detail: genericReadinessDetail('SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY', result),
};
},
};
}
function bigqueryProbeDispatch(
input: QueryHistoryProbeInput,
probe: BigQueryQueryHistoryProbe,
): DispatchedProbe {
return {
label: 'bigquery',
spinnerLabel: `Probing INFORMATION_SCHEMA.JOBS_BY_PROJECT on ${input.connectionId}`,
fastSkipDetail: 'INFORMATION_SCHEMA.JOBS_BY_PROJECT probe skipped (--fast)',
run: async () => {
const result = await probe(input);
return {
status: result.warnings.length > 0 ? 'warn' : 'ok',
detail: genericReadinessDetail('INFORMATION_SCHEMA.JOBS_BY_PROJECT', result),
};
},
};
}
async function buildQueryHistoryStatus(
project: KtxLocalProject,
options: BuildProjectStatusOptions,
): Promise<QueryHistoryStatus[]> {
const targets = Object.entries(project.config.connections)
.filter(([, connection]) => isEnabledPostgresQueryHistory(connection))
.filter(([, connection]) => isQueryHistoryEnabled(connection))
.sort(([left], [right]) => left.localeCompare(right));
const probe = options.postgresQueryHistoryProbe ?? defaultPostgresQueryHistoryProbe;
const postgresProbe = options.postgresQueryHistoryProbe ?? defaultPostgresQueryHistoryProbe;
const snowflakeProbe = options.snowflakeQueryHistoryProbe ?? defaultSnowflakeQueryHistoryProbe;
const bigqueryProbe = options.bigqueryQueryHistoryProbe ?? defaultBigQueryQueryHistoryProbe;
const env = options.env ?? process.env;
const statuses: QueryHistoryStatus[] = [];
for (const [connectionId, connection] of targets) {
if (!isPostgresDriver(connection)) {
const driver = String(connection.driver ?? 'unknown').toLowerCase();
const dialect = queryHistoryDialectForConnection(connection);
if (!dialect) {
statuses.push({
connection: connectionId,
dialect: 'postgres',
driver,
dialect: driver,
status: 'fail',
detail: `connections.${connectionId}.context.queryHistory is enabled but driver is ${String(connection.driver)}`,
fix: `Set connections.${connectionId}.driver to postgres or disable query history for this connection`,
detail: `query history is not supported for driver "${driver}"`,
fix: `Disable connections.${connectionId}.context.queryHistory, or use a postgres, snowflake, or bigquery connection`,
});
continue;
}
const probeInput: QueryHistoryProbeInput = {
projectDir: project.projectDir,
connectionId,
connection,
env,
};
const dispatched =
dialect === 'postgres'
? postgresProbeDispatch(probeInput, postgresProbe)
: dialect === 'snowflake'
? snowflakeProbeDispatch(probeInput, snowflakeProbe)
: bigqueryProbeDispatch(probeInput, bigqueryProbe);
if (options.fast === true) {
statuses.push({
connection: connectionId,
dialect: 'postgres',
driver,
dialect,
status: 'skipped',
detail: 'pg_stat_statements probe skipped (--fast)',
detail: dispatched.fastSkipDetail,
});
continue;
}
try {
const result = await withSpinner(
options.useSpinner === true,
`Probing pg_stat_statements on ${connectionId}`,
() => probe({ projectDir: project.projectDir, connectionId, connection, env }),
);
const outcome = await withSpinner(options.useSpinner === true, dispatched.spinnerLabel, dispatched.run);
statuses.push({
connection: connectionId,
dialect: 'postgres',
status: result.warnings.length > 0 ? 'warn' : 'ok',
detail: readinessDetail(result),
...(result.warnings.length > 0
? {
fix: `Update the Postgres parameter group or config, then rerun \`ktx status --project-dir ${project.projectDir}\``,
}
: {}),
driver,
dialect,
...outcome,
});
} catch (error) {
statuses.push({
connection: connectionId,
dialect: 'postgres',
driver,
dialect,
status: 'fail',
detail: failureDetail(error),
fix: queryHistoryFailureFix(error, connectionId, project.projectDir),
fix: probeFailureFix(error, dispatched.label, connectionId, project.projectDir),
});
}
}
@ -731,6 +883,8 @@ function buildVerdict(
export interface BuildProjectStatusOptions {
env?: NodeJS.ProcessEnv;
postgresQueryHistoryProbe?: PostgresQueryHistoryProbe;
snowflakeQueryHistoryProbe?: SnowflakeQueryHistoryProbe;
bigqueryQueryHistoryProbe?: BigQueryQueryHistoryProbe;
claudeCodeAuthProbe?: ClaudeCodeAuthProbe;
configIssues?: KtxConfigIssue[];
fast?: boolean;