From 5f9c35558a2227a55da4aed64b2011f66f595d68 Mon Sep 17 00:00:00 2001 From: Luca Martial Date: Wed, 13 May 2026 14:18:26 -0700 Subject: [PATCH] feat(cli): add edit flow for context source connections in setup Allow users to edit existing context source connections during setup. Preselects existing values (URLs, credentials, repo details) and offers a "Keep existing credential" option for sensitive fields. Co-Authored-By: Claude Opus 4.6 (1M context) --- packages/cli/src/setup-sources.test.ts | 313 ++++++++++++++++ packages/cli/src/setup-sources.ts | 488 +++++++++++++++++++++---- 2 files changed, 735 insertions(+), 66 deletions(-) diff --git a/packages/cli/src/setup-sources.test.ts b/packages/cli/src/setup-sources.test.ts index 7fe61f76..0a0eab2c 100644 --- a/packages/cli/src/setup-sources.test.ts +++ b/packages/cli/src/setup-sources.test.ts @@ -861,6 +861,7 @@ describe('setup sources step', () => { message: 'Configure dbt', options: [ { value: 'existing:dbt-main', label: 'Use existing dbt connection: dbt-main' }, + { value: 'edit:dbt-main', label: 'Edit existing dbt connection: dbt-main' }, { value: 'new', label: 'Add new dbt connection' }, { value: 'back', label: 'Back' }, ], @@ -988,6 +989,10 @@ describe('setup sources step', () => { value: `existing:${testCase.connectionId}`, label: `Use existing ${testCase.expectedLabel} connection: ${testCase.connectionId}`, }, + { + value: `edit:${testCase.connectionId}`, + label: `Edit existing ${testCase.expectedLabel} connection: ${testCase.connectionId}`, + }, { value: 'new', label: `Add new ${testCase.expectedLabel} connection` }, { value: 'back', label: 'Back' }, ], @@ -996,6 +1001,314 @@ describe('setup sources step', () => { } }); + it('edits an existing Notion source and reopens the page picker with stored pages selected', async () => { + await addPrimarySource(); + await addConnection('notion-main', { + driver: 'notion', + auth_token_ref: 'env:NOTION_TOKEN', + crawl_mode: 'selected_roots', + root_page_ids: ['old-page'], + root_database_ids: [], + root_data_source_ids: [], + }); + const validateNotion = vi.fn(async () => ({ ok: true as const, detail: 'roots=1' })); + const pickNotionRootPages = vi.fn(async () => ({ kind: 'selected' as const, rootPageIds: ['new-page'] })); + const testPrompts = prompts({ + multiselect: [['notion']], + select: ['edit:notion-main', 'keep', 'selected_roots', 'done'], + }); + + await expect( + runKtxSetupSourcesStep( + { projectDir, inputMode: 'auto', runInitialSourceIngest: false, skipSources: false }, + makeIo().io, + { + prompts: testPrompts, + validateNotion, + pickNotionRootPages, + }, + ), + ).resolves.toEqual({ status: 'ready', projectDir, connectionIds: ['notion-main'] }); + + expect(testPrompts.select).toHaveBeenCalledWith({ + message: 'How should KTX find your Notion integration token?', + options: [ + { value: 'keep', label: 'Keep existing credential' }, + { value: 'env', label: 'Use NOTION_TOKEN from the environment' }, + { value: 'paste', label: 'Paste a key and save it as a local secret file' }, + { value: 'back', label: 'Back' }, + ], + }); + expect(pickNotionRootPages).toHaveBeenCalledWith( + { + connectionId: 'notion-main', + connection: expect.objectContaining({ + driver: 'notion', + auth_token_ref: 'env:NOTION_TOKEN', + crawl_mode: 'selected_roots', + root_page_ids: ['old-page'], + }), + }, + expect.anything(), + ); + expect((await readConfig()).connections['notion-main']).toMatchObject({ + driver: 'notion', + auth_token_ref: 'env:NOTION_TOKEN', + crawl_mode: 'selected_roots', + root_page_ids: ['new-page'], + }); + }); + + it('edits an existing Metabase source with the current URL and credential as defaults', async () => { + await addPrimarySource(); + await addConnection('metabase-main', { + driver: 'metabase', + api_url: 'https://metabase-old.example.com', + api_key_ref: 'env:METABASE_API_KEY', // pragma: allowlist secret + mappings: { + databaseMappings: { '1': 'warehouse' }, + syncEnabled: { '1': true }, + syncMode: 'ALL', + }, + }); + const testPrompts = prompts({ + multiselect: [['metabase']], + select: ['edit:metabase-main', 'keep', 'done'], + text: ['https://metabase-new.example.com'], + }); + const discoverMetabaseDatabases = vi.fn(async () => [ + { id: 2, name: 'Analytics', engine: 'postgres', host: 'db.example.com', dbName: 'analytics' }, + ]); + + await expect( + runKtxSetupSourcesStep( + { projectDir, inputMode: 'auto', runInitialSourceIngest: false, skipSources: false }, + makeIo().io, + { + prompts: testPrompts, + discoverMetabaseDatabases, + validateMetabase: vi.fn(async () => ({ ok: true as const, detail: 'mapping validated' })), + runMapping: vi.fn(async () => 0), + }, + ), + ).resolves.toEqual({ status: 'ready', projectDir, connectionIds: ['metabase-main'] }); + + expect(testPrompts.text).toHaveBeenCalledWith({ + message: textInputPrompt('Metabase URL'), + initialValue: 'https://metabase-old.example.com', + }); + expect(testPrompts.select).toHaveBeenCalledWith({ + message: 'How should KTX find your Metabase API key?', + options: [ + { value: 'keep', label: 'Keep existing credential' }, + { value: 'env', label: 'Use METABASE_API_KEY from the environment' }, + { value: 'paste', label: 'Paste a key and save it as a local secret file' }, + { value: 'back', label: 'Back' }, + ], + }); + expect(discoverMetabaseDatabases).toHaveBeenCalledWith({ + sourceUrl: 'https://metabase-new.example.com', + sourceApiKeyRef: 'env:METABASE_API_KEY', + sourceConnectionId: 'metabase-main', + }); + expect((await readConfig()).connections['metabase-main']).toMatchObject({ + driver: 'metabase', + api_url: 'https://metabase-new.example.com', + api_key_ref: 'env:METABASE_API_KEY', + mappings: { + databaseMappings: { '2': 'warehouse' }, + syncEnabled: { '2': true }, + syncMode: 'ALL', + }, + }); + }); + + it('rolls back an edited context source when validation fails', async () => { + await addPrimarySource(); + await addConnection('dbt-main', { + driver: 'dbt', + source_dir: '/repo/existing-dbt', + project_name: 'analytics', + }); + const validateDbt = vi.fn(async () => ({ ok: false as const, message: 'dbt project not found' })); + const testPrompts = prompts({ + multiselect: [['dbt']], + select: ['edit:dbt-main', 'path'], + text: ['/repo/new-dbt', ''], + }); + + await expect( + runKtxSetupSourcesStep( + { projectDir, inputMode: 'auto', runInitialSourceIngest: false, skipSources: false }, + makeIo().io, + { + prompts: testPrompts, + validateDbt, + }, + ), + ).resolves.toEqual({ status: 'failed', projectDir }); + + expect(validateDbt).toHaveBeenCalledWith(expect.objectContaining({ + driver: 'dbt', + source_dir: '/repo/new-dbt', + })); + const config = await readConfig(); + expect(config.connections['dbt-main']).toMatchObject({ + driver: 'dbt', + source_dir: '/repo/existing-dbt', + project_name: 'analytics', + }); + expect(config.ingest.adapters).not.toContain('dbt'); + }); + + it('lets git-backed context source edits keep the existing repo credential', async () => { + await addPrimarySource(); + await addConnection('metricflow-main', { + driver: 'metricflow', + metricflow: { + repoUrl: 'https://github.com/acme/private-metricflow', + branch: 'main', + path: 'metrics', + auth_token_ref: 'env:METRICFLOW_REPO_TOKEN', // pragma: allowlist secret + }, + }); + const testGitRepo = vi.fn(async () => ({ ok: false as const, error: 'authentication required' })); + const testPrompts = prompts({ + multiselect: [['metricflow']], + select: ['edit:metricflow-main', 'git', 'keep', 'done'], + text: ['https://github.com/acme/private-metricflow', 'main', 'metrics'], + }); + + await expect( + runKtxSetupSourcesStep( + { projectDir, inputMode: 'auto', runInitialSourceIngest: false, skipSources: false }, + makeIo().io, + { + prompts: testPrompts, + testGitRepo, + validateMetricflow: vi.fn(async () => ({ ok: true as const, detail: 'metrics=1' })), + }, + ), + ).resolves.toEqual({ status: 'ready', projectDir, connectionIds: ['metricflow-main'] }); + + expect(testPrompts.select).toHaveBeenCalledWith({ + message: 'This MetricFlow repo requires authentication.', + options: [ + { value: 'keep', label: 'Keep existing credential' }, + { value: 'env', label: 'Use GITHUB_TOKEN from the environment' }, + { value: 'paste', label: 'Paste a token and save it as a local secret file' }, + { value: 'skip', label: 'Skip — try without authentication' }, + { value: 'back', label: 'Back' }, + ], + }); + expect((await readConfig()).connections['metricflow-main']).toMatchObject({ + driver: 'metricflow', + metricflow: { + repoUrl: 'https://github.com/acme/private-metricflow', + branch: 'main', + path: 'metrics', + auth_token_ref: 'env:METRICFLOW_REPO_TOKEN', + }, + }); + }); + + it('edits an existing context source from the configured-source follow-up menu', async () => { + await addPrimarySource(); + await addConnection('dbt-main', { + driver: 'dbt', + source_dir: '/repo/existing-dbt', + project_name: 'analytics', + }); + const validateDbt = vi.fn(async () => ({ ok: true as const, detail: 'project=analytics schemas=2' })); + const testPrompts = prompts({ + multiselect: [['dbt']], + select: ['existing:dbt-main', 'edit', 'dbt-main', 'path', 'done'], + text: ['/repo/edited-dbt', ''], + }); + + await expect( + runKtxSetupSourcesStep( + { projectDir, inputMode: 'auto', runInitialSourceIngest: false, skipSources: false }, + makeIo().io, + { + prompts: testPrompts, + validateDbt, + }, + ), + ).resolves.toEqual({ status: 'ready', projectDir, connectionIds: ['dbt-main'] }); + + expect(testPrompts.select).toHaveBeenCalledWith({ + message: '1 context source configured (dbt-main). Add another?', + options: [ + { value: 'done', label: 'Done — continue to context build' }, + { value: 'edit', label: 'Edit an existing context source' }, + { value: 'add', label: 'Add another context source' }, + ], + }); + expect(testPrompts.select).toHaveBeenCalledWith({ + message: 'Context source to edit', + options: [ + { value: 'dbt-main', label: 'dbt-main (dbt)' }, + { value: 'back', label: 'Back' }, + ], + }); + expect(testPrompts.text).toHaveBeenCalledWith({ + message: textInputPrompt('dbt local path'), + initialValue: '/repo/existing-dbt', + }); + expect(validateDbt).toHaveBeenLastCalledWith(expect.objectContaining({ + driver: 'dbt', + source_dir: '/repo/edited-dbt', + })); + expect((await readConfig()).connections['dbt-main']).toMatchObject({ + driver: 'dbt', + source_dir: '/repo/edited-dbt', + project_name: 'analytics', + }); + }); + + it('backs out of editing an existing context source to the source connection menu', async () => { + await addPrimarySource(); + await addConnection('dbt-main', { + driver: 'dbt', + source_dir: '/repo/existing-dbt', + project_name: 'analytics', + }); + const validateDbt = vi.fn(async () => ({ ok: true as const, detail: 'project=analytics schemas=2' })); + const testPrompts = prompts({ + multiselect: [['dbt']], + select: ['edit:dbt-main', 'back', 'existing:dbt-main'], + }); + + await expect( + runKtxSetupSourcesStep( + { projectDir, inputMode: 'auto', runInitialSourceIngest: false, skipSources: false }, + makeIo().io, + { + prompts: testPrompts, + validateDbt, + }, + ), + ).resolves.toEqual({ status: 'ready', projectDir, connectionIds: ['dbt-main'] }); + + expect( + vi + .mocked(testPrompts.select) + .mock.calls.map(([options]) => options.message) + .filter((message) => message === 'Configure dbt'), + ).toHaveLength(2); + expect(validateDbt).toHaveBeenCalledWith({ + driver: 'dbt', + source_dir: '/repo/existing-dbt', + project_name: 'analytics', + }); + expect((await readConfig()).connections['dbt-main']).toMatchObject({ + driver: 'dbt', + source_dir: '/repo/existing-dbt', + project_name: 'analytics', + }); + }); + it('lets Escape from dbt git URL return to source location selection', async () => { await addPrimarySource(); const validateDbt = vi.fn(async () => ({ ok: true as const, detail: 'project=analytics schemas=2' })); diff --git a/packages/cli/src/setup-sources.ts b/packages/cli/src/setup-sources.ts index d18004d9..e0819e4a 100644 --- a/packages/cli/src/setup-sources.ts +++ b/packages/cli/src/setup-sources.ts @@ -224,17 +224,20 @@ async function chooseSourceCredentialRef(input: { label: string; envName: string; secretFileName: string; + existingRef?: string; }): Promise { while (true) { const choice = await input.prompts.select({ message: `How should KTX find your ${input.label}?`, options: [ + ...(input.existingRef ? [{ value: 'keep', label: 'Keep existing credential' }] : []), { value: 'env', label: `Use ${input.envName} from the environment` }, { value: 'paste', label: 'Paste a key and save it as a local secret file' }, { value: 'back', label: 'Back' }, ], }); if (choice === 'back') return 'back'; + if (choice === 'keep' && input.existingRef) return input.existingRef; if (choice === 'paste') { const value = await input.prompts.password({ message: input.label }); if (value === undefined) continue; @@ -256,12 +259,14 @@ async function chooseGitAuthCredentialRef(input: { projectDir: string; source: KtxSetupSourceType; connectionId: string; + existingRef?: string; }): Promise { const label = input.source === 'dbt' ? 'This' : `This ${sourceLabel(input.source)}`; while (true) { const choice = await input.prompts.select({ message: `${label} repo requires authentication.`, options: [ + ...(input.existingRef ? [{ value: 'keep', label: 'Keep existing credential' }] : []), { value: 'env', label: 'Use GITHUB_TOKEN from the environment' }, { value: 'paste', label: 'Paste a token and save it as a local secret file' }, { value: 'skip', label: 'Skip — try without authentication' }, @@ -269,6 +274,7 @@ async function chooseGitAuthCredentialRef(input: { ], }); if (choice === 'back') return 'back'; + if (choice === 'keep' && input.existingRef) return input.existingRef; if (choice === 'skip') return undefined; if (choice === 'paste') { const value = await input.prompts.password({ message: 'Git access token' }); @@ -793,8 +799,14 @@ interface WarehouseConnectionChoice { type InteractiveSourceConnectionChoice = | { kind: 'existing'; connectionId: string; connection: KtxProjectConnectionConfig } | { kind: 'new'; args: KtxSetupSourcesArgs } + | { kind: 'edited'; connectionId: string; args: KtxSetupSourcesArgs } | 'back'; +type SourceSetupChoiceResult = + | { status: 'ready'; connectionId: string } + | { status: 'back' } + | { status: 'failed' }; + async function runSourcePromptSteps( initialState: SourcePromptState, stepsForState: (state: SourcePromptState) => SourcePromptStep[], @@ -828,6 +840,12 @@ function resetRepoLocationFields(state: SourcePromptState): void { delete state.sourceProjectName; } +function sourceLocationFromArgs(args: KtxSetupSourcesArgs): SourceLocationChoice | undefined { + if (args.sourcePath) return 'path'; + if (args.sourceGitUrl) return 'git'; + return undefined; +} + function warehouseConnectionChoices(config: KtxProjectConfig): WarehouseConnectionChoice[] { return Object.entries(config.connections) .filter(([, connection]) => PRIMARY_SOURCE_DRIVERS.has(String(connection.driver ?? '').toLowerCase())) @@ -964,7 +982,7 @@ async function promptForInteractiveSource( testGitRepo: KtxSetupSourcesDeps['testGitRepo'] = testRepoConnection, discoverMetabaseDatabaseList?: KtxSetupSourcesDeps['discoverMetabaseDatabases'], ): Promise { - const initialState: SourcePromptState = { ...args, source }; + const initialState: SourcePromptState = { ...args, source, sourceLocation: sourceLocationFromArgs(args) }; if (args.sourceConnectionId) { initialState.sourceConnectionId = args.sourceConnectionId; } @@ -994,7 +1012,10 @@ async function promptForInteractiveSource( ...(state.sourceLocation === 'path' ? [ async (currentState: SourcePromptState) => { - const sourcePath = await promptText(prompts, { message: `${source} local path` }); + const sourcePath = await promptText(prompts, { + message: `${source} local path`, + ...(currentState.sourcePath ? { initialValue: currentState.sourcePath } : {}), + }); if (sourcePath === undefined) return 'back'; currentState.sourcePath = sourcePath; return 'next'; @@ -1004,13 +1025,19 @@ async function promptForInteractiveSource( ...(state.sourceLocation === 'git' ? [ async (currentState: SourcePromptState) => { - const sourceGitUrl = await promptText(prompts, { message: `${source} git URL` }); + const sourceGitUrl = await promptText(prompts, { + message: `${source} git URL`, + ...(currentState.sourceGitUrl ? { initialValue: currentState.sourceGitUrl } : {}), + }); if (sourceGitUrl === undefined) return 'back'; currentState.sourceGitUrl = sourceGitUrl; return 'next'; }, async (currentState: SourcePromptState) => { - const branch = await promptText(prompts, { message: `${source} git branch`, initialValue: 'main' }); + const branch = await promptText(prompts, { + message: `${source} git branch`, + initialValue: currentState.sourceBranch ?? 'main', + }); if (branch === undefined) return 'back'; currentState.sourceBranch = branch || 'main'; return 'next'; @@ -1031,6 +1058,7 @@ async function promptForInteractiveSource( projectDir: args.projectDir, source, connectionId: currentState.sourceConnectionId ?? `${source}-main`, + existingRef: currentState.sourceAuthTokenRef, }); if (authRef === 'back') return 'back'; if (authRef) { @@ -1104,6 +1132,7 @@ async function promptForInteractiveSource( const subpath = await promptText(prompts, { message: sourceSubpathPrompt(source), placeholder: 'optional', + ...(currentState.sourceSubpath ? { initialValue: currentState.sourceSubpath } : {}), }); if (subpath === undefined) return 'back'; if (subpath) { @@ -1122,7 +1151,10 @@ async function promptForInteractiveSource( return await runSourcePromptSteps(initialState, () => [ ...connectionSteps, async (state) => { - const sourceUrl = await promptText(prompts, { message: 'Metabase URL' }); + const sourceUrl = await promptText(prompts, { + message: 'Metabase URL', + ...(state.sourceUrl ? { initialValue: state.sourceUrl } : {}), + }); if (sourceUrl === undefined) return 'back'; state.sourceUrl = sourceUrl; return 'next'; @@ -1134,6 +1166,7 @@ async function promptForInteractiveSource( label: 'Metabase API key', envName: 'METABASE_API_KEY', secretFileName: `${state.sourceConnectionId ?? 'metabase-main'}-api-key`, + existingRef: state.sourceApiKeyRef, }); if (ref === 'back') return 'back'; state.sourceApiKeyRef = ref; @@ -1165,13 +1198,19 @@ async function promptForInteractiveSource( return await runSourcePromptSteps(initialState, () => [ ...connectionSteps, async (state) => { - const sourceUrl = await promptText(prompts, { message: 'Looker base URL' }); + const sourceUrl = await promptText(prompts, { + message: 'Looker base URL', + ...(state.sourceUrl ? { initialValue: state.sourceUrl } : {}), + }); if (sourceUrl === undefined) return 'back'; state.sourceUrl = sourceUrl; return 'next'; }, async (state) => { - const sourceClientId = await promptText(prompts, { message: 'Looker client id' }); + const sourceClientId = await promptText(prompts, { + message: 'Looker client id', + ...(state.sourceClientId ? { initialValue: state.sourceClientId } : {}), + }); if (sourceClientId === undefined) return 'back'; state.sourceClientId = sourceClientId; return 'next'; @@ -1183,6 +1222,7 @@ async function promptForInteractiveSource( label: 'Looker client secret', envName: 'LOOKER_CLIENT_SECRET', secretFileName: `${state.sourceConnectionId ?? 'looker-main'}-client-secret`, + existingRef: state.sourceClientSecretRef, }); if (ref === 'back') return 'back'; state.sourceClientSecretRef = ref; @@ -1201,6 +1241,7 @@ async function promptForInteractiveSource( const lookerConnectionName = await promptText(prompts, { message: 'Looker connection name', placeholder: 'optional', + ...(state.sourceTarget ? { initialValue: state.sourceTarget } : {}), }); if (lookerConnectionName === undefined) return 'back'; if (lookerConnectionName) { @@ -1222,6 +1263,7 @@ async function promptForInteractiveSource( label: 'Notion integration token', envName: 'NOTION_TOKEN', secretFileName: `${currentState.sourceConnectionId ?? 'notion-main'}-token`, + existingRef: currentState.sourceApiKeyRef, }); if (ref === 'back') return 'back'; currentState.sourceApiKeyRef = ref; @@ -1286,6 +1328,24 @@ function existingConnectionIdsBySource( .sort((left, right) => left.localeCompare(right)); } +function sourceTypeForConnection(connection: KtxProjectConnectionConfig): KtxSetupSourceType | null { + const driver = String(connection.driver ?? '').toLowerCase(); + return SOURCE_OPTIONS.some((option) => option.value === driver) ? (driver as KtxSetupSourceType) : null; +} + +function contextSourceEditTargets(connections: Record): Array<{ + connectionId: string; + source: KtxSetupSourceType; +}> { + return Object.entries(connections) + .map(([connectionId, connection]) => { + const source = sourceTypeForConnection(connection); + return source ? { connectionId, source } : null; + }) + .filter((target): target is { connectionId: string; source: KtxSetupSourceType } => target !== null) + .sort((left, right) => left.connectionId.localeCompare(right.connectionId)); +} + function sourceChecklistForConnections(connections: Record): { options: Array<{ value: KtxSetupSourceType; label: string; hint?: string }>; initialValues: KtxSetupSourceType[]; @@ -1317,6 +1377,180 @@ function defaultConnectionIdForSource( return `${base}-${index}`; } +function firstStringRecordEntry(value: unknown): [string, string] | undefined { + if (!isRecord(value)) return undefined; + for (const [key, raw] of Object.entries(value)) { + if (typeof raw === 'string' && raw.trim().length > 0) { + return [key, raw.trim()]; + } + } + return undefined; +} + +function applyRepoSourceArgs( + args: KtxSetupSourcesArgs, + input: { repoUrl?: string; sourceDir?: string; branch?: string; subpath?: string; authTokenRef?: string }, +): void { + if (input.sourceDir) { + args.sourcePath = input.sourceDir; + } else if (input.repoUrl?.startsWith('file:')) { + args.sourcePath = fileURLToPath(input.repoUrl); + } else if (input.repoUrl) { + args.sourceGitUrl = input.repoUrl; + } + if (input.branch) args.sourceBranch = input.branch; + if (input.subpath) args.sourceSubpath = input.subpath; + if (input.authTokenRef) args.sourceAuthTokenRef = input.authTokenRef; +} + +function sourceArgsFromExistingConnection(input: { + args: KtxSetupSourcesArgs; + source: KtxSetupSourceType; + connectionId: string; + connection: KtxProjectConnectionConfig; +}): KtxSetupSourcesArgs { + const sourceArgs: KtxSetupSourcesArgs = { + projectDir: input.args.projectDir, + inputMode: input.args.inputMode, + source: input.source, + sourceConnectionId: input.connectionId, + runInitialSourceIngest: input.args.runInitialSourceIngest, + skipSources: input.args.skipSources, + }; + + if (input.source === 'dbt') { + applyRepoSourceArgs(sourceArgs, { + sourceDir: stringField(input.connection.source_dir), + repoUrl: stringField(input.connection.repo_url), + branch: stringField(input.connection.branch), + subpath: stringField(input.connection.path), + authTokenRef: stringField(input.connection.auth_token_ref), + }); + const profilesPath = stringField(input.connection.profiles_path); + const target = stringField(input.connection.target); + const projectName = stringField(input.connection.project_name); + if (profilesPath) sourceArgs.sourceProfilesPath = profilesPath; + if (target) sourceArgs.sourceTarget = target; + if (projectName) sourceArgs.sourceProjectName = projectName; + return sourceArgs; + } + + if (input.source === 'metricflow') { + const metricflow = isRecord(input.connection.metricflow) ? input.connection.metricflow : {}; + applyRepoSourceArgs(sourceArgs, { + repoUrl: stringField(metricflow.repoUrl), + branch: stringField(metricflow.branch), + subpath: stringField(metricflow.path), + authTokenRef: stringField(metricflow.auth_token_ref), + }); + return sourceArgs; + } + + if (input.source === 'lookml') { + applyRepoSourceArgs(sourceArgs, { + repoUrl: stringField(input.connection.repoUrl), + branch: stringField(input.connection.branch), + subpath: stringField(input.connection.path), + authTokenRef: stringField(input.connection.auth_token_ref), + }); + const mappings = isRecord(input.connection.mappings) ? input.connection.mappings : {}; + const expectedLookerConnectionName = stringField(mappings.expectedLookerConnectionName); + if (expectedLookerConnectionName) sourceArgs.sourceTarget = expectedLookerConnectionName; + return sourceArgs; + } + + if (input.source === 'metabase') { + sourceArgs.sourceUrl = stringField(input.connection.api_url); + sourceArgs.sourceApiKeyRef = stringField(input.connection.api_key_ref); + const mappings = isRecord(input.connection.mappings) ? input.connection.mappings : {}; + const databaseMapping = firstStringRecordEntry(mappings.databaseMappings); + if (databaseMapping) { + sourceArgs.metabaseDatabaseId = Number.parseInt(databaseMapping[0], 10); + sourceArgs.sourceWarehouseConnectionId = databaseMapping[1]; + } + return sourceArgs; + } + + if (input.source === 'looker') { + sourceArgs.sourceUrl = stringField(input.connection.base_url); + sourceArgs.sourceClientId = stringField(input.connection.client_id); + sourceArgs.sourceClientSecretRef = stringField(input.connection.client_secret_ref); + const mappings = isRecord(input.connection.mappings) ? input.connection.mappings : {}; + const connectionMapping = firstStringRecordEntry(mappings.connectionMappings); + if (connectionMapping) { + sourceArgs.sourceTarget = connectionMapping[0]; + sourceArgs.sourceWarehouseConnectionId = connectionMapping[1]; + } + return sourceArgs; + } + + sourceArgs.sourceApiKeyRef = stringField(input.connection.auth_token_ref); + sourceArgs.notionCrawlMode = + input.connection.crawl_mode === 'all_accessible' ? 'all_accessible' : 'selected_roots'; + if (Array.isArray(input.connection.root_page_ids)) { + sourceArgs.notionRootPageIds = input.connection.root_page_ids.filter( + (pageId): pageId is string => typeof pageId === 'string', + ); + } + return sourceArgs; +} + +async function promptEditedSourceConnection(input: { + args: KtxSetupSourcesArgs; + source: KtxSetupSourceType; + connectionId: string; + connection: KtxProjectConnectionConfig; + prompts: KtxSetupSourcesPromptAdapter; + io: KtxCliIo; + testGitRepo?: KtxSetupSourcesDeps['testGitRepo']; + pickNotionRootPages?: KtxSetupSourcesDeps['pickNotionRootPages']; + discoverMetabaseDatabases?: KtxSetupSourcesDeps['discoverMetabaseDatabases']; +}): Promise | 'back'> { + const sourceArgs = await promptForInteractiveSource( + sourceArgsFromExistingConnection({ + args: input.args, + source: input.source, + connectionId: input.connectionId, + connection: input.connection, + }), + input.source, + input.prompts, + input.io, + { + pickNotionRootPages: input.pickNotionRootPages, + discoverMetabaseDatabases: input.discoverMetabaseDatabases, + }, + input.connectionId, + input.testGitRepo, + input.discoverMetabaseDatabases, + ); + return sourceArgs === 'back' + ? 'back' + : { kind: 'edited', connectionId: input.connectionId, args: sourceArgs }; +} + +async function chooseContextSourceToEdit(input: { + projectDir: string; + prompts: KtxSetupSourcesPromptAdapter; +}): Promise<{ connectionId: string; source: KtxSetupSourceType } | 'back'> { + const project = await loadKtxProject({ projectDir: input.projectDir }); + const targets = contextSourceEditTargets(project.config.connections); + if (targets.length === 0) return 'back'; + const choice = await input.prompts.select({ + message: 'Context source to edit', + options: [ + ...targets.map((target) => ({ + value: target.connectionId, + label: `${target.connectionId} (${sourceLabel(target.source)})`, + })), + { value: 'back', label: 'Back' }, + ], + }); + if (choice === 'back') return 'back'; + const target = targets.find((candidate) => candidate.connectionId === choice); + return target ?? 'back'; +} + async function chooseInteractiveSourceConnection(input: { args: KtxSetupSourcesArgs; source: KtxSetupSourceType; @@ -1356,6 +1590,10 @@ async function chooseInteractiveSourceConnection(input: { value: `existing:${connectionId}`, label: `Use existing ${label} connection: ${connectionId}`, })), + ...existingIds.map((connectionId) => ({ + value: `edit:${connectionId}`, + label: `Edit existing ${label} connection: ${connectionId}`, + })), { value: 'new', label: `Add new ${label} connection` }, { value: 'back', label: 'Back' }, ], @@ -1369,6 +1607,28 @@ async function chooseInteractiveSourceConnection(input: { } continue; } + if (choice.startsWith('edit:')) { + const connectionId = choice.slice('edit:'.length); + const connection = input.connections[connectionId]; + if (!connection) { + continue; + } + const edited = await promptEditedSourceConnection({ + args: input.args, + source: input.source, + connectionId, + connection, + prompts: input.prompts, + io: input.io, + testGitRepo: input.testGitRepo, + pickNotionRootPages: input.pickNotionRootPages, + discoverMetabaseDatabases: input.discoverMetabaseDatabases, + }); + if (edited === 'back') { + continue; + } + return edited; + } const sourceArgs = await promptForInteractiveSource( input.args, input.source, @@ -1433,6 +1693,85 @@ async function validateSource( return await (deps.validateNotion ?? defaultValidateNotion)(args.connection); } +async function saveValidateAndMaybeBuildSource(input: { + args: KtxSetupSourcesArgs; + source: KtxSetupSourceType; + sourceChoice: Exclude; + prompts: KtxSetupSourcesPromptAdapter; + io: KtxCliIo; + deps: KtxSetupSourcesDeps; +}): Promise { + const connectionId = + input.sourceChoice.kind === 'existing' + ? input.sourceChoice.connectionId + : input.sourceChoice.kind === 'edited' + ? input.sourceChoice.connectionId + : (input.sourceChoice.args.sourceConnectionId ?? `${input.source}-main`); + const connection = + input.sourceChoice.kind === 'existing' + ? input.sourceChoice.connection + : buildConnection(input.source, input.sourceChoice.args); + const rollback = + input.sourceChoice.kind === 'existing' + ? undefined + : await writeSourceConnection( + input.args.projectDir, + connectionId, + connection, + sourceAdapter(input.source), + ); + + if (input.sourceChoice.kind === 'existing') { + await ensureSourceAdapterEnabled(input.args.projectDir, input.source); + } + + const validation = await validateSource( + input.source, + { projectDir: input.args.projectDir, connectionId, connection }, + input.deps, + ); + if (!validation.ok) { + await rollback?.(); + input.io.stderr.write(`${validation.message}\n`); + return { status: 'failed' }; + } + + if (input.source === 'metabase' || input.source === 'looker') { + input.prompts.log?.(`Validating ${sourceLabel(input.source)} mapping…`); + const mappingCode = await (input.deps.runMapping ?? defaultRunMapping)( + input.args.projectDir, + connectionId, + createSetupPrefixedIo(input.io), + ); + if (mappingCode !== 0) { + await rollback?.(); + return { status: 'failed' }; + } + } + + if (input.args.runInitialSourceIngest) { + const ingestResult = await runInitialSourceIngestWithRecovery({ + args: input.args, + connectionId, + io: input.io, + prompts: input.prompts, + deps: input.deps, + }); + if (ingestResult === 'failed') { + await rollback?.(); + return { status: 'failed' }; + } + if (ingestResult === 'back') { + await rollback?.(); + return { status: 'back' }; + } + } else { + input.io.stdout.write(`│ Context source ${connectionId} saved. It will be built during the context build step.\n`); + } + + return { status: 'ready', connectionId }; +} + export async function runKtxSetupSourcesStep( args: KtxSetupSourcesArgs, io: KtxCliIo, @@ -1510,62 +1849,27 @@ export async function runKtxSetupSourcesStep( returnToSourceSelection = true; break; } - const connectionId = - sourceChoice.kind === 'existing' - ? sourceChoice.connectionId - : (sourceChoice.args.sourceConnectionId ?? `${source}-main`); - const connection = - sourceChoice.kind === 'existing' ? sourceChoice.connection : buildConnection(source, sourceChoice.args); - const rollback = - sourceChoice.kind === 'existing' - ? undefined - : await writeSourceConnection(args.projectDir, connectionId, connection, sourceAdapter(source)); - if (sourceChoice.kind === 'existing') { - await ensureSourceAdapterEnabled(args.projectDir, source); - } - const validation = await validateSource(source, { projectDir: args.projectDir, connectionId, connection }, deps); - - if (!validation.ok) { - await rollback?.(); - io.stderr.write(`${validation.message}\n`); + const choiceResult = await saveValidateAndMaybeBuildSource({ + args, + source, + sourceChoice, + prompts, + io, + deps, + }); + if (choiceResult.status === 'failed') { return { status: 'failed', projectDir: args.projectDir }; } - if (source === 'metabase' || source === 'looker') { - prompts.log?.(`Validating ${sourceLabel(source)} mapping…`); - const mappingCode = await (deps.runMapping ?? defaultRunMapping)( - args.projectDir, - connectionId, - createSetupPrefixedIo(io), - ); - if (mappingCode !== 0) { - await rollback?.(); - return { status: 'failed', projectDir: args.projectDir }; + if (choiceResult.status === 'back') { + if (args.source) { + return { status: 'back', projectDir: args.projectDir }; } + returnToSourceSelection = true; + break; } - if (args.runInitialSourceIngest) { - const ingestResult = await runInitialSourceIngestWithRecovery({ - args, - connectionId, - io, - prompts, - deps, - }); - if (ingestResult === 'failed') { - await rollback?.(); - return { status: 'failed', projectDir: args.projectDir }; - } - if (ingestResult === 'back') { - await rollback?.(); - if (args.source) { - return { status: 'back', projectDir: args.projectDir }; - } - returnToSourceSelection = true; - break; - } - } else { - io.stdout.write(`│ Context source ${connectionId} saved. It will be built during the context build step.\n`); + if (!readyConnectionIds.includes(choiceResult.connectionId)) { + readyConnectionIds.push(choiceResult.connectionId); } - readyConnectionIds.push(connectionId); } if (returnToSourceSelection) { @@ -1573,14 +1877,66 @@ export async function runKtxSetupSourcesStep( } if (readyConnectionIds.length > 0 && !args.source && args.inputMode !== 'disabled') { - const addMore = await prompts.select({ - message: `${readyConnectionIds.length} context source${readyConnectionIds.length > 1 ? 's' : ''} configured (${readyConnectionIds.join(', ')}). Add another?`, - options: [ - { value: 'done', label: 'Done — continue to context build' }, - { value: 'add', label: 'Add another context source' }, - ], - }); - if (addMore === 'add') { + let restartSourceSelection = false; + while (true) { + const addMore = await prompts.select({ + message: `${readyConnectionIds.length} context source${readyConnectionIds.length > 1 ? 's' : ''} configured (${readyConnectionIds.join(', ')}). Add another?`, + options: [ + { value: 'done', label: 'Done — continue to context build' }, + { value: 'edit', label: 'Edit an existing context source' }, + { value: 'add', label: 'Add another context source' }, + ], + }); + if (addMore === 'add') { + restartSourceSelection = true; + break; + } + if (addMore === 'edit') { + const editTarget = await chooseContextSourceToEdit({ projectDir: args.projectDir, prompts }); + if (editTarget === 'back') { + continue; + } + const projectForEdit = await loadKtxProject({ projectDir: args.projectDir }); + const connection = projectForEdit.config.connections[editTarget.connectionId]; + if (!connection) { + continue; + } + const sourceChoice = await promptEditedSourceConnection({ + args, + source: editTarget.source, + connectionId: editTarget.connectionId, + connection, + prompts, + io, + testGitRepo: deps.testGitRepo, + pickNotionRootPages: deps.pickNotionRootPages, + discoverMetabaseDatabases: deps.discoverMetabaseDatabases, + }); + if (sourceChoice === 'back') { + continue; + } + const choiceResult = await saveValidateAndMaybeBuildSource({ + args, + source: editTarget.source, + sourceChoice, + prompts, + io, + deps, + }); + if (choiceResult.status === 'failed') { + return { status: 'failed', projectDir: args.projectDir }; + } + if (choiceResult.status === 'back') { + continue; + } + if (!readyConnectionIds.includes(choiceResult.connectionId)) { + readyConnectionIds.push(choiceResult.connectionId); + } + continue; + } + break; + } + if (restartSourceSelection) { continue; } }