Skip to content

Commit

Permalink
Merge branch 'main' into add_semi_automated_mount_point_migration
Browse files Browse the repository at this point in the history
  • Loading branch information
william-conti authored Sep 13, 2023
2 parents 9a7876e + 2b568a9 commit 2563ede
Show file tree
Hide file tree
Showing 14 changed files with 167 additions and 37 deletions.
3 changes: 1 addition & 2 deletions src/databricks/labs/ucx/inventory/permissions_inventory.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import logging

import pandas as pd
from databricks.sdk import WorkspaceClient
from pyspark.sql import DataFrame
from pyspark.sql.types import StringType, StructField, StructType
Expand Down Expand Up @@ -39,7 +38,7 @@ def cleanup(self):
def save(self, items: list[PermissionsInventoryItem]):
# TODO: update instead of append
logger.info(f"Saving {len(items)} items to inventory table {self._table}")
serialized_items = pd.DataFrame([item.as_dict() for item in items])
serialized_items = [item.as_dict() for item in items]
df = self.spark.createDataFrame(serialized_items, schema=self._table_schema)
df.write.mode("append").format("delta").saveAsTable(self._table)
logger.info("Successfully saved the items to inventory table")
Expand Down
15 changes: 5 additions & 10 deletions src/databricks/labs/ucx/inventory/verification.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
from typing import Literal

from databricks.sdk import WorkspaceClient
from databricks.sdk.service import workspace

from databricks.labs.ucx.providers.groups_info import GroupMigrationState
from databricks.labs.ucx.support.secrets import SecretScopesSupport


class VerificationManager:
def __init__(self, ws: WorkspaceClient):
def __init__(self, ws: WorkspaceClient, secrets_support: SecretScopesSupport):
self._ws = ws
self._secrets_support = secrets_support

def verify(
self, migration_state: GroupMigrationState, target: Literal["backup", "account"], tuples: list[tuple[str, str]]
Expand Down Expand Up @@ -51,8 +52,8 @@ def verify_applied_scope_acls(
for mi in migration_state.groups:
src_name = getattr(mi, base_attr).display_name
dst_name = getattr(mi, target).display_name
src_permission = self._secret_scope_permission(scope_name, src_name)
dst_permission = self._secret_scope_permission(scope_name, dst_name)
src_permission = self._secrets_support.secret_scope_permission(scope_name, src_name)
dst_permission = self._secrets_support.secret_scope_permission(scope_name, dst_name)
assert src_permission == dst_permission, "Scope ACLs were not applied correctly"

def verify_roles_and_entitlements(self, migration_state: GroupMigrationState, target: Literal["backup", "account"]):
Expand All @@ -65,9 +66,3 @@ def verify_roles_and_entitlements(self, migration_state: GroupMigrationState, ta

assert base_group_info.roles == target_group_info.roles
assert base_group_info.entitlements == target_group_info.entitlements

def _secret_scope_permission(self, scope_name: str, group_name: str) -> workspace.AclPermission | None:
for acl in self._ws.secrets.list_acls(scope=scope_name):
if acl.principal == group_name:
return acl.permission
return None
2 changes: 1 addition & 1 deletion src/databricks/labs/ucx/support/group_level.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ def _crawler_task(self, group: iam.Group, property_name: str):
@sleep_and_retry
@limits(calls=10, period=1)
def _applier_task(self, group_id: str, value: list[iam.ComplexValue], property_name: str):
operations = [iam.Patch(op=iam.PatchOp.ADD, path=property_name, value=value)]
operations = [iam.Patch(op=iam.PatchOp.ADD, path=property_name, value=[e.as_dict() for e in value])]
schemas = [iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP]
self._ws.groups.patch(id=group_id, operations=operations, schemas=schemas)

Expand Down
12 changes: 6 additions & 6 deletions src/databricks/labs/ucx/support/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,11 @@ def __init__(self, ws: WorkspaceClient, num_threads: int, workspace_start_path:
ws=ws,
listings=[
listing_wrapper(ws.clusters.list, "cluster_id", RequestObjectType.CLUSTERS),
listing_wrapper(ws.cluster_policies.list, "cluster_policy_id", RequestObjectType.CLUSTER_POLICIES),
listing_wrapper(ws.cluster_policies.list, "policy_id", RequestObjectType.CLUSTER_POLICIES),
listing_wrapper(ws.instance_pools.list, "instance_pool_id", RequestObjectType.INSTANCE_POOLS),
listing_wrapper(ws.warehouses.list, "id", RequestObjectType.SQL_WAREHOUSES),
listing_wrapper(ws.jobs.list, "job_id", RequestObjectType.JOBS),
listing_wrapper(ws.pipelines.list, "pipeline_id", RequestObjectType.PIPELINES),
listing_wrapper(ws.pipelines.list_pipelines, "pipeline_id", RequestObjectType.PIPELINES),
listing_wrapper(experiments_listing(ws), "experiment_id", RequestObjectType.EXPERIMENTS),
listing_wrapper(models_listing(ws), "id", RequestObjectType.REGISTERED_MODELS),
workspace_listing(ws, num_threads=num_threads, start_path=workspace_start_path),
Expand Down Expand Up @@ -64,13 +64,13 @@ def supports(self) -> dict[str, BaseSupport]:
"roles": self._scim_support,
# generic API
"clusters": self._generic_support,
"cluster_policies": self._generic_support,
"instance_pools": self._generic_support,
"sql_warehouses": self._generic_support,
"cluster-policies": self._generic_support,
"instance-pools": self._generic_support,
"sql/warehouses": self._generic_support,
"jobs": self._generic_support,
"pipelines": self._generic_support,
"experiments": self._generic_support,
"registered_models": self._generic_support,
"registered-models": self._generic_support,
"tokens": self._generic_support,
"passwords": self._generic_support,
# workspace objects
Expand Down
23 changes: 19 additions & 4 deletions src/databricks/labs/ucx/support/permissions.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,12 @@ def _safe_get_permissions(
permissions = ws.permissions.get(request_object_type, object_id)
return permissions
except DatabricksError as e:
if e.error_code in ["RESOURCE_DOES_NOT_EXIST", "RESOURCE_NOT_FOUND", "PERMISSION_DENIED"]:
if e.error_code in [
"RESOURCE_DOES_NOT_EXIST",
"RESOURCE_NOT_FOUND",
"PERMISSION_DENIED",
"FEATURE_DISABLED",
]:
logger.warning(f"Could not get permissions for {request_object_type} {object_id} due to {e.error_code}")
return None
else:
Expand Down Expand Up @@ -79,7 +84,7 @@ def _prepare_new_acl(
def _applier_task(
self, ws: WorkspaceClient, object_id: str, acl: list[iam.AccessControlRequest], request_type: RequestObjectType
):
ws.permissions.update(request_type, object_id, acl)
ws.permissions.update(request_object_type=request_type, request_object_id=object_id, access_control_list=acl)

@sleep_and_retry
@limits(calls=100, period=1)
Expand All @@ -90,10 +95,13 @@ def _crawler_task(
request_type: RequestObjectType,
) -> PermissionsInventoryItem | None:
permissions = self._safe_get_permissions(ws, request_type, object_id)

support = object_id if request_type == RequestObjectType.AUTHORIZATION else request_type.value

if permissions:
return PermissionsInventoryItem(
object_id=object_id,
support=request_type.value,
support=support,
raw_object_permissions=json.dumps(permissions.as_dict()),
)

Expand All @@ -103,10 +111,17 @@ def _get_apply_task(
new_acl = self._prepare_new_acl(
iam.ObjectPermissions.from_dict(json.loads(item.raw_object_permissions)), migration_state, destination
)

request_type = (
RequestObjectType.AUTHORIZATION
if item.support in ("passwords", "tokens")
else RequestObjectType(item.support)
)

return partial(
self._applier_task,
ws=self._ws,
request_type=RequestObjectType(item.support),
request_type=request_type,
acl=new_acl,
object_id=item.object_id,
)
Expand Down
34 changes: 34 additions & 0 deletions src/databricks/labs/ucx/support/secrets.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import json
import random
import time
from functools import partial

from databricks.sdk import WorkspaceClient
Expand Down Expand Up @@ -33,10 +35,42 @@ def is_item_relevant(self, item: PermissionsInventoryItem, migration_state: Grou
mentioned_groups = [acl.principal for acl in acls]
return any(g in mentioned_groups for g in [info.workspace.display_name for info in migration_state.groups])

def secret_scope_permission(self, scope_name: str, group_name: str) -> workspace.AclPermission | None:
for acl in self._ws.secrets.list_acls(scope=scope_name):
if acl.principal == group_name:
return acl.permission
return None

def _inflight_check(
self, scope_name: str, group_name: str, expected_permission: workspace.AclPermission, num_retries: int = 5
):
# in-flight check for the applied permissions
# the api might be inconsistent, therefore we need to check that the permissions were applied
# TODO: add mixin to SDK
retries_left = num_retries
while retries_left > 0:
time.sleep(random.random() * 2)
applied_permission = self.secret_scope_permission(scope_name=scope_name, group_name=group_name)
if applied_permission:
if applied_permission == expected_permission:
return
else:
msg = (
f"Applied permission {applied_permission} is not "
f"equal to expected permission {expected_permission}"
)
raise ValueError(msg)

retries_left -= 1

msg = f"Failed to apply permissions for {group_name} on scope {scope_name} in {num_retries} retries"
raise ValueError(msg)

@sleep_and_retry
@limits(calls=30, period=1)
def _rate_limited_put_acl(self, object_id: str, principal: str, permission: workspace.AclPermission):
self._ws.secrets.put_acl(object_id, principal, permission)
self._inflight_check(scope_name=object_id, group_name=principal, expected_permission=permission)

def _get_apply_task(
self, item: PermissionsInventoryItem, migration_state: GroupMigrationState, destination: Destination
Expand Down
2 changes: 1 addition & 1 deletion src/databricks/labs/ucx/support/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ def _applier_task(self, object_type: sql.ObjectTypePlural, object_id: str, acl:
Please note that we only have SET option (DBSQL Permissions API doesn't support UPDATE operation).
This affects the way how we prepare the new ACL request.
"""
self._ws.dbsql_permissions.set(object_type=object_type, object_id=object_id, acl=acl)
self._ws.dbsql_permissions.set(object_type=object_type, object_id=object_id, access_control_list=acl)

def get_crawler_tasks(self):
for listing in self._listings:
Expand Down
7 changes: 3 additions & 4 deletions src/databricks/labs/ucx/toolkits/group_migration.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,11 @@ def __init__(self, config: MigrationConfig):

self._group_manager = GroupManager(self._ws, config.groups)
self._permissions_inventory = PermissionsInventoryTable(config.inventory_database, self._ws)
self._supports_provider = SupportsProvider(self._ws, self._num_threads, self._workspace_start_path)
self._permissions_manager = PermissionManager(
self._ws,
self._permissions_inventory,
supports_provider=SupportsProvider(self._ws, self._num_threads, self._workspace_start_path),
self._ws, self._permissions_inventory, supports_provider=self._supports_provider
)
self._verification_manager = VerificationManager(self._ws)
self._verification_manager = VerificationManager(self._ws, self._supports_provider.supports["secrets"])

@staticmethod
def _verify_ws_client(w: WorkspaceClient):
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/support/test_group_level.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def test_scim_apply(migration_state):
task()
ws.groups.patch.assert_called_once_with(
id="test-backup",
operations=[iam.Patch(op=iam.PatchOp.ADD, path="roles", value=sample_permissions)],
operations=[iam.Patch(op=iam.PatchOp.ADD, path="roles", value=[p.as_dict() for p in sample_permissions])],
schemas=[iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
)

Expand Down
8 changes: 4 additions & 4 deletions tests/unit/support/test_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,13 @@ def test_supports_provider():
"entitlements",
"roles",
"clusters",
"cluster_policies",
"instance_pools",
"sql_warehouses",
"cluster-policies",
"instance-pools",
"sql/warehouses",
"jobs",
"pipelines",
"experiments",
"registered_models",
"registered-models",
"tokens",
"passwords",
"notebooks",
Expand Down
48 changes: 47 additions & 1 deletion tests/unit/support/test_permissions.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
PermissionsInventoryItem,
RequestObjectType,
)
from databricks.labs.ucx.support.listing import authorization_listing
from databricks.labs.ucx.support.permissions import (
GenericPermissionsSupport,
listing_wrapper,
Expand Down Expand Up @@ -92,7 +93,11 @@ def test_apply(migration_state):
)
]

ws.permissions.update.assert_called_with(RequestObjectType.CLUSTERS, "test", expected_acl_payload)
ws.permissions.update.assert_called_with(
request_object_type=RequestObjectType.CLUSTERS,
request_object_id="test",
access_control_list=expected_acl_payload,
)


def test_relevance():
Expand Down Expand Up @@ -136,3 +141,44 @@ def test_no_permissions():
_task = tasks[0]
item = _task()
assert item is None


def test_passwords_tokens_crawler(migration_state):
ws = MagicMock()

basic_acl = [
iam.AccessControlResponse(
group_name="test",
all_permissions=[iam.Permission(inherited=False, permission_level=iam.PermissionLevel.CAN_USE)],
)
]

ws.permissions.get.side_effect = [
iam.ObjectPermissions(
object_id="passwords", object_type=RequestObjectType.AUTHORIZATION, access_control_list=basic_acl
),
iam.ObjectPermissions(
object_id="tokens", object_type=RequestObjectType.AUTHORIZATION, access_control_list=basic_acl
),
]

sup = GenericPermissionsSupport(ws=ws, listings=[authorization_listing()])
tasks = list(sup.get_crawler_tasks())
assert len(tasks) == 2
auth_items = [task() for task in tasks]
for item in auth_items:
assert item.object_id in ["tokens", "passwords"]
assert item.support in ["tokens", "passwords"]
applier = sup.get_apply_task(item, migration_state, "backup")
new_acl = sup._prepare_new_acl(
permissions=iam.ObjectPermissions.from_dict(json.loads(item.raw_object_permissions)),
migration_state=migration_state,
destination="backup",
)
applier()
ws.permissions.update.assert_called_once_with(
request_object_type=RequestObjectType.AUTHORIZATION,
request_object_id=item.object_id,
access_control_list=new_acl,
)
ws.permissions.update.reset_mock()
42 changes: 42 additions & 0 deletions tests/unit/support/test_secrets.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import json
from unittest.mock import MagicMock, call

import pytest
from databricks.sdk.service import workspace

from databricks.labs.ucx.inventory.types import PermissionsInventoryItem
Expand Down Expand Up @@ -55,6 +56,18 @@ def test_secret_scopes_apply(migration_state: GroupMigrationState):
),
)

# positive case - permissions are applied correctly
ws.secrets.list_acls.return_value = [
workspace.AclItem(
principal="db-temp-test",
permission=workspace.AclPermission.MANAGE,
),
workspace.AclItem(
principal="irrelevant",
permission=workspace.AclPermission.MANAGE,
),
]

task = sup.get_apply_task(item, migration_state, "backup")
task()
assert ws.secrets.put_acl.call_count == 2
Expand All @@ -64,3 +77,32 @@ def test_secret_scopes_apply(migration_state: GroupMigrationState):
call("test", "irrelevant", workspace.AclPermission.MANAGE),
]
ws.secrets.put_acl.assert_has_calls(calls, any_order=False)


def test_secret_scopes_apply_failed():
ws = MagicMock()
sup = SecretScopesSupport(ws=ws)
expected_permission = workspace.AclPermission.MANAGE
with pytest.raises(ValueError) as e:
sup._inflight_check(
group_name="db-temp-test", scope_name="test", expected_permission=expected_permission, num_retries=2
)
assert "Failed to apply permissions" in str(e.value)


def test_secret_scopes_apply_incorrect():
ws = MagicMock()
ws.secrets.list_acls.return_value = [
workspace.AclItem(
principal="db-temp-test",
permission=workspace.AclPermission.READ,
)
]

sup = SecretScopesSupport(ws=ws)
expected_permission = workspace.AclPermission.MANAGE
with pytest.raises(ValueError) as e:
sup._inflight_check(
group_name="db-temp-test", scope_name="test", expected_permission=expected_permission, num_retries=2
)
assert "not equal to expected permission" in str(e.value)
2 changes: 1 addition & 1 deletion tests/unit/support/test_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def test_apply(migration_state):
),
]
ws.dbsql_permissions.set.assert_called_once_with(
object_type=sql.ObjectTypePlural.ALERTS, object_id="test", acl=expected_payload
object_type=sql.ObjectTypePlural.ALERTS, object_id="test", access_control_list=expected_payload
)


Expand Down
Loading

0 comments on commit 2563ede

Please sign in to comment.