From 6a166c115f948fff8955bc3b29475f9aa92a2ce7 Mon Sep 17 00:00:00 2001
From: serge-rider <serge@jkiss.org>
Date: Fri, 17 Jun 2022 09:54:54 +0300
Subject: [PATCH] CB-2174 SQL query execute timeout

---
 .../cloudbeaver/model/session/WebSession.java   |  4 ++++
 .../service/sql/WebSQLConstants.java            |  1 +
 .../service/sql/WebSQLProcessor.java            | 17 +++++++++++++++++
 3 files changed, 22 insertions(+)

diff --git a/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/model/session/WebSession.java b/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/model/session/WebSession.java
index 153564b85c..14d324158e 100644
--- a/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/model/session/WebSession.java
+++ b/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/model/session/WebSession.java
@@ -160,6 +160,10 @@ public class WebSession extends AbstractSessionPersistent implements SMSession,
         return id;
     }
 
+    public WebApplication getApplication() {
+        return application;
+    }
+
     @Override
     public boolean isApplicationSession() {
         return false;
diff --git a/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLConstants.java b/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLConstants.java
index 85d6a6a60f..4aae93d74a 100644
--- a/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLConstants.java
+++ b/server/bundles/io.cloudbeaver.model/src/io/cloudbeaver/service/sql/WebSQLConstants.java
@@ -24,6 +24,7 @@ public class WebSQLConstants {
     public static final String QUOTA_PROP_ROW_LIMIT = "sqlResultSetRowsLimit";
     public static final String QUOTA_PROP_MEMORY_LIMIT = "sqlResultSetMemoryLimit";
     public static final String QUOTA_PROP_QUERY_LIMIT = "sqlMaxRunningQueries";
+    public static final String QUOTA_PROP_SQL_QUERY_TIMEOUT = "sqlQueryTimeout";
     public static final String QUOTA_PROP_TEXT_PREVIEW_MAX_LENGTH = "sqlTextPreviewMaxLength";
     public static final String QUOTA_PROP_BINARY_PREVIEW_MAX_LENGTH = "sqlBinaryPreviewMaxLength";
     public static final String QUOTA_PROP_RM_FILE_SIZE_LIMIT = "resourceManagerFileSizeLimit";
diff --git a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
index 2d9a49baff..ee3a486998 100644
--- a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
+++ b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
@@ -180,6 +180,7 @@ public class WebSQLProcessor implements WebSessionProvider {
                         session.getExecutionContext(),
                         WebSQLProcessor.this,
                         sqlQuery);
+
                     try (DBCStatement dbStat = DBUtils.makeStatement(
                         source,
                         session,
@@ -188,6 +189,22 @@ public class WebSQLProcessor implements WebSessionProvider {
                         webDataFilter.getOffset(),
                         webDataFilter.getLimit()))
                     {
+                        // Set query timeout
+                        int queryTimeout = (int) session.getDataSource().getContainer().getPreferenceStore()
+                            .getDouble(WebSQLConstants.QUOTA_PROP_SQL_QUERY_TIMEOUT);
+                        if (queryTimeout <= 0) {
+                            queryTimeout = CommonUtils.toInt(
+                                getWebSession().getApplication().getAppConfiguration()
+                                    .getResourceQuota(WebSQLConstants.QUOTA_PROP_SQL_QUERY_TIMEOUT));
+                        }
+                        if (queryTimeout > 0) {
+                            try {
+                                dbStat.setStatementTimeout(queryTimeout);
+                            } catch (Throwable e) {
+                                log.debug("Can't set statement timeout:" + e.getMessage());
+                            }
+                        }
+
                         boolean hasResultSet = dbStat.executeStatement();
                         fillQueryResults(contextInfo, dataContainer, dbStat, hasResultSet, executeInfo, webDataFilter, dataFilter, dataFormat);
                     } catch (DBException e) {