add liveliness API and support user defined probe sql for DataSource …

…definition check.
byzer-org · Apr 25, 2022 · cd58f6e · cd58f6e
1 parent 5a0b772
commit cd58f6e
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 8 deletions.
diff --git a/src/main/java/io/kyligence/notebook/console/bean/dto/ConnectionDTO.java b/src/main/java/io/kyligence/notebook/console/bean/dto/ConnectionDTO.java
@@ -2,7 +2,6 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import io.kyligence.notebook.console.bean.entity.ConnectionInfo;
-import io.kyligence.notebook.console.support.EncryptUtils;
 import io.kyligence.notebook.console.util.JacksonUtils;
 import lombok.Data;
 import lombok.NoArgsConstructor;

diff --git a/src/main/java/io/kyligence/notebook/console/controller/SettingsController.java b/src/main/java/io/kyligence/notebook/console/controller/SettingsController.java
@@ -18,12 +18,16 @@
 import io.swagger.annotations.Api;
 import io.swagger.annotations.ApiOperation;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.http.HttpStatus;
+import org.apache.http.protocol.HTTP;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.data.util.Pair;
 import org.springframework.validation.annotation.Validated;
 import org.springframework.web.bind.annotation.*;
 
+import javax.servlet.http.HttpServletResponse;
 import javax.validation.constraints.NotNull;
+import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 
@@ -57,6 +61,17 @@ public Response<VersionInfo> getVersion() {
         return new Response<VersionInfo>().data(systemService.getVersionInfo());
     }
 
+    @ApiOperation("Liveliness Check")
+    @GetMapping("/liveliness")
+    public void checkLiveliness(HttpServletResponse response) throws IOException {
+        response.setContentType("application/json");
+        response.getWriter().write("{\"status\": \"ok\"}");
+        if (!systemService.dbStatusCheck()) {
+            response.setStatus(HttpStatus.SC_BAD_REQUEST);
+            response.getWriter().write("{\"status\": \"ok\", \"msg\":\"meta database unreachable\"}");
+        }
+    }
+
     @ApiOperation("Engine List")
     @GetMapping("/settings/engines")
     @Permission

diff --git a/src/main/java/io/kyligence/notebook/console/service/ConnectionService.java b/src/main/java/io/kyligence/notebook/console/service/ConnectionService.java
@@ -64,11 +64,14 @@ public ConnectionInfo findById(Integer id) {
     public boolean testConnection(ConnectionDTO content) {
         if (content.getName() == null) content.setName("UserConnectionTmp");
         String sql = ConnectionUtils.renderSQL(content);
-        return checkConnection(sql, content.getName(), engineService.getExecutionEngine());
+        String probeSQL = ConnectionUtils.parseProbeSql(content.getParameter());
+        return checkConnection(sql, probeSQL, content.getName(), engineService.getExecutionEngine());
     }
 
     public boolean testConnection(ConnectionInfo info, String engine) {
-        return checkConnection(renderConnectionSQL(info), info.getName(), engine);
+        String probeSQL = ConnectionUtils.parseProbeSql(
+                JacksonUtils.readJsonArray(info.getParameter(), ConnectionDTO.ParameterMap.class));
+        return checkConnection(renderConnectionSQL(info), probeSQL, info.getName(), engine);
     }
 
     public String renderConnectionSQL(ConnectionInfo info) {
@@ -112,7 +115,7 @@ private List<String> showConnectionTables(ConnectionInfo info) {
         return tables;
     }
 
-    private boolean checkConnection(String renderedSQL, String connectionName, String engine) {
+    private boolean checkConnection(String renderedSQL, String probeSQL, String connectionName, String engine) {
         EngineService.RunScriptParams runScriptParams = new EngineService.RunScriptParams();
 
         try {
@@ -124,10 +127,11 @@ private boolean checkConnection(String renderedSQL, String connectionName, Strin
             );
             String testSQL = String.format(
                     "run command as JDBC.`%1$s._` where" +
-                            "`driver-statement-query`=\"select 1\"" +
+                            "`driver-statement-query`=\"%2$s\"" +
                             "and sqlMode=\"query\"" +
-                            "as %1$s_show_tables;",
-                    connectionName
+                            "as %1$s_probe_test;",
+                    connectionName,
+                    probeSQL
             );
             String result = engineService.runScript(
                     runScriptParams.withSql(testSQL)

diff --git a/src/main/java/io/kyligence/notebook/console/service/SystemService.java b/src/main/java/io/kyligence/notebook/console/service/SystemService.java
@@ -7,6 +7,7 @@
 import io.kyligence.notebook.console.dao.SystemConfigRepository;
 import io.kyligence.notebook.console.support.CriteriaQueryBuilder;
 import io.kyligence.notebook.console.support.EncryptUtils;
+import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 import org.springframework.transaction.annotation.Transactional;
@@ -20,6 +21,7 @@
 import java.util.List;
 
 
+@Slf4j
 @Service
 public class SystemService {
 
@@ -31,6 +33,17 @@ public class SystemService {
     @Resource
     private CriteriaQueryBuilder queryBuilder;
 
+    public boolean dbStatusCheck() {
+        // test meta database reachable
+        try {
+            repository.findAll();
+            return true;
+        } catch (Exception ex) {
+            log.error("Can not access meta database table, please check database status!");
+            return false;
+        }
+    }
+
     @Transactional
     public void updateByUser(SystemConfig systemConfig) {
         Query query = queryBuilder.updateNotNullByField(systemConfig, "user");

diff --git a/src/main/java/io/kyligence/notebook/console/util/ConnectionUtils.java b/src/main/java/io/kyligence/notebook/console/util/ConnectionUtils.java
@@ -2,10 +2,13 @@
 
 import io.kyligence.notebook.console.bean.dto.ConnectionDTO;
 import io.kyligence.notebook.console.support.EncryptUtils;
+import org.apache.commons.lang3.StringUtils;
 
 import java.util.List;
+import java.util.Objects;
 
 public class ConnectionUtils {
+    public static final String PROBE_SQL_PARAM_NAME = "probeSQL";
 
     public static String renderSQL(String url, String driver, String user,
                                    String password, String name,
@@ -22,7 +25,9 @@ public static String renderSQL(String url, String driver, String user,
 
         if (parameter != null) {
             parameter.forEach(parameterMap -> {
-                if (parameterMap.getName() != null && !parameterMap.getName().isEmpty()){
+                if (StringUtils.isNotBlank(parameterMap.getName()) &&
+                        !parameterMap.getName().equalsIgnoreCase(PROBE_SQL_PARAM_NAME) &&
+                        StringUtils.isNotBlank(parameterMap.getValue())) {
                     builder.append(
                             String.format("and %1$s=\"%2$s\"\n", parameterMap.getName(), parameterMap.getValue())
                     );
@@ -45,4 +50,17 @@ public static String renderSQL(ConnectionDTO content) {
                 content.getParameter()
         );
     }
+
+    public static String parseProbeSql(List<ConnectionDTO.ParameterMap> userParameters) {
+        if (Objects.nonNull(userParameters)) {
+            for (ConnectionDTO.ParameterMap parameterMap : userParameters) {
+                if (StringUtils.isNotBlank(parameterMap.getName()) &&
+                        parameterMap.getName().equalsIgnoreCase(PROBE_SQL_PARAM_NAME) &&
+                        StringUtils.isNotBlank(parameterMap.getValue())) {
+                    return parameterMap.getValue();
+                }
+            }
+        }
+        return "select 1";
+    }
 }