[LINKIS-5147][EC-JDBC] Add timeout configurations of datasource for j…

…dbc engine (#5148) * [LINKIS-5147] Add timeout configurations of datasource for jdbc engine * fix style
apache · Jul 16, 2024 · 2b55fe6 · 2b55fe6
1 parent cc3a0f6
commit 2b55fe6
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 1 deletion.
diff --git a/...ins/jdbc/src/main/java/org/apache/linkis/manager/engineplugin/jdbc/ConnectionManager.java b/...ins/jdbc/src/main/java/org/apache/linkis/manager/engineplugin/jdbc/ConnectionManager.java
@@ -19,6 +19,7 @@
 
 import org.apache.linkis.common.utils.SecurityUtils;
 import org.apache.linkis.hadoop.common.utils.KerberosUtils;
+import org.apache.linkis.manager.engineplugin.jdbc.conf.JDBCConfiguration$;
 import org.apache.linkis.manager.engineplugin.jdbc.constant.JDBCEngineConnConstant;
 import org.apache.linkis.manager.engineplugin.jdbc.exception.JDBCParamsIllegalException;
 import org.apache.linkis.manager.engineplugin.jdbc.utils.JdbcParamUtils;
@@ -187,6 +188,14 @@ protected DataSource buildDataSource(String dbUrl, Map<String, String> propertie
         JDBCPropertiesParser.getInt(
             properties, JDBCEngineConnConstant.JDBC_POOL_REMOVE_ABANDONED_TIMEOUT, 300);
 
+    int connectionTimeout =
+        JDBCPropertiesParser.getInt(properties, JDBCEngineConnConstant.JDBC_CONNECTION_TIMEOUT, 0);
+    int socketTimeout =
+        JDBCPropertiesParser.getInt(properties, JDBCEngineConnConstant.JDBC_SOCKET_TIMEOUT, 0);
+    int queryTimeout =
+        JDBCPropertiesParser.getInt(
+            properties, JDBCConfiguration$.MODULE$.JDBC_QUERY_TIMEOUT().key(), 0);
+
     DruidDataSource datasource = new DruidDataSource();
     LOG.info("Database connection address information(数据库连接地址信息)=" + dbUrl);
     datasource.setUrl(dbUrl);
@@ -207,6 +216,15 @@ protected DataSource buildDataSource(String dbUrl, Map<String, String> propertie
     datasource.setPoolPreparedStatements(poolPreparedStatements);
     datasource.setRemoveAbandoned(removeAbandoned);
     datasource.setRemoveAbandonedTimeout(removeAbandonedTimeout);
+    if (connectionTimeout > 0) {
+      datasource.setConnectTimeout(connectionTimeout);
+    }
+    if (socketTimeout > 0) {
+      datasource.setSocketTimeout(socketTimeout);
+    }
+    if (queryTimeout > 0) {
+      datasource.setQueryTimeout(queryTimeout);
+    }
     return datasource;
   }
 

diff --git a/...ain/java/org/apache/linkis/manager/engineplugin/jdbc/constant/JDBCEngineConnConstant.java b/...ain/java/org/apache/linkis/manager/engineplugin/jdbc/constant/JDBCEngineConnConstant.java
@@ -34,6 +34,11 @@ private JDBCEngineConnConstant() {}
       "wds.linkis.jdbc.auth.kerberos.proxy.enable";
   public static final String JDBC_PROXY_USER_PROPERTY = "wds.linkis.jdbc.proxy.user.property";
   public static final String JDBC_PROXY_USER = "wds.linkis.jdbc.proxy.user";
+
+  public static final String JDBC_CONNECTION_TIMEOUT = "wds.linkis.jdbc.connection.timeout";
+
+  public static final String JDBC_SOCKET_TIMEOUT = "wds.linkis.jdbc.socket.timeout";
+
   public static final String JDBC_SCRIPTS_EXEC_USER = "execUser";
   public static final String JDBC_ENGINE_RUN_TIME_DS = "wds.linkis.engine.runtime.datasource";
   public static final String JDBC_ENGINE_RUN_TIME_DS_MAX_VERSION_ID =

diff --git a/...n/scala/org/apache/linkis/manager/engineplugin/jdbc/executor/JDBCEngineConnExecutor.scala b/...n/scala/org/apache/linkis/manager/engineplugin/jdbc/executor/JDBCEngineConnExecutor.scala
@@ -146,7 +146,9 @@ class JDBCEngineConnExecutor(override val outputPrintLimit: Int, val id: Int)
 
     try {
       statement = connection.createStatement()
-      statement.setQueryTimeout(JDBCConfiguration.JDBC_QUERY_TIMEOUT.getValue)
+      if (statement.getQueryTimeout == 0) {
+        statement.setQueryTimeout(JDBCConfiguration.JDBC_QUERY_TIMEOUT.getValue)
+      }
       statement.setFetchSize(outputPrintLimit)
       statement.setMaxRows(outputPrintLimit)