From 8ab460f038c09caeafc3cbcd1267ac05ac980655 Mon Sep 17 00:00:00 2001
From: Lily Li <72165460+lilyli9@users.noreply.github.com>
Date: Fri, 24 Feb 2023 11:39:16 -0800
Subject: [PATCH] feat: add descriptive error msgs (#1150)

* feat: add descriptive error msgs

* fix: only get error msg if exec error

* chore: refactor get_datadoc_error_message

* chore: refactor helper func, del data_cell_id arg
---
 querybook/server/tasks/run_datadoc.py | 44 +++++++++++++++++++++++----
 querybook/server/tasks/run_query.py   |  5 ++-
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/querybook/server/tasks/run_datadoc.py b/querybook/server/tasks/run_datadoc.py
index 6d1d11f73..3d07affd6 100644
--- a/querybook/server/tasks/run_datadoc.py
+++ b/querybook/server/tasks/run_datadoc.py
@@ -1,8 +1,11 @@
 from celery import chain
 
-from app.db import DBSession
+from app.db import with_session, DBSession
 from app.flask_app import celery, socketio
 
+from const.db import (
+    description_length,
+)
 from const.query_execution import QueryExecutionStatus, QueryExecutionType
 from const.schedule import TaskRunStatus
 
@@ -98,7 +101,7 @@ def run_datadoc_with_config(
             tasks_to_run.append(
                 _start_query_execution_task.si(
                     **start_query_execution_kwargs,
-                    previous_query_status=QueryExecutionStatus.DONE.value,
+                    previous_query_result=(QueryExecutionStatus.DONE.value, 0),
                 )
                 if index == 0
                 else _start_query_execution_task.s(**start_query_execution_kwargs)
@@ -117,13 +120,14 @@ def run_datadoc_with_config(
 @celery.task(bind=True)
 def _start_query_execution_task(
     self,
-    previous_query_status,
+    previous_query_result,
     cell_id,
     query_execution_params,
     data_doc_id,
 ):
+    previous_query_status, previous_query_execution_id = previous_query_result
     if previous_query_status != QueryExecutionStatus.DONE.value:
-        raise Exception(GENERIC_QUERY_FAILURE_MSG)
+        raise Exception(get_datadoc_error_message(previous_query_execution_id))
 
     with DBSession() as session:
         query_execution = qe_logic.create_query_execution(
@@ -149,14 +153,42 @@ def _start_query_execution_task(
         return query_execution.id
 
 
+@with_session
+def get_datadoc_error_message(query_execution_id, session=None):
+    _, data_cell_id = qe_logic.get_datadoc_id_from_query_execution_id(
+        query_execution_id, session=session
+    )[0]
+    data_cell_name = datadoc_logic.get_data_cell_by_id(
+        data_cell_id, session=session
+    ).meta.get("title", f"Untitled Cell Id [{data_cell_id}]")
+    query_execution_error = qe_logic.get_query_execution_error(
+        query_execution_id, session=session
+    )
+    query_execution_error_message = (
+        query_execution_error.error_message_extracted
+        if query_execution_error.error_message_extracted
+        else query_execution_error.error_message
+    )
+    error_msg = (
+        f'Failure in "{data_cell_name}": {query_execution_error_message}'
+        if query_execution_error_message is not None
+        else GENERIC_QUERY_FAILURE_MSG
+    )[:description_length]
+    return error_msg
+
+
 @celery.task
 def on_datadoc_run_success(
-    last_query_status,
+    last_query_result,
     completion_params,
     **kwargs,
 ):
+    last_query_status, last_query_execution_id = last_query_result
+
     is_success = last_query_status == QueryExecutionStatus.DONE.value
-    error_msg = None if is_success else GENERIC_QUERY_FAILURE_MSG
+    error_msg = (
+        None if is_success else get_datadoc_error_message(last_query_execution_id)
+    )
 
     return on_datadoc_completion(
         is_success=is_success, error_msg=error_msg, **completion_params
diff --git a/querybook/server/tasks/run_query.py b/querybook/server/tasks/run_query.py
index 9c7ff38db..f8d7531bc 100644
--- a/querybook/server/tasks/run_query.py
+++ b/querybook/server/tasks/run_query.py
@@ -68,7 +68,10 @@ def run_query_task(
             if executor and query_execution_status == QueryExecutionStatus.DONE:
                 log_query_per_table_task.delay(query_execution_id)
 
-    return query_execution_status.value if executor is not None else None
+    return (
+        query_execution_status.value if executor is not None else None,
+        query_execution_id,
+    )
 
 
 def run_executor_until_finish(celery_task, executor):