Update for Databricks SDK 0.56+ (#4178)

## Changes

This PR updates the project to work with Databricks SDK 0.56+ which
included breaking changes.

Incidental changes included adding type-hints to ensure that all the
updated functions are linted; some were not, which prevented some of the
breaking changes from being found during linting.

### Linked issues

Supersedes #4144.

This PR is also required for the upstream blueprint, whose CI/CD always
tests UCX with the latest available SDK (even if UCX specifies an
earlier one):

 - databrickslabs/blueprint#243
 - databrickslabs/blueprint#244

### Tests

- existing unit tests
- existing integration tests

Author: asnare

pyproject.toml

@@ -45,7 +45,7 @@ classifiers = [
 ]
 
 
-dependencies = ["databricks-sdk>=0.44.0,<0.54.0",
+dependencies = ["databricks-sdk>=0.56.0,<0.58.0",
                 "databricks-labs-lsql>=0.16.0,<0.17.0",
                 "databricks-labs-blueprint>=0.11.0,<0.12.0",
                 "PyYAML>=6.0.0,<6.1.0",

src/databricks/labs/ucx/hive_metastore/federation.py

@@ -299,11 +299,11 @@ def _add_missing_permissions_if_needed(self, location_name: str, current_user: s
         grants = self._location_grants(location_name)
         if Privilege.CREATE_FOREIGN_SECURABLE not in grants[current_user]:
             change = PermissionsChange(principal=current_user, add=[Privilege.CREATE_FOREIGN_SECURABLE])
-            self._ws.grants.update(SecurableType.EXTERNAL_LOCATION, location_name, changes=[change])
+            self._ws.grants.update(SecurableType.EXTERNAL_LOCATION.value, location_name, changes=[change])
 
     def _location_grants(self, location_name: str) -> dict[str, set[Privilege]]:
         grants: dict[str, set[Privilege]] = collections.defaultdict(set)
-        result = self._ws.grants.get(SecurableType.EXTERNAL_LOCATION, location_name)
+        result = self._ws.grants.get(SecurableType.EXTERNAL_LOCATION.value, location_name)
         if not result.privilege_assignments:
             return grants
         for assignment in result.privilege_assignments:

src/databricks/labs/ucx/hive_metastore/grants.py

@@ -737,11 +737,11 @@ def apply_location_acl(self):
                 self._update_location_permissions(location_name, principals)
         logger.info("Applied all the permission on external location")
 
-    def _update_location_permissions(self, location_name: str, principals: list[str]):
+    def _update_location_permissions(self, location_name: str, principals: list[str]) -> None:
         permissions = [Privilege.CREATE_EXTERNAL_TABLE, Privilege.CREATE_EXTERNAL_VOLUME, Privilege.READ_FILES]
         changes = [PermissionsChange(add=permissions, principal=principal) for principal in principals]
         self._ws.grants.update(
-            SecurableType.EXTERNAL_LOCATION,
+            SecurableType.EXTERNAL_LOCATION.value,
             location_name,
             changes=changes,
         )

src/databricks/labs/ucx/hive_metastore/table_move.py

@@ -192,9 +192,9 @@ def _alias_table(
                 logger.error(f"Failed to alias table {from_table_name}: {err!s}", exc_info=True)
             return False
 
-    def _reapply_grants(self, from_table_name, to_table_name, *, target_view: bool = False):
+    def _reapply_grants(self, from_table_name: str, to_table_name: str, *, target_view: bool = False) -> None:
         try:
-            grants = self._ws.grants.get(SecurableType.TABLE, from_table_name)
+            grants = self._ws.grants.get(SecurableType.TABLE.value, from_table_name)
         except NotFound:
             logger.warning(f"removed on the backend {from_table_name}")
             return
@@ -215,7 +215,7 @@ def _reapply_grants(self, from_table_name, to_table_name, *, target_view: bool =
             if privileges:
                 grants_changes.append(PermissionsChange(list(privileges), permission.principal))
 
-        self._ws.grants.update(SecurableType.TABLE, to_table_name, changes=grants_changes)
+        self._ws.grants.update(SecurableType.TABLE.value, to_table_name, changes=grants_changes)
 
     def _recreate_table(self, from_table_name, to_table_name, *, is_managed: bool, del_table: bool):
         drop_table = f"DROP TABLE {escape_sql_identifier(from_table_name)}"

tests/integration/aws/test_access.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 from databricks.labs.blueprint.tui import MockPrompts
+from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.catalog import AwsIamRoleRequest
 from databricks.sdk.service.compute import DataSecurityMode, AwsAttributes
 from databricks.sdk.service.iam import PermissionLevel
@@ -104,13 +105,13 @@ def test_create_uber_instance_profile(ws, env_or_skip, make_random, make_cluster
 
 def test_create_external_location_validate_acl(
     make_cluster_permissions,
-    ws,
+    ws: WorkspaceClient,
     make_user,
     make_cluster,
     aws_cli_ctx,
     env_or_skip,
     inventory_schema,
-):
+) -> None:
     aws_cli_ctx.workspace_installation.run()
     aws_cli_ctx.with_dummy_resource_permission()
     aws_cli_ctx.sql_backend.save_table(
@@ -149,13 +150,13 @@ def test_create_external_location_validate_acl(
     try:
         external_location_migration.run()
         permissions = ws.grants.get(
-            SecurableType.EXTERNAL_LOCATION, external_location_name, principal=cluster_user.user_name
+            SecurableType.EXTERNAL_LOCATION.value, external_location_name, principal=cluster_user.user_name
         )
         expected_aws_permission = PrivilegeAssignment(
             principal=cluster_user.user_name,
             privileges=[Privilege.CREATE_EXTERNAL_TABLE, Privilege.CREATE_EXTERNAL_VOLUME, Privilege.READ_FILES],
         )
-        assert expected_aws_permission in permissions.privilege_assignments
+        assert permissions.privilege_assignments and expected_aws_permission in permissions.privilege_assignments
     finally:
         remove_aws_permissions = [
             Privilege.CREATE_EXTERNAL_TABLE,
@@ -164,7 +165,7 @@ def test_create_external_location_validate_acl(
         ]
         ws.external_locations.delete(external_location_name)
         ws.grants.update(
-            SecurableType.EXTERNAL_LOCATION,
+            SecurableType.EXTERNAL_LOCATION.value,
             env_or_skip("TEST_A_LOCATION"),
             changes=[PermissionsChange(remove=remove_aws_permissions, principal=cluster_user.user_name)],
         )

tests/integration/azure/test_locations.py

@@ -4,6 +4,7 @@
 import pytest
 from databricks.labs.blueprint.installation import MockInstallation
 from databricks.labs.blueprint.tui import MockPrompts
+from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors.platform import NotFound
 from databricks.sdk.service.iam import PermissionLevel
 from databricks.sdk.service.compute import DataSecurityMode
@@ -217,7 +218,9 @@ def test_overlapping_location(caplog, ws, sql_backend, inventory_schema, az_cli_
         save_delete_location(ws, "uctest_ziyuanqintest_overlap")
 
 
-def test_run_validate_acl(make_cluster_permissions, ws, make_user, make_cluster, az_cli_ctx, env_or_skip):
+def test_run_validate_acl(
+    make_cluster_permissions, ws: WorkspaceClient, make_user, make_cluster, az_cli_ctx, env_or_skip
+) -> None:
     az_cli_ctx.with_dummy_resource_permission()
     az_cli_ctx.save_locations()
     cluster = make_cluster(single_node=True, spark_conf=_SPARK_CONF, data_security_mode=DataSecurityMode.NONE)
@@ -232,21 +235,21 @@ def test_run_validate_acl(make_cluster_permissions, ws, make_user, make_cluster,
     try:
         location_migration.run()
         permissions = ws.grants.get(
-            SecurableType.EXTERNAL_LOCATION, env_or_skip("TEST_A_LOCATION"), principal=user.user_name
+            SecurableType.EXTERNAL_LOCATION.value, env_or_skip("TEST_A_LOCATION"), principal=user.user_name
         )
         expected_azure_permission = PrivilegeAssignment(
             principal=user.user_name,
             privileges=[Privilege.CREATE_EXTERNAL_TABLE, Privilege.CREATE_EXTERNAL_VOLUME, Privilege.READ_FILES],
         )
-        assert expected_azure_permission in permissions.privilege_assignments
+        assert permissions.privilege_assignments and expected_azure_permission in permissions.privilege_assignments
     finally:
         remove_azure_permissions = [
             Privilege.CREATE_EXTERNAL_TABLE,
             Privilege.CREATE_EXTERNAL_VOLUME,
             Privilege.READ_FILES,
         ]
         ws.grants.update(
-            SecurableType.EXTERNAL_LOCATION,
+            SecurableType.EXTERNAL_LOCATION.value,
             env_or_skip("TEST_A_LOCATION"),
             changes=[PermissionsChange(remove=remove_azure_permissions, principal=user.user_name)],
         )

tests/integration/hive_metastore/test_catalog_schema.py

@@ -6,9 +6,8 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
 from databricks.sdk.retries import retried
-from databricks.sdk.service.catalog import PermissionsList
 from databricks.sdk.service.compute import DataSecurityMode, AwsAttributes
-from databricks.sdk.service.catalog import Privilege, SecurableType, PrivilegeAssignment
+from databricks.sdk.service.catalog import GetPermissionsResponse, Privilege, SecurableType, PrivilegeAssignment
 from databricks.sdk.service.iam import PermissionLevel
 
 from databricks.labs.ucx.hive_metastore.grants import Grant
@@ -94,8 +93,8 @@ def test_create_catalog_schema_with_principal_acl_azure(
     mock_prompts = MockPrompts({"Please provide storage location url for catalog: *": ""})
     catalog_schema.create_all_catalogs_schemas(mock_prompts)
 
-    schema_grants = ws.grants.get(SecurableType.SCHEMA, schema_name)
-    catalog_grants = ws.grants.get(SecurableType.CATALOG, catalog_name)
+    schema_grants = ws.grants.get(SecurableType.SCHEMA.value, schema_name)
+    catalog_grants = ws.grants.get(SecurableType.CATALOG.value, catalog_name)
     schema_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_SCHEMA])
     catalog_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_CATALOG])
     assert schema_grant in schema_grants.privilege_assignments
@@ -128,8 +127,8 @@ def test_create_catalog_schema_with_principal_acl_aws(
     mock_prompts = MockPrompts({"Please provide storage location url for catalog: *": ""})
     catalog_schema.create_all_catalogs_schemas(mock_prompts)
 
-    schema_grants = ws.grants.get(SecurableType.SCHEMA, schema_name)
-    catalog_grants = ws.grants.get(SecurableType.CATALOG, catalog_name)
+    schema_grants = ws.grants.get(SecurableType.SCHEMA.value, schema_name)
+    catalog_grants = ws.grants.get(SecurableType.CATALOG.value, catalog_name)
     schema_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_SCHEMA])
     catalog_grant = PrivilegeAssignment(user.user_name, [Privilege.USE_CATALOG])
     assert schema_grant in schema_grants.privilege_assignments
@@ -167,8 +166,8 @@ def test_create_catalog_schema_with_legacy_hive_metastore_privileges(
     runtime_ctx.catalog_schema.create_all_catalogs_schemas(mock_prompts, properties=properties)
 
     @retried(on=[NotFound], timeout=timedelta(seconds=20))
-    def get_schema_permissions_list(full_name: str) -> PermissionsList:
-        return ws.grants.get(SecurableType.SCHEMA, full_name)
+    def get_schema_permissions_list(full_name: str) -> GetPermissionsResponse:
+        return ws.grants.get(SecurableType.SCHEMA.value, full_name)
 
     assert (
         ws.schemas.get(f"{dst_catalog_name}.{dst_schema_name}").owner

tests/integration/hive_metastore/test_migrate.py

@@ -2,12 +2,13 @@
 from datetime import timedelta
 
 import pytest
-from databricks.sdk import AccountClient
+from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
 from databricks.sdk.retries import retried
 from databricks.sdk.service.compute import DataSecurityMode, AwsAttributes
 from databricks.sdk.service.catalog import Privilege, SecurableType, TableInfo, TableType
 from databricks.sdk.service.iam import PermissionLevel
+
 from databricks.labs.ucx.__about__ import __version__
 
 from databricks.labs.ucx.config import WorkspaceConfig
@@ -650,7 +651,9 @@ def test_mapping_reverts_table(ws, sql_backend, runtime_ctx, make_catalog):
 
 
 # @retried(on=[NotFound], timeout=timedelta(minutes=3))
-def test_migrate_managed_tables_with_acl(ws, sql_backend, runtime_ctx, make_catalog, make_user, env_or_skip):
+def test_migrate_managed_tables_with_acl(
+    ws: WorkspaceClient, sql_backend, runtime_ctx, make_catalog, make_user, env_or_skip
+) -> None:
     src_schema = runtime_ctx.make_schema(catalog_name="hive_metastore")
     src_managed_table = runtime_ctx.make_table(catalog_name=src_schema.catalog_name, schema_name=src_schema.name)
     user = make_user()
@@ -687,10 +690,11 @@ def test_migrate_managed_tables_with_acl(ws, sql_backend, runtime_ctx, make_cata
     assert len(target_tables) == 1
 
     target_table_properties = ws.tables.get(f"{dst_schema.full_name}.{src_managed_table.name}").properties
+    assert target_table_properties
     assert target_table_properties["upgraded_from"] == src_managed_table.full_name
     assert target_table_properties[Table.UPGRADED_FROM_WS_PARAM] == str(ws.get_workspace_id())
 
-    target_table_grants = ws.grants.get(SecurableType.TABLE, f"{dst_schema.full_name}.{src_managed_table.name}")
+    target_table_grants = ws.grants.get(SecurableType.TABLE.value, f"{dst_schema.full_name}.{src_managed_table.name}")
     target_principals = [pa for pa in target_table_grants.privilege_assignments or [] if pa.principal == user.user_name]
     assert len(target_principals) == 1, f"Missing grant for user {user.user_name}"
     assert target_principals[0].privileges == [Privilege.MODIFY, Privilege.SELECT]
@@ -721,8 +725,8 @@ def test_migrate_managed_tables_with_acl(ws, sql_backend, runtime_ctx, make_cata
 
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_migrate_external_tables_with_principal_acl_azure(
-    ws, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster, make_ucx_group
-):
+    ws: WorkspaceClient, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster, make_ucx_group
+) -> None:
     if not ws.config.is_azure:
         pytest.skip("only works in azure test env")
     ctx, table_full_name, _, _ = prepared_principal_acl
@@ -741,7 +745,8 @@ def test_migrate_external_tables_with_principal_acl_azure(
     )
     table_migrate.migrate_tables(what=What.EXTERNAL_SYNC)
 
-    target_table_grants = ws.grants.get(SecurableType.TABLE, table_full_name)
+    target_table_grants = ws.grants.get(SecurableType.TABLE.value, table_full_name)
+    assert target_table_grants.privilege_assignments is not None
     match = False
     for _ in target_table_grants.privilege_assignments:
         if _.principal == user_with_cluster_access.user_name and _.privileges == [Privilege.ALL_PRIVILEGES]:
@@ -764,8 +769,8 @@ def test_migrate_external_tables_with_principal_acl_azure(
 
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_migrate_external_tables_with_principal_acl_aws(
-    ws, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster, env_or_skip
-):
+    ws: WorkspaceClient, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster, env_or_skip
+) -> None:
     ctx, table_full_name, _, _ = prepared_principal_acl
     ctx.with_dummy_resource_permission()
     cluster = make_cluster(
@@ -782,7 +787,8 @@ def test_migrate_external_tables_with_principal_acl_aws(
     )
     table_migrate.migrate_tables(what=What.EXTERNAL_SYNC)
 
-    target_table_grants = ws.grants.get(SecurableType.TABLE, table_full_name)
+    target_table_grants = ws.grants.get(SecurableType.TABLE.value, table_full_name)
+    assert target_table_grants.privilege_assignments is not None
     match = False
     for _ in target_table_grants.privilege_assignments:
         if _.principal == user.user_name and _.privileges == [Privilege.ALL_PRIVILEGES]:
@@ -793,8 +799,12 @@ def test_migrate_external_tables_with_principal_acl_aws(
 
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_migrate_external_tables_with_principal_acl_aws_warehouse(
-    ws, make_user, prepared_principal_acl, make_warehouse_permissions, make_warehouse, env_or_skip
-):
+    ws: WorkspaceClient,
+    make_user,
+    prepared_principal_acl,
+    make_warehouse_permissions,
+    make_warehouse,
+) -> None:
     if not ws.config.is_aws:
         pytest.skip("temporary: only works in aws test env")
     ctx, table_full_name, _, _ = prepared_principal_acl
@@ -809,7 +819,8 @@ def test_migrate_external_tables_with_principal_acl_aws_warehouse(
     )
     table_migrate.migrate_tables(what=What.EXTERNAL_SYNC)
 
-    target_table_grants = ws.grants.get(SecurableType.TABLE, table_full_name)
+    target_table_grants = ws.grants.get(SecurableType.TABLE.value, table_full_name)
+    assert target_table_grants.privilege_assignments is not None
     match = False
     for _ in target_table_grants.privilege_assignments:
         if _.principal == user.user_name and _.privileges == [Privilege.ALL_PRIVILEGES]:
@@ -879,8 +890,8 @@ def test_migrate_table_in_mount(
 
 
 def test_migrate_external_tables_with_spn_azure(
-    ws, make_user, prepared_principal_acl, make_cluster_permissions, make_cluster
-):
+    ws: WorkspaceClient, prepared_principal_acl, make_cluster_permissions, make_cluster
+) -> None:
     if not ws.config.is_azure:
         pytest.skip("temporary: only works in azure test env")
     ctx, table_full_name, _, _ = prepared_principal_acl
@@ -897,7 +908,8 @@ def test_migrate_external_tables_with_spn_azure(
     )
     table_migrate.migrate_tables(what=What.EXTERNAL_SYNC)
 
-    target_table_grants = ws.grants.get(SecurableType.TABLE, table_full_name)
+    target_table_grants = ws.grants.get(SecurableType.TABLE.value, table_full_name)
+    assert target_table_grants.privilege_assignments is not None
     match = False
     for _ in target_table_grants.privilege_assignments:
         if _.principal == spn_with_mount_access and _.privileges == [Privilege.ALL_PRIVILEGES]: