byzer-org · chncaesar · May 7, 2022 · Apr 25, 2022
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/bean/dto/SystemStatus.java b/src/main/java/io/kyligence/notebook/console/bean/dto/SystemStatus.java
@@ -0,0 +1,14 @@
+package io.kyligence.notebook.console.bean.dto;
+
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+import java.util.List;
+
+@Data
+@NoArgsConstructor
+public class SystemStatus {
+    private String status;
+    private String msg;
+    private List<EngineStatusDTO> engineStatus;
+}
diff --git a/src/main/java/io/kyligence/notebook/console/controller/SettingsController.java b/src/main/java/io/kyligence/notebook/console/controller/SettingsController.java
@@ -14,18 +14,23 @@
 import io.kyligence.notebook.console.support.DisableInTrial;
 import io.kyligence.notebook.console.support.Permission;
 import io.kyligence.notebook.console.util.EngineStatus;
+import io.kyligence.notebook.console.util.JacksonUtils;
 import io.kyligence.notebook.console.util.WebUtils;
 import io.swagger.annotations.Api;
 import io.swagger.annotations.ApiOperation;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.http.HttpStatus;
 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;
+import java.util.Objects;
 
 @Slf4j
 @Validated
@@ -57,6 +62,44 @@ public Response<VersionInfo> getVersion() {
         return new Response<VersionInfo>().data(systemService.getVersionInfo());
     }
 
+    @ApiOperation("Readiness Check")
+    @GetMapping("/system/readiness")
+    public void readinessCheck(HttpServletResponse response) throws IOException {
+
+        SystemStatus systemStatus = systemStatusCheck();
+        response.setContentType("application/json");
+
+        if (systemStatus.getStatus().equalsIgnoreCase("DOWN")) {
+            response.setStatus(HttpStatus.SC_BAD_REQUEST);
+        }
+        systemStatus.setEngineStatus(EngineListDTO.valueOf(engineService.getEngineStatusMap()).getList());
+        response.getWriter().write(Objects.requireNonNull(JacksonUtils.writeJson(systemStatus)));
+    }
+
+    @ApiOperation("Liveliness Check")
+    @GetMapping("/system/liveliness")
+    public void livelinessCheck(HttpServletResponse response) throws IOException {
+        SystemStatus systemStatus = systemStatusCheck();
+        response.setContentType("application/json");
+
+        if (systemStatus.getStatus().equalsIgnoreCase("DOWN")) {
+            response.setStatus(HttpStatus.SC_BAD_REQUEST);
+        }
+
+        response.getWriter().write(Objects.requireNonNull(JacksonUtils.writeJson(systemStatus)));
+
+    }
+
+    private SystemStatus systemStatusCheck() {
+        SystemStatus systemStatus = new SystemStatus();
+        systemStatus.setStatus("UP");
+        if (!systemService.isMetaDBReachable()) {
+            systemStatus.setStatus("DOWN");
+            systemStatus.setMsg("meta database unreachable");
+        }
+        return systemStatus;
+    }
+
     @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 isMetaDBReachable() {
+        // 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";
+    }
 }
diff --git a/src/test/java/io/kyligence/notebook/console/service/ConnectionServiceTest.java b/src/test/java/io/kyligence/notebook/console/service/ConnectionServiceTest.java
@@ -57,7 +57,11 @@ public void testTestConnection() {
         ).respond(response().withBody(MOCK_CONNECTION_CONTENT));
         Assert.assertTrue(connectionService.testConnection(MOCK_CONNECTION, "default"));
         Assert.assertFalse(connectionService.testConnection(MOCK_CONNECTION, "backup"));
+        Assert.assertTrue(connectionService.testConnection(ConnectionDTO.valueOf(MOCK_CONNECTION)));
 
+        MOCK_CONNECTION.setParameter("[{\"name\":\"probeSQL\",\"value\": \"select 1 from testTB\"}]");
+        Assert.assertTrue(connectionService.testConnection(MOCK_CONNECTION, "default"));
+        Assert.assertFalse(connectionService.testConnection(MOCK_CONNECTION, "backup"));
         Assert.assertTrue(connectionService.testConnection(ConnectionDTO.valueOf(MOCK_CONNECTION)));
     }
 
@@ -71,9 +75,4 @@ public void testRenderSQL() {
                 "as admin-mockConnectionForAdmin;";
         Assert.assertEquals(expectSQL, connectionService.renderConnectionSQL(MOCK_CONNECTION));
     }
-
-    @Test
-    public void testCreateConnection() {
-
-    }
 }
diff --git a/src/test/java/io/kyligence/notebook/console/service/SystemServiceTest.java b/src/test/java/io/kyligence/notebook/console/service/SystemServiceTest.java
@@ -4,19 +4,30 @@
 import io.kyligence.notebook.console.NotebookLauncherBaseTest;
 import io.kyligence.notebook.console.bean.entity.SystemConfig;
 import io.kyligence.notebook.console.dao.SystemConfigRepository;
+import io.kyligence.notebook.console.exception.ByzerException;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.InjectMocks;
+import org.mockito.Mock;
 import org.springframework.beans.factory.annotation.Autowired;
 
 import javax.annotation.PostConstruct;
 
+import static org.mockito.Mockito.when;
+
 public class SystemServiceTest extends NotebookLauncherBaseTest {
 
     private SystemConfig systemConfig;
 
     @Autowired
     private SystemService systemService;
 
+    @InjectMocks
+    private SystemService mockService;
+
+    @Mock
+    private SystemConfigRepository mockRepo;
+
     @Override
     @PostConstruct
     public void mock() {
@@ -38,4 +49,13 @@ public void testUpdateConfig() {
         config = systemService.getConfig("mockConfigUser2");
         Assert.assertEquals("backup", config.getEngine());
     }
+
+    @Test
+    public void testMetaDBReachable () {
+        Assert.assertTrue(systemService.isMetaDBReachable());
+
+        when(mockRepo.findAll()).thenThrow(ByzerException.class);
+
+        Assert.assertFalse(mockService.isMetaDBReachable());
+    }
 }