Merge pull request #1061 from DaanRademaker/add_support_for_volumes

add new permissions support for volumes

.github/workflows/python-package.yml

@@ -5,17 +5,17 @@ name: Python package
 
 on:
   push:
-    branches: [ main ]
+    branches: [main]
   pull_request:
-    branches: [ main ]
+    branches: [main]
 
 jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
 
     steps:
       - uses: actions/checkout@v2
@@ -38,7 +38,7 @@ jobs:
         uses: mikepenz/action-junit-report@v2
         if: always() # always run even if the previous step fails
         with:
-          report_paths: 'aws-lambda/junit/report.xml'
+          report_paths: "aws-lambda/junit/report.xml"
       - name: Upload Coverage to Codecov
         uses: codecov/codecov-action@v1
         with:

aws-lambda/pyproject.toml

@@ -15,13 +15,13 @@ include = ["LICENSE"]
 packages = [{ include = "databricks_cdk", from = "src" }]
 
 [tool.poetry.dependencies]
-python = "^3.7"
+python = "^3.8"
 pydantic = "^1.4.0"
 requests = ">=2.22"
 boto3 = "^1.21.10"
 cfnresponse = "^1.1.2"
 tenacity = "^8.2.2"
-databricks-sdk = "^0.11.0"
+databricks-sdk = "^0.13.0"
 
 [tool.poetry.dev-dependencies]
 coverage = { version = "^6.1.1", extras = ["toml"] }

aws-lambda/src/databricks_cdk/resources/handler.py

@@ -79,6 +79,11 @@
     create_or_update_warehouse_permissions,
     delete_warehouse_permissions,
 )
+from databricks_cdk.resources.permissions.volume_permissions import (
+    VolumePermissionsProperties,
+    create_or_update_volume_permissions,
+    delete_volume_permissions,
+)
 from databricks_cdk.resources.scim.user import UserProperties, create_or_update_user, delete_user
 from databricks_cdk.resources.secrets.secret import SecretProperties, create_or_update_secret, delete_secret
 from databricks_cdk.resources.secrets.secret_scope import (
@@ -197,6 +202,8 @@ def create_or_update_resource(event: DatabricksEvent) -> CnfResponse:
         return create_or_update_registered_model_permissions(
             RegisteredModelPermissionPermissionProperties(**event.ResourceProperties)
         )
+    elif action == "volume-permissions":
+        return create_or_update_volume_permissions(VolumePermissionsProperties(**event.ResourceProperties))
     elif action == "experiment-permission":
         return create_or_update_experiment_permissions(ExperimentPermissionProperties(**event.ResourceProperties))
     elif action == "token":
@@ -292,6 +299,11 @@ def delete_resource(event: DatabricksEvent) -> CnfResponse:
             RegisteredModelPermissionPermissionProperties(**event.ResourceProperties),
             event.PhysicalResourceId,
         )
+    elif action == "volume-permissions":
+        return delete_volume_permissions(
+            VolumePermissionsProperties(**event.ResourceProperties),
+            event.PhysicalResourceId,
+        )
     elif action == "unity-storage-credentials":
         return delete_storage_credential(
             StorageCredentialsProperties(**event.ResourceProperties),

aws-lambda/src/databricks_cdk/resources/permissions/changes.py

@@ -0,0 +1,89 @@
+from typing import Dict, List
+
+from databricks.sdk.service.catalog import PermissionsChange, PermissionsList, Privilege, PrivilegeAssignment
+
+
+def get_assignment_dict_from_permissions_list(permissions_list: PermissionsList) -> Dict[str, List[Privilege]]:
+    """Converts PermissionsList to dict with key of principal and list of associated privileges as value"""
+    privilige_assignments = permissions_list.privilege_assignments
+
+    if privilige_assignments is None:
+        return {}
+
+    return {
+        x.principal: [Privilege(y) for y in x.privileges]
+        for x in privilige_assignments
+        if x.principal is not None and x.privileges is not None
+    }
+
+
+def get_assignment_dict_from_privilege_assignments(
+    privilege_assignments: List[PrivilegeAssignment],
+) -> Dict[str, List[Privilege]]:
+    """Converts list of PrivilegeAssignment to dict with key of principal and list of associated privileges as value"""
+    return {
+        x.principal: x.privileges for x in privilege_assignments if x.principal is not None and x.privileges is not None
+    }
+
+
+def get_permission_changes_principals(
+    assignments_on_databricks_dict: Dict[str, List[Privilege]],
+    assignments_from_properties_dict: Dict[str, List[Privilege]],
+) -> List[PermissionsChange]:
+    """See if there are new principals that need to be added"""
+    permission_changes = []
+
+    principals_on_databricks = set(assignments_on_databricks_dict.keys())
+    principals_from_properties = set(assignments_from_properties_dict.keys())
+
+    principals_to_add = principals_from_properties.difference(principals_on_databricks)
+    principals_to_remove = principals_on_databricks.difference(principals_from_properties)
+
+    for principal in principals_to_add:
+        permission_changes.append(
+            PermissionsChange(principal=principal, add=assignments_from_properties_dict[principal])
+        )
+
+    for principal in principals_to_remove:
+        permission_changes.append(
+            PermissionsChange(principal=principal, remove=assignments_on_databricks_dict[principal])
+        )
+
+    return permission_changes
+
+
+def get_permission_changes_assignments_changed(
+    assignments_on_databricks_dict: Dict[str, List[Privilege]],
+    assignments_from_properties_dict: Dict[str, List[Privilege]],
+) -> List[PermissionsChange]:
+    """See if there are principal assignemnts that need to be updated"""
+    permission_changes = []
+    for principal, privileges in assignments_from_properties_dict.items():
+        privileges_databricks = set(assignments_on_databricks_dict.get(principal, []))
+        privileges_properties = set(privileges)
+
+        if privileges_databricks != privileges_properties and len(privileges_databricks) > 0:
+            to_remove = privileges_databricks.difference(privileges_properties)
+            to_add = privileges_properties.difference(privileges_databricks)
+            permission_changes.append(PermissionsChange(principal=principal, add=list(to_add), remove=list(to_remove)))
+
+    return permission_changes
+
+
+def get_permission_changes(
+    assignments_on_databricks: PermissionsList, assignments_from_properties: List[PrivilegeAssignment]
+) -> List[PermissionsChange]:
+    """Get the changes between the existing grants and the new grants and create PermissionsChange object"""
+    # Convert to dict for easier lookup
+    assignments_on_databricks_dict = get_assignment_dict_from_permissions_list(assignments_on_databricks)
+    assignments_from_properties_dict = get_assignment_dict_from_privilege_assignments(assignments_from_properties)
+
+    permission_changes: List[PermissionsChange] = []
+    permission_changes += get_permission_changes_principals(
+        assignments_on_databricks_dict, assignments_from_properties_dict
+    )
+    permission_changes += get_permission_changes_assignments_changed(
+        assignments_on_databricks_dict, assignments_from_properties_dict
+    )
+
+    return permission_changes

aws-lambda/src/databricks_cdk/resources/permissions/volume_permissions.py

@@ -0,0 +1,52 @@
+from typing import List
+
+from databricks.sdk.service.catalog import PermissionsList, PrivilegeAssignment, SecurableType
+from pydantic import BaseModel
+
+from databricks_cdk.resources.permissions.changes import get_permission_changes
+from databricks_cdk.utils import CnfResponse, get_workspace_client
+
+
+class VolumePermissionsProperties(BaseModel):
+    workspace_url: str
+    volume_name: str
+    privilege_assignments: List[PrivilegeAssignment] = []
+
+
+def create_or_update_volume_permissions(
+    properties: VolumePermissionsProperties,
+) -> CnfResponse:
+    """Create volume permissions on volume at databricks"""
+
+    workspace_client = get_workspace_client(properties.workspace_url)
+    existing_grants: PermissionsList = workspace_client.grants.get(
+        securable_type=SecurableType.VOLUME, full_name=properties.volume_name
+    )
+
+    permission_changes = get_permission_changes(existing_grants, properties.privilege_assignments)
+
+    workspace_client.grants.update(
+        securable_type=SecurableType.VOLUME, full_name=properties.volume_name, changes=permission_changes
+    )
+
+    return CnfResponse(
+        physical_resource_id=f"{properties.volume_name}/permissions",
+    )
+
+
+def delete_volume_permissions(properties: VolumePermissionsProperties, physical_resource_id: str) -> CnfResponse:
+    """Deletes all volume permissions on volume at databricks"""
+    workspace_client = get_workspace_client(properties.workspace_url)
+    existing_grants: PermissionsList = workspace_client.grants.get(
+        securable_type=SecurableType.VOLUME, full_name=properties.volume_name
+    )
+
+    permission_changes = get_permission_changes(existing_grants, [])
+
+    workspace_client.grants.update(
+        securable_type=SecurableType.VOLUME, full_name=properties.volume_name, changes=permission_changes
+    )
+
+    return CnfResponse(
+        physical_resource_id=physical_resource_id,
+    )

aws-lambda/tests/resources/permissions/test_changes.py

@@ -0,0 +1,139 @@
+from databricks.sdk.service.catalog import PermissionsChange, PermissionsList, Privilege, PrivilegeAssignment
+
+from databricks_cdk.resources.permissions.changes import (
+    get_assignment_dict_from_permissions_list,
+    get_assignment_dict_from_privilege_assignments,
+    get_permission_changes,
+    get_permission_changes_assignments_changed,
+    get_permission_changes_principals,
+)
+
+
+def test_get_assignment_dict_from_permissions_list():
+    permissions_list = PermissionsList(
+        **{
+            "privilege_assignments": [
+                PrivilegeAssignment(**{"principal": "principal1", "privileges": ["APPLY_TAG", "WRITE_FILES"]}),
+                PrivilegeAssignment(**{"principal": "principal2", "privileges": ["READ_FILES"]}),
+            ]
+        }
+    )
+
+    result = get_assignment_dict_from_permissions_list(permissions_list)
+
+    assert result == {
+        "principal1": [Privilege.APPLY_TAG, Privilege.WRITE_FILES],
+        "principal2": [Privilege.READ_FILES],
+    }
+
+
+def test_get_assignment_dict_from_permissions_list_empty():
+    permissions_list = PermissionsList(**{"privilege_assignments": None})
+
+    result = get_assignment_dict_from_permissions_list(permissions_list)
+
+    assert result == {}
+
+
+def test_get_assignment_dict_from_privilege_assignments():
+    privilege_assignments = [
+        PrivilegeAssignment(**{"principal": "principal1", "privileges": [Privilege.APPLY_TAG, Privilege.WRITE_FILES]}),
+        PrivilegeAssignment(**{"principal": "principal2", "privileges": [Privilege.READ_FILES]}),
+    ]
+
+    result = get_assignment_dict_from_privilege_assignments(privilege_assignments)
+
+    assert result == {
+        "principal1": [Privilege.APPLY_TAG, Privilege.WRITE_FILES],
+        "principal2": [Privilege.READ_FILES],
+    }
+
+
+def test_get_assignment_dict_from_privilege_assignments_empty():
+    privilege_assignments = [
+        PrivilegeAssignment(**{"principal": "principal1", "privileges": None}),
+        PrivilegeAssignment(**{"principal": None, "privileges": [Privilege.READ_FILES]}),
+    ]
+    result = get_assignment_dict_from_privilege_assignments(privilege_assignments)
+
+    assert result == {}
+
+
+def test_get_permission_changes_assignments_changed():
+    assignments_on_databricks_dict = {"principal": list(tuple([Privilege.APPLY_TAG, Privilege.WRITE_FILES]))}
+    assignments_from_properties_dict = {"principal": list(tuple([Privilege.APPLY_TAG, Privilege.READ_FILES]))}
+
+    result = get_permission_changes_assignments_changed(
+        assignments_on_databricks_dict, assignments_from_properties_dict
+    )
+
+    # READ_FILES should be added WRITE_FILES should be removed
+    assert result == [
+        PermissionsChange(add=[Privilege.READ_FILES], principal="principal", remove=[Privilege.WRITE_FILES])
+    ]
+
+
+def test_get_permission_changes_assignments_changed_no_changes():
+    assignments_on_databricks_dict = {"principal": list(tuple([Privilege.APPLY_TAG, Privilege.WRITE_FILES]))}
+    assignments_from_properties_dict = {"principal": list(tuple([Privilege.APPLY_TAG, Privilege.WRITE_FILES]))}
+
+    result = get_permission_changes_assignments_changed(
+        assignments_on_databricks_dict, assignments_from_properties_dict
+    )
+
+    # No changes
+    assert result == []
+
+
+def test_get_permission_changes_new_principals():
+    assignments_on_databricks_dict = {
+        "principal_to_remove": list(tuple([Privilege.APPLY_TAG])),
+        "principal_to_stay": list(tuple([Privilege.APPLY_TAG])),
+    }
+    assignments_from_properties_dict = {
+        "principal_to_add": list(tuple([Privilege.APPLY_TAG])),
+        "principal_to_stay": list(tuple([Privilege.APPLY_TAG])),
+    }
+
+    result = get_permission_changes_principals(assignments_on_databricks_dict, assignments_from_properties_dict)
+
+    # No new principals
+    assert result == [
+        PermissionsChange(add=[Privilege.APPLY_TAG], principal="principal_to_add", remove=None),
+        PermissionsChange(add=None, principal="principal_to_remove", remove=[Privilege.APPLY_TAG]),
+    ]
+
+
+def test_get_permission_changes_no_principals():
+    assignments_on_databricks_dict = {}
+    assignments_from_properties_dict = {}
+
+    result = get_permission_changes_principals(assignments_on_databricks_dict, assignments_from_properties_dict)
+
+    # No new principals
+    assert result == []
+
+
+def test_get_permission_changes():
+    sample_permissions_list = PermissionsList(
+        privilege_assignments=[
+            PrivilegeAssignment(principal="user1", privileges=[Privilege.APPLY_TAG]),
+            PrivilegeAssignment(principal="user2", privileges=[Privilege.CREATE]),
+            PrivilegeAssignment(principal="userToRemove", privileges=[Privilege.CREATE]),
+        ]
+    )
+    sample_privilege_assignments = [
+        PrivilegeAssignment(principal="user1", privileges=[Privilege.APPLY_TAG]),
+        PrivilegeAssignment(principal="user2", privileges=[Privilege.APPLY_TAG]),
+        PrivilegeAssignment(principal="userToAdd", privileges=[Privilege.CREATE]),
+    ]
+
+    result = get_permission_changes(
+        assignments_on_databricks=sample_permissions_list, assignments_from_properties=sample_privilege_assignments
+    )
+
+    assert result == [
+        PermissionsChange(add=[Privilege.CREATE], principal="userToAdd", remove=None),
+        PermissionsChange(add=None, principal="userToRemove", remove=[Privilege.CREATE]),
+        PermissionsChange(add=[Privilege.APPLY_TAG], principal="user2", remove=[Privilege.CREATE]),
+    ]

typescript/src/resources/deploy-lambda.ts

@@ -7,6 +7,7 @@ import {Workspace, WorkspaceProperties} from "./account/workspace";
 import {InstanceProfile, InstanceProfileProperties} from "./instance-profiles/instance-profile";
 import {Cluster, ClusterProperties} from "./clusters/cluster";
 import {ClusterPermissions, ClusterPermissionsProperties} from "./permissions/cluster-permissions";
+import {VolumePermissions, VolumePermissionsProperties} from "./permissions/volumePermissions";
 import {DbfsFile, DbfsFileProperties} from "./dbfs/dbfs-file";
 import {SecretScope, SecretScopeProperties} from "./secrets/secret-scope";
 import {Job, JobProperties} from "./jobs/job";
@@ -179,6 +180,13 @@ export abstract class IDatabricksDeployLambda extends Construct {
         });
     }
 
+    public createVolumePermissions(scope: Construct, id: string, props: VolumePermissionsProperties): VolumePermissions {
+        return new VolumePermissions(scope, id, {
+            ...props,
+            serviceToken: this.serviceToken
+        });
+    }
+
     public createInstancePool(scope: Construct, id: string, props: InstancePoolProperties): InstancePool {
         return new InstancePool(scope, id, {
             ...props,

typescript/src/resources/permissions/index.ts

@@ -4,3 +4,4 @@ export * from "./job-permissions";
 export * from "./cluster-policy-permissions";
 export * from "./registered-model-permissions";
 export * from "./experiment-permissions";
+export * from "./volumePermissions";