diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py index a8a821d586..9b438ee408 100644 --- a/src/databricks/labs/ucx/cli.py +++ b/src/databricks/labs/ucx/cli.py @@ -98,6 +98,19 @@ def skip(w: WorkspaceClient, schema: str | None = None, table: str | None = None return ctx.table_mapping.skip_schema(schema) +@ucx.command +def unskip(w: WorkspaceClient, schema: str | None = None, table: str | None = None): + """Create a unskip comment on a schema or a table""" + logger.info("Running unskip command") + if not schema: + logger.error("--schema is a required parameter.") + return None + ctx = WorkspaceContext(w) + if table: + return ctx.table_mapping.unskip_table_or_view(schema, table, ctx.tables_crawler.load_one) + return ctx.table_mapping.unskip_schema(schema) + + @ucx.command(is_account=True) def sync_workspace_info(a: AccountClient): """upload workspace config to all workspaces in the account where ucx is installed""" diff --git a/src/databricks/labs/ucx/hive_metastore/mapping.py b/src/databricks/labs/ucx/hive_metastore/mapping.py index 655172e388..94f33a97fe 100644 --- a/src/databricks/labs/ucx/hive_metastore/mapping.py +++ b/src/databricks/labs/ucx/hive_metastore/mapping.py @@ -135,6 +135,25 @@ def skip_table_or_view(self, schema_name: str, table_name: str, load_table: Call except BadRequest as err: logger.error(f"Failed to apply skip marker for Table {schema_name}.{table_name}: {err!s}", exc_info=True) + def unskip_table_or_view(self, schema_name: str, table_name: str, load_table: Callable[[str, str], Table | None]): + # Removes skip mark from the table property + try: + table = load_table(schema_name, table_name) + if table is None: + raise NotFound("[TABLE_OR_VIEW_NOT_FOUND]") + self._sql_backend.execute( + f"ALTER {table.kind} {escape_sql_identifier(schema_name)}.{escape_sql_identifier(table_name)} UNSET TBLPROPERTIES IF EXISTS('{self.UCX_SKIP_PROPERTY}' );" + ) + except NotFound as err: + if "[TABLE_OR_VIEW_NOT_FOUND]" in str(err) or "[DELTA_TABLE_NOT_FOUND]" in str(err): + logger.error(f"Failed to apply skip marker for Table {schema_name}.{table_name}. Table not found.") + else: + logger.error( + f"Failed to apply skip marker for Table {schema_name}.{table_name}: {err!s}", exc_info=True + ) + except BadRequest as err: + logger.error(f"Failed to apply skip marker for Table {schema_name}.{table_name}: {err!s}", exc_info=True) + def skip_schema(self, schema: str): # Marks a schema to be skipped in the migration process by applying a table property try: @@ -149,6 +168,20 @@ def skip_schema(self, schema: str): except BadRequest as err: logger.error(err) + def unskip_schema(self, schema: str): + # Removes skip mark from the schema property + try: + self._sql_backend.execute( + f"ALTER SCHEMA {escape_sql_identifier(schema)} UNSET DBPROPERTIES IF EXISTS('{self.UCX_SKIP_PROPERTY}');" + ) + except NotFound as err: + if "[SCHEMA_NOT_FOUND]" in str(err): + logger.error(f"Failed to remove skip marker for Schema {schema}. Schema not found.") + else: + logger.error(err) + except BadRequest as err: + logger.error(err) + def get_tables_to_migrate(self, tables_crawler: TablesCrawler) -> Collection[TableToMigrate]: rules = self.load() # Getting all the source tables from the rules