Fix: Mysql plugin: add timeout to validation check to avoid connectio…

…n block. (appsmithorg#6555) * add timeout to validation check to avoid connection block.
nnajiabraham · Aug 12, 2021 · 8f99635 · 8f99635
1 parent f19ebba
commit 8f99635
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/app/server/appsmith-plugins/mysqlPlugin/src/main/java/com/external/plugins/MySqlPlugin.java b/app/server/appsmith-plugins/mysqlPlugin/src/main/java/com/external/plugins/MySqlPlugin.java
@@ -44,6 +44,7 @@
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 
+import java.time.Duration;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
@@ -58,6 +59,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -74,6 +76,7 @@ public class MySqlPlugin extends BasePlugin {
     private static final String DATE_COLUMN_TYPE_NAME = "date";
     private static final String DATETIME_COLUMN_TYPE_NAME = "datetime";
     private static final String TIMESTAMP_COLUMN_TYPE_NAME = "timestamp";
+    private static final int VALIDATION_CHECK_TIMEOUT = 4; // seconds
 
     /**
      * Example output for COLUMNS_QUERY:
@@ -235,7 +238,10 @@ public Mono<ActionExecutionResult> executeCommon(Connection connection,
             List<RequestParamDTO> requestParams = List.of(new RequestParamDTO(ACTION_CONFIGURATION_BODY,
                     transformedQuery, null, null, psParams));
 
+            // TODO: need to write a JUnit TC for VALIDATION_CHECK_TIMEOUT
             Flux<Result> resultFlux = Mono.from(connection.validate(ValidationDepth.REMOTE))
+                    .timeout(Duration.ofSeconds(VALIDATION_CHECK_TIMEOUT))
+                    .onErrorMap(TimeoutException.class, error -> new StaleConnectionException())
                     .flatMapMany(isValid -> {
                         if (isValid) {
                             return createAndExecuteQueryFromConnection(finalQuery,
@@ -784,6 +790,8 @@ public Mono<DatasourceStructure> getStructure(Connection connection, DatasourceC
             final Map<String, DatasourceStructure.Key> keyRegistry = new HashMap<>();
 
             return Mono.from(connection.validate(ValidationDepth.REMOTE))
+                    .timeout(Duration.ofSeconds(VALIDATION_CHECK_TIMEOUT))
+                    .onErrorMap(TimeoutException.class, error -> new StaleConnectionException())
                     .flatMapMany(isValid -> {
                         if (isValid) {
                             return connection.createStatement(COLUMNS_QUERY).execute();