Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: snowflake key auth server #34548

Merged
merged 15 commits into from
Jul 1, 2024
Merged

feat: snowflake key auth server #34548

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
1be7356
chore: Add KeyPairAuth as snowflake auth type document
rishabhrathod01 Jun 24, 2024
d7450d1
Merge branch 'release' into chore/addKeyPairAuthSnowFlake
rishabhrathod01 Jun 25, 2024
3ef2735
Add subtype for AuthenticationDTO
rishabhrathod01 Jun 26, 2024
3de7457
Merge branch 'release' of github.com:appsmithorg/appsmith into chore/…
AmanAgarwal041 Jun 26, 2024
5710dfc
chore: key pair validation in snowflake
AmanAgarwal041 Jun 26, 2024
5a26ff4
chore: fix junit test case
AmanAgarwal041 Jun 26, 2024
4025ac2
Merge branch 'release' of github.com:appsmithorg/appsmith into chore/…
AmanAgarwal041 Jun 26, 2024
83a8884
fix: key auth method added in snowflake
Jun 27, 2024
937a77e
fix: removed unnecessary code
Jun 27, 2024
78d700d
Merge branch 'release' into feat/snowflake-key-auth-server
Jun 27, 2024
4267582
Merge branch 'release' into feat/snowflake-key-auth-server
Jun 28, 2024
456f767
Merge branch 'release' into feat/snowflake-key-auth-server
Jul 1, 2024
953d145
fix: code review comments addressed
Jul 1, 2024
219c26c
fix: fixed imports and failing tests
Jul 1, 2024
cb64fff
fix: reverted too many specific imports
Jul 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
211 changes: 141 additions & 70 deletions .../appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
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;
Expand All @@ -15,32 +16,29 @@
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.security.*;
import java.sql.*;
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.util.*;

import static com.appsmith.external.constants.Authentication.BASIC;
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;
Expand Down Expand Up @@ -159,58 +157,31 @@ 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
StringBuilder urlBuilder = new StringBuilder("jdbc:snowflake://"
+ datasourceConfiguration.getUrl() + ".snowflakecomputing.com?");
config.setJdbcUrl(urlBuilder.toString());
sneha122 marked this conversation as resolved.
Show resolved Hide resolved

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());
}
Comment on lines +170 to +173
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Improve exception handling for Hikari pool initialization.

Consider providing more context in the exception message or using a custom exception to improve error handling.

- throw new AppsmithPluginException(
-     AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, e.getMessage());
+ throw new CustomHikariPoolException(
+     AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, "Error initializing Hikari pool: " + e.getMessage());
Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
} catch (HikariPool.PoolInitializationException e) {
throw new AppsmithPluginException(
AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, e.getMessage());
}
} catch (HikariPool.PoolInitializationException e) {
throw new CustomHikariPoolException(
AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, "Error initializing Hikari pool: " + e.getMessage());
}


return datasource;
})
.subscribeOn(scheduler);
return datasource;
})
.subscribeOn(scheduler))
.onErrorMap(
AppsmithPluginException.class,
error -> new AppsmithPluginException(
AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, error.getMessage()));
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
}

@Override
Expand All @@ -234,17 +205,14 @@ 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();
if (authentication instanceof DBAuth) {
DBAuth dbAuth = (DBAuth) authentication;
properties.setProperty("user", dbAuth.getUsername());
properties.setProperty("password", dbAuth.getPassword());
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
}

properties.setProperty(
"warehouse",
String.valueOf(
Expand Down Expand Up @@ -491,5 +459,108 @@ public Mono<DatasourceStructure> getStructure(
})
.subscribeOn(scheduler);
}

private Mono<HikariConfig> getHikariConfig(
DatasourceConfiguration datasourceConfiguration, Properties properties) {
HikariConfig commonConfig = getCommonHikariConfig(properties);
Mono<HikariConfig> configMono = Mono.just(commonConfig);
sneha122 marked this conversation as resolved.
Show resolved Hide resolved

String authenticationType =
datasourceConfiguration.getAuthentication().getAuthenticationType();
if (authenticationType != null) {
switch (authenticationType) {
case BASIC:
configMono = configMono.flatMap(config -> getBasicAuthConfig(config, datasourceConfiguration));
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
break;
case SNOWFLAKE_KEY_PAIR_AUTH:
configMono =
configMono.flatMap(config -> getKeyPairAuthConfig(config, 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 -> {
StringBuilder urlBuilder = new StringBuilder(
"jdbc:snowflake://" + datasourceConfiguration.getUrl() + ".snowflakecomputing.com?");
config.setJdbcUrl(urlBuilder.toString());

SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
ds.setPrivateKey(privateKey);
ds.setUser(keyPairAuthConfig.getUsername());
ds.setUrl(urlBuilder.toString());
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()));
}
sneha122 marked this conversation as resolved.
Show resolved Hide resolved

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;
});
}
sneha122 marked this conversation as resolved.
Show resolved Hide resolved

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 (IllegalArgumentException e) {
return Mono.error(new AppsmithPluginException(
AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
SnowflakeErrorMessages.PASSPHRASE_IS_REQUIRED_FOR_ENCRYPTED_PRIVATE_KEY));
} catch (PKCSException e) {
return Mono.error(new AppsmithPluginException(
AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
SnowflakeErrorMessages.PASSPHRASE_OR_PRIVATE_KEY_INCORRECT));
} catch (Exception e) {
return Mono.error(new AppsmithPluginException(
AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
SnowflakeErrorMessages.UNABLE_TO_CREATE_CONNECTION_ERROR_MSG));
}
}
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
}
}
4 changes: 4 additions & 0 deletions ...snowflakePlugin/src/main/java/com/external/plugins/exceptions/SnowflakeErrorMessages.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 PASSPHRASE_IS_REQUIRED_FOR_ENCRYPTED_PRIVATE_KEY =
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
"Passphrase is required as private key is encrypted";
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
public static final String PASSPHRASE_OR_PRIVATE_KEY_INCORRECT = "Passphrase or private key is incorrect";
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
}
46 changes: 46 additions & 0 deletions .../appsmith-plugins/snowflakePlugin/src/main/java/com/external/utils/SnowflakeKeyUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
package com.external.utils;

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 IllegalArgumentException(
SnowflakeErrorMessages.PASSPHRASE_IS_REQUIRED_FOR_ENCRYPTED_PRIVATE_KEY);
sneha122 marked this conversation as resolved.
Show resolved Hide resolved
}
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);
}
}
Loading