Refine setup table selection flow

This commit is contained in:
Luca Martial 2026-05-12 21:31:11 -07:00
parent 6a5383a398
commit 9a8cb08192
5 changed files with 97 additions and 57 deletions

View file

@ -1,6 +1,7 @@
import { mkdir, mkdtemp, readFile, rm, stat, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { readKtxSetupState } from '@ktx/context/project';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import {
formatInstallSummary,
@ -89,7 +90,7 @@ describe('setup agents', () => {
projectDir: tempDir,
installs: [{ target: 'universal', scope: 'project', mode: 'cli' }],
});
expect(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')).toContain('agents');
expect((await readKtxSetupState(tempDir)).completed_steps).toContain('agents');
expect(io.stderr()).toBe('');
});
@ -143,7 +144,7 @@ describe('setup agents', () => {
await expect(readKtxAgentInstallManifest(tempDir)).resolves.toEqual(null);
});
it('uses prompts in interactive mode and supports Back', async () => {
it('treats cancel as skip in interactive mode', async () => {
const io = makeIo();
const prompts = {
select: vi.fn(async () => 'back'),
@ -165,7 +166,7 @@ describe('setup agents', () => {
io.io,
{ prompts },
),
).resolves.toEqual({ status: 'back', projectDir: tempDir });
).resolves.toEqual({ status: 'skipped', projectDir: tempDir });
});
it('explains how to select multiple agent targets in interactive mode', async () => {

View file

@ -391,10 +391,9 @@ export async function runKtxSetupAgentsStep(
options: [
{ value: 'cli', label: 'CLI tools and skills' },
{ value: 'skip', label: 'Skip' },
{ value: 'back', label: 'Back' },
],
})) as KtxAgentInstallMode | 'skip' | 'back');
if (mode === 'back') return { status: 'back', projectDir: args.projectDir };
if (mode === 'back') return { status: 'skipped', projectDir: args.projectDir };
if (mode === 'skip') return { status: 'skipped', projectDir: args.projectDir };
const targets =

View file

@ -534,7 +534,6 @@ describe('setup databases step', () => {
options: [
{ value: 'continue', label: 'Continue to knowledge sources' },
{ value: 'add', label: 'Add another primary source' },
{ value: 'back', label: 'Back' },
],
});
expect(testConnection).not.toHaveBeenCalled();
@ -585,7 +584,6 @@ describe('setup databases step', () => {
options: [
{ value: 'continue', label: 'Continue to knowledge sources' },
{ value: 'add', label: 'Add another primary source' },
{ value: 'back', label: 'Back' },
],
});
expect(testConnection).toHaveBeenCalledTimes(1);
@ -620,7 +618,6 @@ describe('setup databases step', () => {
options: [
{ value: 'continue', label: 'Continue to knowledge sources' },
{ value: 'add', label: 'Add another primary source' },
{ value: 'back', label: 'Back' },
],
});
const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8'));
@ -655,7 +652,6 @@ describe('setup databases step', () => {
options: [
{ value: 'continue', label: 'Continue to knowledge sources' },
{ value: 'add', label: 'Add another primary source' },
{ value: 'back', label: 'Back' },
],
});
});
@ -698,7 +694,6 @@ describe('setup databases step', () => {
options: [
{ value: 'continue', label: 'Continue to knowledge sources' },
{ value: 'add', label: 'Add another primary source' },
{ value: 'back', label: 'Back' },
],
});
});

View file

@ -366,17 +366,26 @@ async function defaultListSchemas(projectDir: string, connectionId: string): Pro
return [];
}
function configuredSchemas(connection: KtxProjectConnectionConfig | undefined, driver: KtxSetupDatabaseDriver): string[] | undefined {
if (!connection) return undefined;
const spec = SCOPE_DISCOVERY_SPECS[driver];
if (!spec) return undefined;
const values = configuredScopeValues(connection, spec);
return values.length > 0 ? values : undefined;
}
async function defaultListTables(projectDir: string, connectionId: string): Promise<KtxTableListEntry[]> {
const project = await loadKtxProject({ projectDir });
const connection = project.config.connections[connectionId];
const driver = normalizeDriver(connection?.driver);
const schemas = driver ? configuredSchemas(connection, driver) : undefined;
if (driver === 'postgres') {
const { KtxPostgresScanConnector, isKtxPostgresConnectionConfig } = await import('@ktx/connector-postgres');
if (!isKtxPostgresConnectionConfig(connection)) return [];
const connector = new KtxPostgresScanConnector({ connectionId, connection });
try {
return await connector.listTables();
return await connector.listTables(schemas);
} finally {
await connector.cleanup();
}
@ -387,7 +396,7 @@ async function defaultListTables(projectDir: string, connectionId: string): Prom
if (!isKtxMysqlConnectionConfig(connection)) return [];
const connector = new KtxMysqlScanConnector({ connectionId, connection });
try {
return await connector.listTables();
return await connector.listTables(schemas);
} finally {
await connector.cleanup();
}
@ -398,7 +407,7 @@ async function defaultListTables(projectDir: string, connectionId: string): Prom
if (!isKtxSqlServerConnectionConfig(connection)) return [];
const connector = new KtxSqlServerScanConnector({ connectionId, connection });
try {
return await connector.listTables();
return await connector.listTables(schemas);
} finally {
await connector.cleanup();
}
@ -409,7 +418,7 @@ async function defaultListTables(projectDir: string, connectionId: string): Prom
if (!isKtxBigQueryConnectionConfig(connection)) return [];
const connector = new KtxBigQueryScanConnector({ connectionId, connection });
try {
return await connector.listTables();
return await connector.listTables(schemas);
} finally {
await connector.cleanup();
}
@ -420,7 +429,7 @@ async function defaultListTables(projectDir: string, connectionId: string): Prom
if (!isKtxSnowflakeConnectionConfig(connection)) return [];
const connector = new KtxSnowflakeScanConnector({ connectionId, connection });
try {
return await connector.listTables();
return await connector.listTables(schemas);
} finally {
await connector.cleanup();
}
@ -431,7 +440,7 @@ async function defaultListTables(projectDir: string, connectionId: string): Prom
if (!isKtxClickHouseConnectionConfig(connection)) return [];
const connector = new KtxClickHouseScanConnector({ connectionId, connection });
try {
return await connector.listTables();
return await connector.listTables(schemas);
} finally {
await connector.cleanup();
}
@ -476,7 +485,6 @@ function configuredPrimarySourcesPrompt(connectionIds: string[]): {
options: [
{ value: 'continue', label: 'Continue to knowledge sources' },
{ value: 'add', label: 'Add another primary source' },
{ value: 'back', label: 'Back' },
],
};
}
@ -1051,6 +1059,22 @@ async function writeScopeConfig(input: {
});
}
async function clearScopeConfig(projectDir: string, connectionId: string): Promise<void> {
const project = await loadKtxProject({ projectDir });
const connection = project.config.connections[connectionId];
if (!connection) return;
const driver = normalizeDriver(connection.driver);
if (!driver) return;
const spec = SCOPE_DISCOVERY_SPECS[driver];
if (!spec) return;
const cleaned = Object.fromEntries(
Object.entries(connection).filter(
([key]) => key !== spec.configArrayField && key !== spec.configSingleField && key !== 'enabled_tables',
),
) as KtxProjectConnectionConfig;
await writeConnectionConfig({ projectDir, connectionId, connection: cleaned });
}
async function maybeConfigureSchemaScope(input: {
projectDir: string;
connectionId: string;
@ -1204,43 +1228,49 @@ async function maybeConfigureTableScope(input: {
const schemaList = [...bySchema.keys()].sort();
const schemaSummary = schemaList.map((s) => `${s} (${bySchema.get(s)!.length})`).join(', ');
const action = await input.prompts.select({
message: `Tables found in selected schemas\n` +
`${discovered.length} tables across ${schemaList.length} ${schemaList.length === 1 ? 'schema' : 'schemas'}: ${schemaSummary}`,
options: [
{ value: 'all', label: 'Enable all tables' },
{ value: 'customize', label: 'Customize which tables to enable' },
],
});
let selected: string[] | null = null;
if (action === 'back') {
return false;
}
let selected: string[];
if (action === 'all') {
selected = allQualified;
} else {
const choices = await input.prompts.multiselect({
message: withMultiselectNavigation(
`Tables to enable for ${input.connectionId}\n` +
`Deselect any tables agents should not use.`,
),
options: discovered.map((t) => {
const qualified = `${t.schema}.${t.name}`;
const suffix = t.kind === 'view' ? ' (view)' : '';
return { value: qualified, label: `${qualified}${suffix}` };
}),
initialValues: allQualified,
required: true,
while (selected === null) {
const action = await input.prompts.select({
message: `Tables found in selected schemas\n` +
`${discovered.length} tables across ${schemaList.length} ${schemaList.length === 1 ? 'schema' : 'schemas'}: ${schemaSummary}`,
options: [
{ value: 'all', label: 'Enable all tables' },
{ value: 'customize', label: 'Customize which tables to enable' },
{ value: 'back', label: 'Back' },
],
});
if (choices.includes('back')) {
if (action === 'back') {
return false;
}
selected = choices.length > 0 ? choices : allQualified;
if (action === 'all') {
selected = allQualified;
} else {
const choices = await input.prompts.multiselect({
message: withMultiselectNavigation(
`Tables to enable for ${input.connectionId}\n` +
`Deselect any tables agents should not use.`,
),
options: discovered.map((t) => {
const qualified = `${t.schema}.${t.name}`;
const suffix = t.kind === 'view' ? ' (view)' : '';
return { value: qualified, label: `${qualified}${suffix}` };
}),
initialValues: allQualified,
required: true,
});
if (choices.includes('back')) {
continue;
}
if (choices.length === 0) {
input.io.stdout.write('│ KTX needs at least one table enabled. Select a table or press Escape to go back.\n');
continue;
}
selected = choices;
}
}
await writeConnectionConfig({
@ -1383,12 +1413,16 @@ async function validateAndScanConnection(input: {
testLines.push(`Driver: ${driverDisplay}${Number.isFinite(tableCount) ? ` · Tables: ${tableCount}` : ''}`);
writeSetupSection(input.io, `Testing ${input.connectionId}`, testLines);
if (!(await maybeConfigureSchemaScope(input))) {
return false;
}
while (true) {
if (!(await maybeConfigureSchemaScope(input))) {
return false;
}
if (!(await maybeConfigureTableScope(input))) {
return false;
if (await maybeConfigureTableScope(input)) {
break;
}
await clearScopeConfig(input.projectDir, input.connectionId);
}
await maybeRunHistoricSqlSetupProbe({
@ -1559,13 +1593,10 @@ export async function runKtxSetupDatabasesStep(
while (true) {
if (showConfiguredPrimaryMenu) {
const action = await prompts.select(configuredPrimarySourcesPrompt(selectedConnectionIds));
if (action === 'continue') {
if (action === 'continue' || action === 'back') {
await markDatabasesComplete(args.projectDir, selectedConnectionIds);
return { status: 'ready', projectDir: args.projectDir, connectionIds: selectedConnectionIds };
}
if (action === 'back') {
return { status: 'back', projectDir: args.projectDir };
}
}
showConfiguredPrimaryMenu = false;

View file

@ -411,6 +411,20 @@ export async function runLocalScan(options: RunLocalScanOptions): Promise<LocalS
extractedAtFallback: report.createdAt,
});
const structuralSnapshot = enabledTables ? filterSnapshotTables(rawSnapshot, enabledTables) : rawSnapshot;
if (enabledTables && structuralSnapshot.tables.length < rawSnapshot.tables.length) {
const excluded = rawSnapshot.tables.length - structuralSnapshot.tables.length;
let remaining = excluded;
const ds = report.diffSummary;
const subFrom = (field: 'tablesAdded' | 'tablesUnchanged' | 'tablesModified') => {
const take = Math.min(remaining, ds[field]);
ds[field] -= take;
remaining -= take;
};
subFrom('tablesAdded');
subFrom('tablesUnchanged');
subFrom('tablesModified');
await options.progress?.update(0.6, scanChangeSummary(report.diffSummary));
}
const manifestArtifacts = await writeLocalScanManifestShards({
project: options.project,
connectionId: options.connectionId,