handle byzer python script stacktrace

pom.xml

@@ -365,6 +365,24 @@
             <version>1.16.3</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>testcontainers</artifactId>
+            <version>1.16.3</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>junit-jupiter</artifactId>
+            <version>1.16.3</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>mysql</artifactId>
+            <version>1.16.3</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>org.mock-server</groupId>
             <artifactId>mockserver-client-java</artifactId>

src/main/java/io/kyligence/notebook/console/bean/dto/JobInfoDTO.java

@@ -4,8 +4,8 @@
 import io.kyligence.notebook.console.bean.entity.JobInfo;
 import io.kyligence.notebook.console.bean.entity.JobInfoArchive;
 import io.kyligence.notebook.console.bean.model.CurrentJobInfo;
+import io.kyligence.notebook.console.util.EngineExceptionUtils;
 import io.kyligence.notebook.console.util.EntityUtils;
-import io.kyligence.notebook.console.util.ExceptionUtils;
 import lombok.Data;
 import lombok.NoArgsConstructor;
 import org.apache.commons.lang3.StringUtils;
@@ -88,7 +88,7 @@ public static JobInfoDTO valueOf(JobInfo jobInfo, JobProgressDTO jobProgress, Cu
         resp.endTime = EntityUtils.toStr(jobInfo.getFinishTime());
         resp.result = jobInfo.getResult();
         resp.msg = jobInfo.getMsg();
-        resp.rootCause = ExceptionUtils.getRootCause(jobInfo.getMsg());
+        resp.rootCause = EngineExceptionUtils.getRootCause(jobInfo.getMsg());
         resp.notebook = StringUtils.isBlank(jobInfo.getNotebook()) ? "untitled" : jobInfo.getNotebook();
         resp.engine = StringUtils.isBlank(jobInfo.getEngine()) ? "default" : jobInfo.getEngine();
         if (Objects.nonNull(jobProgress)) {

src/main/java/io/kyligence/notebook/console/service/JobService.java

@@ -14,6 +14,7 @@
 import io.kyligence.notebook.console.exception.ByzerException;
 import io.kyligence.notebook.console.exception.ErrorCodeEnum;
 import io.kyligence.notebook.console.support.CriteriaQueryBuilder;
+import io.kyligence.notebook.console.util.EngineExceptionUtils;
 import io.kyligence.notebook.console.util.ExceptionUtils;
 import io.kyligence.notebook.console.util.JacksonUtils;
 import lombok.extern.slf4j.Slf4j;
@@ -259,10 +260,11 @@ public boolean needRetry(Integer status){
     public boolean jobDone(String jobId, Integer status, String result, String msg, Timestamp finishTime) {
         JobInfo jobInfo = new JobInfo();
         jobInfo.setJobId(jobId);
-        jobInfo.setMsg(msg);
         jobInfo.setResult(result);
         jobInfo.setFinishTime(finishTime);
         jobInfo.setStatus(status);
+        jobInfo.setMsg(status == JobInfo.JobStatus.SUCCESS ? msg :
+                EngineExceptionUtils.parseStackTrace(getJobContent(jobId), msg));
 
         try {
             JobProgressDTO jobProgress = getJobProgress(jobId);

src/main/java/io/kyligence/notebook/console/service/SchedulerService.java

@@ -12,6 +12,7 @@
 import io.kyligence.notebook.console.scheduler.SchedulerConfig;
 import io.kyligence.notebook.console.scheduler.SchedulerFactory;
 import io.kyligence.notebook.console.scheduler.dolphin.dto.EntityModification;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.compress.utils.Lists;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
@@ -22,6 +23,7 @@
 import java.util.*;
 import java.util.stream.Collectors;
 
+@Slf4j
 @Service
 public class SchedulerService {
 
@@ -57,8 +59,9 @@ public void initSchedulers() {
 
     public void callback(String token, String scheduleOwner, String entityType, String entityId, String commitId) {
         if (!config.getScheduleCallbackToken().equals(token)) {
-            throw new ByzerException();
+            throw new ByzerException("Scheduler callback token auth failed!");
         }
+        log.info("Receive scheduler callback for {} execute Entity[{}, {}, {}]", scheduleOwner, entityType, entityId, commitId);
         String user = config.getScheduleCallbackUser();
         EngineService.RunScriptParams runScriptParams = new EngineService.RunScriptParams()
                 .withAsync("false")
@@ -87,9 +90,13 @@ public void callback(String token, String scheduleOwner, String entityType, Stri
         // 发送查询
         try {
             engineService.runScript(runScriptParams);
+            log.info("Scheduler callback for {} execute Entity[{}, {}, {}] succeed!",
+                    scheduleOwner, entityType, entityId, commitId);
         } catch (Exception ex) {
             // update job status to FAILED if exception happened
             status = JobInfo.JobStatus.FAILED;
+            log.error("Scheduler callback for {} execute Entity[{}, {}, {}] failed!",
+                    scheduleOwner, entityType, entityId, commitId);
             throw ex;
         } finally {
             jobInfo.setFinishTime(new Timestamp(System.currentTimeMillis()));

src/main/java/io/kyligence/notebook/console/util/EngineExceptionUtils.java

@@ -0,0 +1,102 @@
+package io.kyligence.notebook.console.util;
+
+import org.springframework.data.util.Pair;
+import org.apache.commons.lang3.StringUtils;
+
+import java.util.*;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+public class EngineExceptionUtils {
+
+    private EngineExceptionUtils() {
+        throw new IllegalStateException("Utility class");
+    }
+
+    private static final Pattern pythonLineRegex = Pattern.compile("^(File \"<string>\", line )(\\d+)(,.*)");
+
+    public static Pair<String, Map<Integer, Integer>> parseCodeTypeAndOffset(String script) {
+        String codeType = "byzer";
+        Map<Integer, Integer> offsetTracker = new LinkedHashMap<>();
+        int offset = 0;
+        int lineno = 0;
+
+        boolean prevIsHint = false;
+        for (String line : script.split("\n")) {
+            lineno++;
+            if (line.trim().startsWith("#%")) {
+                String header = StringUtils.stripStart(line.trim(), "#%");
+                codeType = header.contains("=") ? codeType : header;
+                offset++;
+                prevIsHint = true;
+            } else if (prevIsHint) {
+                offsetTracker.put(lineno - offset, offset);
+                prevIsHint = false;
+            }
+        }
+        return Pair.of(codeType, offsetTracker);
+    }
+
+    public static String parseStackTrace(String script, String stackTrace) {
+        if (StringUtils.isBlank(script) || StringUtils.isBlank(stackTrace)) return stackTrace;
+        Pair<String, Map<Integer, Integer>> codeTypeAndOffset = parseCodeTypeAndOffset(script);
+        if (codeTypeAndOffset.getFirst().equalsIgnoreCase("python")) {
+            String result = parsePythonError(stackTrace, codeTypeAndOffset.getSecond());
+            return result.isEmpty() ? stackTrace : result;
+        }
+        return stackTrace;
+    }
+
+    public static String parsePythonError(String stackTrace, Map<Integer, Integer> codeOffset) {
+        List<String> result = new ArrayList<>();
+        boolean start = false;
+        for (String line : stackTrace.split("\n")) {
+            if (start && line.contains("PythonRunner.scala")) break;
+
+            if (line.trim().startsWith("File")) {
+                start = true;
+                result.add(pythonLinenoAlign(line.trim(), codeOffset));
+
+            } else if (start) {
+                result.add(line);
+            }
+        }
+        return StringUtils.join(result, "\n");
+    }
+
+    public static String pythonLinenoAlign(String line, Map<Integer, Integer> codeOffset) {
+        Matcher m = pythonLineRegex.matcher(line);
+        if (!m.find()) return line;
+        Integer lineno = Integer.valueOf(m.group(2));
+        int lineMark = 0;
+        for (Integer mark : codeOffset.keySet()) {
+            if (mark > lineno) break;
+            lineMark = mark;
+        }
+        return m.group(1) + (lineno + codeOffset.get(lineMark)) + m.group(3);
+    }
+
+    public static String getRootCause(String ex) {
+        if (StringUtils.isBlank(ex)) {
+            return "";
+        }
+        // python stacktrace
+        if (ex.trim().startsWith("File")) {
+            String[] lines = ex.split("\n");
+            return lines[lines.length - 1];
+        }
+        List<String> list = Arrays.asList(ex.split("caused by: \n"));
+        List<String> subList = Arrays.asList(list.get(list.size() - 1).split("\n"));
+        String regex = (".*([a-zA-Z]*[.][a-zA-Z]*[:][0-9]*\\))");
+
+        StringBuilder rootCause = new StringBuilder();
+        for (String term : subList) {
+            if (Pattern.matches(regex, term)) break;
+            if (rootCause.toString().isEmpty() || rootCause.toString().startsWith(" ")) {
+                rootCause.append(term);
+            }
+        }
+        return rootCause.toString();
+    }
+
+}

src/main/java/io/kyligence/notebook/console/util/ExceptionUtils.java

@@ -1,9 +1,5 @@
 package io.kyligence.notebook.console.util;
 
-import java.util.Arrays;
-import java.util.List;
-import java.util.regex.Pattern;
-
 public class ExceptionUtils {
 
     public static String getRootCause(Throwable t) {
@@ -23,24 +19,4 @@ public static String getRootCause(Throwable t) {
         }
         return origin.getMessage();
     }
-
-    public static String getRootCause(String ex) {
-        if (ex == null || ex.isEmpty()) {
-            return "";
-        }
-        List<String> list = Arrays.asList(ex.split("caused by: \n"));
-        List<String> subList = Arrays.asList(list.get(list.size() - 1).split("\n"));
-        String regex = (".*([a-zA-Z]*[.][a-zA-Z]*[:][0-9]*\\))");
-
-        StringBuilder rootCause = new StringBuilder();
-        for (String term : subList) {
-            if (Pattern.matches(regex, term)) break;
-            if (rootCause.toString().isEmpty() || rootCause.toString().startsWith(" ")) {
-                rootCause.append(term);
-            }
-        }
-        return rootCause.toString();
-    }
-
-
 }

src/test/java/io/kyligence/notebook/console/NotebookLauncherBaseTest.java

@@ -26,6 +26,7 @@
 import org.springframework.test.web.servlet.request.MockMvcRequestBuilders;
 import org.springframework.web.context.WebApplicationContext;
 import org.testcontainers.containers.MockServerContainer;
+import org.testcontainers.containers.MySQLContainer;
 import org.testcontainers.containers.Network;
 import org.testcontainers.utility.DockerImageName;
 
@@ -61,12 +62,17 @@ public class NotebookLauncherBaseTest extends BaseResourceLoader {
 
     public static final DockerImageName MOCKSERVER_IMAGE = DockerImageName.parse("jamesdbloom/mockserver:mockserver-5.5.4");
 
-    public static MockServerContainer mockServer = new MockServerContainer(MOCKSERVER_IMAGE)
+    public static final MockServerContainer mockServer = new MockServerContainer(MOCKSERVER_IMAGE)
             .withEnv("testcontainers.reuse.enable", "true")
             .withNetworkAliases("notebook-network")
             .withStartupAttempts(3)
             .withReuse(true);
 
+    public static final MySQLContainer mockDatabase = (MySQLContainer) new MySQLContainer("mysql:5.7")
+            .withDatabaseName("notebook")
+            .withUsername("root")
+            .withPassword("root");
+
     /*
       -DNOTEBOOK_HOME=/opt/projects/kyligence/byzer-notebook
       -Dspring.config.name=application,notebook
@@ -95,14 +101,19 @@ public class NotebookLauncherBaseTest extends BaseResourceLoader {
         System.setProperty("PROPERTIES_PATH", testClassPath);
 
         mockServer.start();
-
+        mockDatabase.start();
         client = new MockServerClient(mockServer.getHost(), mockServer.getServerPort());
         assertTrue("Mockserver running", client.isRunning());
 
         // used by !show et;
         client.when(request().withPath("/run/script").withMethod("POST")
         ).respond(response().withBody(etResponse));
-
+        System.setProperty("spring.datasource.url", mockDatabase.getJdbcUrl());
+        System.setProperty("spring.datasource.driver-class-name", mockDatabase.getDriverClassName());
+        System.setProperty("spring.datasource.username", mockDatabase.getUsername());
+        System.setProperty("spring.datasource.password", mockDatabase.getPassword());
+        System.setProperty("notebook.database.ip", mockDatabase.getHost());
+        System.setProperty("notebook.database.port", mockDatabase.getFirstMappedPort().toString());
         System.setProperty("notebook.mlsql.engine-url", String.format("http://%s:%s", mockServer.getHost(), mockServer.getServerPort()));
     }
 

src/test/java/io/kyligence/notebook/console/service/JobServiceTest.java

@@ -20,6 +20,7 @@
 import javax.annotation.PostConstruct;
 import java.sql.Timestamp;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 import static org.mockito.ArgumentMatchers.*;
@@ -226,18 +227,71 @@ public void testJobDone() {
         Assert.assertTrue(resp);
 
         Assert.assertEquals(JobInfo.JobStatus.SUCCESS, (long) jobService.findByJobId(testJobId3).getStatus());
+        String mockPythonCode = "#%python\n" +
+                "#%env=source activate ray1.8.0\n" +
+                "#%schema=st(field(info,string))\n" +
+                "#%input=data\n" +
+                "#%output=contentsTb\n" +
+                "#%dataMode=model\n" +
+                "\n" +
+                "from pyjava.api.mlsql import RayContext, PythonContext\n" +
+                "import pandas as pd\n" +
+                "\n" +
+                "context:PythonContext = context\n" +
+                "ray_context = RayContext.connect(globals(),None)\n" +
+                "data = ray_context.to_pandas()\n" +
+                "\n" +
+                "def test(d):\n" +
+                "    context.log_client.log_to_driver(d)\n" +
+                "    raise RuntimeError()\n" +
+                "for d in data.iterrows():\n" +
+                "    test(d)\n" +
+                "context.build_result([{\"info\":\"data\"}])";
+
+        insertJobInfo("testPythonError-jobId", JobInfo.JobStatus.RUNNING,
+                "testPythonError-user", mockPythonCode);
+        String mockMsg = "Job aborted due to stage failure: Task 0 in stage 34446.0 failed 1 times, most recent failure:" +
+                " Lost task 0.0 in stage 34446.0 (TID 181638) (b77ec5268d1a executor driver): " +
+                "org.apache.spark.SparkException: Traceback (most recent call last):\n" +
+                "File \"/opt/conda/envs/ray1.8.0/lib/python3.6/site-packages/pyjava/worker.py\", line 155, in main\n" +
+                "process()\n" +
+                "File \"/opt/conda/envs/ray1.8.0/lib/python3.6/site-packages/pyjava/worker.py\", line 132, in process\n" +
+                "exec(code, n_local, n_local)\n" +
+                "File \"<string>\", line 14, in <module>\n" +
+                "File \"<string>\", line 12, in test\n" +
+                "RuntimeError\n" +
+                "at tech.mlsql.arrow.python.runner.BasePythonRunner$ReaderIterator.handlePythonException(PythonRunner.scala:324)";
+        jobService.jobDone("testPythonError-jobId", JobInfo.JobStatus.FAILED, "",
+                mockMsg, new Timestamp(System.currentTimeMillis()));
+
+        Assert.assertTrue(jobService.findByJobId("testPythonError-jobId").getMsg().startsWith("File"));
+
+        String mockScript = "select * from ab as output;";
+        insertJobInfo("testByzerError-jobId", JobInfo.JobStatus.RUNNING, "testPythonError-user",
+                mockScript);
+        mockMsg = "Table or view not found: ab; line 1 pos 14;";
+        jobService.jobDone("testByzerError-jobId", JobInfo.JobStatus.FAILED, "",
+                mockMsg, new Timestamp(System.currentTimeMillis()));
+
+        Assert.assertEquals(jobService.findByJobId("testByzerError-jobId").getMsg(), mockMsg);
+
     }
 
     @Test
     public void testGetJobContent() {
-        when(mockJobInfoRepository.getContentByJobId(eq(testJobId))).thenReturn(testJobContent);
-        Assert.assertEquals(testJobContent, mockJobService.getJobContent(testJobId));
+        Assert.assertEquals(testJobContent, jobService.getJobContent(testJobId));
+        Assert.assertNull(jobService.getJobContent("not-exist-job"));
+
+        String archivedContent = "!show resource;";
+        String archivedId = "test-archived";
+        JobInfoArchive archived = new JobInfoArchive(archivedId, JobInfo.JobStatus.SUCCESS);
+        archived.setContent(archivedContent);
+        archived.setUser("testGetContent-user");
+        archived.setName("mock-test-archive");
+        jobInfoArchiveRepository.save(archived);
 
-        when(mockJobInfoRepository.getContentByJobId(eq(testJobId2))).thenReturn(null);
-        Assert.assertNull(mockJobService.getJobContent(testJobId2));
+        Assert.assertEquals(archivedContent, jobService.getJobContent(archivedId));
 
-        when(mockJobInfoArchiveRepository.getContentByJobId(eq(testJobId2))).thenReturn(testJobContent);
-        Assert.assertEquals(testJobContent, mockJobService.getJobContent(testJobId2));
     }
 
     @Test

src/test/java/io/kyligence/notebook/console/service/NotebookServiceTest.java

@@ -83,7 +83,6 @@ public void testSaveCell() {
         Assert.assertNotNull(savedCellInfo);
         Assert.assertEquals(Integer.valueOf(defaultMockNotebookId), savedCellInfo.getNotebookId());
         Assert.assertEquals(content, savedCellInfo.getContent());
-        Assert.assertEquals(currentTimeStamp, savedCellInfo.getUpdateTime().getTime());
     }
 
     @Test