From f05e32a97ef34cfb3471cbaa8be9f1d21d885496 Mon Sep 17 00:00:00 2001
From: Nidhi <nidhi@appsmith.com>
Date: Mon, 16 Aug 2021 16:03:39 +0530
Subject: [PATCH] Throw a stale connection exception when authentication token
 expires in Snowflake (#6612)

* Throw a stale connection on authentication token expired

* Added JUnit test
---
 .../com/external/plugins/SnowflakePlugin.java |  4 ++
 .../external/plugins/SnowflakePluginTest.java | 37 ++++++++++++++++++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java b/app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java
index f38970de91..399fc4850f 100644
--- a/app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java
+++ b/app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java
@@ -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 {
diff --git a/app/server/appsmith-plugins/snowflakePlugin/src/test/java/com/external/plugins/SnowflakePluginTest.java b/app/server/appsmith-plugins/snowflakePlugin/src/test/java/com/external/plugins/SnowflakePluginTest.java
index ba164a4749..28758ba135 100644
--- a/app/server/appsmith-plugins/snowflakePlugin/src/test/java/com/external/plugins/SnowflakePluginTest.java
+++ b/app/server/appsmith-plugins/snowflakePlugin/src/test/java/com/external/plugins/SnowflakePluginTest.java
@@ -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();
+    }
+
 }