@kaelio/ktx 0.8.0 → 0.9.0

package/dist/setup-databases.js

@@ -4,15 +4,23 @@ import { delimiter, dirname, join } from 'node:path';
 import { fileURLToPath } from 'node:url';
 import { promisify } from 'node:util';
 import { getDriverRegistration } from './context/connections/drivers.js';
+import { createLocalKtxLlmRuntimeFromConfig } from './context/llm/local-config.js';
 import { queryHistoryDialectForConnection } from './context/ingest/adapters/historic-sql/connection-dialect.js';
+import { proposeQueryHistoryServiceAccountFilters, } from './context/ingest/adapters/historic-sql/query-history-filter-picker.js';
+import { resolveQueryHistoryScopeFloor } from './context/ingest/adapters/historic-sql/scope-floor.js';
 import { runHistoricSqlReadinessProbe, } from './context/ingest/historic-sql-probes.js';
 import { serializeKtxProjectConfig } from './context/project/config.js';
 import { loadKtxProject } from './context/project/project.js';
 import { markKtxSetupStateStepComplete, setKtxSetupDatabaseConnectionIds } from './context/project/setup-config.js';
+import { getKtxCliPackageInfo } from './cli-runtime.js';
 import { errorMessage, flushPrefixedBufferedCommandOutput, writePrefixedLines, } from './clack.js';
 import { runKtxConnection } from './connection.js';
+import { createBufferedCommandIo } from './io/buffered-command-io.js';
+import { runConnectionSetupWithRecovery, } from './connection-recovery.js';
 import { pickDatabaseScope as defaultPickDatabaseScope, } from './database-tree-picker.js';
 import { withMultiselectNavigation, withTextInputNavigation } from './prompt-navigation.js';
+import { createKtxCliHistoricSqlRuntime } from './local-adapters.js';
+import { queryHistoryPullConfig } from './public-ingest.js';
 import { runKtxScan } from './scan.js';
 import { writeProjectLocalSecretReference } from './setup-secrets.js';
 import { isDemoConnection } from './telemetry/demo-detect.js';
@@ -671,10 +679,13 @@ async function maybeApplyHistoricSqlConfig(input) {
     if (!enabled) {
         return withQueryHistoryConfig(input.connection, { ...existing, enabled: false });
     }
+    const existingFilters = existing.filters && typeof existing.filters === 'object' && !Array.isArray(existing.filters)
+        ? existing.filters
+        : {};
     const common = {
         ...existing,
         enabled: true,
-        filters: historicSqlFiltersForSetup(input.args.queryHistoryServiceAccountPatterns),
+        filters: historicSqlFiltersForSetup(input.args.queryHistoryServiceAccountPatterns, existingFilters),
     };
     if (dialect === 'postgres') {
         return withQueryHistoryConfig(input.connection, {
@@ -688,9 +699,10 @@ async function maybeApplyHistoricSqlConfig(input) {
         redactionPatterns: input.args.queryHistoryRedactionPatterns ?? [],
     });
 }
-function historicSqlFiltersForSetup(patterns) {
+function historicSqlFiltersForSetup(patterns, existingFilters = {}) {
     const serviceAccountPatterns = patterns ?? [];
     return {
+        ...existingFilters,
         dropTrivialProbes: true,
         ...(serviceAccountPatterns.length > 0
             ? {
@@ -715,29 +727,6 @@ async function defaultScanConnection(projectDir, connectionId, io) {
         dryRun: false,
     }, io);
 }
-function createBufferedCommandIo() {
-    let stdout = '';
-    let stderr = '';
-    return {
-        stdout: {
-            isTTY: false,
-            write(chunk) {
-                stdout += chunk;
-            },
-        },
-        stderr: {
-            write(chunk) {
-                stderr += chunk;
-            },
-        },
-        stdoutText() {
-            return stdout;
-        },
-        stderrText() {
-            return stderr;
-        },
-    };
-}
 function envWithCurrentNodeFirst(env = process.env) {
     return {
         ...env,
@@ -889,6 +878,27 @@ async function disableConnectionQueryHistory(projectDir, connectionId) {
         connection: withQueryHistoryConfig(connection, { ...existing, enabled: false }),
     });
 }
+function okValidateResult() {
+    return { status: 'ok' };
+}
+function backValidateResult() {
+    return { status: 'back' };
+}
+function failedValidateResult() {
+    return { status: 'failed' };
+}
+function queryHistoryUnavailableResult(projectDir, connectionId) {
+    return {
+        status: 'failed',
+        extraActions: [
+            {
+                value: 'disable-query-history',
+                label: 'Disable query history and retry',
+                run: () => disableConnectionQueryHistory(projectDir, connectionId),
+            },
+        ],
+    };
+}
 async function createConnectionConfigRollback(projectDir, connectionId) {
     const project = await loadKtxProject({ projectDir });
     const previousConnection = project.config.connections[connectionId];
@@ -983,14 +993,14 @@ async function maybeConfigureDatabaseScope(input) {
     const connection = project.config.connections[input.connectionId];
     const driver = normalizeDriver(connection?.driver);
     if (!driver || driver === 'sqlite')
-        return 'ready';
+        return okValidateResult();
     const spec = SCOPE_DISCOVERY_SPECS[driver];
     const existingTables = connection?.enabled_tables;
     const hasExistingTables = Array.isArray(existingTables) && existingTables.length > 0;
     const existingScope = spec ? configuredScopeValues(connection, spec) : [];
     const hasExistingScope = !spec || existingScope.length > 0;
     if (hasExistingTables && hasExistingScope && input.forcePrompt !== true) {
-        return 'ready';
+        return okValidateResult();
     }
     const cliSchemas = input.args.databaseSchemas;
     if (input.args.inputMode === 'disabled') {
@@ -1003,7 +1013,7 @@ async function maybeConfigureDatabaseScope(input) {
                 catch (error) {
                     const detail = error instanceof Error ? error.message : String(error);
                     input.io.stderr.write(`Could not discover ${spec.promptLabel.toLowerCase()} for ${input.connectionId}; ${detail}\n`);
-                    return 'ready';
+                    return okValidateResult();
                 }
             }
             if (scopeToWrite.length > 0) {
@@ -1019,7 +1029,7 @@ async function maybeConfigureDatabaseScope(input) {
                 ]);
             }
         }
-        return 'ready';
+        return okValidateResult();
     }
     if (spec && cliSchemas.length > 0) {
         await writeScopeConfig({
@@ -1050,7 +1060,7 @@ async function maybeConfigureDatabaseScope(input) {
                 spec,
             });
             if (typed === undefined)
-                return 'back';
+                return backValidateResult();
             effectiveCliSchemas = typed;
             listedSchemas = typed;
             if (typed.length > 0) {
@@ -1065,7 +1075,7 @@ async function maybeConfigureDatabaseScope(input) {
     }
     const schemas = unique(listedSchemas);
     if (spec && schemas.length === 0) {
-        return 'ready';
+        return okValidateResult();
     }
     const schemaSuggestion = effectiveCliSchemas.length > 0
         ? { excluded: new Set(), suggested: new Set(effectiveCliSchemas) }
@@ -1093,10 +1103,10 @@ async function maybeConfigureDatabaseScope(input) {
         writePrefixedLines((chunk) => input.io.stderr.write(chunk), input.forcePrompt === true
             ? `Could not discover tables for ${input.connectionId}; edit was not saved. ${detail}`
             : `Could not discover tables for ${input.connectionId}; continuing without table filter. ${detail}`);
-        return input.forcePrompt === true ? 'failed' : 'ready';
+        return input.forcePrompt === true ? failedValidateResult() : okValidateResult();
     }
     if (pickResult.kind === 'back') {
-        return 'back';
+        return backValidateResult();
     }
     const enabledTables = pickResult.enabledTables;
     const activeSchemas = pickResult.activeSchemas;
@@ -1111,7 +1121,7 @@ async function maybeConfigureDatabaseScope(input) {
     const refreshedProject = await loadKtxProject({ projectDir: input.projectDir });
     const currentConnection = refreshedProject.config.connections[input.connectionId];
     if (!currentConnection)
-        return 'ready';
+        return okValidateResult();
     await writeConnectionConfig({
         projectDir: input.projectDir,
         connectionId: input.connectionId,
@@ -1127,7 +1137,7 @@ async function maybeConfigureDatabaseScope(input) {
     writeSetupSection(input.io, `Tables enabled for ${input.connectionId}`, [
         `✓ ${enabledTables.length} tables enabled`,
     ]);
-    return 'ready';
+    return okValidateResult();
 }
 async function ensureHistoricSqlIngestDefaults(projectDir) {
     const project = await loadKtxProject({ projectDir });
@@ -1188,6 +1198,158 @@ async function maybeRunHistoricSqlSetupProbe(input) {
     }
     return result.ok;
 }
+function hasServiceAccountsBlock(connection) {
+    const queryHistory = queryHistoryConfigRecord(connection);
+    const filters = queryHistory?.filters;
+    if (!filters || typeof filters !== 'object' || Array.isArray(filters)) {
+        return false;
+    }
+    return 'serviceAccounts' in filters;
+}
+function printQueryHistoryFilterProposal(io, proposal) {
+    if (proposal.excludedRoles.length === 0) {
+        if (proposal.skipped?.reason === 'no-llm') {
+            io.stdout.write('│  Query-history filter picker skipped: no LLM is configured.\n');
+        }
+        else if (proposal.skipped?.reason === 'no-daemon') {
+            io.stdout.write('│  Query-history filter picker skipped: SQL analysis is unavailable.\n');
+        }
+        else if (proposal.skipped?.reason === 'no-in-scope-history') {
+            io.stdout.write('│  Query-history filter picker found no in-scope service-account exclusions.\n');
+        }
+        for (const warning of proposal.warnings) {
+            io.stdout.write(`│  ! ${warning}\n`);
+        }
+        return;
+    }
+    io.stdout.write('│  Proposed query-history service-account filters:\n');
+    for (const excluded of proposal.excludedRoles) {
+        io.stdout.write(`│  - ${excluded.role}: ${excluded.reason}\n`);
+    }
+}
+async function shouldApplyQueryHistoryFilterProposal(input) {
+    if (input.proposal.excludedRoles.length === 0 || input.proposal.skipped?.reason === 'user-block-present') {
+        return false;
+    }
+    if (input.args.yes === true || input.args.inputMode === 'disabled') {
+        return true;
+    }
+    const choice = await input.prompts.select({
+        message: `Apply ${input.proposal.excludedRoles.length} derived query-history service-account exclusion${input.proposal.excludedRoles.length === 1 ? '' : 's'}?`,
+        options: [
+            { value: 'apply', label: 'Apply derived filters (recommended)' },
+            { value: 'skip', label: 'Leave query history filters unchanged' },
+        ],
+    });
+    return choice === 'apply';
+}
+function createSetupQueryHistoryLlmRuntime(input) {
+    try {
+        return (input.deps.createQueryHistoryLlmRuntime?.(input.projectDir, input.project) ??
+            createLocalKtxLlmRuntimeFromConfig(input.project.config.llm, {
+                projectDir: input.projectDir,
+            }));
+    }
+    catch {
+        return null;
+    }
+}
+/** @internal */
+export function managedDaemonOptionsForSetupQueryHistoryPicker(input) {
+    return {
+        cliVersion: input.args.cliVersion ?? getKtxCliPackageInfo().version,
+        projectDir: input.projectDir,
+        installPolicy: input.args.runtimeInstallPolicy ?? (input.args.inputMode === 'disabled' ? 'never' : 'prompt'),
+        io: input.io,
+    };
+}
+async function maybeProposeQueryHistoryFilters(input) {
+    const project = await loadKtxProject({ projectDir: input.projectDir });
+    const connection = project.config.connections[input.connectionId];
+    const queryHistory = queryHistoryConfigRecord(connection);
+    if (!connection || queryHistory?.enabled !== true) {
+        return;
+    }
+    const dialect = queryHistoryDialectForConnection(connection);
+    if (!dialect) {
+        return;
+    }
+    const picker = input.deps.queryHistoryFilterPicker ?? proposeQueryHistoryServiceAccountFilters;
+    const llmRuntime = createSetupQueryHistoryLlmRuntime({
+        projectDir: input.projectDir,
+        project,
+        deps: input.deps,
+    });
+    if (!llmRuntime && !input.deps.queryHistoryFilterPicker) {
+        printQueryHistoryFilterProposal(input.io, {
+            excludedRoles: [],
+            consideredRoleCount: 0,
+            skipped: { reason: 'no-llm' },
+            warnings: [],
+        });
+        return;
+    }
+    const runtime = createKtxCliHistoricSqlRuntime(project, input.connectionId, {
+        managedDaemon: managedDaemonOptionsForSetupQueryHistoryPicker({
+            projectDir: input.projectDir,
+            args: input.args,
+            io: input.io,
+        }),
+    });
+    if (!runtime) {
+        return;
+    }
+    const userServiceAccountsPresent = hasServiceAccountsBlock(connection);
+    const scopeFloor = await resolveQueryHistoryScopeFloor({
+        projectDir: input.projectDir,
+        connectionId: input.connectionId,
+        driver: String(connection.driver ?? ''),
+        connection: connection,
+        storedQueryHistory: queryHistory,
+    });
+    const pullConfig = queryHistoryPullConfig({
+        stored: queryHistory,
+        dialect,
+        enabledTables: scopeFloor.enabledTables,
+        enabledSchemas: scopeFloor.enabledSchemas,
+        modeledTableCatalog: scopeFloor.modeledTableCatalog,
+        scopeFloorWarnings: scopeFloor.warnings,
+    });
+    const proposal = await picker({
+        connectionId: input.connectionId,
+        dialect,
+        queryClient: runtime.queryClient,
+        reader: runtime.reader,
+        sqlAnalysis: runtime.sqlAnalysis,
+        llmRuntime,
+        pullConfig,
+        userServiceAccountsPresent,
+    });
+    printQueryHistoryFilterProposal(input.io, proposal);
+    if (proposal.skipped?.reason === 'user-block-present') {
+        input.io.stdout.write('│  Existing query-history service-account filters left unchanged.\n');
+        return;
+    }
+    if (!(await shouldApplyQueryHistoryFilterProposal({ args: input.args, prompts: input.prompts, proposal }))) {
+        return;
+    }
+    await writeConnectionConfig({
+        projectDir: input.projectDir,
+        connectionId: input.connectionId,
+        connection: withQueryHistoryConfig(connection, {
+            ...queryHistory,
+            filters: {
+                ...(queryHistory.filters && typeof queryHistory.filters === 'object' && !Array.isArray(queryHistory.filters)
+                    ? queryHistory.filters
+                    : {}),
+                serviceAccounts: {
+                    mode: 'exclude',
+                    patterns: proposal.excludedRoles.map((role) => role.pattern),
+                },
+            },
+        }),
+    });
+}
 async function applyHistoricSqlConfigToExistingConnection(input) {
     if (input.args.inputMode === 'disabled' &&
         input.args.enableQueryHistory !== true &&
@@ -1225,7 +1387,7 @@ async function validateAndScanConnection(input) {
     if (testCode !== 0) {
         flushPrefixedBufferedCommandOutput(input.io, testIo);
         writePrefixedLines((chunk) => input.io.stderr.write(chunk), `Connection test failed for ${input.connectionId}.`);
-        return 'failed';
+        return failedValidateResult();
     }
     const testOutput = testIo.stdoutText();
     const outputDriver = normalizeDriver(readOutputValue(testOutput, 'Driver'));
@@ -1233,7 +1395,7 @@ async function validateAndScanConnection(input) {
     const testLines = ['✓ Connection test passed', `Driver: ${driverDisplay}`];
     writeSetupSection(input.io, `Testing ${input.connectionId}`, testLines);
     const scopeStatus = await maybeConfigureDatabaseScope({ ...input, forcePrompt: input.forceScopeAndTables });
-    if (scopeStatus !== 'ready') {
+    if (scopeStatus.status !== 'ok') {
         return scopeStatus;
     }
     const queryHistoryAvailable = await maybeRunHistoricSqlSetupProbe({
@@ -1282,15 +1444,27 @@ async function validateAndScanConnection(input) {
             ].join('\n'));
         }
         if (scanCode !== 0) {
-            return queryHistoryAvailable ? 'failed' : 'failed-query-history-unavailable';
+            return queryHistoryAvailable
+                ? failedValidateResult()
+                : queryHistoryUnavailableResult(input.projectDir, input.connectionId);
         }
     }
     const scanOutput = scanIo.stdoutText();
     writeSetupSection(input.io, `Schema context complete for ${input.connectionId}`, [`Changes: ${summarizeScanChanges(scanOutput)}`]);
+    if (queryHistoryAvailable) {
+        await maybeProposeQueryHistoryFilters({
+            projectDir: input.projectDir,
+            connectionId: input.connectionId,
+            io: input.io,
+            deps: input.deps,
+            args: input.args,
+            prompts: input.prompts,
+        });
+    }
     writeSetupSection(input.io, 'Database ready', [
         `${input.connectionId} · ${driverDisplay} · schema context complete`,
     ]);
-    return 'ready';
+    return okValidateResult();
 }
 async function chooseDrivers(args, io, prompts, options) {
     if (args.databaseDrivers && args.databaseDrivers.length > 0) {
@@ -1394,64 +1568,137 @@ async function choosePrimarySourceToEdit(input) {
     });
     return choice === 'back' ? 'back' : choice;
 }
-async function runPrimarySourceFullEdit(input) {
+async function configureDatabaseConnection(input) {
     const project = await loadKtxProject({ projectDir: input.projectDir });
-    const existing = project.config.connections[input.connectionId];
-    const driver = normalizeDriver(existing?.driver);
-    if (!existing || !driver) {
-        writePrefixedLines((chunk) => input.io.stderr.write(chunk), `Connection "${input.connectionId}" is not a configured database.`);
-        return 'failed';
-    }
-    const rollback = await createConnectionConfigRollback(input.projectDir, input.connectionId);
-    const replacement = await buildConnectionConfig({
-        driver,
+    const latestConnection = project.config.connections[input.connectionId];
+    let connection = await buildConnectionConfig({
+        driver: input.driver,
         connectionId: input.connectionId,
         args: input.args,
         prompts: input.prompts,
-        existingConnection: existing,
+        existingConnection: latestConnection,
     });
-    if (replacement === 'back') {
-        await rollback();
+    while (!connection && input.args.inputMode !== 'disabled') {
+        const action = await input.prompts.select(missingConnectionDetailsPrompt(driverLabel(input.driver), input.canReturnToDriverSelection));
+        if (action === 'back') {
+            return 'back';
+        }
+        connection = await buildConnectionConfig({
+            driver: input.driver,
+            connectionId: input.connectionId,
+            args: input.args,
+            prompts: input.prompts,
+            existingConnection: latestConnection,
+        });
+    }
+    if (connection === 'back') {
         return 'back';
     }
-    if (!replacement) {
-        await rollback();
-        return 'failed';
+    if (!connection) {
+        input.io.stderr.write(`Missing connection details for ${driverLabel(input.driver)}.\n`);
+        return 'cancelled';
     }
     const withHistoricSql = await maybeApplyHistoricSqlConfig({
-        connection: replacement,
-        driver,
+        connection,
+        driver: input.driver,
         args: input.args,
         prompts: input.prompts,
     });
     if (withHistoricSql === 'back') {
-        await rollback();
         return 'back';
     }
     await writeConnectionConfig({
         projectDir: input.projectDir,
         connectionId: input.connectionId,
-        connection: withExistingPrimaryEditPromptDefaults({
-            previous: existing,
-            next: withHistoricSql,
-            driver,
-        }),
+        connection: input.editBaseline
+            ? withExistingPrimaryEditPromptDefaults({
+                previous: input.editBaseline,
+                next: withHistoricSql,
+                driver: input.driver,
+            })
+            : withHistoricSql,
+        io: input.io,
+    });
+    return 'configured';
+}
+async function runDatabaseConnectionSetupWithRecovery(input) {
+    let configureCalls = 0;
+    // `configureDatabaseConnection` returns 'cancelled' only when required
+    // connection details are absent in non-interactive mode. The recovery
+    // primitive collapses that into 'failed', so we track it here to restore the
+    // distinct 'missing-input' outcome the surrounding step reports for
+    // incomplete flags (vs. a real connection/probe failure).
+    let sawMissingInput = false;
+    const outcome = await runConnectionSetupWithRecovery({
+        label: input.connectionId,
+        interactive: input.interactive ?? input.args.inputMode !== 'disabled',
+        allowSkip: input.allowSkip,
         io: input.io,
+        prompts: input.prompts,
+        snapshot: () => createConnectionConfigRollback(input.projectDir, input.connectionId),
+        configure: async () => {
+            configureCalls += 1;
+            if (input.reuseExistingOnFirstConfigure && configureCalls === 1) {
+                const historicSqlResult = await applyHistoricSqlConfigToExistingConnection({
+                    projectDir: input.projectDir,
+                    connectionId: input.connectionId,
+                    args: input.args,
+                    prompts: input.prompts,
+                });
+                return historicSqlResult === 'back' ? 'back' : 'configured';
+            }
+            const configured = await configureDatabaseConnection({
+                projectDir: input.projectDir,
+                connectionId: input.connectionId,
+                driver: input.driver,
+                args: input.args,
+                prompts: input.prompts,
+                io: input.io,
+                canReturnToDriverSelection: input.canReturnToDriverSelection,
+                editBaseline: input.editBaseline,
+            });
+            if (configured === 'cancelled') {
+                sawMissingInput = true;
+            }
+            return configured;
+        },
+        validate: () => validateAndScanConnection({
+            projectDir: input.projectDir,
+            connectionId: input.connectionId,
+            io: input.io,
+            deps: input.deps,
+            args: input.args,
+            prompts: input.prompts,
+            forceScopeAndTables: input.forceScopeAndTables,
+        }),
     });
-    const validated = await validateAndScanConnection({
+    if (outcome === 'failed' && sawMissingInput) {
+        return 'missing-input';
+    }
+    return outcome;
+}
+async function runPrimarySourceFullEdit(input) {
+    const project = await loadKtxProject({ projectDir: input.projectDir });
+    const existing = project.config.connections[input.connectionId];
+    const driver = normalizeDriver(existing?.driver);
+    if (!existing || !driver) {
+        writePrefixedLines((chunk) => input.io.stderr.write(chunk), `Connection "${input.connectionId}" is not a configured database.`);
+        return 'failed';
+    }
+    const outcome = await runDatabaseConnectionSetupWithRecovery({
         projectDir: input.projectDir,
         connectionId: input.connectionId,
-        io: input.io,
-        deps: input.deps,
+        driver,
         args: input.args,
         prompts: input.prompts,
+        io: input.io,
+        deps: input.deps,
+        canReturnToDriverSelection: true,
+        allowSkip: false,
         forceScopeAndTables: true,
+        editBaseline: existing,
     });
-    if (validated !== 'ready') {
-        await rollback();
-        return validated === 'failed-query-history-unavailable' ? 'failed' : validated;
-    }
-    return 'ready';
+    return outcome === 'skip' ? 'back' : outcome;
 }
 export async function runKtxSetupDatabasesStep(args, io, deps = {}) {
     if (args.skipDatabases) {
@@ -1462,29 +1709,37 @@ export async function runKtxSetupDatabasesStep(args, io, deps = {}) {
     if (args.databaseConnectionIds && args.databaseConnectionIds.length > 0) {
         const selectedConnectionIds = [];
         for (const connectionId of unique(args.databaseConnectionIds)) {
-            const historicSqlResult = await applyHistoricSqlConfigToExistingConnection({
+            const project = await loadKtxProject({ projectDir: args.projectDir });
+            const driver = normalizeDriver(project.config.connections[connectionId]?.driver);
+            if (!driver) {
+                writePrefixedLines((chunk) => io.stderr.write(chunk), `Connection "${connectionId}" is not configured.`);
+                return { status: 'failed', projectDir: args.projectDir };
+            }
+            const setupOutcome = await runDatabaseConnectionSetupWithRecovery({
                 projectDir: args.projectDir,
                 connectionId,
+                driver,
                 args,
                 prompts,
-            });
-            if (historicSqlResult === 'back')
-                return { status: 'back', projectDir: args.projectDir };
-            const setupStatus = await validateAndScanConnection({
-                projectDir: args.projectDir,
-                connectionId,
                 io,
                 deps,
-                args,
-                prompts,
+                canReturnToDriverSelection: false,
+                allowSkip: false,
+                interactive: false,
+                reuseExistingOnFirstConfigure: true,
             });
-            if (setupStatus === 'back') {
+            if (setupOutcome === 'back') {
                 return { status: 'back', projectDir: args.projectDir };
             }
-            if (setupStatus === 'failed') {
+            if (setupOutcome === 'missing-input') {
+                return { status: 'missing-input', projectDir: args.projectDir };
+            }
+            if (setupOutcome === 'failed') {
                 return { status: 'failed', projectDir: args.projectDir };
             }
-            selectedConnectionIds.push(connectionId);
+            if (setupOutcome === 'ready') {
+                selectedConnectionIds.push(connectionId);
+            }
         }
         await markDatabasesComplete(args.projectDir, selectedConnectionIds);
         return { status: 'ready', projectDir: args.projectDir, connectionIds: selectedConnectionIds };
@@ -1533,6 +1788,9 @@ export async function runKtxSetupDatabasesStep(args, io, deps = {}) {
                     showConfiguredPrimaryMenu = true;
                     continue;
                 }
+                if (editResult === 'missing-input') {
+                    return { status: 'missing-input', projectDir: args.projectDir };
+                }
                 if (editResult === 'failed') {
                     return { status: 'failed', projectDir: args.projectDir };
                 }
@@ -1587,7 +1845,6 @@ export async function runKtxSetupDatabasesStep(args, io, deps = {}) {
                 io.stderr.write('Missing database connection id: pass --database-connection-id.\n');
                 return { status: 'missing-input', projectDir: args.projectDir };
             }
-            let connectionAlreadyValidated = false;
             if (connectionChoice.kind === 'edit') {
                 const editResult = await runPrimarySourceFullEdit({
                     projectDir: args.projectDir,
@@ -1603,192 +1860,44 @@ export async function runKtxSetupDatabasesStep(args, io, deps = {}) {
                     returnToDriverSelection = true;
                     break;
                 }
-                if (editResult === 'failed') {
-                    return { status: 'failed', projectDir: args.projectDir };
-                }
-                connectionAlreadyValidated = true;
-            }
-            else if (connectionChoice.kind === 'new') {
-                let connection = await buildConnectionConfig({
-                    driver,
-                    connectionId: connectionChoice.connectionId,
-                    args,
-                    prompts,
-                });
-                if (connection === 'back') {
-                    if (!canReturnToDriverSelection)
-                        return { status: 'back', projectDir: args.projectDir };
-                    returnToDriverSelection = true;
-                    break;
-                }
-                while (!connection && args.inputMode !== 'disabled') {
-                    const label = driverLabel(driver);
-                    const action = await prompts.select(missingConnectionDetailsPrompt(label, canReturnToDriverSelection));
-                    if (action === 'back') {
-                        if (!canReturnToDriverSelection)
-                            return { status: 'back', projectDir: args.projectDir };
-                        returnToDriverSelection = true;
-                        break;
-                    }
-                    connection = await buildConnectionConfig({
-                        driver,
-                        connectionId: connectionChoice.connectionId,
-                        args,
-                        prompts,
-                    });
-                    if (connection === 'back') {
-                        if (!canReturnToDriverSelection)
-                            return { status: 'back', projectDir: args.projectDir };
-                        returnToDriverSelection = true;
-                        break;
-                    }
-                }
-                if (returnToDriverSelection) {
-                    break;
-                }
-                if (connection === 'back') {
-                    break;
-                }
-                if (!connection) {
-                    io.stderr.write(`Missing connection details for ${driverLabel(driver)}.\n`);
+                if (editResult === 'missing-input') {
                     return { status: 'missing-input', projectDir: args.projectDir };
                 }
-                const withHistoricSql = await maybeApplyHistoricSqlConfig({ connection, driver, args, prompts });
-                if (withHistoricSql === 'back') {
-                    if (!canReturnToDriverSelection)
-                        return { status: 'back', projectDir: args.projectDir };
-                    returnToDriverSelection = true;
-                    break;
+                if (editResult === 'failed') {
+                    return { status: 'failed', projectDir: args.projectDir };
                 }
-                await writeConnectionConfig({
-                    projectDir: args.projectDir,
-                    connectionId: connectionChoice.connectionId,
-                    connection: withHistoricSql,
-                    io,
-                });
             }
             else {
-                const existing = project.config.connections[connectionChoice.connectionId];
-                const withHistoricSql = await maybeApplyHistoricSqlConfig({ connection: existing, driver, args, prompts });
-                if (withHistoricSql === 'back') {
-                    if (!canReturnToDriverSelection)
-                        return { status: 'back', projectDir: args.projectDir };
-                    returnToDriverSelection = true;
-                    break;
-                }
-                await writeConnectionConfig({
-                    projectDir: args.projectDir,
-                    connectionId: connectionChoice.connectionId,
-                    connection: withHistoricSql,
-                    io,
-                });
-            }
-            let connectionSkipped = false;
-            let setupStatus = connectionAlreadyValidated
-                ? 'ready'
-                : await validateAndScanConnection({
+                const setupOutcome = await runDatabaseConnectionSetupWithRecovery({
                     projectDir: args.projectDir,
                     connectionId: connectionChoice.connectionId,
-                    io,
-                    deps,
+                    driver,
                     args,
                     prompts,
+                    io,
+                    deps,
+                    canReturnToDriverSelection,
+                    allowSkip: true,
+                    reuseExistingOnFirstConfigure: connectionChoice.kind === 'existing',
                 });
-            while (!connectionAlreadyValidated && setupStatus !== 'ready') {
-                if (setupStatus === 'back') {
-                    if (!canReturnToDriverSelection)
-                        return { status: 'back', projectDir: args.projectDir };
-                    returnToDriverSelection = true;
-                    break;
-                }
-                if (args.inputMode === 'disabled')
-                    return { status: 'failed', projectDir: args.projectDir };
-                const failureOptions = [
-                    { value: 'retry', label: 'Retry connection test' },
-                    { value: 're-enter', label: 'Re-enter connection details' },
-                    ...(setupStatus === 'failed-query-history-unavailable'
-                        ? [{ value: 'disable-query-history', label: 'Disable query history and retry' }]
-                        : []),
-                    { value: 'skip', label: 'Skip this database' },
-                    { value: 'back', label: 'Back' },
-                ];
-                const action = await prompts.select({
-                    message: `Database setup failed for ${connectionChoice.connectionId}`,
-                    options: failureOptions,
-                });
-                if (action === 'back') {
+                if (setupOutcome === 'back') {
                     if (!canReturnToDriverSelection)
                         return { status: 'back', projectDir: args.projectDir };
                     returnToDriverSelection = true;
                     break;
                 }
-                if (action === 'skip') {
-                    connectionSkipped = true;
-                    break;
-                }
-                if (action === 'retry') {
-                    setupStatus = await validateAndScanConnection({
-                        projectDir: args.projectDir,
-                        connectionId: connectionChoice.connectionId,
-                        io,
-                        deps,
-                        args,
-                        prompts,
-                    });
+                if (setupOutcome === 'missing-input') {
+                    return { status: 'missing-input', projectDir: args.projectDir };
                 }
-                else if (action === 'disable-query-history') {
-                    await disableConnectionQueryHistory(args.projectDir, connectionChoice.connectionId);
-                    setupStatus = await validateAndScanConnection({
-                        projectDir: args.projectDir,
-                        connectionId: connectionChoice.connectionId,
-                        io,
-                        deps,
-                        args,
-                        prompts,
-                    });
+                if (setupOutcome === 'failed') {
+                    return { status: 'failed', projectDir: args.projectDir };
                 }
-                else if (action === 're-enter') {
-                    const connection = await buildConnectionConfig({
-                        driver,
-                        connectionId: connectionChoice.connectionId,
-                        args,
-                        prompts,
-                    });
-                    if (connection === 'back') {
-                        if (!canReturnToDriverSelection)
-                            return { status: 'back', projectDir: args.projectDir };
-                        returnToDriverSelection = true;
-                        break;
-                    }
-                    if (!connection)
-                        continue;
-                    const withHistoricSql = await maybeApplyHistoricSqlConfig({ connection, driver, args, prompts });
-                    if (withHistoricSql === 'back') {
-                        if (!canReturnToDriverSelection)
-                            return { status: 'back', projectDir: args.projectDir };
-                        returnToDriverSelection = true;
-                        break;
-                    }
-                    await writeConnectionConfig({
-                        projectDir: args.projectDir,
-                        connectionId: connectionChoice.connectionId,
-                        connection: withHistoricSql,
-                        io,
-                    });
-                    setupStatus = await validateAndScanConnection({
-                        projectDir: args.projectDir,
-                        connectionId: connectionChoice.connectionId,
-                        io,
-                        deps,
-                        args,
-                        prompts,
-                    });
+                if (setupOutcome === 'skip') {
+                    continue;
                 }
             }
             if (returnToDriverSelection)
                 break;
-            if (connectionSkipped)
-                continue;
             pushUniqueConnectionId(selectedConnectionIds, connectionChoice.connectionId);
         }
         if (returnToDriverSelection) {