feat: snowflake key auth server (#34548)

## Description
This PR adds new authentication method for snowflake: Use of public and
private keys to authorise a snowflake database.

Snowflake provides a couple of authentication mechanisms in order to
authenticate the DB. Currently appsmith only provides a way to authorise
using username and password. This PR adds support for private key
authentication, where by user can set a public key on their DB and they
can use corresponding private key in appsmith to authorise the
datasource.

In snowflake DB form, we have added a new field for authentication type
which has two options: Basic and Key pair. Basic will allows us to
authenticate using username and password. Key pair will provide us a way
to upload the private key along with input field for adding passphrase
(In case of encrypted private key).

More info: https://docs.snowflake.com/en/user-guide/key-pair-auth



https://github.com/appsmithorg/appsmith/assets/30018882/99774925-12a3-4cc0-af0a-614c3574cdc3


Fixes #`Issue Number`  
_or_  
Fixes `Issue URL`
> [!WARNING]  
> _If no issue exists, please create an issue first, and check with the
maintainers if the issue is valid._

## Automation

/ok-to-test tags="@tag.Datasource"

### 🔍 Cypress test results
<!-- This is an auto-generated comment: Cypress test results  -->
> [!TIP]
> 🟢 🟢 🟢 All cypress tests have passed! 🎉 🎉 🎉
> Workflow run:
<https://github.com/appsmithorg/appsmith/actions/runs/9743830603>
> Commit: cb64fff
> <a
href="https://internal.appsmith.com/app/cypress-dashboard/rundetails-65890b3c81d7400d08fa9ee5?branch=master&workflowId=9743830603&attempt=1"
target="_blank">Cypress dashboard</a>.
> Tags: `@tag.Datasource`

<!-- end of auto-generated comment: Cypress test results  -->




## Communication
Should the DevRel and Marketing teams inform users about this change?
- [ ] Yes
- [ ] No


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced Snowflake integration with improved connection creation and
authentication handling (key pair and basic auth types).

- **Bug Fixes**
- Fixed issues related to handling authentication objects within
datasource configuration to ensure more reliable connections.

- **Documentation**
- Added new error messages for passphrase requirements and incorrect
passphrase or private key for encrypted private keys.

- **Tests**
- Updated tests to verify the correct setting of authentication types in
the Snowflake plugin.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Rishabh-Rathod <rishabh.rathod@appsmith.com>
Co-authored-by: Aman Agarwal <aman@appsmith.com>
Co-authored-by: “sneha122” <“sneha@appsmith.com”>

app/client/src/api/DatasourcesApi.ts

@@ -143,6 +143,9 @@ class DatasourcesApi extends API {
       toastMessage: undefined,
       datasourceConfiguration: datasourceConfig.datasourceConfiguration && {
         ...datasourceConfig.datasourceConfiguration,
+        authentication: DatasourcesApi.cleanAuthenticationObject(
+          datasourceConfig.datasourceConfiguration.authentication,
+        ),
         connection: datasourceConfig.datasourceConfiguration.connection && {
           ...datasourceConfig.datasourceConfiguration.connection,
           ssl: {

app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java

@@ -1,12 +1,12 @@
 package com.external.plugins;
 
-import com.appsmith.external.constants.Authentication;
 import com.appsmith.external.exceptions.pluginExceptions.AppsmithPluginError;
 import com.appsmith.external.exceptions.pluginExceptions.AppsmithPluginException;
 import com.appsmith.external.exceptions.pluginExceptions.StaleConnectionException;
 import com.appsmith.external.models.ActionConfiguration;
 import com.appsmith.external.models.ActionExecutionRequest;
 import com.appsmith.external.models.ActionExecutionResult;
+import com.appsmith.external.models.AuthenticationDTO;
 import com.appsmith.external.models.DBAuth;
 import com.appsmith.external.models.DatasourceConfiguration;
 import com.appsmith.external.models.DatasourceStructure;
@@ -15,32 +15,28 @@
 import com.appsmith.external.plugins.BasePlugin;
 import com.appsmith.external.plugins.PluginExecutor;
 import com.external.plugins.exceptions.SnowflakeErrorMessages;
+import com.external.utils.SnowflakeKeyUtils;
 import com.external.utils.SqlUtils;
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.HikariPoolMXBean;
 import com.zaxxer.hikari.pool.HikariPool;
 import lombok.extern.slf4j.Slf4j;
+import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
+import org.bouncycastle.pkcs.PKCSException;
 import org.pf4j.Extension;
 import org.pf4j.PluginWrapper;
 import org.springframework.util.StringUtils;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 
-import java.sql.Connection;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.Statement;
-import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedHashMap;
-import java.util.Map;
-import java.util.Properties;
-import java.util.Set;
+import java.security.PrivateKey;
+import java.sql.*;
+import java.util.*;
 
+import static com.appsmith.external.constants.Authentication.DB_AUTH;
+import static com.appsmith.external.constants.Authentication.SNOWFLAKE_KEY_PAIR_AUTH;
 import static com.appsmith.external.constants.PluginConstants.PluginName.SNOWFLAKE_PLUGIN_NAME;
 import static com.external.utils.ExecutionUtils.getRowsFromQueryResult;
 import static com.external.utils.SnowflakeDatasourceUtils.getConnectionFromHikariConnectionPool;
@@ -159,58 +155,30 @@ public Mono<ActionExecutionResult> execute(
         @Override
         public Mono<HikariDataSource> createConnectionClient(
                 DatasourceConfiguration datasourceConfiguration, Properties properties) {
-            return Mono.fromCallable(() -> {
-                        HikariConfig config = new HikariConfig();
-
-                        config.setDriverClassName(properties.getProperty("driver_name"));
-
-                        config.setMinimumIdle(
-                                Integer.parseInt(properties.get("minimumIdle").toString()));
-                        config.setMaximumPoolSize(Integer.parseInt(
-                                properties.get("maximunPoolSize").toString()));
-
-                        config.setInitializationFailTimeout(Long.parseLong(
-                                properties.get("initializationFailTimeout").toString()));
-                        config.setConnectionTimeout(Long.parseLong(
-                                properties.get("connectionTimeoutMillis").toString()));
-
-                        if (Authentication.SNOWFLAKE_KEY_PAIR_AUTH.equals(
-                                datasourceConfiguration.getAuthentication().getAuthenticationType())) {
-                            KeyPairAuth authentication = (KeyPairAuth) datasourceConfiguration.getAuthentication();
-                            // Set authentication properties
-                            if (authentication.getUsername() != null) {
-                                config.setUsername(authentication.getUsername());
-                            }
-                        } else {
-                            // Set authentication properties
-                            DBAuth authentication = (DBAuth) datasourceConfiguration.getAuthentication();
-                            if (authentication.getUsername() != null) {
-                                config.setUsername(authentication.getUsername());
-                            }
-                            if (authentication.getPassword() != null) {
-                                config.setPassword(authentication.getPassword());
-                            }
-                        }
-
-                        // Set up the connection URL
-                        StringBuilder urlBuilder = new StringBuilder(
-                                "jdbc:snowflake://" + datasourceConfiguration.getUrl() + ".snowflakecomputing.com?");
-                        config.setJdbcUrl(urlBuilder.toString());
+            return getHikariConfig(datasourceConfiguration, properties)
+                    .flatMap(config -> Mono.fromCallable(() -> {
+                                // Set up the connection URL
+                                String jdbcUrl = getJDBCUrl(datasourceConfiguration);
+                                config.setJdbcUrl(jdbcUrl);
 
-                        config.setDataSourceProperties(properties);
+                                config.setDataSourceProperties(properties);
 
-                        // Now create the connection pool from the configuration
-                        HikariDataSource datasource = null;
-                        try {
-                            datasource = new HikariDataSource(config);
-                        } catch (HikariPool.PoolInitializationException e) {
-                            throw new AppsmithPluginException(
-                                    AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, e.getMessage());
-                        }
+                                // Now create the connection pool from the configuration
+                                HikariDataSource datasource = null;
+                                try {
+                                    datasource = new HikariDataSource(config);
+                                } catch (HikariPool.PoolInitializationException e) {
+                                    throw new AppsmithPluginException(
+                                            AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, e.getMessage());
+                                }
 
-                        return datasource;
-                    })
-                    .subscribeOn(scheduler);
+                                return datasource;
+                            })
+                            .subscribeOn(scheduler))
+                    .onErrorMap(
+                            AppsmithPluginException.class,
+                            error -> new AppsmithPluginException(
+                                    AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, error.getMessage()));
         }
 
         @Override
@@ -234,17 +202,16 @@ public Properties addPluginSpecificProperties(
         @Override
         public Properties addAuthParamsToConnectionConfig(
                 DatasourceConfiguration datasourceConfiguration, Properties properties) {
-
-            if (Authentication.SNOWFLAKE_KEY_PAIR_AUTH.equals(
-                    datasourceConfiguration.getAuthentication().getAuthenticationType())) {
-                KeyPairAuth authentication = (KeyPairAuth) datasourceConfiguration.getAuthentication();
-                properties.setProperty("user", authentication.getUsername());
-            } else {
-                DBAuth authentication = (DBAuth) datasourceConfiguration.getAuthentication();
-                properties.setProperty("user", authentication.getUsername());
-                properties.setProperty("password", authentication.getPassword());
+            // Only for username password auth, we need to set these properties, for others
+            // like key-pair auth, authentication specific properties need to be set on config itself
+            AuthenticationDTO authentication = datasourceConfiguration.getAuthentication();
+            switch (authentication.getAuthenticationType()) {
+                case DB_AUTH:
+                    DBAuth dbAuth = (DBAuth) authentication;
+                    properties.setProperty("user", dbAuth.getUsername());
+                    properties.setProperty("password", dbAuth.getPassword());
+                    break;
             }
-
             properties.setProperty(
                     "warehouse",
                     String.valueOf(
@@ -317,7 +284,7 @@ public Set<String> validateDatasource(DatasourceConfiguration datasourceConfigur
             if (datasourceConfiguration.getAuthentication() == null) {
                 invalids.add(SnowflakeErrorMessages.DS_MISSING_AUTHENTICATION_DETAILS_ERROR_MSG);
             } else {
-                if (Authentication.SNOWFLAKE_KEY_PAIR_AUTH.equals(
+                if (SNOWFLAKE_KEY_PAIR_AUTH.equals(
                         datasourceConfiguration.getAuthentication().getAuthenticationType())) {
                     KeyPairAuth authentication = (KeyPairAuth) datasourceConfiguration.getAuthentication();
                     if (StringUtils.isEmpty(authentication.getUsername())) {
@@ -491,5 +458,112 @@ public Mono<DatasourceStructure> getStructure(
                     })
                     .subscribeOn(scheduler);
         }
+
+        private Mono<HikariConfig> getHikariConfig(
+                DatasourceConfiguration datasourceConfiguration, Properties properties) {
+            HikariConfig commonConfig = getCommonHikariConfig(properties);
+            Mono<HikariConfig> configMono = Mono.empty();
+
+            String authenticationType =
+                    datasourceConfiguration.getAuthentication().getAuthenticationType();
+            if (authenticationType != null) {
+                switch (authenticationType) {
+                    case DB_AUTH:
+                        configMono = getBasicAuthConfig(commonConfig, datasourceConfiguration);
+                        break;
+                    case SNOWFLAKE_KEY_PAIR_AUTH:
+                        configMono = getKeyPairAuthConfig(commonConfig, datasourceConfiguration);
+                        break;
+                    default:
+                        break;
+                }
+            }
+            return configMono;
+        }
+
+        private Mono<HikariConfig> getKeyPairAuthConfig(
+                HikariConfig config, DatasourceConfiguration datasourceConfiguration) {
+            KeyPairAuth keyPairAuthConfig = (KeyPairAuth) datasourceConfiguration.getAuthentication();
+            byte[] keyBytes = keyPairAuthConfig.getPrivateKey().getDecodedContent();
+            String passphrase = keyPairAuthConfig.getPassphrase();
+            return getPrivateKeyFromBase64(keyBytes, passphrase)
+                    .flatMap(privateKey -> {
+                        String jdbcUrl = getJDBCUrl(datasourceConfiguration);
+
+                        // Prepare datasource object to be passed to hikariConfig
+                        SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
+                        ds.setPrivateKey(privateKey);
+                        ds.setUser(keyPairAuthConfig.getUsername());
+                        ds.setUrl(jdbcUrl);
+                        ds.setWarehouse(String.valueOf(
+                                datasourceConfiguration.getProperties().get(0).getValue()));
+                        ds.setDatabaseName(String.valueOf(
+                                datasourceConfiguration.getProperties().get(1).getValue()));
+                        ds.setRole(String.valueOf(
+                                datasourceConfiguration.getProperties().get(3).getValue()));
+                        ds.setSchema(String.valueOf(
+                                datasourceConfiguration.getProperties().get(2).getValue()));
+                        config.setDataSource(ds);
+
+                        return Mono.just(config);
+                    })
+                    .onErrorMap(
+                            AppsmithPluginException.class,
+                            error -> new AppsmithPluginException(
+                                    AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, error.getMessage()));
+        }
+
+        private Mono<HikariConfig> getBasicAuthConfig(
+                HikariConfig config, DatasourceConfiguration datasourceConfiguration) {
+            return Mono.fromCallable(() -> {
+                DBAuth authentication = (DBAuth) datasourceConfiguration.getAuthentication();
+                if (authentication.getUsername() != null) {
+                    config.setUsername(authentication.getUsername());
+                }
+                if (authentication.getPassword() != null) {
+                    config.setPassword(authentication.getPassword());
+                }
+                return config;
+            });
+        }
+
+        private HikariConfig getCommonHikariConfig(Properties properties) {
+            HikariConfig config = new HikariConfig();
+            config.setDriverClassName(properties.getProperty("driver_name"));
+
+            config.setMinimumIdle(Integer.parseInt(properties.get("minimumIdle").toString()));
+            config.setMaximumPoolSize(
+                    Integer.parseInt(properties.get("maximunPoolSize").toString()));
+
+            config.setInitializationFailTimeout(
+                    Long.parseLong(properties.get("initializationFailTimeout").toString()));
+            config.setConnectionTimeout(
+                    Long.parseLong(properties.get("connectionTimeoutMillis").toString()));
+            return config;
+        }
+
+        private Mono<PrivateKey> getPrivateKeyFromBase64(byte[] keyBytes, String passphrase) {
+            try {
+                return Mono.just(SnowflakeKeyUtils.readEncryptedPrivateKey(keyBytes, passphrase));
+            } catch (AppsmithPluginException e) {
+                return Mono.error(new AppsmithPluginException(
+                        AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
+                        SnowflakeErrorMessages.DS_MISSING_PASSPHRASE_FOR_ENCRYPTED_PRIVATE_KEY));
+            } catch (PKCSException e) {
+                return Mono.error(new AppsmithPluginException(
+                        AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
+                        SnowflakeErrorMessages.DS_INCORRECT_PASSPHRASE_OR_PRIVATE_KEY));
+            } catch (Exception e) {
+                return Mono.error(new AppsmithPluginException(
+                        AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
+                        SnowflakeErrorMessages.UNABLE_TO_CREATE_CONNECTION_ERROR_MSG));
+            }
+        }
+
+        private String getJDBCUrl(DatasourceConfiguration dsConfig) {
+            StringBuilder urlBuilder =
+                    new StringBuilder("jdbc:snowflake://" + dsConfig.getUrl() + ".snowflakecomputing.com?");
+            return urlBuilder.toString();
+        }
     }
 }

app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/exceptions/SnowflakeErrorMessages.java

@@ -36,4 +36,8 @@ public class SnowflakeErrorMessages extends BasePluginErrorMessages {
     public static final String DS_MISSING_PASSWORD_ERROR_MSG = "Missing password for authentication.";
 
     public static final String DS_MISSING_PRIVATE_KEY_ERROR_MSG = "Missing private key for authentication.";
+
+    public static final String DS_MISSING_PASSPHRASE_FOR_ENCRYPTED_PRIVATE_KEY =
+            "Passphrase is required as private key is encrypted";
+    public static final String DS_INCORRECT_PASSPHRASE_OR_PRIVATE_KEY = "Passphrase or private key is incorrect";
 }

app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/utils/SnowflakeKeyUtils.java

@@ -0,0 +1,49 @@
+package com.external.utils;
+
+import com.appsmith.external.exceptions.pluginExceptions.AppsmithPluginError;
+import com.appsmith.external.exceptions.pluginExceptions.AppsmithPluginException;
+import com.external.plugins.exceptions.SnowflakeErrorMessages;
+import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
+import org.bouncycastle.jce.provider.BouncyCastleProvider;
+import org.bouncycastle.openssl.PEMParser;
+import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
+import org.bouncycastle.openssl.jcajce.JceOpenSSLPKCS8DecryptorProviderBuilder;
+import org.bouncycastle.operator.InputDecryptorProvider;
+import org.bouncycastle.pkcs.PKCS8EncryptedPrivateKeyInfo;
+
+import java.io.StringReader;
+import java.security.PrivateKey;
+import java.security.Security;
+
+import static org.apache.commons.lang.StringUtils.isEmpty;
+
+public class SnowflakeKeyUtils {
+    static {
+        Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
+    }
+
+    public static PrivateKey readEncryptedPrivateKey(byte[] keyBytes, String passphrase) throws Exception {
+        PrivateKeyInfo privateKeyInfo = null;
+        String privateKeyPEM = new String(keyBytes);
+        PEMParser pemParser = new PEMParser(new StringReader(privateKeyPEM));
+        Object pemObject = pemParser.readObject();
+        if (pemObject instanceof PKCS8EncryptedPrivateKeyInfo) {
+            // Handle the case where the private key is encrypted.
+            if (isEmpty(passphrase)) {
+                throw new AppsmithPluginException(
+                        AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
+                        SnowflakeErrorMessages.DS_MISSING_PASSPHRASE_FOR_ENCRYPTED_PRIVATE_KEY);
+            }
+            PKCS8EncryptedPrivateKeyInfo encryptedPrivateKeyInfo = (PKCS8EncryptedPrivateKeyInfo) pemObject;
+            InputDecryptorProvider pkcs8Prov =
+                    new JceOpenSSLPKCS8DecryptorProviderBuilder().build(passphrase.toCharArray());
+            privateKeyInfo = encryptedPrivateKeyInfo.decryptPrivateKeyInfo(pkcs8Prov);
+        } else if (pemObject instanceof PrivateKeyInfo) {
+            // Handle the case where the private key is unencrypted.
+            privateKeyInfo = (PrivateKeyInfo) pemObject;
+        }
+        pemParser.close();
+        JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(BouncyCastleProvider.PROVIDER_NAME);
+        return converter.getPrivateKey(privateKeyInfo);
+    }
+}

app/server/appsmith-plugins/snowflakePlugin/src/test/java/com/external/plugins/SnowflakePluginTest.java

@@ -102,6 +102,7 @@ public void testDatasourceWithInvalidUrl() {
         DBAuth auth = new DBAuth();
         auth.setUsername("test");
         auth.setPassword("test");
+        auth.setAuthenticationType("dbAuth");
         datasourceConfiguration.setAuthentication(auth);
         List<Property> properties = new ArrayList<>();
         properties.add(new Property("warehouse", "warehouse"));