Skip to content

Commit

Permalink
Add grant progress encoder (#3079)
Browse files Browse the repository at this point in the history
## Changes

Add grant progress encoder

### Linked issues

Resolves #3058

### Functionality

- [x] modified existing workflow: `migration-progress`

### Tests

- [x] added unit tests
  • Loading branch information
JCZuurmond authored Oct 25, 2024
1 parent dfeb12d commit 69ba7e7
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 0 deletions.
20 changes: 20 additions & 0 deletions src/databricks/labs/ucx/progress/grants.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
from dataclasses import replace

from databricks.labs.ucx.hive_metastore.grants import Grant
from databricks.labs.ucx.progress.history import ProgressEncoder
from databricks.labs.ucx.progress.install import Historical


class GrantsProgressEncoder(ProgressEncoder[Grant]):
"""Encoder class:Grant to class:History.
A failure for a grants implies it cannot be mapped to Unity Catalog.
"""

def _encode_record_as_historical(self, record: Grant) -> Historical:
historical = super()._encode_record_as_historical(record)
failures = []
if not record.uc_grant_sql():
type_, key = record.this_type_and_key()
failures = [f"Action '{record.action_type}' on {type_} '{key}' unmappable to Unity Catalog"]
return replace(historical, failures=historical.failures + failures)
70 changes: 70 additions & 0 deletions tests/unit/progress/test_grants.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
from unittest.mock import create_autospec

import pytest

from databricks.labs.ucx.framework.owners import Ownership
from databricks.labs.ucx.framework.utils import escape_sql_identifier
from databricks.labs.ucx.hive_metastore.grants import Grant
from databricks.labs.ucx.progress.grants import GrantsProgressEncoder


@pytest.mark.parametrize(
"grant",
[
Grant("principal", "SELECT", "hive_metastore", "schema", udf="function"),
Grant("principal", "OWN", "hive_metastore", "schema", "table"),
Grant("principal", "SELECT", "hive_metastore", "schema"),
Grant("principal", "USAGE", "catalog"),
],
)
def test_grants_progress_encoder_no_failures(mock_backend, grant: Grant) -> None:
ownership = create_autospec(Ownership)
ownership.owner_of.return_value = "user"
encoder = GrantsProgressEncoder(
mock_backend,
ownership,
Grant,
run_id=1,
workspace_id=123456789,
catalog="test",
)

encoder.append_inventory_snapshot([grant])

rows = mock_backend.rows_written_for(escape_sql_identifier(encoder.full_name), "append")
assert len(rows) > 0, f"No rows written for: {encoder.full_name}"
assert len(rows[0].failures) == 0
ownership.owner_of.assert_called_once()


@pytest.mark.parametrize(
"grant, failure",
[
(
Grant("principal", "DENY", "hive_metastore", "schema"),
"Action 'DENY' on DATABASE 'hive_metastore.schema' unmappable to Unity Catalog",
),
(
Grant("principal", "READ_METADATA", "hive_metastore", "schema", "table"),
"Action 'READ_METADATA' on TABLE 'hive_metastore.schema.table' unmappable to Unity Catalog",
),
],
)
def test_grants_progress_encoder_failures(mock_backend, grant: Grant, failure: str) -> None:
ownership = create_autospec(Ownership)
ownership.owner_of.return_value = "user"
encoder = GrantsProgressEncoder(
mock_backend,
ownership,
Grant,
run_id=1,
workspace_id=123456789,
catalog="test",
)

encoder.append_inventory_snapshot([grant])

rows = mock_backend.rows_written_for(escape_sql_identifier(encoder.full_name), "append")
assert len(rows) > 0, f"No rows written for: {encoder.full_name}"
assert rows[0].failures == [failure]
ownership.owner_of.assert_called_once()

0 comments on commit 69ba7e7

Please sign in to comment.