diff --git a/packages/cli/src/setup-databases.test.ts b/packages/cli/src/setup-databases.test.ts index d010a908..fa4ca3f2 100644 --- a/packages/cli/src/setup-databases.test.ts +++ b/packages/cli/src/setup-databases.test.ts @@ -240,8 +240,9 @@ describe('setup databases step', () => { expect(prompts.select).toHaveBeenCalledWith({ message: 'Configure PostgreSQL', options: [ - { value: 'existing:warehouse', label: 'Use existing PostgreSQL connection: warehouse' }, - { value: 'new', label: 'Add new PostgreSQL connection' }, + { value: 'existing:warehouse', label: 'Keep existing PostgreSQL connection: warehouse' }, + { value: 'edit:warehouse', label: 'Edit PostgreSQL connection: warehouse' }, + { value: 'new', label: 'Add another PostgreSQL connection' }, { value: 'back', label: 'Back' }, ], }); @@ -564,7 +565,8 @@ describe('setup databases step', () => { message: 'Primary sources already configured: warehouse\nWhat would you like to do?', options: [ { value: 'continue', label: 'Continue to knowledge sources' }, - { value: 'add', label: 'Add another primary source' }, + { value: 'edit', label: 'Edit an existing primary source' }, + { value: 'add', label: 'Add additional primary sources' }, ], }); expect(testConnection).not.toHaveBeenCalled(); @@ -608,11 +610,16 @@ describe('setup databases step', () => { connectionIds: ['warehouse', 'mysql-warehouse'], }); expect(prompts.multiselect).toHaveBeenCalledTimes(1); + expect(prompts.multiselect).toHaveBeenCalledWith(expect.objectContaining({ + initialValues: ['postgres'], + required: true, + })); expect(prompts.select).toHaveBeenCalledWith({ message: 'Primary sources already configured: warehouse\nWhat would you like to do?', options: [ { value: 'continue', label: 'Continue to knowledge sources' }, - { value: 'add', label: 'Add another primary source' }, + { value: 'edit', label: 'Edit an existing primary source' }, + { value: 'add', label: 'Add additional primary sources' }, ], }); expect(testConnection).toHaveBeenCalledTimes(1); @@ -642,11 +649,16 @@ describe('setup databases step', () => { connectionIds: ['postgres-warehouse', 'mysql-warehouse'], }); expect(prompts.multiselect).toHaveBeenCalledTimes(2); + expect(prompts.multiselect).toHaveBeenNthCalledWith(2, expect.objectContaining({ + initialValues: ['postgres'], + required: true, + })); expect(prompts.select).toHaveBeenCalledWith({ message: 'Primary sources already configured: postgres-warehouse\nWhat would you like to do?', options: [ { value: 'continue', label: 'Continue to knowledge sources' }, - { value: 'add', label: 'Add another primary source' }, + { value: 'edit', label: 'Edit an existing primary source' }, + { value: 'add', label: 'Add additional primary sources' }, ], }); const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); @@ -675,12 +687,17 @@ describe('setup databases step', () => { connectionIds: ['postgres-warehouse'], }); expect(prompts.multiselect).toHaveBeenCalledTimes(2); + expect(prompts.multiselect).toHaveBeenNthCalledWith(2, expect.objectContaining({ + initialValues: ['postgres'], + required: true, + })); expect(io.stdout()).not.toContain('KTX cannot work without at least one primary source'); expect(prompts.select).toHaveBeenNthCalledWith(2, { message: 'Primary sources already configured: postgres-warehouse\nWhat would you like to do?', options: [ { value: 'continue', label: 'Continue to knowledge sources' }, - { value: 'add', label: 'Add another primary source' }, + { value: 'edit', label: 'Edit an existing primary source' }, + { value: 'add', label: 'Add additional primary sources' }, ], }); }); @@ -715,16 +732,389 @@ describe('setup databases step', () => { ); expect(result).toEqual({ status: 'ready', projectDir: tempDir, connectionIds: ['warehouse'] }); + expect(prompts.multiselect).toHaveBeenCalledWith(expect.objectContaining({ + initialValues: ['postgres'], + required: true, + })); expect(io.stdout()).not.toContain('KTX cannot work without at least one primary source'); expect(prompts.select).toHaveBeenNthCalledWith(2, { message: 'Primary sources already configured: warehouse\nWhat would you like to do?', options: [ { value: 'continue', label: 'Continue to knowledge sources' }, - { value: 'add', label: 'Add another primary source' }, + { value: 'edit', label: 'Edit an existing primary source' }, + { value: 'add', label: 'Add additional primary sources' }, ], }); }); + it('returns from primary source edit selection back to the configured source menu', async () => { + await writeFile( + join(tempDir, 'ktx.yaml'), + [ + 'project: warehouse', + 'connections:', + ' warehouse:', + ' driver: postgres', + ' url: env:DATABASE_URL', + 'setup:', + ' database_connection_ids:', + ' - warehouse', + '', + ].join('\n'), + 'utf-8', + ); + await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); + const prompts = makePromptAdapter({ + selectValues: ['edit', 'back', 'continue'], + }); + const testConnection = vi.fn(async () => 0); + const scanConnection = vi.fn(async () => 0); + + const result = await runKtxSetupDatabasesStep( + { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, + makeIo().io, + { prompts, testConnection, scanConnection }, + ); + + expect(result).toEqual({ status: 'ready', projectDir: tempDir, connectionIds: ['warehouse'] }); + expect(prompts.select).toHaveBeenNthCalledWith(2, { + message: 'Primary source to edit', + options: [ + { value: 'warehouse', label: 'warehouse (PostgreSQL)' }, + { value: 'back', label: 'Back' }, + ], + }); + expect(prompts.select).toHaveBeenNthCalledWith(3, { + message: 'Primary sources already configured: warehouse\nWhat would you like to do?', + options: [ + { value: 'continue', label: 'Continue to knowledge sources' }, + { value: 'edit', label: 'Edit an existing primary source' }, + { value: 'add', label: 'Add additional primary sources' }, + ], + }); + expect(testConnection).not.toHaveBeenCalled(); + expect(scanConnection).not.toHaveBeenCalled(); + }); + + it('reruns table selection after editing schema scope so stale enabled tables are removed', async () => { + await writeFile( + join(tempDir, 'ktx.yaml'), + [ + 'project: warehouse', + 'connections:', + ' warehouse:', + ' driver: postgres', + ' url: env:DATABASE_URL', + ' schemas:', + ' - public', + ' enabled_tables:', + ' - public.orders', + 'setup:', + ' database_connection_ids:', + ' - warehouse', + '', + ].join('\n'), + 'utf-8', + ); + await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); + const prompts = makePromptAdapter({ + textValues: ['env:DATABASE_URL'], + multiselectValues: [['analytics']], + }); + let primaryMenuCount = 0; + vi.mocked(prompts.select).mockImplementation(async (options) => { + if (options.message === 'Primary sources already configured: warehouse\nWhat would you like to do?') { + primaryMenuCount += 1; + return primaryMenuCount === 1 ? 'edit' : 'continue'; + } + if (options.message === 'Primary source to edit') return 'warehouse'; + if (options.message === 'How do you want to connect to PostgreSQL?') return 'url'; + return 'back'; + }); + const testConnection = vi.fn(async () => 0); + const scanConnection = vi.fn(async () => 0); + const listSchemas = vi.fn(async () => ['analytics', 'public']); + const listTables = vi.fn(async () => [{ schema: 'analytics', name: 'customers', kind: 'table' as const }]); + + const result = await runKtxSetupDatabasesStep( + { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, + makeIo().io, + { prompts, testConnection, scanConnection, listSchemas, listTables }, + ); + + expect(result).toEqual({ status: 'ready', projectDir: tempDir, connectionIds: ['warehouse'] }); + expect(prompts.text).toHaveBeenCalledWith({ + message: textInputPrompt('PostgreSQL connection URL'), + placeholder: 'env:DATABASE_URL', + initialValue: 'env:DATABASE_URL', + }); + expect(listTables).toHaveBeenCalledWith(tempDir, 'warehouse'); + expect(testConnection).toHaveBeenCalledWith(tempDir, 'warehouse', expect.anything()); + expect(scanConnection).toHaveBeenCalledWith(tempDir, 'warehouse', expect.anything()); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.warehouse).toMatchObject({ + schemas: ['analytics'], + enabled_tables: ['analytics.customers'], + }); + }); + + it('preselects existing schema and table choices when editing a primary source', async () => { + await writeFile( + join(tempDir, 'ktx.yaml'), + [ + 'project: warehouse', + 'connections:', + ' warehouse:', + ' driver: postgres', + ' url: env:DATABASE_URL', + ' schemas:', + ' - public', + ' enabled_tables:', + ' - public.customers', + ' - public.orders', + 'setup:', + ' database_connection_ids:', + ' - warehouse', + '', + ].join('\n'), + 'utf-8', + ); + await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); + const prompts = makePromptAdapter({ + textValues: ['env:DATABASE_URL'], + multiselectValues: [['public'], ['public.customers', 'public.orders']], + }); + let primaryMenuCount = 0; + vi.mocked(prompts.select).mockImplementation(async (options) => { + if (options.message === 'Primary sources already configured: warehouse\nWhat would you like to do?') { + primaryMenuCount += 1; + return primaryMenuCount === 1 ? 'edit' : 'continue'; + } + if (options.message === 'Primary source to edit') return 'warehouse'; + if (options.message === 'How do you want to connect to PostgreSQL?') return 'url'; + if (options.message.startsWith('Tables found in selected schemas')) return 'customize'; + return 'back'; + }); + const listSchemas = vi.fn(async () => ['orbit_analytics', 'orbit_raw', 'public']); + const listTables = vi.fn(async () => [ + { schema: 'public', name: 'customers', kind: 'table' as const }, + { schema: 'public', name: 'orders', kind: 'table' as const }, + { schema: 'public', name: 'products', kind: 'table' as const }, + ]); + + const result = await runKtxSetupDatabasesStep( + { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, + makeIo().io, + { + prompts, + testConnection: vi.fn(async () => 0), + scanConnection: vi.fn(async () => 0), + listSchemas, + listTables, + }, + ); + + expect(result).toEqual({ status: 'ready', projectDir: tempDir, connectionIds: ['warehouse'] }); + expect(prompts.multiselect).toHaveBeenNthCalledWith(1, { + message: expect.stringContaining('PostgreSQL schemas to scan'), + options: [ + { value: 'orbit_analytics', label: 'orbit_analytics' }, + { value: 'orbit_raw', label: 'orbit_raw' }, + { value: 'public', label: 'public' }, + ], + initialValues: ['public'], + required: true, + }); + expect(prompts.multiselect).toHaveBeenNthCalledWith(2, { + message: expect.stringContaining('Tables to enable for warehouse'), + options: [ + { value: 'public.customers', label: 'public.customers' }, + { value: 'public.orders', label: 'public.orders' }, + { value: 'public.products', label: 'public.products' }, + ], + initialValues: ['public.customers', 'public.orders'], + required: true, + }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.warehouse).toMatchObject({ + schemas: ['public'], + enabled_tables: ['public.customers', 'public.orders'], + }); + }); + + it('returns to the configured primary menu when backing out of schema review during edit', async () => { + await writeFile( + join(tempDir, 'ktx.yaml'), + [ + 'project: warehouse', + 'connections:', + ' warehouse:', + ' driver: postgres', + ' url: env:DATABASE_URL', + ' schemas:', + ' - public', + ' enabled_tables:', + ' - public.orders', + 'setup:', + ' database_connection_ids:', + ' - warehouse', + '', + ].join('\n'), + 'utf-8', + ); + await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); + const prompts = makePromptAdapter({ + textValues: ['env:DATABASE_URL'], + multiselectValues: [['back']], + }); + let primaryMenuCount = 0; + vi.mocked(prompts.select).mockImplementation(async (options) => { + if (options.message === 'Primary sources already configured: warehouse\nWhat would you like to do?') { + primaryMenuCount += 1; + return primaryMenuCount === 1 ? 'edit' : 'continue'; + } + if (options.message === 'Primary source to edit') return 'warehouse'; + if (options.message === 'How do you want to connect to PostgreSQL?') return 'url'; + return 'back'; + }); + const testConnection = vi.fn(async () => 0); + const scanConnection = vi.fn(async () => 0); + const listSchemas = vi.fn(async () => ['analytics', 'public']); + const listTables = vi.fn(async () => [{ schema: 'analytics', name: 'customers', kind: 'table' as const }]); + + const result = await runKtxSetupDatabasesStep( + { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, + makeIo().io, + { prompts, testConnection, scanConnection, listSchemas, listTables }, + ); + + expect(result).toEqual({ status: 'ready', projectDir: tempDir, connectionIds: ['warehouse'] }); + expect(primaryMenuCount).toBe(2); + expect(testConnection).toHaveBeenCalledWith(tempDir, 'warehouse', expect.anything()); + expect(scanConnection).not.toHaveBeenCalled(); + expect(listTables).not.toHaveBeenCalled(); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.warehouse).toMatchObject({ + url: 'env:DATABASE_URL', + schemas: ['public'], + enabled_tables: ['public.orders'], + }); + }); + + it('returns to the configured primary menu when backing out of table review during edit', async () => { + await writeFile( + join(tempDir, 'ktx.yaml'), + [ + 'project: warehouse', + 'connections:', + ' warehouse:', + ' driver: postgres', + ' url: env:DATABASE_URL', + ' schemas:', + ' - public', + ' enabled_tables:', + ' - public.orders', + 'setup:', + ' database_connection_ids:', + ' - warehouse', + '', + ].join('\n'), + 'utf-8', + ); + await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); + const prompts = makePromptAdapter({ textValues: ['env:DATABASE_URL'] }); + let primaryMenuCount = 0; + vi.mocked(prompts.select).mockImplementation(async (options) => { + if (options.message === 'Primary sources already configured: warehouse\nWhat would you like to do?') { + primaryMenuCount += 1; + return primaryMenuCount === 1 ? 'edit' : 'continue'; + } + if (options.message === 'Primary source to edit') return 'warehouse'; + if (options.message === 'How do you want to connect to PostgreSQL?') return 'url'; + if (options.message.startsWith('Tables found in selected schemas')) return 'back'; + return 'back'; + }); + const testConnection = vi.fn(async () => 0); + const scanConnection = vi.fn(async () => 0); + const listSchemas = vi.fn(async () => ['public']); + const listTables = vi.fn(async () => [ + { schema: 'public', name: 'customers', kind: 'table' as const }, + { schema: 'public', name: 'orders', kind: 'table' as const }, + ]); + + const result = await runKtxSetupDatabasesStep( + { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, + makeIo().io, + { prompts, testConnection, scanConnection, listSchemas, listTables }, + ); + + expect(result).toEqual({ status: 'ready', projectDir: tempDir, connectionIds: ['warehouse'] }); + expect(primaryMenuCount).toBe(2); + expect(listTables).toHaveBeenCalledWith(tempDir, 'warehouse'); + expect(scanConnection).not.toHaveBeenCalled(); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.warehouse).toMatchObject({ + url: 'env:DATABASE_URL', + schemas: ['public'], + enabled_tables: ['public.orders'], + }); + }); + + it('restores an existing primary source edit when the follow-up scan fails', async () => { + await writeFile( + join(tempDir, 'ktx.yaml'), + [ + 'project: warehouse', + 'connections:', + ' warehouse:', + ' driver: postgres', + ' url: env:DATABASE_URL', + ' schemas:', + ' - public', + ' enabled_tables:', + ' - public.orders', + 'setup:', + ' database_connection_ids:', + ' - warehouse', + '', + ].join('\n'), + 'utf-8', + ); + await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); + const prompts = makePromptAdapter({ + textValues: ['env:DATABASE_URL'], + multiselectValues: [['public']], + }); + vi.mocked(prompts.select).mockImplementation(async (options) => { + if (options.message === 'Primary sources already configured: warehouse\nWhat would you like to do?') return 'edit'; + if (options.message === 'Primary source to edit') return 'warehouse'; + if (options.message === 'How do you want to connect to PostgreSQL?') return 'url'; + if (options.message.startsWith('Tables found in selected schemas')) return 'all'; + return 'back'; + }); + const listTables = vi.fn(async () => [ + { schema: 'public', name: 'customers', kind: 'table' as const }, + { schema: 'public', name: 'orders', kind: 'table' as const }, + ]); + + const result = await runKtxSetupDatabasesStep( + { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, + makeIo().io, + { + prompts, + testConnection: vi.fn(async () => 0), + scanConnection: vi.fn(async () => 1), + listTables, + }, + ); + + expect(result).toEqual({ status: 'failed', projectDir: tempDir }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.warehouse).toMatchObject({ + enabled_tables: ['public.orders'], + }); + }); + it('lets Escape from connection fields return to connection method selection', async () => { const prompts = makePromptAdapter({ selectValues: ['fields', 'url'], diff --git a/packages/cli/src/setup-databases.ts b/packages/cli/src/setup-databases.ts index 5b5b5f8a..9db80689 100644 --- a/packages/cli/src/setup-databases.ts +++ b/packages/cli/src/setup-databases.ts @@ -176,6 +176,7 @@ const SCOPE_DISCOVERY_SPECS: Partial; +type ConnectionSetupStatus = 'ready' | 'back' | 'failed'; const DRIVER_CONNECTION_DEFAULTS: Record = { postgres: { port: '5432' }, @@ -227,6 +228,16 @@ function unique(values: string[]): string[] { return [...new Set(values.filter((value) => value.trim().length > 0))]; } +function stringConfigField(connection: KtxProjectConnectionConfig | undefined, field: string): string | undefined { + const value = connection?.[field]; + return typeof value === 'string' && value.trim().length > 0 ? value.trim() : undefined; +} + +function numberConfigField(connection: KtxProjectConnectionConfig | undefined, field: string): number | undefined { + const value = connection?.[field]; + return typeof value === 'number' && Number.isFinite(value) ? value : undefined; +} + function historicSqlConfigRecord(connection: KtxProjectConnectionConfig | undefined): Record | null { const historicSql = connection?.historicSql; return historicSql && typeof historicSql === 'object' && !Array.isArray(historicSql) @@ -454,6 +465,18 @@ function configuredPrimaryConnectionIds( .sort((left, right) => left.localeCompare(right)); } +function configuredPrimaryDrivers( + connections: Record, + connectionIds: string[], +): KtxSetupDatabaseDriver[] { + const configured = new Set( + connectionIds + .map((connectionId) => normalizeDriver(connections[connectionId]?.driver)) + .filter((driver): driver is KtxSetupDatabaseDriver => driver !== null), + ); + return DRIVER_OPTIONS.map((option) => option.value).filter((driver) => configured.has(driver)); +} + function configuredPrimarySourcesPrompt(connectionIds: string[]): { message: string; options: Array<{ value: string; label: string }>; @@ -462,7 +485,8 @@ function configuredPrimarySourcesPrompt(connectionIds: string[]): { message: `Primary sources already configured: ${connectionIds.join(', ')}\nWhat would you like to do?`, options: [ { value: 'continue', label: 'Continue to knowledge sources' }, - { value: 'add', label: 'Add another primary source' }, + { value: 'edit', label: 'Edit an existing primary source' }, + { value: 'add', label: 'Add additional primary sources' }, ], }; } @@ -552,23 +576,40 @@ async function buildFieldsConnectionConfig(input: { connectionId: string; args: KtxSetupDatabasesArgs; prompts: KtxSetupDatabasesPromptAdapter; + existingConnection?: KtxProjectConnectionConfig; }): Promise { const label = driverLabel(input.driver); const defaults = DRIVER_CONNECTION_DEFAULTS[input.driver]; - const host = await promptText(input.prompts, `${label} host`, 'localhost'); + const host = await promptText( + input.prompts, + `${label} host`, + stringConfigField(input.existingConnection, 'host') ?? 'localhost', + ); if (host === undefined) return 'back'; if (!host) return null; - const portStr = await promptText(input.prompts, `${label} port`, defaults.port); + const portStr = await promptText( + input.prompts, + `${label} port`, + String(numberConfigField(input.existingConnection, 'port') ?? defaults.port), + ); if (portStr === undefined) return 'back'; const port = Number(portStr || defaults.port); - const database = await promptText(input.prompts, `${label} database name`); + const database = await promptText( + input.prompts, + `${label} database name`, + stringConfigField(input.existingConnection, 'database'), + ); if (database === undefined) return 'back'; if (!database) return null; - const username = await promptText(input.prompts, `${label} username`); + const username = await promptText( + input.prompts, + `${label} username`, + stringConfigField(input.existingConnection, 'username'), + ); if (username === undefined) return 'back'; if (!username) return null; @@ -583,6 +624,7 @@ async function buildFieldsConnectionConfig(input: { }); if (credentialResult === 'back') return 'back'; if (credentialResult) passwordRef = credentialResult; + if (!credentialResult) passwordRef = stringConfigField(input.existingConnection, 'password'); } return { @@ -601,9 +643,14 @@ async function buildPastedUrlConnectionConfig(input: { connectionId: string; args: KtxSetupDatabasesArgs; prompts: KtxSetupDatabasesPromptAdapter; + existingConnection?: KtxProjectConnectionConfig; }): Promise { const label = driverLabel(input.driver); - const rawUrl = await promptText(input.prompts, `${label} connection URL`); + const rawUrl = await promptText( + input.prompts, + `${label} connection URL`, + stringConfigField(input.existingConnection, 'url'), + ); if (rawUrl === undefined) return 'back'; if (!rawUrl) return null; @@ -642,6 +689,7 @@ async function buildUrlConnectionConfig(input: { connectionId: string; args: KtxSetupDatabasesArgs; prompts: KtxSetupDatabasesPromptAdapter; + existingConnection?: KtxProjectConnectionConfig; }): Promise { if (input.args.inputMode === 'disabled' && !input.args.databaseUrl) return null; @@ -689,6 +737,7 @@ async function buildConnectionConfig(input: { connectionId: string; args: KtxSetupDatabasesArgs; prompts: KtxSetupDatabasesPromptAdapter; + existingConnection?: KtxProjectConnectionConfig; }): Promise { const { driver, args, prompts } = input; if (driver === 'sqlite') { @@ -698,22 +747,37 @@ async function buildConnectionConfig(input: { (await promptText( prompts, 'SQLite database file\nEnter a relative or absolute path, for example ./warehouse.sqlite.', + stringConfigField(input.existingConnection, 'path'), )); if (path === undefined) return 'back'; return path ? { driver: 'sqlite', path } : null; } if (driver === 'postgres' || driver === 'mysql' || driver === 'clickhouse' || driver === 'sqlserver') { - return await buildUrlConnectionConfig({ driver, connectionId: input.connectionId, args, prompts }); + return await buildUrlConnectionConfig({ + driver, + connectionId: input.connectionId, + args, + prompts, + existingConnection: input.existingConnection, + }); } if (driver === 'bigquery') { - const datasetId = await promptText(prompts, 'BigQuery dataset\nFor example analytics.'); + const datasetId = await promptText( + prompts, + 'BigQuery dataset\nFor example analytics.', + stringConfigField(input.existingConnection, 'dataset_id'), + ); if (datasetId === undefined) return 'back'; - const credentialsPath = await promptText(prompts, 'Path to service account JSON file'); + const credentialsPath = await promptText( + prompts, + 'Path to service account JSON file', + stringConfigField(input.existingConnection, 'credentials_json'), + ); if (credentialsPath === undefined) return 'back'; const location = await promptText( prompts, 'BigQuery location\nPress Enter for US, or enter a location like EU.', - 'US', + stringConfigField(input.existingConnection, 'location') ?? 'US', ); if (location === undefined) return 'back'; if (!datasetId || !credentialsPath) return null; @@ -725,19 +789,35 @@ async function buildConnectionConfig(input: { }; } if (driver === 'snowflake') { - const account = await promptText(prompts, 'Snowflake account identifier'); + const account = await promptText( + prompts, + 'Snowflake account identifier', + stringConfigField(input.existingConnection, 'account'), + ); if (account === undefined) return 'back'; - const warehouse = await promptText(prompts, 'Snowflake warehouse\nFor example ANALYTICS_WH.'); + const warehouse = await promptText( + prompts, + 'Snowflake warehouse\nFor example ANALYTICS_WH.', + stringConfigField(input.existingConnection, 'warehouse'), + ); if (warehouse === undefined) return 'back'; - const database = await promptText(prompts, 'Snowflake database name'); + const database = await promptText( + prompts, + 'Snowflake database name', + stringConfigField(input.existingConnection, 'database'), + ); if (database === undefined) return 'back'; const schemaName = await promptText( prompts, 'Snowflake schema\nPress Enter for PUBLIC, or enter a schema name.', - 'PUBLIC', + stringConfigField(input.existingConnection, 'schema_name') ?? 'PUBLIC', ); if (schemaName === undefined) return 'back'; - const username = await promptText(prompts, 'Snowflake username'); + const username = await promptText( + prompts, + 'Snowflake username', + stringConfigField(input.existingConnection, 'username'), + ); if (username === undefined) return 'back'; const passwordRef = await promptCredential({ prompts, @@ -747,9 +827,14 @@ async function buildConnectionConfig(input: { secretName: 'password', // pragma: allowlist secret }); if (passwordRef === 'back') return 'back'; // pragma: allowlist secret - const role = await promptText(prompts, 'Snowflake role (optional)\nPress Enter to skip.'); + const role = await promptText( + prompts, + 'Snowflake role (optional)\nPress Enter to skip.', + stringConfigField(input.existingConnection, 'role'), + ); if (role === undefined) return 'back'; - if (!account || !warehouse || !database || !schemaName || !username || !passwordRef) return null; + const resolvedPasswordRef = passwordRef ?? stringConfigField(input.existingConnection, 'password'); + if (!account || !warehouse || !database || !schemaName || !username || !resolvedPasswordRef) return null; return { driver: 'snowflake', authMethod: 'password', @@ -758,7 +843,7 @@ async function buildConnectionConfig(input: { database, schema_name: schemaName, username, - password: passwordRef, + password: resolvedPasswordRef, ...(role ? { role } : {}), }; } @@ -1096,6 +1181,59 @@ async function writeConnectionConfig(input: { } } +async function createConnectionConfigRollback(projectDir: string, connectionId: string): Promise<() => Promise> { + const project = await loadKtxProject({ projectDir }); + const previousConnection = project.config.connections[connectionId]; + const hadPreviousConnection = previousConnection !== undefined; + return async () => { + const latest = await loadKtxProject({ projectDir }); + const connections = { ...latest.config.connections }; + if (hadPreviousConnection) { + connections[connectionId] = previousConnection; + } else { + delete connections[connectionId]; + } + await writeFile( + latest.configPath, + serializeKtxProjectConfig({ + ...latest.config, + connections, + }), + 'utf-8', + ); + }; +} + +function withExistingPrimaryEditPromptDefaults(input: { + previous: KtxProjectConnectionConfig; + next: KtxProjectConnectionConfig; + driver: KtxSetupDatabaseDriver; +}): KtxProjectConnectionConfig { + const merged: KtxProjectConnectionConfig = { ...input.next }; + const spec = SCOPE_DISCOVERY_SPECS[input.driver]; + if (spec) { + const nextArray = input.next[spec.configArrayField]; + const previousArray = input.previous[spec.configArrayField]; + if ( + !(Array.isArray(nextArray) && nextArray.length > 0) && + Array.isArray(previousArray) && + previousArray.length > 0 + ) { + delete merged[spec.configSingleField]; + merged[spec.configArrayField] = previousArray; + } else if (!Object.hasOwn(input.next, spec.configArrayField) && !Object.hasOwn(input.next, spec.configSingleField)) { + const previousSingle = input.previous[spec.configSingleField]; + if (typeof previousSingle === 'string' && previousSingle.trim().length > 0) { + merged[spec.configSingleField] = previousSingle; + } + } + } + if (!Object.hasOwn(input.next, 'enabled_tables') && Array.isArray(input.previous.enabled_tables)) { + merged.enabled_tables = input.previous.enabled_tables; + } + return merged; +} + function configuredScopeValues( connection: KtxProjectConnectionConfig | undefined, spec: ScopeDiscoverySpec, @@ -1156,18 +1294,19 @@ async function maybeConfigureSchemaScope(input: { prompts: KtxSetupDatabasesPromptAdapter; deps: KtxSetupDatabasesDeps; io: KtxCliIo; -}): Promise { + forcePrompt?: boolean; +}): Promise { const project = await loadKtxProject({ projectDir: input.projectDir }); const connection = project.config.connections[input.connectionId]; const driver = normalizeDriver(connection?.driver); - if (!driver) return true; + if (!driver) return 'ready'; const spec = SCOPE_DISCOVERY_SPECS[driver]; - if (!spec) return true; + if (!spec) return 'ready'; const arrayVal = connection?.[spec.configArrayField]; - if (Array.isArray(arrayVal) && arrayVal.length > 0) { - return true; + if (Array.isArray(arrayVal) && arrayVal.length > 0 && input.forcePrompt !== true) { + return 'ready'; } if (input.args.databaseSchemas.length > 0) { @@ -1177,7 +1316,7 @@ async function maybeConfigureSchemaScope(input: { values: input.args.databaseSchemas, spec, }); - return true; + return 'ready'; } writeSetupSection(input.io, `Discovering ${spec.promptLabel.toLowerCase()}`, [ @@ -1190,14 +1329,18 @@ async function maybeConfigureSchemaScope(input: { await (input.deps.listSchemas ?? defaultListSchemas)(input.projectDir, input.connectionId), ); } catch (error) { + const detail = error instanceof Error ? error.message : String(error); input.io.stderr.write( - `Could not discover ${spec.promptLabel.toLowerCase()} for ${input.connectionId}; continuing with existing ${spec.noun} scope. ` + - `Pass --database-schema to set it explicitly. ${error instanceof Error ? error.message : String(error)}\n`, + input.forcePrompt === true + ? `Could not discover ${spec.promptLabel.toLowerCase()} for ${input.connectionId}; edit was not saved. ` + + `Pass --database-schema to set it explicitly. ${detail}\n` + : `Could not discover ${spec.promptLabel.toLowerCase()} for ${input.connectionId}; continuing with existing ${spec.noun} scope. ` + + `Pass --database-schema to set it explicitly. ${detail}\n`, ); - return true; + return input.forcePrompt === true ? 'failed' : 'ready'; } if (discovered.length === 0) { - return true; + return 'ready'; } let selected: string[]; @@ -1217,7 +1360,7 @@ async function maybeConfigureSchemaScope(input: { required: true, }); if (choices.includes('back')) { - return false; + return 'back'; } selected = choices.length > 0 ? choices : initialValues; } @@ -1232,7 +1375,7 @@ async function maybeConfigureSchemaScope(input: { writeSetupSection(input.io, `${capitalNounPlural} saved for ${input.connectionId}`, [ `✓ ${selected.join(', ')}`, ]); - return true; + return 'ready'; } async function maybeConfigureTableScope(input: { @@ -1242,19 +1385,20 @@ async function maybeConfigureTableScope(input: { prompts: KtxSetupDatabasesPromptAdapter; io: KtxCliIo; deps: KtxSetupDatabasesDeps; -}): Promise { + forcePrompt?: boolean; +}): Promise { const project = await loadKtxProject({ projectDir: input.projectDir }); const connection = project.config.connections[input.connectionId]; const driver = normalizeDriver(connection?.driver); - if (!driver || driver === 'sqlite') return true; + if (!driver || driver === 'sqlite') return 'ready'; const existingTables = connection?.enabled_tables; - if (Array.isArray(existingTables) && existingTables.length > 0) { - return true; + if (Array.isArray(existingTables) && existingTables.length > 0 && input.forcePrompt !== true) { + return 'ready'; } if (input.args.inputMode === 'disabled') { - return true; + return 'ready'; } writeSetupSection(input.io, 'Discovering tables', [ @@ -1268,15 +1412,20 @@ async function maybeConfigureTableScope(input: { input.connectionId, ); } catch (error) { + const detail = error instanceof Error ? error.message : String(error); input.io.stderr.write( - `Could not discover tables for ${input.connectionId}; continuing without table filter. ` + - `${error instanceof Error ? error.message : String(error)}\n`, + input.forcePrompt === true + ? `Could not discover tables for ${input.connectionId}; edit was not saved. ${detail}\n` + : `Could not discover tables for ${input.connectionId}; continuing without table filter. ${detail}\n`, ); - return true; + return input.forcePrompt === true ? 'failed' : 'ready'; } if (discovered.length === 0) { - return true; + if (input.forcePrompt === true) { + input.io.stderr.write(`No tables discovered for ${input.connectionId}; edit was not saved.\n`); + } + return input.forcePrompt === true ? 'failed' : 'ready'; } const allQualified = discovered.map((t) => `${t.schema}.${t.name}`); @@ -1290,7 +1439,7 @@ async function maybeConfigureTableScope(input: { writeSetupSection(input.io, `Tables enabled for ${input.connectionId}`, [ `✓ ${allQualified[0]}`, ]); - return true; + return 'ready'; } const bySchema = new Map(); @@ -1316,7 +1465,7 @@ async function maybeConfigureTableScope(input: { }); if (action === 'back') { - return false; + return 'back'; } if (action === 'all') { @@ -1332,7 +1481,10 @@ async function maybeConfigureTableScope(input: { const suffix = t.kind === 'view' ? ' (view)' : ''; return { value: qualified, label: `${qualified}${suffix}` }; }), - initialValues: allQualified, + initialValues: + Array.isArray(existingTables) && input.forcePrompt === true + ? existingTables.filter((table): table is string => typeof table === 'string' && allQualified.includes(table)) + : allQualified, required: true, }); @@ -1356,7 +1508,7 @@ async function maybeConfigureTableScope(input: { writeSetupSection(input.io, `Tables enabled for ${input.connectionId}`, [ `✓ ${selected.length}/${discovered.length} tables enabled`, ]); - return true; + return 'ready'; } async function ensureHistoricSqlIngestDefaults(projectDir: string): Promise { @@ -1466,7 +1618,8 @@ async function validateAndScanConnection(input: { deps: KtxSetupDatabasesDeps; args: KtxSetupDatabasesArgs; prompts: KtxSetupDatabasesPromptAdapter; -}): Promise { + forceScopeAndTables?: boolean; +}): Promise { const testConnection = input.deps.testConnection ?? defaultTestConnection; const scanConnection = input.deps.scanConnection ?? defaultScanConnection; const project = await loadKtxProject({ projectDir: input.projectDir }); @@ -1477,7 +1630,7 @@ async function validateAndScanConnection(input: { if (testCode !== 0) { flushBufferedCommandOutput(input.io, testIo); input.io.stderr.write(`Connection test failed for ${input.connectionId}.\n`); - return false; + return 'failed'; } const testOutput = testIo.stdoutText(); const outputDriver = normalizeDriver(readOutputValue(testOutput, 'Driver')); @@ -1486,14 +1639,24 @@ async function validateAndScanConnection(input: { writeSetupSection(input.io, `Testing ${input.connectionId}`, testLines); while (true) { - if (!(await maybeConfigureSchemaScope(input))) { - return false; + const schemaStatus = await maybeConfigureSchemaScope({ ...input, forcePrompt: input.forceScopeAndTables }); + if (schemaStatus !== 'ready') { + return schemaStatus; } - if (await maybeConfigureTableScope(input)) { + const tableStatus = await maybeConfigureTableScope({ ...input, forcePrompt: input.forceScopeAndTables }); + if (tableStatus === 'ready') { break; } + if (input.forceScopeAndTables) { + return tableStatus; + } + + if (tableStatus === 'failed') { + return 'failed'; + } + await clearScopeConfig(input.projectDir, input.connectionId); } @@ -1554,7 +1717,7 @@ async function validateAndScanConnection(input: { ); } if (scanCode !== 0) { - return false; + return 'failed'; } } const scanOutput = scanIo.stdoutText(); @@ -1570,14 +1733,14 @@ async function validateAndScanConnection(input: { writeSetupSection(input.io, 'Primary source ready', [ `${input.connectionId} · ${driverDisplay} · structural scan complete`, ]); - return true; + return 'ready'; } async function chooseDrivers( args: KtxSetupDatabasesArgs, io: KtxCliIo, prompts: KtxSetupDatabasesPromptAdapter, - options?: { hasPrimarySources?: boolean }, + options?: { hasPrimarySources?: boolean; initialDrivers?: KtxSetupDatabaseDriver[] }, ): Promise { if (args.databaseDrivers && args.databaseDrivers.length > 0) { return [...new Set(args.databaseDrivers)]; @@ -1592,10 +1755,12 @@ async function chooseDrivers( return 'missing-input'; } while (true) { + const initialValues = unique(options?.initialDrivers ?? []); const choices = await prompts.multiselect({ message: withMultiselectNavigation('Which primary sources should KTX connect to?'), options: [...DRIVER_OPTIONS], - required: false, + ...(initialValues.length > 0 ? { initialValues } : {}), + required: options?.hasPrimarySources === true, }); if (choices.includes('back')) { return 'back'; @@ -1617,7 +1782,7 @@ async function chooseConnectionIdForDriver(input: { connections: Record; args: KtxSetupDatabasesArgs; prompts: KtxSetupDatabasesPromptAdapter; -}): Promise<{ kind: 'existing' | 'new'; connectionId: string } | 'back' | 'missing-input'> { +}): Promise<{ kind: 'existing' | 'new' | 'edit'; connectionId: string } | 'back' | 'missing-input'> { if (input.args.databaseConnectionId) { return { kind: 'new', connectionId: input.args.databaseConnectionId }; } @@ -1647,14 +1812,19 @@ async function chooseConnectionIdForDriver(input: { options: [ ...existingIds.map((connectionId) => ({ value: `existing:${connectionId}`, - label: `Use existing ${label} connection: ${connectionId}`, + label: `Keep existing ${label} connection: ${connectionId}`, })), - { value: 'new', label: `Add new ${label} connection` }, + ...existingIds.map((connectionId) => ({ + value: `edit:${connectionId}`, + label: `Edit ${label} connection: ${connectionId}`, + })), + { value: 'new', label: `Add another ${label} connection` }, { value: 'back', label: 'Back' }, ], }); if (choice === 'back') return 'back'; if (choice.startsWith('existing:')) return { kind: 'existing', connectionId: choice.slice('existing:'.length) }; + if (choice.startsWith('edit:')) return { kind: 'edit', connectionId: choice.slice('edit:'.length) }; const entered = await input.prompts.text({ message: withTextInputNavigation(connectionNamePrompt(label)), placeholder: defaultId, @@ -1666,6 +1836,102 @@ async function chooseConnectionIdForDriver(input: { } } +async function choosePrimarySourceToEdit(input: { + projectDir: string; + connectionIds: string[]; + prompts: KtxSetupDatabasesPromptAdapter; +}): Promise { + const project = await loadKtxProject({ projectDir: input.projectDir }); + const options = input.connectionIds + .map((connectionId) => { + const driver = normalizeDriver(project.config.connections[connectionId]?.driver); + if (!driver) return null; + return { value: connectionId, label: `${connectionId} (${driverLabel(driver)})` }; + }) + .filter((option): option is { value: string; label: string } => option !== null); + if (options.length === 0) return 'back'; + const choice = await input.prompts.select({ + message: 'Primary source to edit', + options: [...options, { value: 'back', label: 'Back' }], + }); + return choice === 'back' ? 'back' : choice; +} + +async function runPrimarySourceFullEdit(input: { + projectDir: string; + connectionId: string; + args: KtxSetupDatabasesArgs; + prompts: KtxSetupDatabasesPromptAdapter; + io: KtxCliIo; + deps: KtxSetupDatabasesDeps; +}): Promise<'ready' | 'back' | 'failed'> { + const project = await loadKtxProject({ projectDir: input.projectDir }); + const existing = project.config.connections[input.connectionId]; + const driver = normalizeDriver(existing?.driver); + if (!existing || !driver) { + input.io.stderr.write(`Connection "${input.connectionId}" is not a configured primary source.\n`); + return 'failed'; + } + + const rollback = await createConnectionConfigRollback(input.projectDir, input.connectionId); + const replacement = await buildConnectionConfig({ + driver, + connectionId: input.connectionId, + args: input.args, + prompts: input.prompts, + existingConnection: existing, + }); + if (replacement === 'back') { + await rollback(); + return 'back'; + } + if (!replacement) { + await rollback(); + return 'failed'; + } + + const withHistoricSql = await maybeApplyHistoricSqlConfig({ + connection: replacement, + 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, + ...(!Object.hasOwn(withHistoricSql, 'historicSql') && existing.historicSql !== undefined + ? { historicSql: existing.historicSql } + : {}), + }, + driver, + }), + }); + + const validated = await validateAndScanConnection({ + projectDir: input.projectDir, + connectionId: input.connectionId, + io: input.io, + deps: input.deps, + args: input.args, + prompts: input.prompts, + forceScopeAndTables: true, + }); + if (validated !== 'ready') { + await rollback(); + return validated; + } + return 'ready'; +} + export async function runKtxSetupDatabasesStep( args: KtxSetupDatabasesArgs, io: KtxCliIo, @@ -1688,7 +1954,18 @@ export async function runKtxSetupDatabasesStep( prompts, }); if (historicSqlResult === 'back') return { status: 'back', projectDir: args.projectDir }; - if (!(await validateAndScanConnection({ projectDir: args.projectDir, connectionId, io, deps, args, prompts }))) { + const setupStatus = await validateAndScanConnection({ + projectDir: args.projectDir, + connectionId, + io, + deps, + args, + prompts, + }); + if (setupStatus === 'back') { + return { status: 'back', projectDir: args.projectDir }; + } + if (setupStatus === 'failed') { return { status: 'failed', projectDir: args.projectDir }; } selectedConnectionIds.push(connectionId); @@ -1712,10 +1989,43 @@ export async function runKtxSetupDatabasesStep( await markDatabasesComplete(args.projectDir, selectedConnectionIds); return { status: 'ready', projectDir: args.projectDir, connectionIds: selectedConnectionIds }; } + if (action === 'edit') { + const connectionId = await choosePrimarySourceToEdit({ + projectDir: args.projectDir, + connectionIds: selectedConnectionIds, + prompts, + }); + if (connectionId === 'back') { + showConfiguredPrimaryMenu = true; + continue; + } + const editResult = await runPrimarySourceFullEdit({ + projectDir: args.projectDir, + connectionId, + args, + prompts, + io, + deps, + }); + if (editResult === 'back') { + showConfiguredPrimaryMenu = true; + continue; + } + if (editResult === 'failed') { + return { status: 'failed', projectDir: args.projectDir }; + } + pushUniqueConnectionId(selectedConnectionIds, connectionId); + showConfiguredPrimaryMenu = true; + continue; + } } showConfiguredPrimaryMenu = false; - const drivers = await chooseDrivers(args, io, prompts, { hasPrimarySources: selectedConnectionIds.length > 0 }); + const driverProject = await loadKtxProject({ projectDir: args.projectDir }); + const drivers = await chooseDrivers(args, io, prompts, { + hasPrimarySources: selectedConnectionIds.length > 0, + initialDrivers: configuredPrimaryDrivers(driverProject.config.connections, selectedConnectionIds), + }); if (drivers === 'back') { if (selectedConnectionIds.length > 0 && canReturnToDriverSelection && args.inputMode !== 'disabled') { showConfiguredPrimaryMenu = true; @@ -1750,7 +2060,26 @@ export async function runKtxSetupDatabasesStep( return { status: 'missing-input', projectDir: args.projectDir }; } - if (connectionChoice.kind === 'new') { + let connectionAlreadyValidated = false; + if (connectionChoice.kind === 'edit') { + const editResult = await runPrimarySourceFullEdit({ + projectDir: args.projectDir, + connectionId: connectionChoice.connectionId, + args, + prompts, + io, + deps, + }); + if (editResult === 'back') { + if (!canReturnToDriverSelection) return { status: 'back', projectDir: args.projectDir }; + 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, @@ -1819,16 +2148,22 @@ export async function runKtxSetupDatabasesStep( } let connectionSkipped = false; - while ( - !(await validateAndScanConnection({ - projectDir: args.projectDir, - connectionId: connectionChoice.connectionId, - io, - deps, - args, - prompts, - })) - ) { + let setupStatus: ConnectionSetupStatus = connectionAlreadyValidated + ? 'ready' + : await validateAndScanConnection({ + projectDir: args.projectDir, + connectionId: connectionChoice.connectionId, + io, + deps, + args, + prompts, + }); + 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 action = await prompts.select({ message: `Primary source setup failed for ${connectionChoice.connectionId}`, @@ -1848,7 +2183,16 @@ export async function runKtxSetupDatabasesStep( connectionSkipped = true; break; } - if (action === 're-enter') { + if (action === 'retry') { + setupStatus = await validateAndScanConnection({ + projectDir: args.projectDir, + connectionId: connectionChoice.connectionId, + io, + deps, + args, + prompts, + }); + } else if (action === 're-enter') { const connection = await buildConnectionConfig({ driver, connectionId: connectionChoice.connectionId, @@ -1872,6 +2216,14 @@ export async function runKtxSetupDatabasesStep( connectionId: connectionChoice.connectionId, connection: withHistoricSql, }); + setupStatus = await validateAndScanConnection({ + projectDir: args.projectDir, + connectionId: connectionChoice.connectionId, + io, + deps, + args, + prompts, + }); } } if (returnToDriverSelection) break; 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; } }