feat: add error message to SQL Exception

emdgroup · Oct 25, 2024 · 25e5e2b · 25e5e2b
1 parent 4c5ad68
commit 25e5e2b
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 1 deletion.
diff --git a/libs/foundry-dev-tools/src/foundry_dev_tools/errors/meta.py b/libs/foundry-dev-tools/src/foundry_dev_tools/errors/meta.py
@@ -59,7 +59,9 @@ def __init__(self, response: requests.Response | None = None, info: str | None =
         term_size = shutil.get_terminal_size().columns
         msg = self.message
         if self.info:
-            msg += "\n{self.info}"
+            msg += f"\n{self.info}\n"
+        else:
+            msg += "\n"
         request_sep = "-" * int((term_size - 7) / 2)
         msg += request_sep + "REQUEST" + request_sep + "\n"
         if self.response is not None:

diff --git a/libs/foundry-dev-tools/src/foundry_dev_tools/errors/sql.py b/libs/foundry-dev-tools/src/foundry_dev_tools/errors/sql.py
@@ -2,14 +2,23 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from foundry_dev_tools.errors.meta import FoundryAPIError
 
+if TYPE_CHECKING:
+    import requests
+
 
 class FoundrySqlQueryFailedError(FoundryAPIError):
     """Exception is thrown when SQL Query execution failed."""
 
     message = "Foundry SQL Query Failed."
 
+    def __init__(self, response: requests.Response):
+        self.error_message = response.json().get("status", {}).get("failed", {}).get("errorMessage", "")
+        super().__init__(response=response, info=self.error_message)
+
 
 class FoundrySqlQueryClientTimedOutError(FoundryAPIError):
     """Exception is thrown when the Query execution time exceeded the client timeout value."""

diff --git a/tests/unit/clients/test_foundry_sql_server.py b/tests/unit/clients/test_foundry_sql_server.py
@@ -5,6 +5,7 @@
 
 from foundry_dev_tools.errors.sql import (
     FoundrySqlQueryClientTimedOutError,
+    FoundrySqlQueryFailedError,
     FoundrySqlSerializationFormatNotImplementedError,
 )
 from foundry_dev_tools.utils.clients import build_api_url
@@ -60,3 +61,63 @@ def test_read_arrow_optional_polars(test_context_mock):
     pa_table = arrow_stream.read_all()
     df = pl.from_arrow(pa_table)
     assert df.shape == (1, 1)
+
+
+def test_exception(mocker, test_context_mock):
+    mocker.patch("time.sleep")  # we do not want to wait 500 ms in a test ;)
+
+    test_context_mock.mock_adapter.register_uri(
+        "POST",
+        build_api_url(TEST_HOST.url, "foundry-sql-server", "queries/execute"),
+        json={"queryId": "6941e90b-7a18-4ee9-8e82-63b64d7a7bde", "status": {"type": "running", "running": {}}},
+    )
+
+    test_context_mock.mock_adapter.register_uri(
+        "GET",
+        re.compile(re.escape(build_api_url(TEST_HOST.url, "foundry-sql-server", "")) + "queries/[^/]*/status"),
+        json={
+            "status": {
+                "type": "failed",
+                "failed": {
+                    "errorMessage": "org.apache.spark.sql.AnalysisException: ",
+                    "failureReason": "FAILED_TO_EXECUTE",
+                },
+            }
+        },
+    )
+    with pytest.raises(FoundrySqlQueryFailedError) as exception:
+        test_context_mock.foundry_sql_server.query_foundry_sql(
+            "SELECT notExistingColumn FROM `ri.foundry.main.dataset.249d2faf-4dcb-490b-9de5-3ae1160e8fbf`",
+            timeout=0.001,
+        )
+    assert exception.value.error_message == "org.apache.spark.sql.AnalysisException: "
+
+
+def test_exception_unknown_json(mocker, test_context_mock):
+    mocker.patch("time.sleep")  # we do not want to wait 500 ms in a test ;)
+
+    test_context_mock.mock_adapter.register_uri(
+        "POST",
+        build_api_url(TEST_HOST.url, "foundry-sql-server", "queries/execute"),
+        json={"queryId": "6941e90b-7a18-4ee9-8e82-63b64d7a7bde", "status": {"type": "running", "running": {}}},
+    )
+
+    test_context_mock.mock_adapter.register_uri(
+        "GET",
+        re.compile(re.escape(build_api_url(TEST_HOST.url, "foundry-sql-server", "")) + "queries/[^/]*/status"),
+        json={
+            "status": {
+                "type": "failed",
+                "unknownKey": {
+                    "errorMessage": "org.apache.spark.sql.AnalysisException: ",
+                    "failureReason": "FAILED_TO_EXECUTE",
+                },
+            }
+        },
+    )
+    with pytest.raises(FoundrySqlQueryFailedError) as exception:
+        test_context_mock.foundry_sql_server.query_foundry_sql(
+            "SELECT notExistingColumn FROM `ri.foundry.main.dataset.249d2faf-4dcb-490b-9de5-3ae1160e8fbf`",
+            timeout=0.001,
+        )
+    assert exception.value.error_message == ""