From f8ef73b8a66dc79557557708caa5352cc9405a20 Mon Sep 17 00:00:00 2001
From: Trisha Anand <trisha@appsmith.com>
Date: Thu, 28 Dec 2023 12:10:58 +0530
Subject: [PATCH] fix: [Databricks] Adding stale connection exception handling
 in case of cluster restart (#29890)

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

## Summary by CodeRabbit

- **Bug Fixes**
- Improved error handling in the database connection process by
identifying and throwing an exception for stale connections.
- Enhanced error messages to include specific SQL states for better
troubleshooting.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 .../java/com/external/plugins/DatabricksPlugin.java  | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/app/server/appsmith-plugins/databricksPlugin/src/main/java/com/external/plugins/DatabricksPlugin.java b/app/server/appsmith-plugins/databricksPlugin/src/main/java/com/external/plugins/DatabricksPlugin.java
index fd82863f2fd7..ad724be113bd 100644
--- a/app/server/appsmith-plugins/databricksPlugin/src/main/java/com/external/plugins/DatabricksPlugin.java
+++ b/app/server/appsmith-plugins/databricksPlugin/src/main/java/com/external/plugins/DatabricksPlugin.java
@@ -138,11 +138,21 @@ public Mono<ActionExecutionResult> execute(
                             }
 
                         } catch (SQLException e) {
+
+                            String sqlState = e.getSQLState();
+                            // Databricks returns true on isValid check even if the connection is stale.
+                            // This scenario in particular happens when the connection was established before
+                            // the cluster restarts. The sql state here corresponds to bad connection link
+                            // and hence the correct action is to throw a StaleConnectionException.
+                            if (sqlState != null && sqlState.equals("08S01")) {
+                                return Mono.error(new StaleConnectionException(CONNECTION_CLOSED_ERROR_MSG));
+                            }
+
                             return Mono.error(new AppsmithPluginException(
                                     QUERY_EXECUTION_FAILED,
                                     QUERY_EXECUTION_FAILED_ERROR_MSG,
                                     e.getMessage(),
-                                    "SQLSTATE: " + e.getSQLState()));
+                                    "SQLSTATE: " + sqlState));
                         }
 
                         ActionExecutionResult result = new ActionExecutionResult();