Added permission migration support for feature tables and the root permissions for models and feature tables (#997)

Author: william-conti

src/databricks/labs/ucx/mixins/fixtures.py

@@ -283,6 +283,12 @@ def _path(ws, path):
             [PermissionLevel.CAN_VIEW, PermissionLevel.CAN_MANAGE],
             _simple,
         ),
+        (
+            "feature_table",
+            "feature-tables",
+            [PermissionLevel.CAN_VIEW_METADATA, PermissionLevel.CAN_EDIT_METADATA, PermissionLevel.CAN_MANAGE],
+            _simple,
+        ),
     ]
 
 
@@ -1126,3 +1132,23 @@ def remove(endpoint_name: str):
             logger.info(f"Can't remove endpoint {e}")
 
     yield from factory("Serving endpoint", create, remove)
+
+
+@pytest.fixture
+def make_feature_table(ws, make_random):
+    def create():
+        feature_table_name = make_random(6) + "." + make_random(6)
+        table = ws.api_client.do(
+            "POST",
+            "/api/2.0/feature-store/feature-tables/create",
+            body={"name": feature_table_name, "primary_keys": [{"name": "pk", "data_type": "string"}]},
+        )
+        return table['feature_table']
+
+    def remove(table: dict):
+        try:
+            ws.api_client.do("DELETE", "/api/2.0/feature-store/feature-tables/delete", body={"name": table["name"]})
+        except RuntimeError as e:
+            logger.info(f"Can't remove feature table {e}")
+
+    yield from factory("Feature table", create, remove)

src/databricks/labs/ucx/workspace_access/generic.py

@@ -410,6 +410,34 @@ def inner() -> Iterator[ml.Experiment]:
     return inner
 
 
+def feature_store_listing(ws: WorkspaceClient):
+    def inner() -> list[GenericPermissionsInfo]:
+        feature_tables = []
+        token = None
+        while True:
+            result = ws.api_client.do(
+                "GET", "/api/2.0/feature-store/feature-tables/search", query={"page_token": token, "max_results": 200}
+            )
+            for table in result.get("feature_tables", []):
+                feature_tables.append(GenericPermissionsInfo(table["id"], "feature-tables"))
+
+            if "next_page_token" not in result:
+                break
+            token = result["next_page_token"]  # type: ignore[index]
+
+        return feature_tables
+
+    return inner
+
+
+def feature_tables_root_page():
+    return [GenericPermissionsInfo("/root", "feature-tables")]
+
+
+def models_root_page():
+    return [GenericPermissionsInfo("/root", "registered-models")]
+
+
 def tokens_and_passwords():
     for _value in ("tokens", "passwords"):
         yield GenericPermissionsInfo(_value, "authorization")

src/databricks/labs/ucx/workspace_access/manager.py

@@ -54,7 +54,10 @@ def factory(
             generic.Listing(ws.serving_endpoints.list, "id", "serving-endpoints"),
             generic.Listing(generic.experiments_listing(ws), "experiment_id", "experiments"),
             generic.Listing(generic.models_listing(ws, num_threads), "id", "registered-models"),
+            generic.Listing(generic.models_root_page, "object_id", "registered-models"),
             generic.Listing(generic.tokens_and_passwords, "object_id", "authorization"),
+            generic.Listing(generic.feature_store_listing(ws), "object_id", "feature-tables"),
+            generic.Listing(generic.feature_tables_root_page, "object_id", "feature-tables"),
             generic.WorkspaceListing(
                 ws,
                 sql_backend=sql_backend,

tests/integration/workspace_access/test_generic.py

@@ -1,18 +1,22 @@
 import json
 from datetime import timedelta
 
+from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import BadRequest, NotFound
 from databricks.sdk.retries import retried
 from databricks.sdk.service import iam
-from databricks.sdk.service.iam import PermissionLevel
+from databricks.sdk.service.iam import AccessControlRequest, PermissionLevel
 
 from databricks.labs.ucx.workspace_access.base import Permissions
 from databricks.labs.ucx.workspace_access.generic import (
     GenericPermissionsSupport,
     Listing,
     WorkspaceListing,
     experiments_listing,
+    feature_store_listing,
+    feature_tables_root_page,
     models_listing,
+    models_root_page,
     tokens_and_passwords,
 )
 from databricks.labs.ucx.workspace_access.groups import MigratedGroup
@@ -439,3 +443,86 @@ def test_endpoints(
 
     after = generic_permissions.load_as_dict("serving-endpoints", endpoint.response.id)
     assert after[group_b.display_name] == PermissionLevel.CAN_MANAGE
+
+
+def test_feature_tables(ws: WorkspaceClient, make_feature_table, make_group, make_feature_table_permissions):
+    group_a = make_group()
+    group_b = make_group()
+    feature_table = make_feature_table()
+    make_feature_table_permissions(
+        object_id=feature_table["id"],
+        permission_level=PermissionLevel.CAN_EDIT_METADATA,
+        group_name=group_a.display_name,
+    )
+
+    generic_permissions = GenericPermissionsSupport(
+        ws, [Listing(feature_store_listing(ws), "object_id", "feature-tables")]
+    )
+    before = generic_permissions.load_as_dict("feature-tables", feature_table["id"])
+    assert before[group_a.display_name] == PermissionLevel.CAN_EDIT_METADATA
+
+    apply_tasks(
+        generic_permissions,
+        [
+            MigratedGroup.partial_info(group_a, group_b),
+        ],
+    )
+
+    after = generic_permissions.load_as_dict("feature-tables", feature_table["id"])
+    assert after[group_b.display_name] == PermissionLevel.CAN_EDIT_METADATA
+
+
+def test_feature_store_root_page(ws: WorkspaceClient, make_group):
+    group_a = make_group()
+    group_b = make_group()
+    ws.permissions.update(
+        "feature-tables",
+        "/root",
+        access_control_list=[
+            AccessControlRequest(group_name=group_a.display_name, permission_level=PermissionLevel.CAN_EDIT_METADATA)
+        ],
+    )
+
+    generic_permissions = GenericPermissionsSupport(
+        ws, [Listing(feature_tables_root_page, "object_id", "feature-tables")]
+    )
+    before = generic_permissions.load_as_dict("feature-tables", "/root")
+    assert before[group_a.display_name] == PermissionLevel.CAN_EDIT_METADATA
+
+    apply_tasks(
+        generic_permissions,
+        [
+            MigratedGroup.partial_info(group_a, group_b),
+        ],
+    )
+
+    after = generic_permissions.load_as_dict("feature-tables", "/root")
+    assert after[group_b.display_name] == PermissionLevel.CAN_EDIT_METADATA
+
+
+def test_models_root_page(ws: WorkspaceClient, make_group):
+    group_a = make_group()
+    group_b = make_group()
+    ws.permissions.update(
+        "registered-models",
+        "/root",
+        access_control_list=[
+            AccessControlRequest(
+                group_name=group_a.display_name, permission_level=PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS
+            )
+        ],
+    )
+
+    generic_permissions = GenericPermissionsSupport(ws, [Listing(models_root_page, "object_id", "registered-models")])
+    before = generic_permissions.load_as_dict("registered-models", "/root")
+    assert before[group_a.display_name] == PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS
+
+    apply_tasks(
+        generic_permissions,
+        [
+            MigratedGroup.partial_info(group_a, group_b),
+        ],
+    )
+
+    after = generic_permissions.load_as_dict("registered-models", "/root")
+    assert after[group_b.display_name] == PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS

tests/unit/workspace_access/test_generic.py

@@ -26,7 +26,10 @@
     WorkspaceListing,
     WorkspaceObjectInfo,
     experiments_listing,
+    feature_store_listing,
+    feature_tables_root_page,
     models_listing,
+    models_root_page,
     tokens_and_passwords,
 )
 from databricks.labs.ucx.workspace_access.groups import MigrationState
@@ -849,3 +852,67 @@ def test_verify_task_should_fail_if_acls_missing():
 
     with pytest.raises(ValueError):
         sup.get_verify_task(item)
+
+
+def test_feature_tables_listing():
+    ws = MagicMock()
+
+    def do_api_side_effect(*_, query):
+        if not query["page_token"]:
+            return {"feature_tables": [{"id": "table1"}, {"id": "table2"}], "next_page_token": "token"}
+        return {"feature_tables": [{"id": "table3"}, {"id": "table4"}]}
+
+    ws.api_client.do.side_effect = do_api_side_effect
+
+    wrapped = Listing(feature_store_listing(ws), id_attribute="object_id", object_type="feature-tables")
+    result = list(wrapped)
+
+    assert len(result) == 4
+    assert result[0].object_id == "table1"
+    assert result[0].request_type == "feature-tables"
+
+
+def test_root_page_listing():
+    ws = MagicMock()
+
+    basic_acl = [
+        iam.AccessControlResponse(
+            group_name="test",
+            all_permissions=[iam.Permission(inherited=False, permission_level=iam.PermissionLevel.CAN_EDIT_METADATA)],
+        )
+    ]
+
+    ws.permissions.get.side_effect = [
+        iam.ObjectPermissions(object_id="/root", object_type="feature-tables", access_control_list=basic_acl),
+    ]
+
+    sup = GenericPermissionsSupport(ws=ws, listings=[Listing(feature_tables_root_page, "object_id", "feature-tables")])
+    tasks = list(sup.get_crawler_tasks())
+    assert len(tasks) == 1
+    auth_items = [task() for task in tasks]
+    for item in auth_items:
+        assert item.object_id == "/root"
+        assert item.object_type == "feature-tables"
+
+
+def test_models_page_listing():
+    ws = MagicMock()
+
+    basic_acl = [
+        iam.AccessControlResponse(
+            group_name="test",
+            all_permissions=[iam.Permission(inherited=False, permission_level=iam.PermissionLevel.CAN_EDIT_METADATA)],
+        )
+    ]
+
+    ws.permissions.get.side_effect = [
+        iam.ObjectPermissions(object_id="/root", object_type="registered-models", access_control_list=basic_acl),
+    ]
+
+    sup = GenericPermissionsSupport(ws=ws, listings=[Listing(models_root_page, "object_id", "registered-models")])
+    tasks = list(sup.get_crawler_tasks())
+    assert len(tasks) == 1
+    auth_items = [task() for task in tasks]
+    for item in auth_items:
+        assert item.object_id == "/root"
+        assert item.object_type == "registered-models"

tests/unit/workspace_access/test_manager.py

@@ -215,6 +215,7 @@ def test_factory(mocker):
             "entitlements",
             "roles",
             'serving-endpoints',
+            "feature-tables",
             "ANY FILE",
             "FUNCTION",
             "ANONYMOUS FUNCTION",