From 9a8cb081924d47fdc9be09be9003b60e02bdbd40 Mon Sep 17 00:00:00 2001 From: Luca Martial Date: Tue, 12 May 2026 21:31:11 -0700 Subject: [PATCH] Refine setup table selection flow --- packages/cli/src/setup-agents.test.ts | 7 +- packages/cli/src/setup-agents.ts | 3 +- packages/cli/src/setup-databases.test.ts | 5 - packages/cli/src/setup-databases.ts | 125 ++++++++++++++--------- packages/context/src/scan/local-scan.ts | 14 +++ 5 files changed, 97 insertions(+), 57 deletions(-) diff --git a/packages/cli/src/setup-agents.test.ts b/packages/cli/src/setup-agents.test.ts index 9a984352..a5a065e8 100644 --- a/packages/cli/src/setup-agents.test.ts +++ b/packages/cli/src/setup-agents.test.ts @@ -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 () => { diff --git a/packages/cli/src/setup-agents.ts b/packages/cli/src/setup-agents.ts index 6a9721b9..97d8e610 100644 --- a/packages/cli/src/setup-agents.ts +++ b/packages/cli/src/setup-agents.ts @@ -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 = diff --git a/packages/cli/src/setup-databases.test.ts b/packages/cli/src/setup-databases.test.ts index 4c2abfcd..231aae84 100644 --- a/packages/cli/src/setup-databases.test.ts +++ b/packages/cli/src/setup-databases.test.ts @@ -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' }, ], }); }); diff --git a/packages/cli/src/setup-databases.ts b/packages/cli/src/setup-databases.ts index 53c6a7ab..b90ad9c5 100644 --- a/packages/cli/src/setup-databases.ts +++ b/packages/cli/src/setup-databases.ts @@ -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 { 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 { + 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; diff --git a/packages/context/src/scan/local-scan.ts b/packages/context/src/scan/local-scan.ts index 8cb50126..7f3c00a0 100644 --- a/packages/context/src/scan/local-scan.ts +++ b/packages/context/src/scan/local-scan.ts @@ -411,6 +411,20 @@ export async function runLocalScan(options: RunLocalScanOptions): Promise { + 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,