From f6d9841599cb8577b7258f44bdc2a288d2e0e867 Mon Sep 17 00:00:00 2001
From: Liran Bareket <lbareket@gmail.com>
Date: Mon, 27 Jan 2025 02:49:47 -0500
Subject: [PATCH] Fixed Skip/Unskip schema functionality  (#3567)

Closes #3494
---
 src/databricks/labs/ucx/hive_metastore/mapping.py | 9 +++++----
 tests/unit/hive_metastore/test_mapping.py         | 6 ++++--
 tests/unit/test_cli.py                            | 4 ++--
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/databricks/labs/ucx/hive_metastore/mapping.py b/src/databricks/labs/ucx/hive_metastore/mapping.py
index bef402d73d..78237ea507 100644
--- a/src/databricks/labs/ucx/hive_metastore/mapping.py
+++ b/src/databricks/labs/ucx/hive_metastore/mapping.py
@@ -172,7 +172,7 @@ def skip_schema(self, schema: str):
         # Marks a schema to be skipped in the migration process by applying a table property
         try:
             self._sql_backend.execute(
-                f"ALTER SCHEMA {escape_sql_identifier(schema)} SET DBPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)"
+                f"ALTER SCHEMA hive_metastore.{escape_sql_identifier(schema)} SET DBPROPERTIES('{self.UCX_SKIP_PROPERTY}' = true)"
             )
         except NotFound as err:
             if "[SCHEMA_NOT_FOUND]" in str(err):
@@ -190,7 +190,7 @@ def unskip_schema(self, schema: str) -> None:
         """
         try:
             self._sql_backend.execute(
-                f"ALTER SCHEMA hive_metastore.{escape_sql_identifier(schema)} UNSET DBPROPERTIES IF EXISTS('{self.UCX_SKIP_PROPERTY}');"
+                f"ALTER SCHEMA hive_metastore.{escape_sql_identifier(schema)} SET DBPROPERTIES('{self.UCX_SKIP_PROPERTY}' = false);"
             )
         except (NotFound, BadRequest) as e:
             logger.error(f"Failed to remove skip marker from schema: {schema}.", exc_info=e)
@@ -251,7 +251,8 @@ def _get_database_in_scope_task(self, database: str) -> str | None:
         properties = describe.get("Properties", "")
         if not properties:
             return database
-        if self.UCX_SKIP_PROPERTY in TablesCrawler.parse_database_props(properties.lower()):
+        tbl_props = TablesCrawler.parse_database_props(properties.lower())
+        if tbl_props.get(self.UCX_SKIP_PROPERTY, "false") == "true":
             logger.info(f"Database {database} is marked to be skipped")
             return None
         return database
@@ -268,7 +269,7 @@ def _get_table_in_scope_task(self, table_to_migrate: TableToMigrate, check_uc_ta
             return None
 
         for value in properties:
-            if value["key"] == self.UCX_SKIP_PROPERTY:
+            if value["key"] == self.UCX_SKIP_PROPERTY and value["value"] == "true":
                 logger.info(f"{table.key} is marked to be skipped")
                 return None
             if value["key"] == "upgraded_to":
diff --git a/tests/unit/hive_metastore/test_mapping.py b/tests/unit/hive_metastore/test_mapping.py
index f07399287d..c9751c6561 100644
--- a/tests/unit/hive_metastore/test_mapping.py
+++ b/tests/unit/hive_metastore/test_mapping.py
@@ -220,7 +220,9 @@ def test_skip_happy_path(caplog):
     )
     assert len(caplog.records) == 0
     mapping.skip_schema(schema="schema")
-    sbe.execute.assert_called_with(f"ALTER SCHEMA `schema` SET DBPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = true)")
+    sbe.execute.assert_called_with(
+        f"ALTER SCHEMA hive_metastore.`schema` SET DBPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = true)"
+    )
     assert len(caplog.records) == 0
 
 
@@ -262,7 +264,7 @@ def test_unskip_on_schema() -> None:
     mapping.unskip_schema(schema="schema")
     ws.tables.get.assert_not_called()
     assert (
-        f"ALTER SCHEMA hive_metastore.`schema` UNSET DBPROPERTIES IF EXISTS('{mapping.UCX_SKIP_PROPERTY}');"
+        f"ALTER SCHEMA hive_metastore.`schema` SET DBPROPERTIES('{mapping.UCX_SKIP_PROPERTY}' = false);"
         in mock_backend.queries
     )
 
diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py
index d6574590a1..969a65d8a8 100644
--- a/tests/unit/test_cli.py
+++ b/tests/unit/test_cli.py
@@ -204,7 +204,7 @@ def test_skip_with_schema(ws) -> None:
 
     ws.statement_execution.execute_statement.assert_called_with(
         warehouse_id='test',
-        statement="ALTER SCHEMA `schema` SET DBPROPERTIES('databricks.labs.ucx.skip' = true)",
+        statement="ALTER SCHEMA hive_metastore.`schema` SET DBPROPERTIES('databricks.labs.ucx.skip' = true)",
         byte_limit=None,
         catalog=None,
         schema=None,
@@ -261,7 +261,7 @@ def test_unskip_with_schema(ws) -> None:
 
     ws.statement_execution.execute_statement.assert_called_with(
         warehouse_id='test',
-        statement="ALTER SCHEMA hive_metastore.`schema` UNSET DBPROPERTIES IF EXISTS('databricks.labs.ucx.skip');",
+        statement="ALTER SCHEMA hive_metastore.`schema` SET DBPROPERTIES('databricks.labs.ucx.skip' = false);",
         byte_limit=None,
         catalog=None,
         schema=None,