Converted RuntimeBackend query executions to SDK exceptions (#769)

Convert the query execution exceptions to SDK exceptions like NotFound,
PermissionDenied and SyntaxError (currently used BadRequest to map
syntax errors).

Integration and unit tests added.

Fixes #726

---------

Co-authored-by: Serge Smertin <[email protected]>

Author: pritishpai

src/databricks/labs/ucx/framework/crawlers.py

@@ -8,7 +8,7 @@
 from typing import Any, ClassVar, Generic, Protocol, TypeVar
 
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import NotFound
+from databricks.sdk.errors import BadRequest, NotFound, PermissionDenied, Unknown
 
 from databricks.labs.ucx.mixins.sql import Row, StatementExecutionExt
 
@@ -151,11 +151,21 @@ def __init__(self):
 
     def execute(self, sql):
         logger.debug(f"[spark][execute] {sql}")
-        self._spark.sql(sql)
+        try:
+            immediate_response = self._spark.sql(sql)
+        except Exception as e:
+            error_message = str(e)
+            self._raise_spark_sql_exceptions(error_message)
+        return immediate_response
 
     def fetch(self, sql) -> Iterator[Row]:
         logger.debug(f"[spark][fetch] {sql}")
-        return self._spark.sql(sql).collect()
+        try:
+            fetch_query_response = self._spark.sql(sql).collect()
+        except Exception as e:
+            error_message = str(e)
+            self._raise_spark_sql_exceptions(error_message)
+        return fetch_query_response
 
     def save_table(self, full_name: str, rows: Sequence[DataclassInstance], klass: Dataclass, mode: str = "append"):
         rows = self._filter_none_rows(rows, klass)
@@ -167,6 +177,19 @@ def save_table(self, full_name: str, rows: Sequence[DataclassInstance], klass: D
         df = self._spark.createDataFrame(rows, self._schema_for(klass))
         df.write.saveAsTable(full_name, mode=mode)
 
+    @staticmethod
+    def _raise_spark_sql_exceptions(error_message: str):
+        if "SCHEMA_NOT_FOUND" in error_message:
+            raise NotFound(error_message) from None
+        elif "TABLE_OR_VIEW_NOT_FOUND" in error_message:
+            raise NotFound(error_message) from None
+        elif "PARSE_SYNTAX_ERROR" in error_message:
+            raise BadRequest(error_message) from None
+        elif "Operation not allowed" in error_message:
+            raise PermissionDenied(error_message) from None
+        else:
+            raise Unknown(error_message) from None
+
 
 class CrawlerBase(Generic[Result]):
     def __init__(self, backend: SqlBackend, catalog: str, schema: str, table: str, klass: type[Result]):

tests/integration/framework/test_crawlers.py

@@ -1,3 +1,5 @@
+from databricks.labs.blueprint.commands import CommandExecutor
+
 from databricks.labs.ucx.framework.crawlers import SchemaDeployer
 from databricks.labs.ucx.hive_metastore.grants import Grant
 
@@ -9,3 +11,166 @@ def test_deploys_database(sql_backend, inventory_schema):
     deployer.deploy_schema()
     deployer.deploy_table("grants", Grant)
     deployer.deploy_view("grant_detail", "queries/views/grant_detail.sql")
+
+
+def test_runtime_backend_incorrect_schema_and_table_handled(ws, wsfs_wheel, make_random):
+    commands = CommandExecutor(ws.clusters, ws.command_execution, lambda: ws.config.cluster_id)
+
+    commands.install_notebook_library(f"/Workspace{wsfs_wheel}")
+    query_response_incorrect_schema_execute = commands.run(
+        f"""
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import NotFound
+        backend = RuntimeBackend()
+        try:
+            backend.execute("USE {make_random()}")
+            return "FAILED"
+        except NotFound as e:
+            return "PASSED"
+        """
+    )
+    assert query_response_incorrect_schema_execute == "PASSED"
+
+    query_response_incorrect_table_execute = commands.run(
+        f"""
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import NotFound
+        backend = RuntimeBackend()
+        try:
+            backend.execute("SELECT * FROM default.{make_random()}")
+            return "FAILED"
+        except NotFound as e:
+            return "PASSED"
+        """
+    )
+    assert query_response_incorrect_table_execute == "PASSED"
+
+    query_response_incorrect_schema_fetch = commands.run(
+        f"""
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import NotFound
+        backend = RuntimeBackend()
+        try:
+            query_response = backend.fetch("DESCRIBE {make_random()}")
+            return "FAILED"
+        except NotFound as e:
+            return "PASSED"
+        """
+    )
+    assert query_response_incorrect_schema_fetch == "PASSED"
+
+    query_response_incorrect_table_fetch = commands.run(
+        f"""
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import NotFound
+        backend = RuntimeBackend()
+        try:
+            query_response = backend.fetch("SELECT * FROM default.{make_random()}")
+            return "FAILED"
+        except NotFound as e:
+            return "PASSED"
+        """
+    )
+    assert query_response_incorrect_table_fetch == "PASSED"
+
+
+def test_runtime_backend_incorrect_syntax_handled(ws, wsfs_wheel):
+    commands = CommandExecutor(ws.clusters, ws.command_execution, lambda: ws.config.cluster_id)
+
+    commands.install_notebook_library(f"/Workspace{wsfs_wheel}")
+    query_response_incorrect_syntax_execute = commands.run(
+        """
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import BadRequest
+        backend = RuntimeBackend()
+        try:
+            backend.execute("SHWO DTABASES")
+            return "FAILED"
+        except BadRequest:
+            return "PASSED"
+        """
+    )
+    assert query_response_incorrect_syntax_execute == "PASSED"
+
+    query_response_incorrect_syntax_fetch = commands.run(
+        """
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import BadRequest
+        backend = RuntimeBackend()
+        try:
+            query_response = backend.fetch("SHWO DTABASES")
+            return "FAILED"
+        except BadRequest:
+            return "PASSED"
+        """
+    )
+    assert query_response_incorrect_syntax_fetch == "PASSED"
+
+
+def test_runtime_backend_permission_denied_handled(ws, wsfs_wheel):
+    commands = CommandExecutor(ws.clusters, ws.command_execution, lambda: ws.config.cluster_id)
+
+    commands.install_notebook_library(f"/Workspace{wsfs_wheel}")
+    query_response_permission_denied_execute = commands.run(
+        """
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import PermissionDenied
+        backend = RuntimeBackend()
+        try:
+            current_user = backend.fetch("SELECT current_user()")
+            backend.execute(f"GRANT CREATE EXTERNAL LOCATION ON METASTORE TO {current_user}")
+            return "FAILED"
+        except PermissionDenied:
+            return "PASSED"
+        """
+    )
+    assert query_response_permission_denied_execute == "PASSED"
+
+    query_response_permission_denied_fetch = commands.run(
+        """
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import PermissionDenied
+        backend = RuntimeBackend()
+        try:
+            current_user = backend.fetch(f"SELECT current_user()")
+            grants = backend.fetch(f"GRANT CREATE EXTERNAL LOCATION ON METASTORE TO {current_user}")
+            return "FAILED"
+        except PermissionDenied:
+            return "PASSED"
+        """
+    )
+    assert query_response_permission_denied_fetch == "PASSED"
+
+
+def test_runtime_backend_unknown_error_handled(ws, wsfs_wheel):
+    commands = CommandExecutor(ws.clusters, ws.command_execution, lambda: ws.config.cluster_id)
+
+    commands.install_notebook_library(f"/Workspace{wsfs_wheel}")
+
+    query_response_unknown_execute = commands.run(
+        """
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import Unknown
+        backend = RuntimeBackend()
+        try:
+            backend.execute("SHOW GRANTS ON METASTORE")
+            return "FAILED"
+        except Unknown:
+            return "PASSED"
+        """
+    )
+    assert query_response_unknown_execute == "PASSED"
+
+    query_response_unknown_fetch = commands.run(
+        """
+        from databricks.labs.ucx.framework.crawlers import RuntimeBackend
+        from databricks.sdk.errors import Unknown
+        backend = RuntimeBackend()
+        try:
+            grants = backend.fetch("SHOW GRANTS ON METASTORE")
+            return "FAILED"
+        except Unknown:
+            return "PASSED"
+        """
+    )
+    assert query_response_unknown_fetch == "PASSED"

tests/unit/framework/test_crawlers.py

@@ -4,7 +4,7 @@
 from unittest import mock
 
 import pytest
-from databricks.sdk.errors import NotFound
+from databricks.sdk.errors import BadRequest, NotFound, PermissionDenied, Unknown
 from databricks.sdk.service import sql
 
 from databricks.labs.ucx.framework.crawlers import (
@@ -244,3 +244,108 @@ def test_save_table_with_not_null_constraint_violated(mocker):
         assert (
             str(exc_info.value) == "Not null constraint violated for column key, row = {'key': None, 'value': 'value'}"
         )
+
+
+def test_raise_spark_sql_exceptions(mocker):
+    with mock.patch.dict(os.environ, {"DATABRICKS_RUNTIME_VERSION": "14.0"}):
+        pyspark_sql_session = mocker.Mock()
+        sys.modules["pyspark.sql.session"] = pyspark_sql_session
+
+        rb = RuntimeBackend
+        error_message_invalid_schema = "SCHEMA_NOT_FOUND foo schema does not exist"
+        with pytest.raises(NotFound):
+            rb._raise_spark_sql_exceptions(error_message_invalid_schema)
+
+        error_message_invalid_table = "TABLE_OR_VIEW_NOT_FOUND foo table does not exist"
+        with pytest.raises(NotFound):
+            rb._raise_spark_sql_exceptions(error_message_invalid_table)
+
+        error_message_invalid_syntax = "PARSE_SYNTAX_ERROR foo"
+        with pytest.raises(BadRequest):
+            rb._raise_spark_sql_exceptions(error_message_invalid_syntax)
+
+        error_message_permission_denied = "foo Operation not allowed"
+        with pytest.raises(PermissionDenied):
+            rb._raise_spark_sql_exceptions(error_message_permission_denied)
+
+        error_message_invalid_schema = "foo error failure"
+        with pytest.raises(Unknown):
+            rb._raise_spark_sql_exceptions(error_message_invalid_schema)
+
+
+def test_execute(mocker):
+    with mock.patch.dict(os.environ, {"DATABRICKS_RUNTIME_VERSION": "14.0"}):
+        pyspark_sql_session = mocker.Mock()
+        sys.modules["pyspark.sql.session"] = pyspark_sql_session
+        rb = RuntimeBackend()
+
+        sql_query = "SELECT * from bar"
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "SCHEMA_NOT_FOUND"
+        )
+        with pytest.raises(NotFound):
+            rb.execute(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "TABLE_OR_VIEW_NOT_FOUND"
+        )
+        with pytest.raises(NotFound):
+            rb.execute(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "PARSE_SYNTAX_ERROR"
+        )
+        with pytest.raises(BadRequest):
+            rb.execute(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "Operation not allowed"
+        )
+        with pytest.raises(PermissionDenied):
+            rb.execute(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "foo error occurred"
+        )
+        with pytest.raises(Unknown):
+            rb.execute(sql_query)
+
+
+def test_fetch(mocker):
+    with mock.patch.dict(os.environ, {"DATABRICKS_RUNTIME_VERSION": "14.0"}):
+        pyspark_sql_session = mocker.Mock()
+        sys.modules["pyspark.sql.session"] = pyspark_sql_session
+        rb = RuntimeBackend()
+
+        sql_query = "SELECT * from bar"
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "SCHEMA_NOT_FOUND"
+        )
+        with pytest.raises(NotFound):
+            rb.fetch(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "TABLE_OR_VIEW_NOT_FOUND"
+        )
+        with pytest.raises(NotFound):
+            rb.fetch(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "PARSE_SYNTAX_ERROR"
+        )
+        with pytest.raises(BadRequest):
+            rb.fetch(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "Operation not allowed"
+        )
+        with pytest.raises(PermissionDenied):
+            rb.fetch(sql_query)
+
+        pyspark_sql_session.SparkSession.builder.getOrCreate.return_value.sql.side_effect = Exception(
+            "foo error occurred"
+        )
+        with pytest.raises(Unknown):
+            rb.fetch(sql_query)