From d111128a234a549034112ca4b097c65cf82e1989 Mon Sep 17 00:00:00 2001 From: mattsenicksigma Date: Sat, 27 Jun 2026 00:07:12 -0700 Subject: [PATCH 1/4] feat(snowflake): add OAuth token authentication Adds authMethod: oauth to the Snowflake connector alongside the existing password and rsa methods. An OAuth access token is supplied via token_ref (env: or file: reference) or a literal token field, and is passed to the Snowflake SDK as authenticator: OAUTH. username is optional for OAuth since the token carries identity, but is accepted when the account requires it. Docs and five new config-resolution tests included. Co-Authored-By: Claude Sonnet 4.6 --- .../docs/integrations/primary-sources.mdx | 21 ++++- .../cli/src/connectors/snowflake/connector.ts | 36 ++++++-- .../connectors/snowflake/connector.test.ts | 83 +++++++++++++++++++ 3 files changed, 130 insertions(+), 10 deletions(-) diff --git a/docs-site/content/docs/integrations/primary-sources.mdx b/docs-site/content/docs/integrations/primary-sources.mdx index 6cb2d26fe..6a14f103b 100644 --- a/docs-site/content/docs/integrations/primary-sources.mdx +++ b/docs-site/content/docs/integrations/primary-sources.mdx @@ -118,7 +118,7 @@ This helps **ktx** understand how your team actually queries the data. ## Snowflake -Connects via the Snowflake SDK. Supports multi-schema scanning, RSA key authentication, and query-history configuration for Snowflake query history. +Connects via the Snowflake SDK. Supports multi-schema scanning, password, RSA key pair, and OAuth authentication, and query-history configuration for Snowflake query history. ### Connection config @@ -146,8 +146,25 @@ single schema, `schema_name: PUBLIC` is accepted as an equivalent shorthand. | Method | Config | |--------|--------| -| Password | `password: env:SNOWFLAKE_PASSWORD` | +| Password (default) | `password: env:SNOWFLAKE_PASSWORD` | | RSA key pair | `authMethod: rsa`, `privateKey: file:~/.ssh/snowflake_key.pem`, optional `passphrase` | +| OAuth | `authMethod: oauth`, `token_ref: env:SNOWFLAKE_OAUTH_TOKEN` | + +For OAuth, `username` is optional — the token carries identity. If your Snowflake account requires it alongside the token (e.g. for auditing), add `username` as normal. + +```yaml title="ktx.yaml — OAuth example" +connections: + my-snowflake: + driver: snowflake + authMethod: oauth + account: xy12345 + warehouse: ANALYTICS_WH + database: PROD + token_ref: env:SNOWFLAKE_OAUTH_TOKEN + role: ANALYST +``` + +`token_ref` follows the same reference syntax as other secrets: `env:VAR` reads from an environment variable, `file:/path/to/token` reads from a file (supports `~` expansion). A literal token value can be supplied directly with `token:`, but `token_ref` is preferred. ### Features diff --git a/packages/cli/src/connectors/snowflake/connector.ts b/packages/cli/src/connectors/snowflake/connector.ts index 5f0166757..c95cae4de 100644 --- a/packages/cli/src/connectors/snowflake/connector.ts +++ b/packages/cli/src/connectors/snowflake/connector.ts @@ -33,7 +33,7 @@ import { configureSnowflakeSdkLogger } from './sdk-logger.js'; export interface KtxSnowflakeConnectionConfig { driver?: string; - authMethod?: 'password' | 'rsa'; + authMethod?: 'password' | 'rsa' | 'oauth'; account?: string; warehouse?: string; database?: string; @@ -43,21 +43,24 @@ export interface KtxSnowflakeConnectionConfig { password?: string; privateKey?: string; passphrase?: string; + token?: string; + token_ref?: string; role?: string; maxConnections?: number; [key: string]: unknown; } export interface KtxSnowflakeResolvedConnectionConfig { - authMethod: 'password' | 'rsa'; + authMethod: 'password' | 'rsa' | 'oauth'; account: string; warehouse: string; database: string; schemas: string[]; - username: string; + username?: string; password?: string; privateKey?: string; passphrase?: string; + token?: string; role?: string; maxConnections: number; } @@ -260,7 +263,8 @@ export function snowflakeConnectionConfigFromConfig(input: { if (!database) { throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.database`); } - if (!username) { + // username is required for password/rsa; optional for oauth (token carries identity). + if (authMethod !== 'oauth' && !username) { throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.username`); } assertSafeSnowflakeIdentifier(warehouse, 'warehouse'); @@ -275,7 +279,7 @@ export function snowflakeConnectionConfigFromConfig(input: { warehouse, database, schemas: resolvedSchemas, - username, + ...(username ? { username } : {}), maxConnections: positiveIntegerConfigValue({ connection: input.connection, key: 'maxConnections', @@ -297,6 +301,18 @@ export function snowflakeConnectionConfigFromConfig(input: { if (!resolved.privateKey) { throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.privateKey for RSA auth`); } + } else if (authMethod === 'oauth') { + // token_ref is the preferred form; token accepts a literal value. + const tokenRef = input.connection?.token_ref; + const token = tokenRef + ? resolveStringReference(tokenRef as string, env) + : stringConfigValue(input.connection, 'token', env); + if (!token) { + throw new Error( + `Native Snowflake connector requires connections.${input.connectionId}.token_ref (or token) for OAuth auth`, + ); + } + resolved.token = token; } else { resolved.password = stringConfigValue(input.connection, 'password', env); if (!resolved.password) { @@ -488,9 +504,13 @@ class SnowflakeSdkDriver implements KtxSnowflakeDriver { clientSessionKeepAliveHeartbeatFrequency: 900, ...patch?.sdkOptions, }; - return this.resolved.authMethod === 'rsa' - ? { ...baseConfig, authenticator: 'SNOWFLAKE_JWT', privateKey: this.decryptPrivateKey() } - : { ...baseConfig, password: this.resolved.password }; + if (this.resolved.authMethod === 'rsa') { + return { ...baseConfig, authenticator: 'SNOWFLAKE_JWT', privateKey: this.decryptPrivateKey() }; + } + if (this.resolved.authMethod === 'oauth') { + return { ...baseConfig, authenticator: 'OAUTH', token: this.resolved.token }; + } + return { ...baseConfig, password: this.resolved.password }; } private async executeSnowflakeQuery( diff --git a/packages/cli/test/connectors/snowflake/connector.test.ts b/packages/cli/test/connectors/snowflake/connector.test.ts index 1b00061bb..e96457e43 100644 --- a/packages/cli/test/connectors/snowflake/connector.test.ts +++ b/packages/cli/test/connectors/snowflake/connector.test.ts @@ -194,6 +194,89 @@ describe('KtxSnowflakeScanConnector', () => { } }); + it('resolves OAuth config with token_ref and does not require username', () => { + const resolved = snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'oauth', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + token_ref: 'env:SNOWFLAKE_OAUTH_TOKEN', + }, + env: { SNOWFLAKE_OAUTH_TOKEN: 'test-token-value' }, // pragma: allowlist secret + }); + expect(resolved.authMethod).toBe('oauth'); + expect(resolved.token).toBe('test-token-value'); + expect(resolved.username).toBeUndefined(); + }); + + it('resolves OAuth config with inline token', () => { + const resolved = snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'oauth', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + token: 'inline-token', // pragma: allowlist secret + }, + }); + expect(resolved.authMethod).toBe('oauth'); + expect(resolved.token).toBe('inline-token'); + }); + + it('resolves OAuth config with optional username', () => { + const resolved = snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'oauth', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + username: 'svc_account', + token_ref: 'env:SNOWFLAKE_OAUTH_TOKEN', + }, + env: { SNOWFLAKE_OAUTH_TOKEN: 'test-token-value' }, // pragma: allowlist secret + }); + expect(resolved.username).toBe('svc_account'); + expect(resolved.token).toBe('test-token-value'); + }); + + it('throws when OAuth config is missing token and token_ref', () => { + expect(() => + snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'oauth', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + }, + }), + ).toThrow('connections.warehouse.token_ref'); + }); + + it('throws when password auth is missing username', () => { + expect(() => + snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'password', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + password: 'fixture-pass', // pragma: allowlist secret + }, + }), + ).toThrow('connections.warehouse.username'); + }); + it('rejects stale Snowflake pool config key', () => { const baseConnection: KtxSnowflakeConnectionConfig = { driver: 'snowflake', From 61f6c6e8e61b97afd8c42cf8a95307414bd237a2 Mon Sep 17 00:00:00 2001 From: mattsenicksigma Date: Sat, 27 Jun 2026 00:33:11 -0700 Subject: [PATCH 2/4] feat(snowflake): add PAT and externalbrowser auth methods MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the Snowflake connector with two additional authentication paths: - `authMethod: pat` — programmatic access token via PROGRAMMATIC_ACCESS_TOKEN authenticator, suitable for headless and CI environments - `authMethod: externalbrowser` — SSO browser popup via EXTERNALBROWSER authenticator, username now optional so IDP resolves identity Username is now only required for password and RSA auth; oauth, pat, and externalbrowser all treat it as optional. Co-Authored-By: Claude Sonnet 4.6 --- .../docs/integrations/primary-sources.mdx | 12 +++-- .../cli/src/connectors/snowflake/connector.ts | 25 +++++---- .../connectors/snowflake/connector.test.ts | 52 +++++++++++++++++++ 3 files changed, 76 insertions(+), 13 deletions(-) diff --git a/docs-site/content/docs/integrations/primary-sources.mdx b/docs-site/content/docs/integrations/primary-sources.mdx index 6a14f103b..c78f9af6e 100644 --- a/docs-site/content/docs/integrations/primary-sources.mdx +++ b/docs-site/content/docs/integrations/primary-sources.mdx @@ -118,7 +118,7 @@ This helps **ktx** understand how your team actually queries the data. ## Snowflake -Connects via the Snowflake SDK. Supports multi-schema scanning, password, RSA key pair, and OAuth authentication, and query-history configuration for Snowflake query history. +Connects via the Snowflake SDK. Supports multi-schema scanning, password, RSA key pair, OAuth token, programmatic access token, and browser-based SSO authentication, and query-history configuration for Snowflake query history. ### Connection config @@ -148,9 +148,11 @@ single schema, `schema_name: PUBLIC` is accepted as an equivalent shorthand. |--------|--------| | Password (default) | `password: env:SNOWFLAKE_PASSWORD` | | RSA key pair | `authMethod: rsa`, `privateKey: file:~/.ssh/snowflake_key.pem`, optional `passphrase` | -| OAuth | `authMethod: oauth`, `token_ref: env:SNOWFLAKE_OAUTH_TOKEN` | +| OAuth token | `authMethod: oauth`, `token_ref: env:SNOWFLAKE_OAUTH_TOKEN` | +| Programmatic access token | `authMethod: pat`, `token_ref: env:SNOWFLAKE_PAT` | +| Browser SSO | `authMethod: externalbrowser` — opens a browser window for SSO/MFA login | -For OAuth, `username` is optional — the token carries identity. If your Snowflake account requires it alongside the token (e.g. for auditing), add `username` as normal. +For `oauth`, `pat`, and `externalbrowser`, `username` is optional — identity is carried by the token or resolved by the IDP. ```yaml title="ktx.yaml — OAuth example" connections: @@ -164,7 +166,9 @@ connections: role: ANALYST ``` -`token_ref` follows the same reference syntax as other secrets: `env:VAR` reads from an environment variable, `file:/path/to/token` reads from a file (supports `~` expansion). A literal token value can be supplied directly with `token:`, but `token_ref` is preferred. +`token_ref` follows the same reference syntax as other secrets: `env:VAR` reads from an environment variable, `file:/path/to/token` reads from a file. A literal value can be supplied directly with `token:`, but `token_ref` is preferred. + +Browser SSO authentication is interactive — running `ktx connection test` or `ktx setup` will open a browser window to complete the Snowflake SSO flow. Not suitable for headless or CI environments. ### Features diff --git a/packages/cli/src/connectors/snowflake/connector.ts b/packages/cli/src/connectors/snowflake/connector.ts index c95cae4de..da2312c8c 100644 --- a/packages/cli/src/connectors/snowflake/connector.ts +++ b/packages/cli/src/connectors/snowflake/connector.ts @@ -33,7 +33,7 @@ import { configureSnowflakeSdkLogger } from './sdk-logger.js'; export interface KtxSnowflakeConnectionConfig { driver?: string; - authMethod?: 'password' | 'rsa' | 'oauth'; + authMethod?: 'password' | 'rsa' | 'oauth' | 'pat' | 'externalbrowser'; account?: string; warehouse?: string; database?: string; @@ -51,7 +51,7 @@ export interface KtxSnowflakeConnectionConfig { } export interface KtxSnowflakeResolvedConnectionConfig { - authMethod: 'password' | 'rsa' | 'oauth'; + authMethod: 'password' | 'rsa' | 'oauth' | 'pat' | 'externalbrowser'; account: string; warehouse: string; database: string; @@ -263,9 +263,11 @@ export function snowflakeConnectionConfigFromConfig(input: { if (!database) { throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.database`); } - // username is required for password/rsa; optional for oauth (token carries identity). - if (authMethod !== 'oauth' && !username) { - throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.username`); + // username is required for password/rsa; optional for oauth/pat/externalbrowser. + if (authMethod === 'password' || authMethod === 'rsa') { + if (!username) { + throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.username`); + } } assertSafeSnowflakeIdentifier(warehouse, 'warehouse'); assertSafeSnowflakeIdentifier(database, 'database'); @@ -301,19 +303,18 @@ export function snowflakeConnectionConfigFromConfig(input: { if (!resolved.privateKey) { throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.privateKey for RSA auth`); } - } else if (authMethod === 'oauth') { - // token_ref is the preferred form; token accepts a literal value. + } else if (authMethod === 'oauth' || authMethod === 'pat') { const tokenRef = input.connection?.token_ref; const token = tokenRef ? resolveStringReference(tokenRef as string, env) : stringConfigValue(input.connection, 'token', env); if (!token) { throw new Error( - `Native Snowflake connector requires connections.${input.connectionId}.token_ref (or token) for OAuth auth`, + `Native Snowflake connector requires connections.${input.connectionId}.token_ref (or token) for ${authMethod} auth`, ); } resolved.token = token; - } else { + } else if (authMethod !== 'externalbrowser') { resolved.password = stringConfigValue(input.connection, 'password', env); if (!resolved.password) { throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.password`); @@ -510,6 +511,12 @@ class SnowflakeSdkDriver implements KtxSnowflakeDriver { if (this.resolved.authMethod === 'oauth') { return { ...baseConfig, authenticator: 'OAUTH', token: this.resolved.token }; } + if (this.resolved.authMethod === 'pat') { + return { ...baseConfig, authenticator: 'PROGRAMMATIC_ACCESS_TOKEN', token: this.resolved.token }; + } + if (this.resolved.authMethod === 'externalbrowser') { + return { ...baseConfig, authenticator: 'EXTERNALBROWSER' }; + } return { ...baseConfig, password: this.resolved.password }; } diff --git a/packages/cli/test/connectors/snowflake/connector.test.ts b/packages/cli/test/connectors/snowflake/connector.test.ts index e96457e43..115e0d54e 100644 --- a/packages/cli/test/connectors/snowflake/connector.test.ts +++ b/packages/cli/test/connectors/snowflake/connector.test.ts @@ -212,6 +212,41 @@ describe('KtxSnowflakeScanConnector', () => { expect(resolved.username).toBeUndefined(); }); + it('resolves PAT config with token_ref', () => { + const resolved = snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'pat', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + username: 'svc_user', + token_ref: 'env:SNOWFLAKE_PAT', + }, + env: { SNOWFLAKE_PAT: 'pat-token-value' }, // pragma: allowlist secret + }); + expect(resolved.authMethod).toBe('pat'); + expect(resolved.token).toBe('pat-token-value'); + expect(resolved.username).toBe('svc_user'); + }); + + it('throws when PAT config is missing token and token_ref', () => { + expect(() => + snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'pat', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + username: 'svc_user', + }, + }), + ).toThrow('connections.warehouse.token_ref'); + }); + it('resolves OAuth config with inline token', () => { const resolved = snowflakeConnectionConfigFromConfig({ connectionId: 'warehouse', @@ -261,6 +296,23 @@ describe('KtxSnowflakeScanConnector', () => { ).toThrow('connections.warehouse.token_ref'); }); + it('resolves externalbrowser config without username or credentials', () => { + const resolved = snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'externalbrowser', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + }, + }); + expect(resolved.authMethod).toBe('externalbrowser'); + expect(resolved.username).toBeUndefined(); + expect(resolved.password).toBeUndefined(); + expect(resolved.privateKey).toBeUndefined(); + }); + it('throws when password auth is missing username', () => { expect(() => snowflakeConnectionConfigFromConfig({ From ab4fecdf0d7a947b87a94fb40c219f1741cb3540 Mon Sep 17 00:00:00 2001 From: mattsenicksigma Date: Thu, 2 Jul 2026 14:58:07 -0700 Subject: [PATCH 3/4] update --- .../docs/integrations/primary-sources.mdx | 10 +- .../cli/src/connectors/snowflake/connector.ts | 72 +++++++--- packages/cli/src/setup-databases.ts | 78 ++++++++--- .../connectors/snowflake/connector.test.ts | 53 +++++-- packages/cli/test/setup-databases.test.ts | 131 ++++++++++++++++++ 5 files changed, 291 insertions(+), 53 deletions(-) diff --git a/docs-site/content/docs/integrations/primary-sources.mdx b/docs-site/content/docs/integrations/primary-sources.mdx index a2814bde9..87eae3a5c 100644 --- a/docs-site/content/docs/integrations/primary-sources.mdx +++ b/docs-site/content/docs/integrations/primary-sources.mdx @@ -156,8 +156,8 @@ single schema, `schema_name: PUBLIC` is accepted as an equivalent shorthand. |--------|--------| | Password (default) | `password: env:SNOWFLAKE_PASSWORD` | | RSA key pair | `authMethod: rsa`, `privateKey: file:~/.ssh/snowflake_key.pem`, optional `passphrase` | -| OAuth token | `authMethod: oauth`, `token_ref: env:SNOWFLAKE_OAUTH_TOKEN` | -| Programmatic access token | `authMethod: pat`, `token_ref: env:SNOWFLAKE_PAT` | +| OAuth token | `authMethod: oauth`, `token: env:SNOWFLAKE_OAUTH_TOKEN` | +| Programmatic access token | `authMethod: pat`, `token: env:SNOWFLAKE_PAT` | | Browser SSO | `authMethod: externalbrowser` — opens a browser window for SSO/MFA login | For `oauth`, `pat`, and `externalbrowser`, `username` is optional — identity is carried by the token or resolved by the IDP. @@ -170,13 +170,13 @@ connections: account: xy12345 warehouse: ANALYTICS_WH database: PROD - token_ref: env:SNOWFLAKE_OAUTH_TOKEN + token: env:SNOWFLAKE_OAUTH_TOKEN role: ANALYST ``` -`token_ref` follows the same reference syntax as other secrets: `env:VAR` reads from an environment variable, `file:/path/to/token` reads from a file. A literal value can be supplied directly with `token:`, but `token_ref` is preferred. +`token` follows the same reference syntax as other secrets: `env:VAR` reads from an environment variable, `file:/path/to/token` reads from a file, or a literal value can be supplied directly. -Browser SSO authentication is interactive — running `ktx connection test` or `ktx setup` will open a browser window to complete the Snowflake SSO flow. Not suitable for headless or CI environments. +Browser SSO authentication is interactive — running `ktx connection test` or `ktx setup` will open a browser window to complete the Snowflake SSO flow. Not suitable for headless or CI environments. `ktx setup`'s Snowflake wizard supports all five auth methods, both when creating a new connection and when editing an existing one. ### Features diff --git a/packages/cli/src/connectors/snowflake/connector.ts b/packages/cli/src/connectors/snowflake/connector.ts index 6697ef84d..b05859a03 100644 --- a/packages/cli/src/connectors/snowflake/connector.ts +++ b/packages/cli/src/connectors/snowflake/connector.ts @@ -45,7 +45,6 @@ export interface KtxSnowflakeConnectionConfig { privateKey?: string; passphrase?: string; token?: string; - token_ref?: string; role?: string; maxConnections?: number; [key: string]: unknown; @@ -62,6 +61,8 @@ export interface KtxSnowflakeResolvedConnectionConfig { privateKey?: string; passphrase?: string; token?: string; + /** Re-reads the token source (env var or file) so a rotated short-lived token is picked up on each new pooled connection. */ + refreshToken?: () => string; role?: string; maxConnections: number; deadlineMs: number; @@ -268,6 +269,12 @@ export function snowflakeConnectionConfigFromConfig(input: { } const env = input.env ?? process.env; const authMethod = input.connection?.authMethod ?? 'password'; + const knownAuthMethods = ['password', 'rsa', 'oauth', 'pat', 'externalbrowser']; + if (!knownAuthMethods.includes(authMethod)) { + throw new Error( + `connections.${input.connectionId}.authMethod "${authMethod}" is not supported; use one of ${knownAuthMethods.join(', ')}`, + ); + } const account = stringConfigValue(input.connection, 'account', env); const warehouse = stringConfigValue(input.connection, 'warehouse', env); const database = stringConfigValue(input.connection, 'database', env); @@ -323,16 +330,17 @@ export function snowflakeConnectionConfigFromConfig(input: { throw new Error(`Native Snowflake connector requires connections.${input.connectionId}.privateKey for RSA auth`); } } else if (authMethod === 'oauth' || authMethod === 'pat') { - const tokenRef = input.connection?.token_ref; - const token = tokenRef - ? resolveStringReference(tokenRef as string, env) - : stringConfigValue(input.connection, 'token', env); - if (!token) { - throw new Error( - `Native Snowflake connector requires connections.${input.connectionId}.token_ref (or token) for ${authMethod} auth`, - ); - } - resolved.token = token; + const connection = input.connection; + const connectionId = input.connectionId; + const refreshToken = (): string => { + const token = stringConfigValue(connection, 'token', env); + if (!token) { + throw new Error(`Native Snowflake connector requires connections.${connectionId}.token for ${authMethod} auth`); + } + return token; + }; + resolved.token = refreshToken(); + resolved.refreshToken = refreshToken; } else if (authMethod !== 'externalbrowser') { resolved.password = stringConfigValue(input.connection, 'password', env); if (!resolved.password) { @@ -504,16 +512,27 @@ class SnowflakeSdkDriver implements KtxSnowflakeDriver { private async getPool(): Promise> { if (!this.pool) { - this.pool = snowflake.createPool(await this.resolveConnectionOptions(), { - min: 0, - max: this.resolved.maxConnections, - evictionRunIntervalMillis: 30_000, - acquireTimeoutMillis: 60_000, - }); + this.pool = snowflake.createPool(await this.resolveConnectionOptions(), this.poolPolicy()); } return this.pool; } + private poolPolicy(): { min: number; max: number; evictionRunIntervalMillis: number; idleTimeoutMillis?: number; acquireTimeoutMillis: number } { + if (this.resolved.authMethod !== 'externalbrowser') { + return { min: 0, max: this.resolved.maxConnections, evictionRunIntervalMillis: 30_000, acquireTimeoutMillis: 60_000 }; + } + // externalbrowser opens a system browser for SSO/MFA (default 120s per attempt, see + // WAIT_FOR_BROWSER_ACTION_TIMEOUT in the SDK). Keep at least one connection warm and + // avoid idle eviction so a scan doesn't force a repeat login mid-run, and give acquire + // more headroom than the SDK's own browser timeout so the pool never gives up first. + return { + min: 1, + max: this.resolved.maxConnections, + evictionRunIntervalMillis: 0, + acquireTimeoutMillis: 150_000, + }; + } + private async resolveConnectionOptions(): Promise { const patch = await this.sdkOptionsProvider?.resolve({ account: this.resolved.account, @@ -537,11 +556,20 @@ class SnowflakeSdkDriver implements KtxSnowflakeDriver { if (this.resolved.authMethod === 'rsa') { return { ...baseConfig, authenticator: 'SNOWFLAKE_JWT', privateKey: this.decryptPrivateKey() }; } - if (this.resolved.authMethod === 'oauth') { - return { ...baseConfig, authenticator: 'OAUTH', token: this.resolved.token }; - } - if (this.resolved.authMethod === 'pat') { - return { ...baseConfig, authenticator: 'PROGRAMMATIC_ACCESS_TOKEN', token: this.resolved.token }; + if (this.resolved.authMethod === 'oauth' || this.resolved.authMethod === 'pat') { + const authenticator = this.resolved.authMethod === 'oauth' ? 'OAUTH' : 'PROGRAMMATIC_ACCESS_TOKEN'; + const refreshToken = this.resolved.refreshToken; + const fallbackToken = this.resolved.token; + return { + ...baseConfig, + authenticator, + // The pool reuses this options object for every new physical connection it opens + // (min: 0 plus idle eviction means that can happen well after construction), so + // re-read a short-lived token each time rather than baking in the value from setup. + get token() { + return refreshToken ? refreshToken() : fallbackToken; + }, + }; } if (this.resolved.authMethod === 'externalbrowser') { return { ...baseConfig, authenticator: 'EXTERNALBROWSER' }; diff --git a/packages/cli/src/setup-databases.ts b/packages/cli/src/setup-databases.ts index efeba2092..c354e28b4 100644 --- a/packages/cli/src/setup-databases.ts +++ b/packages/cli/src/setup-databases.ts @@ -895,25 +895,32 @@ async function buildConnectionConfig(input: { stringConfigField(input.existingConnection, 'database'), ); if (database === undefined) return 'back'; - const username = await promptText( - prompts, - 'Snowflake username', - stringConfigField(input.existingConnection, 'username'), - ); - if (username === undefined) return 'back'; const authChoice = await prompts.select({ message: 'Snowflake authentication method', options: [ { value: 'password', label: 'Password' }, { value: 'rsa', label: 'Key-pair (RSA / JWT)' }, + { value: 'oauth', label: 'OAuth token' }, + { value: 'pat', label: 'Programmatic access token' }, + { value: 'externalbrowser', label: 'Browser SSO' }, { value: 'back', label: 'Back' }, ], }); if (authChoice === 'back') return 'back'; - const authMethod: 'password' | 'rsa' = authChoice === 'rsa' ? 'rsa' : 'password'; + const authMethod = authChoice as 'password' | 'rsa' | 'oauth' | 'pat' | 'externalbrowser'; + const usernameRequired = authMethod === 'password' || authMethod === 'rsa'; + const username = await promptText( + prompts, + usernameRequired + ? 'Snowflake username' + : 'Snowflake username (optional)\nPress Enter to skip — identity is carried by the token or resolved by the IDP.', + stringConfigField(input.existingConnection, 'username'), + ); + if (username === undefined) return 'back'; let passwordRef: string | null = null; let privateKeyInput: string | undefined; let passphraseRef: string | null = null; + let tokenRef: string | null = null; if (authMethod === 'password') { const ref = await promptCredential({ prompts, @@ -924,7 +931,7 @@ async function buildConnectionConfig(input: { }); if (ref === 'back') return 'back'; // pragma: allowlist secret passwordRef = ref; - } else { + } else if (authMethod === 'rsa') { 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.', @@ -940,7 +947,18 @@ async function buildConnectionConfig(input: { }); if (phr === 'back') return 'back'; passphraseRef = phr; + } else if (authMethod === 'oauth' || authMethod === 'pat') { + const ref = await promptCredential({ + prompts, + message: authMethod === 'oauth' ? 'Snowflake OAuth token' : 'Snowflake programmatic access token (PAT)', + projectDir: args.projectDir, + connectionId: input.connectionId, + secretName: 'token', // pragma: allowlist secret + }); + if (ref === 'back') return 'back'; + tokenRef = ref; } + // externalbrowser needs no credential prompt — the SDK opens a browser for SSO/MFA. const role = await promptText( prompts, 'Snowflake role (optional)\nPress Enter to skip.', @@ -961,20 +979,46 @@ async function buildConnectionConfig(input: { ...(role ? { role } : {}), }; } - const resolvedPrivateKey = privateKeyInput - ? normalizeFileReference(privateKeyInput) - : stringConfigField(input.existingConnection, 'privateKey'); - if (!account || !warehouse || !database || !username || !resolvedPrivateKey) return null; - const resolvedPassphrase = passphraseRef ?? stringConfigField(input.existingConnection, 'passphrase'); + if (authMethod === 'rsa') { + const resolvedPrivateKey = privateKeyInput + ? normalizeFileReference(privateKeyInput) + : stringConfigField(input.existingConnection, 'privateKey'); + if (!account || !warehouse || !database || !username || !resolvedPrivateKey) return null; + const resolvedPassphrase = passphraseRef ?? stringConfigField(input.existingConnection, 'passphrase'); + return { + driver: 'snowflake', + authMethod: 'rsa', + account, + warehouse, + database, + username, + privateKey: resolvedPrivateKey, + ...(resolvedPassphrase ? { passphrase: resolvedPassphrase } : {}), + ...(role ? { role } : {}), + }; + } + if (authMethod === 'oauth' || authMethod === 'pat') { + const resolvedToken = tokenRef ?? stringConfigField(input.existingConnection, 'token'); + if (!account || !warehouse || !database || !resolvedToken) return null; + return { + driver: 'snowflake', + authMethod, + account, + warehouse, + database, + ...(username ? { username } : {}), + token: resolvedToken, + ...(role ? { role } : {}), + }; + } + if (!account || !warehouse || !database) return null; return { driver: 'snowflake', - authMethod: 'rsa', + authMethod: 'externalbrowser', account, warehouse, database, - username, - privateKey: resolvedPrivateKey, - ...(resolvedPassphrase ? { passphrase: resolvedPassphrase } : {}), + ...(username ? { username } : {}), ...(role ? { role } : {}), }; } diff --git a/packages/cli/test/connectors/snowflake/connector.test.ts b/packages/cli/test/connectors/snowflake/connector.test.ts index 3d71a6c90..0fc7d3a58 100644 --- a/packages/cli/test/connectors/snowflake/connector.test.ts +++ b/packages/cli/test/connectors/snowflake/connector.test.ts @@ -195,7 +195,7 @@ describe('KtxSnowflakeScanConnector', () => { } }); - it('resolves OAuth config with token_ref and does not require username', () => { + it('resolves OAuth config with token env reference and does not require username', () => { const resolved = snowflakeConnectionConfigFromConfig({ connectionId: 'warehouse', connection: { @@ -204,7 +204,7 @@ describe('KtxSnowflakeScanConnector', () => { account: 'acct', warehouse: 'WH', database: 'ANALYTICS', - token_ref: 'env:SNOWFLAKE_OAUTH_TOKEN', + token: 'env:SNOWFLAKE_OAUTH_TOKEN', }, env: { SNOWFLAKE_OAUTH_TOKEN: 'test-token-value' }, // pragma: allowlist secret }); @@ -213,7 +213,7 @@ describe('KtxSnowflakeScanConnector', () => { expect(resolved.username).toBeUndefined(); }); - it('resolves PAT config with token_ref', () => { + it('resolves PAT config with token env reference', () => { const resolved = snowflakeConnectionConfigFromConfig({ connectionId: 'warehouse', connection: { @@ -223,7 +223,7 @@ describe('KtxSnowflakeScanConnector', () => { warehouse: 'WH', database: 'ANALYTICS', username: 'svc_user', - token_ref: 'env:SNOWFLAKE_PAT', + token: 'env:SNOWFLAKE_PAT', }, env: { SNOWFLAKE_PAT: 'pat-token-value' }, // pragma: allowlist secret }); @@ -232,7 +232,7 @@ describe('KtxSnowflakeScanConnector', () => { expect(resolved.username).toBe('svc_user'); }); - it('throws when PAT config is missing token and token_ref', () => { + it('throws when PAT config is missing token', () => { expect(() => snowflakeConnectionConfigFromConfig({ connectionId: 'warehouse', @@ -245,7 +245,42 @@ describe('KtxSnowflakeScanConnector', () => { username: 'svc_user', }, }), - ).toThrow('connections.warehouse.token_ref'); + ).toThrow('connections.warehouse.token'); + }); + + it('re-reads the token source on each call so a rotated token is picked up', () => { + const env: Record = { SNOWFLAKE_OAUTH_TOKEN: 'first-token' }; // pragma: allowlist secret + const resolved = snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'oauth', + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + token: 'env:SNOWFLAKE_OAUTH_TOKEN', + }, + env, + }); + expect(resolved.token).toBe('first-token'); + env.SNOWFLAKE_OAUTH_TOKEN = 'rotated-token'; + expect(resolved.refreshToken?.()).toBe('rotated-token'); + }); + + it('rejects an unknown authMethod instead of falling through to password auth', () => { + expect(() => + snowflakeConnectionConfigFromConfig({ + connectionId: 'warehouse', + connection: { + driver: 'snowflake', + authMethod: 'oauth2' as never, + account: 'acct', + warehouse: 'WH', + database: 'ANALYTICS', + password: 'stale-password', // pragma: allowlist secret + }, + }), + ).toThrow('connections.warehouse.authMethod "oauth2" is not supported'); }); it('resolves OAuth config with inline token', () => { @@ -274,7 +309,7 @@ describe('KtxSnowflakeScanConnector', () => { warehouse: 'WH', database: 'ANALYTICS', username: 'svc_account', - token_ref: 'env:SNOWFLAKE_OAUTH_TOKEN', + token: 'env:SNOWFLAKE_OAUTH_TOKEN', }, env: { SNOWFLAKE_OAUTH_TOKEN: 'test-token-value' }, // pragma: allowlist secret }); @@ -282,7 +317,7 @@ describe('KtxSnowflakeScanConnector', () => { expect(resolved.token).toBe('test-token-value'); }); - it('throws when OAuth config is missing token and token_ref', () => { + it('throws when OAuth config is missing token', () => { expect(() => snowflakeConnectionConfigFromConfig({ connectionId: 'warehouse', @@ -294,7 +329,7 @@ describe('KtxSnowflakeScanConnector', () => { database: 'ANALYTICS', }, }), - ).toThrow('connections.warehouse.token_ref'); + ).toThrow('connections.warehouse.token'); }); it('resolves externalbrowser config without username or credentials', () => { diff --git a/packages/cli/test/setup-databases.test.ts b/packages/cli/test/setup-databases.test.ts index feff6613b..25691c12a 100644 --- a/packages/cli/test/setup-databases.test.ts +++ b/packages/cli/test/setup-databases.test.ts @@ -2705,6 +2705,137 @@ describe('setup databases step', () => { expect(config.connections.snowflake.password).toBeUndefined(); }); + it('configures Snowflake with OAuth token auth via setup wizard', async () => { + const io = makeIo(); + const result = await runKtxSetupDatabasesStep( + { + projectDir: tempDir, + inputMode: 'disabled', + databaseDrivers: ['snowflake'], + databaseConnectionId: 'snowflake', + databaseSchemas: ['PUBLIC'], + skipDatabases: false, + }, + io.io, + { + testConnection: vi.fn(async () => 0), + scanConnection: vi.fn(async () => 0), + prompts: makePromptAdapter({ + selectValues: ['oauth'], + textValues: ['env:SNOWFLAKE_ACCOUNT', 'WH', 'ANALYTICS', 'svc_account', ''], + passwordValues: ['env:SNOWFLAKE_OAUTH_TOKEN'], + }), + }, + ); + + expect(result.status).toBe('ready'); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.snowflake).toMatchObject({ + driver: 'snowflake', + authMethod: 'oauth', + account: 'env:SNOWFLAKE_ACCOUNT', + warehouse: 'WH', + database: 'ANALYTICS', + username: 'svc_account', + token: 'env:SNOWFLAKE_OAUTH_TOKEN', // pragma: allowlist secret + }); + expect(config.connections.snowflake.password).toBeUndefined(); + }); + + it('configures Snowflake with browser SSO auth via setup wizard, without prompting for credentials', async () => { + const io = makeIo(); + const result = await runKtxSetupDatabasesStep( + { + projectDir: tempDir, + inputMode: 'disabled', + databaseDrivers: ['snowflake'], + databaseConnectionId: 'snowflake', + databaseSchemas: ['PUBLIC'], + skipDatabases: false, + }, + io.io, + { + testConnection: vi.fn(async () => 0), + scanConnection: vi.fn(async () => 0), + prompts: makePromptAdapter({ + selectValues: ['externalbrowser'], + textValues: ['env:SNOWFLAKE_ACCOUNT', 'WH', 'ANALYTICS', '', ''], + }), + }, + ); + + expect(result.status).toBe('ready'); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.snowflake).toMatchObject({ + driver: 'snowflake', + authMethod: 'externalbrowser', + account: 'env:SNOWFLAKE_ACCOUNT', + warehouse: 'WH', + database: 'ANALYTICS', + }); + expect(config.connections.snowflake.username).toBeUndefined(); + expect(config.connections.snowflake.password).toBeUndefined(); + expect(config.connections.snowflake.token).toBeUndefined(); + }); + + it('round-trips an existing OAuth Snowflake connection when re-editing via setup wizard', async () => { + await writeFile( + join(tempDir, 'ktx.yaml'), + [ + 'connections:', + ' warehouse:', + ' driver: snowflake', + ' authMethod: oauth', + ' account: env:SNOWFLAKE_ACCOUNT', + ' warehouse: WH', + ' database: ANALYTICS', + ' username: svc_account', + ' token: env:SNOWFLAKE_OAUTH_TOKEN', + 'setup:', + ' database_connection_ids:', + ' - warehouse', + '', + ].join('\n'), + 'utf-8', + ); + await writeKtxSetupState(tempDir, { completed_steps: ['databases'] }); + const prompts = makePromptAdapter({ textValues: [], passwordValues: [] }); + let primaryMenuCount = 0; + vi.mocked(prompts.select).mockImplementation(async (options) => { + if (options.message === 'Databases configured: warehouse\nWhat would you like to do?') { + primaryMenuCount += 1; + return primaryMenuCount === 1 ? 'edit' : 'continue'; + } + if (options.message === 'Database to edit') return 'warehouse'; + if (options.message === 'Snowflake authentication method') return 'oauth'; + if (options.message.startsWith('Enable query-history ingest')) return 'no'; + 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 () => [{ catalog: 'ANALYTICS', schema: 'PUBLIC', name: 'ORDERS', kind: 'table' as const }]); + const pickers = makePickerStubs({ scopes: [{ schemas: ['PUBLIC'], tables: [] }] }); + + const result = await runKtxSetupDatabasesStep( + { projectDir: tempDir, inputMode: 'auto', skipDatabases: false, databaseSchemas: [] }, + makeIo().io, + { prompts, testConnection, scanConnection, listSchemas, listTables, pickDatabaseScope: pickers.pickDatabaseScope }, + ); + + expect(result).toEqual({ status: 'ready', projectDir: tempDir, connectionIds: ['warehouse'] }); + const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8')); + expect(config.connections.warehouse).toMatchObject({ + driver: 'snowflake', + authMethod: 'oauth', + account: 'env:SNOWFLAKE_ACCOUNT', + warehouse: 'WH', + database: 'ANALYTICS', + username: 'svc_account', + token: 'env:SNOWFLAKE_OAUTH_TOKEN', // pragma: allowlist secret + }); + }); + it('writes Postgres query history config with minExecutions and ignores window/redaction output', async () => { const io = makeIo(); const result = await runKtxSetupDatabasesStep( From 7b2b94d004f66e2a9bdabe1816b650bdcacbe941 Mon Sep 17 00:00:00 2001 From: mattsenicksigma Date: Thu, 2 Jul 2026 16:06:56 -0700 Subject: [PATCH 4/4] add allowlist to false secret concern --- docs-site/content/docs/integrations/primary-sources.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs-site/content/docs/integrations/primary-sources.mdx b/docs-site/content/docs/integrations/primary-sources.mdx index 87eae3a5c..9c66caa53 100644 --- a/docs-site/content/docs/integrations/primary-sources.mdx +++ b/docs-site/content/docs/integrations/primary-sources.mdx @@ -174,7 +174,7 @@ connections: role: ANALYST ``` -`token` follows the same reference syntax as other secrets: `env:VAR` reads from an environment variable, `file:/path/to/token` reads from a file, or a literal value can be supplied directly. +`token` follows the same reference syntax as other secrets: `env:VAR` reads from an environment variable, `file:/path/to/token` reads from a file, or a literal value can be supplied directly. Browser SSO authentication is interactive — running `ktx connection test` or `ktx setup` will open a browser window to complete the Snowflake SSO flow. Not suitable for headless or CI environments. `ktx setup`'s Snowflake wizard supports all five auth methods, both when creating a new connection and when editing an existing one.