fix: show custom errors in SQL Lab (#14959)

* fix: show custom errors in SQL Lab

* Fix lint

* Fix test

* Update superset/views/base.py

Co-authored-by: ʈᵃᵢ <tai@apache.org>

* Fix lint

* Debug failing test

* Remove print()

* Debug flaky tests

* Fix test

Co-authored-by: ʈᵃᵢ <tai@apache.org>

superset/db_engine_specs/postgres.py

@@ -80,6 +80,10 @@ CONNECTION_HOST_DOWN_REGEX = re.compile(
 CONNECTION_UNKNOWN_DATABASE_REGEX = re.compile(
     'database "(?P<database>.*?)" does not exist'
 )
+COLUMN_DOES_NOT_EXIST_REGEX = re.compile(
+    r'postgresql error: column "(?P<column_name>.+?)" '
+    r"does not exist\s+LINE (?P<location>\d+?)"
+)
 
 
 class PostgresBaseEngineSpec(BaseEngineSpec):
@@ -100,7 +104,7 @@ class PostgresBaseEngineSpec(BaseEngineSpec):
         "P1Y": "DATE_TRUNC('year', {col})",
     }
 
-    custom_errors = {
+    custom_errors: Dict[Pattern[str], Tuple[str, SupersetErrorType, Dict[str, Any]]] = {
         CONNECTION_INVALID_USERNAME_REGEX: (
             __('The username "%(username)s" does not exist.'),
             SupersetErrorType.CONNECTION_INVALID_USERNAME_ERROR,
@@ -139,6 +143,14 @@ class PostgresBaseEngineSpec(BaseEngineSpec):
             SupersetErrorType.CONNECTION_UNKNOWN_DATABASE_ERROR,
             {"invalid": ["database"]},
         ),
+        COLUMN_DOES_NOT_EXIST_REGEX: (
+            __(
+                'We can\'t seem to resolve the column "%(column_name)s" at '
+                "line %(location)s.",
+            ),
+            SupersetErrorType.COLUMN_DOES_NOT_EXIST_ERROR,
+            {},
+        ),
     }
 
     @classmethod

superset/views/base.py

@@ -207,6 +207,9 @@ def handle_api_exception(
             return json_errors_response(
                 errors=[ex.error], status=ex.status, payload=ex.payload
             )
+        except SupersetErrorsException as ex:
+            logger.warning(ex, exc_info=True)
+            return json_errors_response(errors=ex.errors, status=ex.status)
         except SupersetErrorException as ex:
             logger.warning(ex)
             return json_errors_response(errors=[ex.error], status=ex.status)

superset/views/core.py

@@ -78,6 +78,7 @@ from superset.exceptions import (
     CertificateException,
     DatabaseNotFound,
     SerializationError,
+    SupersetErrorsException,
     SupersetException,
     SupersetGenericDBErrorException,
     SupersetSecurityException,
@@ -2426,7 +2427,15 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
             raise SupersetGenericDBErrorException(utils.error_msg_from_exception(ex))
 
         if data.get("status") == QueryStatus.FAILED:
-            raise SupersetGenericDBErrorException(data["error"])
+            msg = data["error"]
+            query = _session.query(Query).filter_by(id=query_id).one()
+            database = query.database
+            db_engine_spec = database.db_engine_spec
+            errors = db_engine_spec.extract_errors(msg)
+            _session.close()
+            if errors:
+                raise SupersetErrorsException(errors)
+            raise SupersetGenericDBErrorException(msg)
         return json_success(payload)
 
     @has_access_api

tests/celery_tests.py

@@ -141,6 +141,8 @@ def get_select_star(table: str, schema: Optional[str] = None):
 
 @pytest.mark.parametrize("ctas_method", [CtasMethod.TABLE, CtasMethod.VIEW])
 def test_run_sync_query_dont_exist(setup_sqllab, ctas_method):
+    examples_db = get_example_database()
+    engine_name = examples_db.db_engine_spec.engine_name
     sql_dont_exist = "SELECT name FROM table_dont_exist"
     result = run_sql(sql_dont_exist, cta=True, ctas_method=ctas_method)
     if backend() == "sqlite" and ctas_method == CtasMethod.VIEW:
@@ -157,7 +159,8 @@ def test_run_sync_query_dont_exist(setup_sqllab, ctas_method):
                     "code": 1002,
                     "message": "Issue 1002 - The database returned an unexpected error.",
                 }
-            ]
+            ],
+            "engine_name": engine_name,
         }
 
 

tests/sqllab_tests.py

@@ -75,6 +75,9 @@ class TestSqlLab(SupersetTestCase):
 
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_sql_json(self):
+        examples_db = get_example_database()
+        engine_name = examples_db.db_engine_spec.engine_name
+
         self.login("admin")
 
         data = self.run_sql("SELECT * FROM birth_names LIMIT 10", "1")
@@ -91,7 +94,8 @@ class TestSqlLab(SupersetTestCase):
                     "code": 1002,
                     "message": "Issue 1002 - The database returned an unexpected error.",
                 }
-            ]
+            ],
+            "engine_name": engine_name,
         }
 
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")