Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add unskip CLI command to undo a skip on schema or a table #2734

Merged
merged 19 commits into from
Sep 26, 2024
Merged
Show file tree
Hide file tree
Changes from 14 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ See [contributing instructions](CONTRIBUTING.md) to help improve this project.
* [`migrate-locations` command](#migrate-locations-command)
* [`create-table-mapping` command](#create-table-mapping-command)
* [`skip` command](#skip-command)
* [`unskip` command](#unskip-command)
* [`create-catalogs-schemas` command](#create-catalogs-schemas-command)
* [`migrate-tables` command](#migrate-tables-command)
* [`revert-migrated-tables` command](#revert-migrated-tables-command)
Expand Down Expand Up @@ -1422,6 +1423,15 @@ Once you're done with table migration, proceed to the [code migration](#code-mig

[[back to top](#databricks-labs-ucx)]

## `unskip` command

```commandline
databricks labs ucx unskip --schema X [--table Y]
```
This command removes the mark set by the [`skip` command](#skip-command) on the given schema or table.

[[back to top](#databricks-labs-ucx)]

## `create-catalogs-schemas` command

```text
Expand Down
2 changes: 1 addition & 1 deletion src/databricks/labs/ucx/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ def skip(w: WorkspaceClient, schema: str | None = None, table: str | None = None

@ucx.command
def unskip(w: WorkspaceClient, schema: str | None = None, table: str | None = None):
"""Create a unskip comment on a schema or a table"""
"""Unset the skip mark from a schema or a table"""
logger.info("Running unskip command")
if not schema:
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
logger.error("--schema is a required parameter.")
Expand Down
50 changes: 31 additions & 19 deletions src/databricks/labs/ucx/hive_metastore/mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,24 +135,32 @@ 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
def unskip_table_or_view(
self, schema_name: str, table_name: str, load_table: Callable[[str, str], Table | None]
) -> None:
"""Removes skip mark from the table property.

Args:
schema_name (str): The schema name of the table to be unskipped.
table_name (str): The table name of the table to be unskipped.
load_table (Callable[[str, str], Table | None]): A function that loads a table from the metastore.
"""
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}' );"
f"ALTER {table.kind} {escape_sql_identifier(table.full_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:
except NotFound as e:
if "[TABLE_OR_VIEW_NOT_FOUND]" in str(e) or "[DELTA_TABLE_NOT_FOUND]" in str(e):
logger.error(
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
f"Failed to apply skip marker for Table {schema_name}.{table_name}: {err!s}", exc_info=True
f"Failed to remove skip marker from table: {schema_name}.{table_name}. Table not found.", exc_info=e
)
except BadRequest as err:
logger.error(f"Failed to apply skip marker for Table {schema_name}.{table_name}: {err!s}", exc_info=True)
else:
logger.error(f"Failed to remove skip marker from table: {schema_name}.{table_name}", exc_info=e)
except BadRequest as e:
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
logger.error(f"Failed to remove skip marker from table: {schema_name}.{table_name}: {e!s}", exc_info=e)
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved

def skip_schema(self, schema: str):
# Marks a schema to be skipped in the migration process by applying a table property
Expand All @@ -168,19 +176,23 @@ 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
def unskip_schema(self, schema: str) -> None:
"""Removes skip mark from the schema property.

Args:
schema (str): The schema name of the table to be unskipped.
"""
try:
self._sql_backend.execute(
f"ALTER SCHEMA {escape_sql_identifier(schema)} UNSET DBPROPERTIES IF EXISTS('{self.UCX_SKIP_PROPERTY}');"
f"ALTER SCHEMA hive_metastore.{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.")
except NotFound as e:
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
if "[SCHEMA_NOT_FOUND]" in str(e):
logger.error(f"Failed to remove skip marker from schema: {schema}. Schema not found.", exc_info=e)
else:
logger.error(err)
except BadRequest as err:
logger.error(err)
logger.error(f"Failed to remove skip marker from schema: {schema}.", exc_info=e)
except BadRequest as e:
logger.error(f"Failed to remove skip marker from schema: {schema}.", exc_info=e)

def get_tables_to_migrate(self, tables_crawler: TablesCrawler) -> Collection[TableToMigrate]:
rules = self.load()
Expand Down
54 changes: 54 additions & 0 deletions tests/unit/hive_metastore/test_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,60 @@ def test_skip_happy_path(caplog):
assert len(caplog.records) == 0


def test_unskip_on_table():
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
ws = create_autospec(WorkspaceClient)
mock_backend = MockBackend()
installation = MockInstallation()
mapping = TableMapping(installation, ws, mock_backend)
table = Table(catalog="catalog", database="schema", name="table", object_type="table", table_format="csv")
mapping.unskip_table_or_view(schema_name="schema", table_name="table", load_table=lambda _schema, _table: table)
ws.tables.get.assert_not_called()
assert (
f"ALTER TABLE `catalog`.`schema`.`table` UNSET TBLPROPERTIES IF EXISTS('{mapping.UCX_SKIP_PROPERTY}');"
in mock_backend.queries
)


def test_unskip_on_view():
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
ws = create_autospec(WorkspaceClient)
mock_backend = MockBackend()
installation = MockInstallation()
mapping = TableMapping(installation, ws, mock_backend)
view = Table(
catalog="catalog", database="schema", name="view", object_type="table", table_format="csv", view_text="stuff"
)
mapping.unskip_table_or_view(schema_name="schema", table_name="view", load_table=lambda _schema, _table: view)
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
ws.tables.get.assert_not_called()
assert (
f"ALTER VIEW `catalog`.`schema`.`view` UNSET TBLPROPERTIES IF EXISTS('{mapping.UCX_SKIP_PROPERTY}');"
in mock_backend.queries
)


def test_unskip_on_schema():
aminmovahed-db marked this conversation as resolved.
Show resolved Hide resolved
ws = create_autospec(WorkspaceClient)
mock_backend = MockBackend()
installation = MockInstallation()
mapping = TableMapping(installation, ws, mock_backend)
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}');"
in mock_backend.queries
)


def test_unskip_missing_table(caplog):
ws = create_autospec(WorkspaceClient)
sbe = create_autospec(SqlBackend)
installation = MockInstallation()
sbe.execute.side_effect = NotFound("[TABLE_OR_VIEW_NOT_FOUND]")
mapping = TableMapping(installation, ws, sbe)
mapping.unskip_table_or_view(schema_name='foo', table_name="table", load_table=lambda schema, table: None)
ws.tables.get.assert_not_called()
assert [rec.message for rec in caplog.records if "table not found" in rec.message.lower()]


def test_skip_missing_schema(caplog):
ws = create_autospec(WorkspaceClient)
sbe = create_autospec(SqlBackend)
Expand Down
Loading