From 6fd40a3025e21e9457e9d4f8510897a6f5e2b77d Mon Sep 17 00:00:00 2001
From: Elizabeth Thompson <eschutho@gmail.com>
Date: Thu, 19 Jan 2023 11:48:36 -0800
Subject: [PATCH] fix logic to extract errors on databricks

---
 superset/db_engine_specs/databricks.py        | 13 +++--
 .../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 98fc44e77fc10..df82b6d2f93e9 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -218,12 +218,15 @@ def extract_errors(
         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 1962f4af36c0d..86ffbc613c110 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.",
+                    }
+                ],
+            },
+        )
+    ]