feat: Add Snowflake SQL integration support (#121)

Author: jankuca

src/messageTypes.ts

@@ -182,6 +182,7 @@ export type LocalizedMessages = {
     // Integration type labels
     integrationsPostgresTypeLabel: string;
     integrationsBigQueryTypeLabel: string;
+    integrationsSnowflakeTypeLabel: string;
     // PostgreSQL form strings
     integrationsPostgresNameLabel: string;
     integrationsPostgresNamePlaceholder: string;
@@ -204,9 +205,36 @@ export type LocalizedMessages = {
     integrationsBigQueryCredentialsLabel: string;
     integrationsBigQueryCredentialsPlaceholder: string;
     integrationsBigQueryCredentialsRequired: string;
+    // Snowflake form strings
+    integrationsSnowflakeNameLabel: string;
+    integrationsSnowflakeNamePlaceholder: string;
+    integrationsSnowflakeAccountLabel: string;
+    integrationsSnowflakeAccountPlaceholder: string;
+    integrationsSnowflakeAuthMethodLabel: string;
+    integrationsSnowflakeAuthMethodSubLabel: string;
+    integrationsSnowflakeAuthMethodUsernamePassword: string;
+    integrationsSnowflakeAuthMethodKeyPair: string;
+    integrationsSnowflakeUnsupportedAuthMethod: string;
+    integrationsSnowflakeUsernameLabel: string;
+    integrationsSnowflakePasswordLabel: string;
+    integrationsSnowflakePasswordPlaceholder: string;
+    integrationsSnowflakeServiceAccountUsernameLabel: string;
+    integrationsSnowflakeServiceAccountUsernameHelp: string;
+    integrationsSnowflakePrivateKeyLabel: string;
+    integrationsSnowflakePrivateKeyHelp: string;
+    integrationsSnowflakePrivateKeyPlaceholder: string;
+    integrationsSnowflakePrivateKeyPassphraseLabel: string;
+    integrationsSnowflakePrivateKeyPassphraseHelp: string;
+    integrationsSnowflakeDatabaseLabel: string;
+    integrationsSnowflakeDatabasePlaceholder: string;
+    integrationsSnowflakeRoleLabel: string;
+    integrationsSnowflakeRolePlaceholder: string;
+    integrationsSnowflakeWarehouseLabel: string;
+    integrationsSnowflakeWarehousePlaceholder: string;
     // Common form strings
     integrationsRequiredField: string;
     integrationsOptionalField: string;
+    integrationsUnnamedIntegration: string;
 };
 // Map all messages to specific payloads
 export class IInteractiveWindowMapping {

src/notebooks/deepnote/integrations/integrationManager.ts

@@ -5,8 +5,15 @@ import { IExtensionContext } from '../../../platform/common/types';
 import { Commands } from '../../../platform/common/constants';
 import { logger } from '../../../platform/logging';
 import { IIntegrationDetector, IIntegrationManager, IIntegrationStorage, IIntegrationWebviewProvider } from './types';
-import { IntegrationStatus, IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes';
+import {
+    DEEPNOTE_TO_INTEGRATION_TYPE,
+    IntegrationStatus,
+    IntegrationType,
+    IntegrationWithStatus,
+    RawIntegrationType
+} from '../../../platform/notebooks/deepnote/integrationTypes';
 import { BlockWithIntegration, scanBlocksForIntegrations } from './integrationUtils';
+import { IDeepnoteNotebookManager } from '../../types';
 
 /**
  * Manages integration UI and commands for Deepnote notebooks
@@ -21,7 +28,8 @@ export class IntegrationManager implements IIntegrationManager {
         @inject(IExtensionContext) private readonly extensionContext: IExtensionContext,
         @inject(IIntegrationDetector) private readonly integrationDetector: IIntegrationDetector,
         @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage,
-        @inject(IIntegrationWebviewProvider) private readonly webviewProvider: IIntegrationWebviewProvider
+        @inject(IIntegrationWebviewProvider) private readonly webviewProvider: IIntegrationWebviewProvider,
+        @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager
     ) {}
 
     public activate(): void {
@@ -150,9 +158,33 @@ export class IntegrationManager implements IIntegrationManager {
         if (selectedIntegrationId && !integrations.has(selectedIntegrationId)) {
             logger.debug(`IntegrationManager: Adding requested integration ${selectedIntegrationId} to the map`);
             const config = await this.integrationStorage.getIntegrationConfig(selectedIntegrationId);
+
+            // Try to get integration metadata from the project
+            const project = this.notebookManager.getOriginalProject(projectId);
+            const projectIntegration = project?.project.integrations?.find((i) => i.id === selectedIntegrationId);
+
+            let integrationName: string | undefined;
+            let integrationType: IntegrationType | undefined;
+
+            if (projectIntegration) {
+                integrationName = projectIntegration.name;
+
+                // Validate that projectIntegration.type exists in the mapping before lookup
+                if (projectIntegration.type in DEEPNOTE_TO_INTEGRATION_TYPE) {
+                    // Map the Deepnote integration type to our IntegrationType
+                    integrationType = DEEPNOTE_TO_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType];
+                } else {
+                    logger.warn(
+                        `IntegrationManager: Unknown integration type '${projectIntegration.type}' for integration ID '${selectedIntegrationId}' in project '${projectId}'. Integration type will be undefined.`
+                    );
+                }
+            }
+
             integrations.set(selectedIntegrationId, {
                 config: config || null,
-                status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected
+                status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected,
+                integrationName,
+                integrationType
             });
         }
 

src/notebooks/deepnote/integrations/integrationWebview.ts

@@ -131,6 +131,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider {
             integrationsConfigureTitle: localize.Integrations.configureTitle,
             integrationsPostgresTypeLabel: localize.Integrations.postgresTypeLabel,
             integrationsBigQueryTypeLabel: localize.Integrations.bigQueryTypeLabel,
+            integrationsSnowflakeTypeLabel: localize.Integrations.snowflakeTypeLabel,
             integrationsCancel: localize.Integrations.cancel,
             integrationsSave: localize.Integrations.save,
             integrationsRequiredField: localize.Integrations.requiredField,
@@ -154,7 +155,34 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider {
             integrationsBigQueryProjectIdPlaceholder: localize.Integrations.bigQueryProjectIdPlaceholder,
             integrationsBigQueryCredentialsLabel: localize.Integrations.bigQueryCredentialsLabel,
             integrationsBigQueryCredentialsPlaceholder: localize.Integrations.bigQueryCredentialsPlaceholder,
-            integrationsBigQueryCredentialsRequired: localize.Integrations.bigQueryCredentialsRequired
+            integrationsBigQueryCredentialsRequired: localize.Integrations.bigQueryCredentialsRequired,
+            integrationsSnowflakeNameLabel: localize.Integrations.snowflakeNameLabel,
+            integrationsSnowflakeNamePlaceholder: localize.Integrations.snowflakeNamePlaceholder,
+            integrationsSnowflakeAccountLabel: localize.Integrations.snowflakeAccountLabel,
+            integrationsSnowflakeAccountPlaceholder: localize.Integrations.snowflakeAccountPlaceholder,
+            integrationsSnowflakeAuthMethodLabel: localize.Integrations.snowflakeAuthMethodLabel,
+            integrationsSnowflakeAuthMethodSubLabel: localize.Integrations.snowflakeAuthMethodSubLabel,
+            integrationsSnowflakeAuthMethodUsernamePassword: localize.Integrations.snowflakeAuthMethodUsernamePassword,
+            integrationsSnowflakeAuthMethodKeyPair: localize.Integrations.snowflakeAuthMethodKeyPair,
+            integrationsSnowflakeUnsupportedAuthMethod: localize.Integrations.snowflakeUnsupportedAuthMethod,
+            integrationsSnowflakeUsernameLabel: localize.Integrations.snowflakeUsernameLabel,
+            integrationsSnowflakePasswordLabel: localize.Integrations.snowflakePasswordLabel,
+            integrationsSnowflakePasswordPlaceholder: localize.Integrations.snowflakePasswordPlaceholder,
+            integrationsSnowflakeServiceAccountUsernameLabel:
+                localize.Integrations.snowflakeServiceAccountUsernameLabel,
+            integrationsSnowflakeServiceAccountUsernameHelp: localize.Integrations.snowflakeServiceAccountUsernameHelp,
+            integrationsSnowflakePrivateKeyLabel: localize.Integrations.snowflakePrivateKeyLabel,
+            integrationsSnowflakePrivateKeyHelp: localize.Integrations.snowflakePrivateKeyHelp,
+            integrationsSnowflakePrivateKeyPlaceholder: localize.Integrations.snowflakePrivateKeyPlaceholder,
+            integrationsSnowflakePrivateKeyPassphraseLabel: localize.Integrations.snowflakePrivateKeyPassphraseLabel,
+            integrationsSnowflakePrivateKeyPassphraseHelp: localize.Integrations.snowflakePrivateKeyPassphraseHelp,
+            integrationsSnowflakeDatabaseLabel: localize.Integrations.snowflakeDatabaseLabel,
+            integrationsSnowflakeDatabasePlaceholder: localize.Integrations.snowflakeDatabasePlaceholder,
+            integrationsSnowflakeRoleLabel: localize.Integrations.snowflakeRoleLabel,
+            integrationsSnowflakeRolePlaceholder: localize.Integrations.snowflakeRolePlaceholder,
+            integrationsSnowflakeWarehouseLabel: localize.Integrations.snowflakeWarehouseLabel,
+            integrationsSnowflakeWarehousePlaceholder: localize.Integrations.snowflakeWarehousePlaceholder,
+            integrationsUnnamedIntegration: localize.Integrations.unnamedIntegration('{0}')
         };
 
         await this.currentPanel.webview.postMessage({

src/notebooks/deepnote/sqlCellStatusBarProvider.ts

@@ -437,6 +437,8 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
                 return l10n.t('PostgreSQL');
             case IntegrationType.BigQuery:
                 return l10n.t('BigQuery');
+            case IntegrationType.Snowflake:
+                return l10n.t('Snowflake');
             default:
                 return String(type);
         }

src/platform/common/utils/localize.ts

@@ -830,10 +830,12 @@ export namespace Integrations {
     export const save = l10n.t('Save');
     export const requiredField = l10n.t('*');
     export const optionalField = l10n.t('(optional)');
+    export const unnamedIntegration = (id: string) => l10n.t('Unnamed Integration ({0})', id);
 
     // Integration type labels
     export const postgresTypeLabel = l10n.t('PostgreSQL');
     export const bigQueryTypeLabel = l10n.t('BigQuery');
+    export const snowflakeTypeLabel = l10n.t('Snowflake');
 
     // PostgreSQL form strings
     export const postgresNameLabel = l10n.t('Name (optional)');
@@ -849,7 +851,6 @@ export namespace Integrations {
     export const postgresPasswordLabel = l10n.t('Password');
     export const postgresPasswordPlaceholder = l10n.t('••••••••');
     export const postgresSslLabel = l10n.t('Use SSL');
-    export const postgresUnnamedIntegration = (id: string) => l10n.t('Unnamed PostgreSQL Integration ({0})', id);
 
     // BigQuery form strings
     export const bigQueryNameLabel = l10n.t('Name (optional)');
@@ -860,7 +861,43 @@ export namespace Integrations {
     export const bigQueryCredentialsPlaceholder = l10n.t('{"type": "service_account", ...}');
     export const bigQueryCredentialsRequired = l10n.t('Credentials are required');
     export const bigQueryInvalidJson = (message: string) => l10n.t('Invalid JSON: {0}', message);
-    export const bigQueryUnnamedIntegration = (id: string) => l10n.t('Unnamed BigQuery Integration ({0})', id);
+
+    // Snowflake form strings
+    export const snowflakeNameLabel = l10n.t('Name (optional)');
+    export const snowflakeNamePlaceholder = l10n.t('My Snowflake Database');
+    export const snowflakeAccountLabel = l10n.t('Account name');
+    export const snowflakeAccountPlaceholder = l10n.t('ptb34938.us-east-1');
+    export const snowflakeAuthMethodLabel = l10n.t('Authentication');
+    export const snowflakeAuthMethodSubLabel = l10n.t('Method');
+    export const snowflakeAuthMethodUsernamePassword = l10n.t('Username & password');
+    export const snowflakeAuthMethodKeyPair = l10n.t('Key-pair (service account)');
+    export const snowflakeUnsupportedAuthMethod = l10n.t(
+        'This Snowflake integration uses an authentication method that is not supported in VS Code. You can view the integration details but cannot edit or use it.'
+    );
+    export const snowflakeUsernameLabel = l10n.t('Username');
+    export const snowflakeUsernamePlaceholder = l10n.t('user');
+    export const snowflakePasswordLabel = l10n.t('Password');
+    export const snowflakePasswordPlaceholder = l10n.t('••••••••');
+    export const snowflakeServiceAccountUsernameLabel = l10n.t('Service Account Username');
+    export const snowflakeServiceAccountUsernameHelp = l10n.t(
+        'The username of the service account that will be used to connect to Snowflake'
+    );
+    export const snowflakePrivateKeyLabel = l10n.t('Private Key');
+    export const snowflakePrivateKeyHelp = l10n.t(
+        'The private key in PEM format. Make sure to include the entire key, including BEGIN and END markers.'
+    );
+    export const snowflakePrivateKeyPlaceholder = l10n.t("Begins with '-----BEGIN PRIVATE KEY-----'");
+    export const snowflakePrivateKeyPassphraseLabel = l10n.t('Private Key Passphrase (optional)');
+    export const snowflakePrivateKeyPassphraseHelp = l10n.t(
+        'If the private key is encrypted, provide the passphrase to decrypt it.'
+    );
+    export const snowflakePrivateKeyPassphrasePlaceholder = l10n.t('');
+    export const snowflakeDatabaseLabel = l10n.t('Database (optional)');
+    export const snowflakeDatabasePlaceholder = l10n.t('');
+    export const snowflakeRoleLabel = l10n.t('Role (optional)');
+    export const snowflakeRolePlaceholder = l10n.t('');
+    export const snowflakeWarehouseLabel = l10n.t('Warehouse (optional)');
+    export const snowflakeWarehousePlaceholder = l10n.t('');
 }
 
 export namespace Deprecated {

src/platform/errors/types.ts

@@ -105,7 +105,8 @@ export type ErrorCategory =
     | 'unknownProduct'
     | 'invalidInterpreter'
     | 'pythonAPINotInitialized'
-    | 'deepnoteserver';
+    | 'deepnoteserver'
+    | 'unsupported_integration';
 
 // If there are errors, then the are added to the telementry properties.
 export type TelemetryErrorProperties = {

src/platform/errors/unsupportedIntegrationError.ts

@@ -0,0 +1,18 @@
+import { BaseError } from './types';
+
+/**
+ * Error thrown when an unsupported integration type is encountered.
+ *
+ * Cause:
+ * An integration configuration has a type that is not supported by the SQL integration system,
+ * or an integration uses an authentication method that is not supported in VSCode.
+ *
+ * Handled by:
+ * Callers should handle this error and inform the user that the integration type or
+ * authentication method is not supported.
+ */
+export class UnsupportedIntegrationError extends BaseError {
+    constructor(message: string) {
+        super('unsupported_integration', message);
+    }
+}

src/platform/notebooks/deepnote/integrationTypes.ts

@@ -9,15 +9,17 @@ export const DATAFRAME_SQL_INTEGRATION_ID = 'deepnote-dataframe-sql';
  */
 export enum IntegrationType {
     Postgres = 'postgres',
-    BigQuery = 'bigquery'
+    BigQuery = 'bigquery',
+    Snowflake = 'snowflake'
 }
 
 /**
  * Map our IntegrationType enum to Deepnote integration type strings
  */
 export const INTEGRATION_TYPE_TO_DEEPNOTE = {
     [IntegrationType.Postgres]: 'pgsql',
-    [IntegrationType.BigQuery]: 'big-query'
+    [IntegrationType.BigQuery]: 'big-query',
+    [IntegrationType.Snowflake]: 'snowflake'
 } as const satisfies { [type in IntegrationType]: string };
 
 export type RawIntegrationType = (typeof INTEGRATION_TYPE_TO_DEEPNOTE)[keyof typeof INTEGRATION_TYPE_TO_DEEPNOTE];
@@ -27,7 +29,8 @@ export type RawIntegrationType = (typeof INTEGRATION_TYPE_TO_DEEPNOTE)[keyof typ
  */
 export const DEEPNOTE_TO_INTEGRATION_TYPE: Record<RawIntegrationType, IntegrationType> = {
     pgsql: IntegrationType.Postgres,
-    'big-query': IntegrationType.BigQuery
+    'big-query': IntegrationType.BigQuery,
+    snowflake: IntegrationType.Snowflake
 };
 
 /**
@@ -61,10 +64,62 @@ export interface BigQueryIntegrationConfig extends BaseIntegrationConfig {
     credentials: string; // JSON string of service account credentials
 }
 
+// Import and re-export Snowflake auth constants from shared module
+import {
+    type SnowflakeAuthMethod,
+    SnowflakeAuthMethods,
+    SUPPORTED_SNOWFLAKE_AUTH_METHODS,
+    isSupportedSnowflakeAuthMethod
+} from './snowflakeAuthConstants';
+export {
+    type SnowflakeAuthMethod,
+    SnowflakeAuthMethods,
+    SUPPORTED_SNOWFLAKE_AUTH_METHODS,
+    isSupportedSnowflakeAuthMethod
+};
+
+/**
+ * Base Snowflake configuration with common fields
+ */
+interface BaseSnowflakeConfig extends BaseIntegrationConfig {
+    type: IntegrationType.Snowflake;
+    account: string;
+    warehouse?: string;
+    database?: string;
+    role?: string;
+}
+
+/**
+ * Snowflake integration configuration (discriminated union)
+ */
+export type SnowflakeIntegrationConfig = BaseSnowflakeConfig &
+    (
+        | {
+              authMethod: typeof SnowflakeAuthMethods.PASSWORD | null;
+              username: string;
+              password: string;
+          }
+        | {
+              authMethod: typeof SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR;
+              username: string;
+              privateKey: string;
+              privateKeyPassphrase?: string;
+          }
+        | {
+              // Unsupported auth methods - we store them but don't allow editing
+              authMethod:
+                  | typeof SnowflakeAuthMethods.OKTA
+                  | typeof SnowflakeAuthMethods.NATIVE_SNOWFLAKE
+                  | typeof SnowflakeAuthMethods.AZURE_AD
+                  | typeof SnowflakeAuthMethods.KEY_PAIR;
+              [key: string]: unknown; // Allow any additional fields for unsupported methods
+          }
+    );
+
 /**
  * Union type of all integration configurations
  */
-export type IntegrationConfig = PostgresIntegrationConfig | BigQueryIntegrationConfig;
+export type IntegrationConfig = PostgresIntegrationConfig | BigQueryIntegrationConfig | SnowflakeIntegrationConfig;
 
 /**
  * Integration connection status

src/platform/notebooks/deepnote/snowflakeAuthConstants.ts

@@ -0,0 +1,33 @@
+/**
+ * Snowflake authentication methods
+ */
+export const SnowflakeAuthMethods = {
+    PASSWORD: 'PASSWORD',
+    OKTA: 'OKTA',
+    NATIVE_SNOWFLAKE: 'NATIVE_SNOWFLAKE',
+    AZURE_AD: 'AZURE_AD',
+    KEY_PAIR: 'KEY_PAIR',
+    SERVICE_ACCOUNT_KEY_PAIR: 'SERVICE_ACCOUNT_KEY_PAIR'
+} as const;
+
+export type SnowflakeAuthMethod = (typeof SnowflakeAuthMethods)[keyof typeof SnowflakeAuthMethods];
+
+/**
+ * Supported auth methods that we can configure in VSCode
+ */
+export const SUPPORTED_SNOWFLAKE_AUTH_METHODS = [
+    null, // Legacy username+password (no authMethod field)
+    SnowflakeAuthMethods.PASSWORD,
+    SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR
+] as const;
+
+export type SupportedSnowflakeAuthMethod = (typeof SUPPORTED_SNOWFLAKE_AUTH_METHODS)[number];
+
+/**
+ * Type guard to check if a value is a supported Snowflake auth method
+ * @param value The value to check
+ * @returns true if the value is one of the supported auth methods
+ */
+export function isSupportedSnowflakeAuthMethod(value: unknown): value is SupportedSnowflakeAuthMethod {
+    return (SUPPORTED_SNOWFLAKE_AUTH_METHODS as readonly unknown[]).includes(value);
+}