Snowflake: key-pair auth from a Malloy config file

packages/malloy-db-snowflake/src/index.ts

@@ -104,6 +104,15 @@ registerConnectionType('snowflake', {
       type: 'password',
       optional: true,
     },
+    {
+      name: 'authenticator',
+      displayName: 'Authenticator',
+      type: 'string',
+      optional: true,
+      advanced: true,
+      description:
+        'snowflake-sdk authenticator, e.g. SNOWFLAKE_JWT, EXTERNALBROWSER, OAUTH. Defaults to SNOWFLAKE_JWT when a private key is set, even if a password is also set; otherwise password authentication.',
+    },
     {
       name: 'timeoutMs',
       displayName: 'Timeout (ms)',

packages/malloy-db-snowflake/src/snowflake_config.unit.spec.ts

@@ -0,0 +1,116 @@
+/*
+ * Copyright Contributors to the Malloy project
+ * SPDX-License-Identifier: MIT
+ */
+
+import {MalloyConfig} from '@malloydata/malloy';
+import snowflake from 'snowflake-sdk';
+import type {ConnectionOptions} from 'snowflake-sdk';
+import {
+  SnowflakeConnection,
+  withKeyPairAuthenticator,
+} from './snowflake_connection';
+// Registers the 'snowflake' connection type for the MalloyConfig tests.
+import './index';
+
+const PEM =
+  '-----BEGIN PRIVATE KEY-----\nMIIBVQ==\n-----END PRIVATE KEY-----\n';
+
+describe('withKeyPairAuthenticator', () => {
+  it('selects SNOWFLAKE_JWT for a private key with no authenticator', () => {
+    expect(
+      withKeyPairAuthenticator({account: 'a', privateKey: PEM}).authenticator
+    ).toBe('SNOWFLAKE_JWT');
+  });
+
+  it('selects SNOWFLAKE_JWT for a private key path with no authenticator', () => {
+    expect(
+      withKeyPairAuthenticator({account: 'a', privateKeyPath: '/k.p8'})
+        .authenticator
+    ).toBe('SNOWFLAKE_JWT');
+  });
+
+  it('leaves an explicit authenticator alone', () => {
+    expect(
+      withKeyPairAuthenticator({
+        account: 'a',
+        privateKey: PEM,
+        authenticator: 'OAUTH',
+      }).authenticator
+    ).toBe('OAUTH');
+  });
+
+  it('leaves password auth alone', () => {
+    expect(
+      withKeyPairAuthenticator({account: 'a', password: 'p'}).authenticator
+    ).toBeUndefined();
+  });
+
+  it('does not take an empty key or key path as a key', () => {
+    for (const key of [{privateKey: ''}, {privateKeyPath: ''}]) {
+      expect(
+        withKeyPairAuthenticator({account: 'a', password: 'p', ...key})
+          .authenticator
+      ).toBeUndefined();
+    }
+  });
+});
+
+describe('Snowflake key-pair auth from a Malloy config', () => {
+  const KEY_ENV = 'TEST_SNOWFLAKE_PRIVATE_KEY';
+  const PASS_ENV = 'TEST_SNOWFLAKE_KEY_PASSPHRASE';
+  let createPool: jest.SpiedFunction<typeof snowflake.createPool>;
+  let conn: SnowflakeConnection | undefined;
+
+  beforeEach(() => {
+    process.env[KEY_ENV] = PEM;
+    process.env[PASS_ENV] = 'hunter2';
+    createPool = jest.spyOn(snowflake, 'createPool');
+  });
+
+  afterEach(async () => {
+    await conn?.close();
+    conn = undefined;
+    createPool.mockRestore();
+    delete process.env[KEY_ENV];
+    delete process.env[PASS_ENV];
+  });
+
+  async function connect(
+    entry: Record<string, unknown>
+  ): Promise<ConnectionOptions> {
+    const config = new MalloyConfig({
+      connections: {
+        // poolMin 0 so the pool opens no connection to the fake account.
+        sf: {is: 'snowflake', account: 'acme', poolMin: 0, ...entry},
+      },
+    });
+    expect(config.log).toEqual([]);
+    const found = await config.connections.lookupConnection('sf');
+    expect(found).toBeInstanceOf(SnowflakeConnection);
+    if (found instanceof SnowflakeConnection) conn = found;
+    expect(createPool).toHaveBeenCalledTimes(1);
+    return createPool.mock.calls[0][0];
+  }
+
+  it('reaches the driver as key-pair auth with no authenticator written', async () => {
+    const opts = await connect({
+      username: 'malloy',
+      privateKey: {env: KEY_ENV},
+      privateKeyPass: {env: PASS_ENV},
+    });
+    expect(opts).toMatchObject({
+      authenticator: 'SNOWFLAKE_JWT',
+      privateKey: PEM,
+      privateKeyPass: 'hunter2',
+    });
+  });
+
+  it('passes a written authenticator through to the driver', async () => {
+    const opts = await connect({
+      username: 'malloy',
+      authenticator: 'EXTERNALBROWSER',
+    });
+    expect(opts.authenticator).toBe('EXTERNALBROWSER');
+  });
+});

packages/malloy-db-snowflake/src/snowflake_connection.ts

@@ -120,6 +120,23 @@ export async function sampleVariantBlocks(
   return undefined;
 }
 
+/**
+ * snowflake-sdk uses a private key only when `authenticator` is
+ * `SNOWFLAKE_JWT`; with no authenticator it falls back to password auth and
+ * fails asking for a password. A key supplied with no authenticator means
+ * key-pair auth, even when a password is also present; the password is then
+ * unused. An empty string is not a key: the driver ignores an empty
+ * `privateKeyPath`, so it must not switch a password login to key-pair. An
+ * explicit authenticator is left alone.
+ */
+export function withKeyPairAuthenticator(
+  opts: ConnectionOptions
+): ConnectionOptions {
+  if (opts.authenticator !== undefined) return opts;
+  if (!opts.privateKey && !opts.privateKeyPath) return opts;
+  return {...opts, authenticator: 'SNOWFLAKE_JWT'};
+}
+
 export interface SnowflakeConnectionOptions {
   // snowflake sdk connection options
   connOptions?: ConnectionOptions;
@@ -190,7 +207,8 @@ export class SnowflakeConnection
       // try to get connection options from ~/.snowflake/connections.toml
       connOptions = SnowflakeExecutor.getConnectionOptionsFromToml();
     }
-    this.connOptions = connOptions ?? {};
+    connOptions = withKeyPairAuthenticator(connOptions);
+    this.connOptions = connOptions;
     this.setupSQL = options?.setupSQL;
     this.executor = new SnowflakeExecutor(
       connOptions,

packages/malloy/src/connection/CONTEXT.md

@@ -174,8 +174,8 @@ A service account key can arrive three ways, in this order of precedence: `servi
 `host` (string), `port` (number), `username` (string), `password` (password), `databaseName` (string), `connectionString` (string, advanced), `setupSQL` (text, advanced), `ssl` (json, advanced)
 
 **Snowflake** (`displayName: "Snowflake"`):
-`account` (string, required), `username` (string), `password` (password), `role` (string), `warehouse` (string), `database` (string), `schema` (string), `privateKeyPath` (file), `privateKeyPass` (password), `timeoutMs` (number, advanced), `schemaSampleTimeoutMs` (number, advanced), `schemaSampleRowLimit` (number, advanced), `schemaSampleFullScanMaxBytes` (number, advanced), `setupSQL` (text, advanced), `poolMin` (number, advanced), `poolMax` (number, advanced), `poolTestOnBorrow` (boolean, advanced)
-Factory extracts `name`, `setupSQL`, `timeoutMs`, and the three pool fields; passes remaining properties as snowflake-sdk `ConnectionOptions`. The pool fields are assembled into a `generic-pool` options object via `buildPoolOptions()` and shallow-merged with `SnowflakeExecutor.defaultPoolOptions_` (`{min: 1, max: 1, testOnBorrow: true, testOnReturn: true}`); omitting all three preserves the defaults.
+`account` (string, required), `username` (string), `password` (password), `role` (string), `warehouse` (string), `database` (string), `schema` (string), `privateKeyPath` (file), `privateKey` (password), `privateKeyPass` (password), `authenticator` (string, advanced), `timeoutMs` (number, advanced), `schemaSampleTimeoutMs` (number, advanced), `schemaSampleRowLimit` (number, advanced), `schemaSampleFullScanMaxBytes` (number, advanced), `setupSQL` (text, advanced), `poolMin` (number, advanced), `poolMax` (number, advanced), `poolTestOnBorrow` (boolean, advanced)
+Factory extracts `name`, `setupSQL`, `timeoutMs`, and the three pool fields; passes remaining properties as snowflake-sdk `ConnectionOptions`. The pool fields are assembled into a `generic-pool` options object via `buildPoolOptions()` and shallow-merged with `SnowflakeExecutor.defaultPoolOptions_` (`{min: 1, max: 1, testOnBorrow: true, testOnReturn: true}`); omitting all three preserves the defaults. The connection sets `authenticator: "SNOWFLAKE_JWT"` when a private key is given without an authenticator, even if a password is also set, since snowflake-sdk otherwise ignores the key and demands a password (`withKeyPairAuthenticator()`).
 
 **Trino** (`displayName: "Trino"`):
 `server` (string), `port` (number), `catalog` (string), `schema` (string), `user` (string), `password` (password), `setupSQL` (text, advanced), `source` (string, advanced), `ssl` (json, advanced), `session` (json, advanced), `extraCredential` (json, advanced), `extraHeaders` (json, advanced)