Skip to content

Commit d664998

Browse files
jankucajankuca
committed
Refactor Snowflake integration to use discriminated union types
- Change auth method from enum to const object with PASSWORD, OKTA, NATIVE_SNOWFLAKE, AZURE_AD, KEY_PAIR, SERVICE_ACCOUNT_KEY_PAIR - Implement discriminated union for SnowflakeIntegrationConfig based on authMethod - Support PASSWORD and SERVICE_ACCOUNT_KEY_PAIR auth methods (non-user-specific) - Mark OKTA, NATIVE_SNOWFLAKE, AZURE_AD, KEY_PAIR as unsupported - Add type-safe connection string generation with proper type narrowing - Show error message and disable form for unsupported auth methods - Add localization for unsupported auth method warning - Update form submission to create correct discriminated union types
1 parent 795d710 commit d664998

File tree

7 files changed

+362
-178
lines changed

7 files changed

+362
-178
lines changed

‎src/messageTypes.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,7 @@ export type LocalizedMessages = {
214214
integrationsSnowflakeAuthMethodSubLabel: string;
215215
integrationsSnowflakeAuthMethodUsernamePassword: string;
216216
integrationsSnowflakeAuthMethodKeyPair: string;
217+
integrationsSnowflakeUnsupportedAuthMethod: string;
217218
integrationsSnowflakeUsernameLabel: string;
218219
integrationsSnowflakeUsernamePlaceholder: string;
219220
integrationsSnowflakePasswordLabel: string;

‎src/notebooks/deepnote/integrations/integrationWebview.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider {
164164
integrationsSnowflakeAuthMethodSubLabel: localize.Integrations.snowflakeAuthMethodSubLabel,
165165
integrationsSnowflakeAuthMethodUsernamePassword: localize.Integrations.snowflakeAuthMethodUsernamePassword,
166166
integrationsSnowflakeAuthMethodKeyPair: localize.Integrations.snowflakeAuthMethodKeyPair,
167+
integrationsSnowflakeUnsupportedAuthMethod: localize.Integrations.snowflakeUnsupportedAuthMethod,
167168
integrationsSnowflakeUsernameLabel: localize.Integrations.snowflakeUsernameLabel,
168169
integrationsSnowflakeUsernamePlaceholder: localize.Integrations.snowflakeUsernamePlaceholder,
169170
integrationsSnowflakePasswordLabel: localize.Integrations.snowflakePasswordLabel,

‎src/platform/common/utils/localize.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -872,6 +872,9 @@ export namespace Integrations {
872872
export const snowflakeAuthMethodSubLabel = l10n.t('Method');
873873
export const snowflakeAuthMethodUsernamePassword = l10n.t('Username & password');
874874
export const snowflakeAuthMethodKeyPair = l10n.t('Key-pair (service account)');
875+
export const snowflakeUnsupportedAuthMethod = l10n.t(
876+
'This Snowflake integration uses an authentication method that is not supported in VSCode. You can view the integration details but cannot edit or use it.'
877+
);
875878
export const snowflakeUsernameLabel = l10n.t('Username');
876879
export const snowflakeUsernamePlaceholder = l10n.t('WEBSITE_ANALYTICS_USER');
877880
export const snowflakePasswordLabel = l10n.t('Password');

‎src/platform/notebooks/deepnote/integrationTypes.ts‎

Lines changed: 55 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -65,32 +65,71 @@ export interface BigQueryIntegrationConfig extends BaseIntegrationConfig {
6565
}
6666

6767
/**
68-
* Snowflake authentication method
68+
* Snowflake authentication methods
6969
*/
70-
export enum SnowflakeAuthMethod {
71-
UsernamePassword = 'username_password',
72-
KeyPair = 'key_pair'
73-
}
70+
export const SnowflakeAuthMethods = {
71+
PASSWORD: 'PASSWORD',
72+
OKTA: 'OKTA',
73+
NATIVE_SNOWFLAKE: 'NATIVE_SNOWFLAKE',
74+
AZURE_AD: 'AZURE_AD',
75+
KEY_PAIR: 'KEY_PAIR',
76+
SERVICE_ACCOUNT_KEY_PAIR: 'SERVICE_ACCOUNT_KEY_PAIR'
77+
} as const;
78+
79+
export type SnowflakeAuthMethod = (typeof SnowflakeAuthMethods)[keyof typeof SnowflakeAuthMethods] | null;
80+
81+
/**
82+
* Supported auth methods that we can configure in VSCode
83+
*/
84+
export const SUPPORTED_SNOWFLAKE_AUTH_METHODS = [
85+
null, // Legacy username+password (no authMethod field)
86+
SnowflakeAuthMethods.PASSWORD,
87+
SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR
88+
] as const;
7489

7590
/**
76-
* Snowflake integration configuration
91+
* Base Snowflake configuration with common fields
7792
*/
78-
export interface SnowflakeIntegrationConfig extends BaseIntegrationConfig {
93+
interface BaseSnowflakeConfig extends BaseIntegrationConfig {
7994
type: IntegrationType.Snowflake;
8095
account: string;
81-
authMethod: SnowflakeAuthMethod;
82-
username: string;
83-
// For username+password auth
84-
password?: string;
85-
// For key-pair auth
86-
privateKey?: string;
87-
privateKeyPassphrase?: string;
88-
// Optional fields
89-
database?: string;
9096
warehouse?: string;
97+
database?: string;
9198
role?: string;
9299
}
93100

101+
/**
102+
* Snowflake integration configuration (discriminated union)
103+
*/
104+
export type SnowflakeIntegrationConfig = BaseSnowflakeConfig &
105+
(
106+
| {
107+
authMethod: null;
108+
username: string;
109+
password: string;
110+
}
111+
| {
112+
authMethod: typeof SnowflakeAuthMethods.PASSWORD;
113+
username: string;
114+
password: string;
115+
}
116+
| {
117+
authMethod: typeof SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR;
118+
username: string;
119+
privateKey: string;
120+
privateKeyPassphrase?: string;
121+
}
122+
| {
123+
// Unsupported auth methods - we store them but don't allow editing
124+
authMethod:
125+
| typeof SnowflakeAuthMethods.OKTA
126+
| typeof SnowflakeAuthMethods.NATIVE_SNOWFLAKE
127+
| typeof SnowflakeAuthMethods.AZURE_AD
128+
| typeof SnowflakeAuthMethods.KEY_PAIR;
129+
[key: string]: unknown; // Allow any additional fields for unsupported methods
130+
}
131+
);
132+
94133
/**
95134
* Union type of all integration configurations
96135
*/

‎src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import {
1010
DATAFRAME_SQL_INTEGRATION_ID,
1111
IntegrationConfig,
1212
IntegrationType,
13-
SnowflakeAuthMethod
13+
SnowflakeAuthMethods
1414
} from './integrationTypes';
1515

1616
/**
@@ -82,16 +82,17 @@ function convertIntegrationConfigToJson(config: IntegrationConfig): string {
8282
// Build Snowflake connection URL
8383
// Format depends on auth method:
8484
// Username+password: snowflake://{username}:{password}@{account}/{database}?warehouse={warehouse}&role={role}&application=YourApp
85-
// Key-pair: snowflake://{username}@{account}/{database}?warehouse={warehouse}&role={role}&authenticator=snowflake_jwt&application=YourApp
86-
const encodedUsername = encodeURIComponent(config.username);
85+
// Service account key-pair: snowflake://{username}@{account}/{database}?warehouse={warehouse}&role={role}&authenticator=snowflake_jwt&application=YourApp
8786
const encodedAccount = encodeURIComponent(config.account);
8887

8988
let url: string;
9089
const params: Record<string, unknown> = {};
9190

92-
if (config.authMethod === SnowflakeAuthMethod.UsernamePassword) {
91+
// Check if this is a supported auth method
92+
if (config.authMethod === null || config.authMethod === SnowflakeAuthMethods.PASSWORD) {
9393
// Username+password authentication
94-
const encodedPassword = encodeURIComponent(config.password || '');
94+
const encodedUsername = encodeURIComponent(config.username);
95+
const encodedPassword = encodeURIComponent(config.password);
9596
const database = config.database ? `/${encodeURIComponent(config.database)}` : '';
9697
url = `snowflake://${encodedUsername}:${encodedPassword}@${encodedAccount}${database}`;
9798

@@ -107,8 +108,9 @@ function convertIntegrationConfigToJson(config: IntegrationConfig): string {
107108
if (queryParams.length > 0) {
108109
url += `?${queryParams.join('&')}`;
109110
}
110-
} else {
111-
// Key-pair authentication
111+
} else if (config.authMethod === SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR) {
112+
// Service account key-pair authentication
113+
const encodedUsername = encodeURIComponent(config.username);
112114
const database = config.database ? `/${encodeURIComponent(config.database)}` : '';
113115
url = `snowflake://${encodedUsername}@${encodedAccount}${database}`;
114116

@@ -127,10 +129,15 @@ function convertIntegrationConfigToJson(config: IntegrationConfig): string {
127129
}
128130

129131
// For key-pair auth, pass the private key and passphrase as params
130-
params.private_key = config.privateKey || '';
132+
params.private_key = config.privateKey;
131133
if (config.privateKeyPassphrase) {
132134
params.private_key_passphrase = config.privateKeyPassphrase;
133135
}
136+
} else {
137+
// Unsupported auth method
138+
throw new UnsupportedIntegrationError(
139+
`Snowflake integration with auth method '${config.authMethod}' is not supported in VSCode`
140+
);
134141
}
135142

136143
return JSON.stringify({

0 commit comments

Comments
 (0)