Protei/PromucFlow_constructor
2
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Throw a stale connection exception when authentication token expires in Snowflake (#6612)

Browse Source 
* Throw a stale connection on authentication token expired

* Added JUnit test
This commit is contained in:
Nidhi 2021-08-16 16:03:39 +05:30 committed by GitHub
parent b2ade8e17e
commit f05e32a97e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 40 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java vendored
View File

@ -14,6 +14,7 @@ import com.appsmith.external.plugins.BasePlugin;
import com.appsmith.external.plugins.PluginExecutor;
import com.external.utils.SqlUtils;
import lombok.extern.slf4j.Slf4j;
import net.snowflake.client.jdbc.SnowflakeReauthenticationRequest;
import org.pf4j.Extension;
import org.pf4j.PluginWrapper;
import org.springframework.util.StringUtils;
@ -88,6 +89,9 @@ public class SnowflakePlugin extends BasePlugin {
rowsList.add(row);
}
} catch (SQLException e) {
if (e instanceof SnowflakeReauthenticationRequest) {
throw new StaleConnectionException();
}
e.printStackTrace();
throw new AppsmithPluginException(AppsmithPluginError.PLUGIN_EXECUTE_ARGUMENT_ERROR, e.getMessage());
} finally {

37
app/server/appsmith-plugins/snowflakePlugin/src/test/java/com/external/plugins/SnowflakePluginTest.java vendored
View File

@ -1,11 +1,21 @@
package com.external.plugins;
import com.appsmith.external.exceptions.pluginExceptions.StaleConnectionException;
import com.appsmith.external.models.ActionConfiguration;
import com.appsmith.external.models.ActionExecutionResult;
import com.appsmith.external.models.DBAuth;
import com.appsmith.external.models.DatasourceConfiguration;
import com.appsmith.external.models.Property;
import lombok.extern.log4j.Log4j;
import net.snowflake.client.jdbc.SnowflakeReauthenticationRequest;
import org.junit.Test;
import org.mockito.Mockito;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.List;
import java.util.Set;
@ -17,7 +27,7 @@ public class SnowflakePluginTest {
SnowflakePlugin.SnowflakePluginExecutor pluginExecutor = new SnowflakePlugin.SnowflakePluginExecutor();
@Test
public void testValidateDatasource_InvalidCredentials_returnsInvalids() {
public void testValidateDatasource_withInvalidCredentials_returnsInvalids() {
DatasourceConfiguration datasourceConfiguration = new DatasourceConfiguration();
DBAuth auth = new DBAuth();
auth.setUsername(null);
@ -32,4 +42,29 @@ public class SnowflakePluginTest {
assertTrue(output.contains("Missing database name."));
}
@Test
public void testExecute_authenticationTimeout_returnsStaleConnectionException() throws SQLException {
final String testQuery = "testQuery";
final Connection connection = Mockito.mock(Connection.class);
Mockito.when(connection.isValid(30))
.thenReturn(true);
final Statement statement = Mockito.mock(Statement.class);
Mockito.when(connection.createStatement())
.thenReturn(statement);
Mockito.when(statement.executeQuery(testQuery))
.thenThrow(new SnowflakeReauthenticationRequest(
"1",
"Authentication token expired",
"",
0));
final ActionConfiguration actionConfiguration = new ActionConfiguration();
actionConfiguration.setBody(testQuery);
final Mono<ActionExecutionResult> actionExecutionResultMono =
pluginExecutor.execute(connection, new DatasourceConfiguration(), actionConfiguration);
StepVerifier.create(actionExecutionResultMono)
.expectErrorMatches(e -> e instanceof StaleConnectionException)
.verify();
}
}