Skip to content

Commit

Permalink
Use make_secret_scope fixture (#163)
Browse files Browse the repository at this point in the history
  • Loading branch information
nfx authored Sep 6, 2023
1 parent 0ac7c58 commit c00cda7
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 70 deletions.
30 changes: 24 additions & 6 deletions src/databricks/labs/ucx/inventory/permissions.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from typing import Literal

from databricks.sdk import WorkspaceClient
from databricks.sdk.service import workspace
from databricks.sdk.service.iam import AccessControlRequest, Group, ObjectPermissions
from databricks.sdk.service.workspace import AclItem as SdkAclItem
from ratelimit import limits, sleep_and_retry
Expand All @@ -23,7 +24,7 @@
RolesAndEntitlements,
)
from databricks.labs.ucx.providers.groups_info import GroupMigrationState
from databricks.labs.ucx.utils import ThreadedExecution, safe_get_acls
from databricks.labs.ucx.utils import ThreadedExecution

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -189,13 +190,13 @@ def _scope_permissions_applicator(self, request_payload: SecretsPermissionReques
# the api might be inconsistent, therefore we need to check that the permissions were applied
for _ in range(3):
time.sleep(random.random() * 2)
applied_acls = safe_get_acls(
self._ws, scope_name=request_payload.object_id, group_name=_acl_item.principal
applied_permission = self._secret_scope_permission(
scope_name=request_payload.object_id, group_name=_acl_item.principal
)
assert applied_acls, f"Failed to apply permissions for {_acl_item.principal}"
assert applied_acls.permission == _acl_item.permission, (
assert applied_permission, f"Failed to apply permissions for {_acl_item.principal}"
assert applied_permission == _acl_item.permission, (
f"Failed to apply permissions for {_acl_item.principal}. "
f"Expected: {_acl_item.permission}. Actual: {applied_acls.permission}"
f"Expected: {_acl_item.permission}. Actual: {applied_permission}"
)

@sleep_and_retry
Expand Down Expand Up @@ -323,3 +324,20 @@ def verify_applied_permissions(
assert [t.all_permissions for t in dst_permissions] == [
s.all_permissions for s in src_permissions
], f"Target permissions were not applied correctly for {object_type}/{object_id}"

def verify_applied_scope_acls(
self, scope_name: str, migration_state: GroupMigrationState, target: Literal["backup", "account"]
):
base_attr = "workspace" if target == "backup" else "backup"
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)
assert src_permission == dst_permission, "Scope ACLs were not applied correctly"

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
10 changes: 0 additions & 10 deletions src/databricks/labs/ucx/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,6 @@
from concurrent.futures import ALL_COMPLETED, ThreadPoolExecutor
from typing import Generic, TypeVar

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

from databricks.labs.ucx.generic import StrEnum

ExecutableResult = TypeVar("ExecutableResult")
Expand Down Expand Up @@ -83,10 +80,3 @@ class WorkspaceLevelEntitlement(StrEnum):
DATABRICKS_SQL_ACCESS = "databricks-sql-access"
ALLOW_CLUSTER_CREATE = "allow-cluster-create"
ALLOW_INSTANCE_POOL_CREATE = "allow-instance-pool-create"


def safe_get_acls(ws: WorkspaceClient, scope_name: str, group_name: str) -> AclItem | None:
all_acls = ws.secrets.list_acls(scope=scope_name)
for acl in all_acls:
if acl.principal == group_name:
return acl
32 changes: 1 addition & 31 deletions tests/integration/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,7 @@
CreateWarehouseRequestWarehouseType,
GetWarehouseResponse,
)
from databricks.sdk.service.workspace import (
AclPermission,
ObjectInfo,
ObjectType,
SecretScope,
)
from databricks.sdk.service.workspace import ObjectInfo, ObjectType

from databricks.labs.ucx.config import InventoryTable
from databricks.labs.ucx.inventory.types import RequestObjectType
Expand All @@ -53,15 +48,13 @@
NUM_TEST_GROUPS = int(os.environ.get("NUM_TEST_GROUPS", 5))
NUM_TEST_INSTANCE_PROFILES = int(os.environ.get("NUM_TEST_INSTANCE_PROFILES", 3))
NUM_TEST_CLUSTERS = int(os.environ.get("NUM_TEST_CLUSTERS", 3))
NUM_TEST_INSTANCE_POOLS = int(os.environ.get("NUM_TEST_INSTANCE_POOLS", 3))
NUM_TEST_CLUSTER_POLICIES = int(os.environ.get("NUM_TEST_CLUSTER_POLICIES", 3))
NUM_TEST_PIPELINES = int(os.environ.get("NUM_TEST_PIPELINES", 3))
NUM_TEST_JOBS = int(os.environ.get("NUM_TEST_JOBS", 3))
NUM_TEST_EXPERIMENTS = int(os.environ.get("NUM_TEST_EXPERIMENTS", 3))
NUM_TEST_MODELS = int(os.environ.get("NUM_TEST_MODELS", 3))
NUM_TEST_WAREHOUSES = int(os.environ.get("NUM_TEST_WAREHOUSES", 3))
NUM_TEST_TOKENS = int(os.environ.get("NUM_TEST_TOKENS", 3))
NUM_TEST_SECRET_SCOPES = int(os.environ.get("NUM_TEST_SECRET_SCOPES", 10))

NUM_THREADS = int(os.environ.get("NUM_TEST_THREADS", 20))
UCX_TESTING_PREFIX = os.environ.get("UCX_TESTING_PREFIX", "ucx")
Expand Down Expand Up @@ -515,27 +508,6 @@ def tokens(ws: WorkspaceClient, env: EnvironmentInfo) -> list[AccessControlReque
yield token_permissions


@pytest.fixture
def secret_scopes(ws: WorkspaceClient, env: EnvironmentInfo) -> list[SecretScope]:
logger.debug("Creating test secret scopes")

for i in range(NUM_TEST_SECRET_SCOPES):
ws.secrets.create_scope(f"{env.test_uid}-test-{i}")

test_secret_scopes = [s for s in ws.secrets.list_scopes() if s.name.startswith(env.test_uid)]

for scope in test_secret_scopes:
random_permission = random.choice(list(AclPermission))
random_ws_group, _ = random.choice(env.groups)
ws.secrets.put_acl(scope.name, random_ws_group.display_name, random_permission)

yield test_secret_scopes

logger.debug("Deleting test secret scopes")
executables = [partial(ws.secrets.delete_scope, s.name) for s in test_secret_scopes]
Threader(executables).run()


@pytest.fixture
def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObjects:
logger.info(f"Creating test workspace objects under /{env.test_uid}")
Expand Down Expand Up @@ -601,12 +573,10 @@ def verifiable_objects(
models,
warehouses,
tokens,
secret_scopes,
workspace_objects,
) -> list[tuple[list, str, RequestObjectType | None]]:
_verifiable_objects = [
(workspace_objects, "workspace_objects", None),
(secret_scopes, "secret_scopes", None),
(tokens, "tokens", RequestObjectType.AUTHORIZATION),
(clusters, "cluster_id", RequestObjectType.CLUSTERS),
(cluster_policies, "policy_id", RequestObjectType.CLUSTER_POLICIES),
Expand Down
34 changes: 11 additions & 23 deletions tests/integration/test_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@

import pytest
from databricks.sdk import WorkspaceClient
from databricks.sdk.service import workspace
from databricks.sdk.service.iam import (
AccessControlRequest,
AccessControlResponse,
Permission,
PermissionLevel,
)
from databricks.sdk.service.workspace import SecretScope
from pyspark.errors import AnalysisException

from databricks.labs.ucx.config import (
Expand All @@ -24,7 +24,6 @@
from databricks.labs.ucx.inventory.types import RequestObjectType
from databricks.labs.ucx.providers.groups_info import GroupMigrationState
from databricks.labs.ucx.toolkits.group_migration import GroupMigrationToolkit
from databricks.labs.ucx.utils import safe_get_acls

from .utils import EnvironmentInfo, WorkspaceObjects

Expand Down Expand Up @@ -68,27 +67,10 @@ def _verify_group_permissions(
assert len(base_acls) == len(target_acls)

elif id_attribute == "secret_scopes":
_scopes: list[SecretScope] = objects
comparison_base = [
getattr(mi, "workspace" if target == "backup" else "backup")
for mi in toolkit.group_manager.migration_groups_provider.groups
]

comparison_target = [getattr(mi, target) for mi in toolkit.group_manager.migration_groups_provider.groups]

for scope in _scopes:
for base_group, target_group in zip(comparison_base, comparison_target, strict=True):
base_acl = safe_get_acls(ws, scope.name, base_group.display_name)
target_acl = safe_get_acls(ws, scope.name, target_group.display_name)

if base_acl:
if not target_acl:
msg = "Target ACL is empty, while base ACL is not"
raise AssertionError(msg)

assert (
base_acl.permission == target_acl.permission
), f"Target permissions were not applied correctly for scope {scope.name}"
for scope_name in objects:
toolkit.permissions_manager.verify_applied_scope_acls(
scope_name, toolkit.group_manager.migration_groups_provider, target
)

elif id_attribute in ("tokens", "passwords"):
_typed_objects: list[AccessControlRequest] = objects
Expand Down Expand Up @@ -149,6 +131,8 @@ def test_e2e(
verifiable_objects: list[tuple[list, str, RequestObjectType | None]],
make_instance_pool,
make_instance_pool_permissions,
make_secret_scope,
make_secret_scope_acl,
):
logger.debug(f"Test environment: {env.test_uid}")
ws_group = env.groups[0][0]
Expand All @@ -161,6 +145,10 @@ def test_e2e(
)
verifiable_objects.append(([pool], "instance_pool_id", RequestObjectType.INSTANCE_POOLS))

scope = make_secret_scope()
make_secret_scope_acl(scope=scope, principal=ws_group.display_name, permission=workspace.AclPermission.WRITE)
verifiable_objects.append(([scope], "secret_scopes", None))

config = MigrationConfig(
connect=ConnectConfig.from_databricks_config(ws.config),
inventory=InventoryConfig(table=inventory_table),
Expand Down

0 comments on commit c00cda7

Please sign in to comment.