From f1a275144f0577bc3466512b40f05a1fde2c647c Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Fri, 22 May 2026 16:44:17 +0200 Subject: [PATCH] feat(cli): add RSA key-pair auth option to Snowflake setup wizard Extends the interactive Snowflake setup flow with an authentication-method prompt (password vs RSA/JWT key-pair). The RSA branch collects a private-key path (env/file/absolute) and an optional passphrase; the resulting connection config records `authMethod: 'rsa'` with `privateKey` and `passphrase` instead of `password`. --- packages/cli/src/setup-databases.test.ts | 50 +++++++++++++++- packages/cli/src/setup-databases.ts | 73 ++++++++++++++++++++---- 2 files changed, 111 insertions(+), 12 deletions(-) diff --git a/packages/cli/src/setup-databases.test.ts b/packages/cli/src/setup-databases.test.ts index 28c9e937..24b3a450 100644 --- a/packages/cli/src/setup-databases.test.ts +++ b/packages/cli/src/setup-databases.test.ts @@ -516,7 +516,7 @@ describe('setup databases step', () => { }, { driver: 'snowflake', - selectValues: ['no'], + selectValues: ['password', 'no'], textValues: ['', 'env:SNOWFLAKE_ACCOUNT', 'ANALYTICS_WH', 'ANALYTICS', '', 'env:SNOWFLAKE_USER', ''], passwordValues: ['env:SNOWFLAKE_PASSWORD'], expectedTextPrompts: [ @@ -2004,6 +2004,7 @@ describe('setup databases step', () => { testConnection: vi.fn(async () => 0), scanConnection: vi.fn(async () => 0), prompts: makePromptAdapter({ + selectValues: ['password'], textValues: ['env:SNOWFLAKE_ACCOUNT', 'WH', 'ANALYTICS', 'PUBLIC', 'reader', ''], passwordValues: ['env:SNOWFLAKE_PASSWORD'], }), @@ -2038,6 +2039,53 @@ describe('setup databases step', () => { expect(config.ingest.adapters).toEqual([]); }); + it('configures Snowflake with RSA key-pair auth via setup wizard', async () => { + const io = makeIo(); + const result = await runKtxSetupDatabasesStep( + { + projectDir: tempDir, + inputMode: 'disabled', + databaseDrivers: ['snowflake'], + databaseConnectionId: 'snowflake', + databaseSchemas: [], + skipDatabases: false, + }, + io.io, + { + testConnection: vi.fn(async () => 0), + scanConnection: vi.fn(async () => 0), + prompts: makePromptAdapter({ + selectValues: ['rsa'], + textValues: [ + 'env:SNOWFLAKE_ACCOUNT', + 'WH', + 'ANALYTICS', + 'PUBLIC', + 'reader', + '~/.ssh/snowflake_rsa_key.p8', + '', + ], + passwordValues: ['env:SNOWFLAKE_KEY_PASS'], + }), + }, + ); + + expect(result.status).toBe('ready'); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.snowflake).toMatchObject({ + driver: 'snowflake', + authMethod: 'rsa', + account: 'env:SNOWFLAKE_ACCOUNT', + warehouse: 'WH', + database: 'ANALYTICS', + schema_name: 'PUBLIC', + username: 'reader', + privateKey: 'file:~/.ssh/snowflake_rsa_key.p8', // pragma: allowlist secret + passphrase: 'env:SNOWFLAKE_KEY_PASS', // pragma: allowlist secret + }); + expect(config.connections.snowflake.password).toBeUndefined(); + }); + it('writes Postgres query history config with minExecutions and ignores window/redaction output', async () => { const io = makeIo(); const result = await runKtxSetupDatabasesStep( diff --git a/packages/cli/src/setup-databases.ts b/packages/cli/src/setup-databases.ts index acdebeec..e721ca57 100644 --- a/packages/cli/src/setup-databases.ts +++ b/packages/cli/src/setup-databases.ts @@ -964,31 +964,82 @@ async function buildConnectionConfig(input: { stringConfigField(input.existingConnection, 'username'), ); if (username === undefined) return 'back'; - const passwordRef = await promptCredential({ - prompts, - message: 'Snowflake password', - projectDir: args.projectDir, - connectionId: input.connectionId, - secretName: 'password', // pragma: allowlist secret + const authChoice = await prompts.select({ + message: 'Snowflake authentication method', + options: [ + { value: 'password', label: 'Password' }, + { value: 'rsa', label: 'Key-pair (RSA / JWT)' }, + { value: 'back', label: 'Back' }, + ], }); - if (passwordRef === 'back') return 'back'; // pragma: allowlist secret + if (authChoice === 'back') return 'back'; + const authMethod: 'password' | 'rsa' = authChoice === 'rsa' ? 'rsa' : 'password'; + let passwordRef: string | null = null; + let privateKeyInput: string | undefined; + let passphraseRef: string | null = null; + if (authMethod === 'password') { + const ref = await promptCredential({ + prompts, + message: 'Snowflake password', + projectDir: args.projectDir, + connectionId: input.connectionId, + secretName: 'password', // pragma: allowlist secret + }); + if (ref === 'back') return 'back'; // pragma: allowlist secret + passwordRef = ref; + } else { + privateKeyInput = await promptText( + prompts, + 'Path to Snowflake private key (PEM)\nFor example ~/.ssh/snowflake_rsa_key.p8, or $ENV_VAR / env:NAME / file:/abs/path.', + displayFileReference(stringConfigField(input.existingConnection, 'privateKey')), + ); + if (privateKeyInput === undefined) return 'back'; + const phr = await promptCredential({ + prompts, + message: 'Private key passphrase (optional)\nPress Enter to skip.', + projectDir: args.projectDir, + connectionId: input.connectionId, + secretName: 'snowflake-passphrase', // pragma: allowlist secret + }); + if (phr === 'back') return 'back'; + passphraseRef = phr; + } const role = await promptText( prompts, 'Snowflake role (optional)\nPress Enter to skip.', stringConfigField(input.existingConnection, 'role'), ); if (role === undefined) return 'back'; - const resolvedPasswordRef = passwordRef ?? stringConfigField(input.existingConnection, 'password'); - if (!account || !warehouse || !database || !schemaName || !username || !resolvedPasswordRef) return null; + if (authMethod === 'password') { + const resolvedPasswordRef = passwordRef ?? stringConfigField(input.existingConnection, 'password'); + if (!account || !warehouse || !database || !schemaName || !username || !resolvedPasswordRef) return null; + return { + driver: 'snowflake', + authMethod: 'password', + account, + warehouse, + database, + schema_name: schemaName, + username, + password: resolvedPasswordRef, + ...(role ? { role } : {}), + }; + } + const resolvedPrivateKey = privateKeyInput + ? normalizeFileReference(privateKeyInput) + : stringConfigField(input.existingConnection, 'privateKey'); + if (!account || !warehouse || !database || !schemaName || !username || !resolvedPrivateKey) return null; + const resolvedPassphrase = passphraseRef ?? stringConfigField(input.existingConnection, 'passphrase'); return { driver: 'snowflake', - authMethod: 'password', + authMethod: 'rsa', account, warehouse, database, schema_name: schemaName, username, - password: resolvedPasswordRef, + privateKey: resolvedPrivateKey, + ...(resolvedPassphrase ? { passphrase: resolvedPassphrase } : {}), ...(role ? { role } : {}), }; }