[superset] branch master updated: fix: better logic to extract errors on databricks (#22792)

[el...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

elizabeth pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new d091a68909 fix: better logic to extract errors on databricks (#22792)
d091a68909 is described below

commit d091a6890996997080c7a1d10e2937157393d8ac
Author: Elizabeth Thompson <es...@gmail.com>
AuthorDate: Fri Jan 20 15:35:09 2023 -0800

    fix: better logic to extract errors on databricks (#22792)
---
 superset/db_engine_specs/databricks.py             | 13 +++--
 .../unit_tests/db_engine_specs/test_databricks.py  | 56 ++++++++++++++++++++++
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py
index 98fc44e77f..df82b6d2f9 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -218,12 +218,15 @@ class DatabricksNativeEngineSpec(DatabricksODBCEngineSpec, BasicParametersMixin)
         raw_message = cls._extract_error_message(ex)
 
         context = context or {}
+        # access_token isn't currently parseable from the
+        # databricks error response, but adding it in here
+        # for reference if their error message changes
         context = {
-            "host": context["hostname"],
-            "access_token": context["password"],
-            "port": context["port"],
-            "username": context["username"],
-            "database": context["database"],
+            "host": context.get("hostname"),
+            "access_token": context.get("password"),
+            "port": context.get("port"),
+            "username": context.get("username"),
+            "database": context.get("database"),
         }
         for regex, (message, error_type, extra) in cls.custom_errors.items():
             match = regex.search(raw_message)
diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py b/tests/unit_tests/db_engine_specs/test_databricks.py
index 1962f4af36..86ffbc613c 100644
--- a/tests/unit_tests/db_engine_specs/test_databricks.py
+++ b/tests/unit_tests/db_engine_specs/test_databricks.py
@@ -18,8 +18,11 @@
 
 import json
 
+import pytest
 from pytest_mock import MockerFixture
 
+from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.utils.core import GenericDataType
 
 
@@ -197,3 +200,56 @@ def test_get_extra_params(mocker: MockerFixture) -> None:
             }
         }
     }
+
+
+def test_extract_errors() -> None:
+    """
+    Test that custom error messages are extracted correctly.
+    """
+
+    msg = ": mismatched input 'fromm'. Expecting: "
+    result = DatabricksNativeEngineSpec.extract_errors(Exception(msg))
+
+    assert result == [
+        SupersetError(
+            message=": mismatched input 'fromm'. Expecting: ",
+            error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
+            level=ErrorLevel.ERROR,
+            extra={
+                "engine_name": "Databricks",
+                "issue_codes": [
+                    {
+                        "code": 1002,
+                        "message": "Issue 1002 - The database returned an unexpected error.",
+                    }
+                ],
+            },
+        )
+    ]
+
+
+def test_extract_errors_with_context() -> None:
+    """
+    Test that custom error messages are extracted correctly with context.
+    """
+
+    msg = ": mismatched input 'fromm'. Expecting: "
+    context = {"hostname": "foo"}
+    result = DatabricksNativeEngineSpec.extract_errors(Exception(msg), context)
+
+    assert result == [
+        SupersetError(
+            message=": mismatched input 'fromm'. Expecting: ",
+            error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
+            level=ErrorLevel.ERROR,
+            extra={
+                "engine_name": "Databricks",
+                "issue_codes": [
+                    {
+                        "code": 1002,
+                        "message": "Issue 1002 - The database returned an unexpected error.",
+                    }
+                ],
+            },
+        )
+    ]