Fix missing security_policy when updating SQL warehouse config (#2409)

1. Fix `InvalidParameterValue: Endpoint security policy is required and
must be one of NONE, DATA_ACCESS_CONTROL, PASSTHROUGH.` error in #2405
2. Print new SQL warehouse data access config in the log when
`databricks.sdk.errors.platform.InvalidParameterValue:
enable_serverless_compute is required.` error in #2405 occurs. With that
users can manually copy paste the config in UI to finish setting up the
uber principal,

Author: qziyuan

src/databricks/labs/ucx/aws/access.py

@@ -12,6 +12,7 @@
 from databricks.sdk.errors import NotFound, ResourceDoesNotExist
 from databricks.sdk.service.catalog import Privilege
 from databricks.sdk.service.compute import Policy
+from databricks.sdk.service.sql import SetWorkspaceWarehouseConfigRequestSecurityPolicy
 
 from databricks.labs.ucx.assessment.aws import (
     AWSInstanceProfile,
@@ -266,10 +267,16 @@ def _update_sql_dac_with_instance_profile(self, iam_instance_profile: AWSInstanc
                 f"workspace warehouse config. Do you want UCX to to update it with the uber instance profile?"
             ):
                 return
+        security_policy = (
+            SetWorkspaceWarehouseConfigRequestSecurityPolicy(warehouse_config.security_policy.value)
+            if warehouse_config.security_policy
+            else SetWorkspaceWarehouseConfigRequestSecurityPolicy.NONE
+        )
         self._ws.warehouses.set_workspace_warehouse_config(
             data_access_config=warehouse_config.data_access_config,
             sql_configuration_parameters=warehouse_config.sql_configuration_parameters,
             instance_profile_arn=iam_instance_profile.instance_profile_arn,
+            security_policy=security_policy,
         )
 
     def get_instance_profile(self, instance_profile_name: str) -> AWSInstanceProfile | None:

src/databricks/labs/ucx/azure/access.py

@@ -10,9 +10,9 @@
 from databricks.labs.blueprint.parallel import ManyError, Threads
 from databricks.labs.blueprint.tui import Prompts
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import NotFound, ResourceAlreadyExists
+from databricks.sdk.errors import InvalidParameterValue, NotFound, ResourceAlreadyExists
 from databricks.sdk.service.catalog import Privilege
-from databricks.sdk.service.sql import EndpointConfPair
+from databricks.sdk.service.sql import EndpointConfPair, SetWorkspaceWarehouseConfigRequestSecurityPolicy
 
 from databricks.labs.ucx.azure.resources import (
     AccessConnector,
@@ -240,10 +240,25 @@ def _update_sql_dac_with_spn(
                     ),
                 ]
             )
-        self._ws.warehouses.set_workspace_warehouse_config(
-            data_access_config=sql_dac,
-            sql_configuration_parameters=warehouse_config.sql_configuration_parameters,
+        security_policy = (
+            SetWorkspaceWarehouseConfigRequestSecurityPolicy(warehouse_config.security_policy.value)
+            if warehouse_config.security_policy
+            else SetWorkspaceWarehouseConfigRequestSecurityPolicy.NONE
         )
+        try:
+            self._ws.warehouses.set_workspace_warehouse_config(
+                data_access_config=sql_dac,
+                sql_configuration_parameters=warehouse_config.sql_configuration_parameters,
+                security_policy=security_policy,
+            )
+        # TODO: Remove following try except once https://github.com/databricks/databricks-sdk-py/issues/305 is fixed
+        except InvalidParameterValue as error:
+            sql_dac_log_msg = "\n".join(f"{config_pair.key} {config_pair.value}" for config_pair in sql_dac)
+            logger.error(
+                f'Adding uber principal to SQL warehouse Data Access Properties is failed using Python SDK with error "{error}". '
+                f'Please try applying the following configs manually in the worksapce admin UI:\n{sql_dac_log_msg}'
+            )
+            raise error
 
     def create_uber_principal(self, prompts: Prompts):
         config = self._installation.load(WorkspaceConfig)

tests/unit/aws/test_access.py

@@ -11,7 +11,12 @@
 from databricks.sdk.service import iam
 from databricks.sdk.service.catalog import AwsIamRoleResponse, ExternalLocationInfo, StorageCredentialInfo
 from databricks.sdk.service.compute import InstanceProfile, Policy
-from databricks.sdk.service.sql import EndpointConfPair, GetWorkspaceWarehouseConfigResponse
+from databricks.sdk.service.sql import (
+    EndpointConfPair,
+    GetWorkspaceWarehouseConfigResponse,
+    SetWorkspaceWarehouseConfigRequestSecurityPolicy,
+    GetWorkspaceWarehouseConfigResponseSecurityPolicy,
+)
 
 from databricks.labs.ucx.assessment.aws import AWSPolicyAction, AWSResources, AWSRole, AWSRoleAction
 from databricks.labs.ucx.aws.access import AWSResourcePermissions
@@ -252,6 +257,7 @@ def test_create_uber_principal_existing_role(mock_ws, mock_installation, backend
         data_access_config=None,
         instance_profile_arn='arn:aws:iam::12345:instance-profile/role1',
         sql_configuration_parameters=None,
+        security_policy=SetWorkspaceWarehouseConfigRequestSecurityPolicy.NONE,
     )
 
 
@@ -287,9 +293,43 @@ def test_create_uber_principal_no_existing_role(mock_ws, mock_installation, back
         data_access_config=[EndpointConfPair("jdbc", "jdbc:sqlserver://localhost:1433;databaseName=master")],
         instance_profile_arn='arn:aws:iam::12345:instance-profile/role1',
         sql_configuration_parameters=None,
+        security_policy=SetWorkspaceWarehouseConfigRequestSecurityPolicy.NONE,
     )
 
 
+@pytest.mark.parametrize(
+    "get_security_policy, set_security_policy",
+    [
+        (
+            GetWorkspaceWarehouseConfigResponseSecurityPolicy.DATA_ACCESS_CONTROL,
+            SetWorkspaceWarehouseConfigRequestSecurityPolicy.DATA_ACCESS_CONTROL,
+        ),
+        (GetWorkspaceWarehouseConfigResponseSecurityPolicy.NONE, SetWorkspaceWarehouseConfigRequestSecurityPolicy.NONE),
+    ],
+)
+def test_create_uber_principal_set_warehouse_config_security_policy(
+    mock_ws, mock_installation, backend, get_security_policy, set_security_policy
+):
+    mock_ws.cluster_policies.get.return_value = Policy(definition=json.dumps({"foo": "bar"}))
+    mock_ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        security_policy=get_security_policy
+    )
+
+    instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
+    aws = create_autospec(AWSResources)
+    aws.get_instance_profile_arn.return_value = instance_profile_arn
+
+    aws_resource_permissions = AWSResourcePermissions(
+        mock_installation,
+        mock_ws,
+        aws,
+        ExternalLocations(mock_ws, backend, "ucx"),
+    )
+    aws_resource_permissions.create_uber_principal(MockPrompts({".*": "yes"}))
+
+    assert mock_ws.warehouses.set_workspace_warehouse_config.call_args.kwargs["security_policy"] == set_security_policy
+
+
 def test_create_uber_principal_no_storage(mock_ws, mock_installation, locations):
     cluster_policy = Policy(
         policy_id="foo", name="Unity Catalog Migration (ucx) ([email protected])", definition=json.dumps({"foo": "bar"})

tests/unit/azure/test_access.py

@@ -9,7 +9,12 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
 from databricks.sdk.service.compute import Policy
-from databricks.sdk.service.sql import GetWorkspaceWarehouseConfigResponse, EndpointConfPair
+from databricks.sdk.service.sql import (
+    GetWorkspaceWarehouseConfigResponse,
+    EndpointConfPair,
+    GetWorkspaceWarehouseConfigResponseSecurityPolicy,
+    SetWorkspaceWarehouseConfigRequestSecurityPolicy,
+)
 from databricks.sdk.service.workspace import GetSecretResponse
 
 from databricks.labs.ucx.azure.access import AzureResourcePermissions
@@ -24,6 +29,7 @@
 from databricks.labs.ucx.hive_metastore import ExternalLocations
 
 from . import azure_api_client
+from .. import DEFAULT_CONFIG
 
 
 def test_save_spn_permissions_no_external_table(caplog):
@@ -744,9 +750,43 @@ def test_create_global_spn():
             ),
         ],
         sql_configuration_parameters=None,
+        security_policy=SetWorkspaceWarehouseConfigRequestSecurityPolicy.NONE,
     )
 
 
+@pytest.mark.parametrize(
+    "get_security_policy, set_security_policy",
+    [
+        (
+            GetWorkspaceWarehouseConfigResponseSecurityPolicy.DATA_ACCESS_CONTROL,
+            SetWorkspaceWarehouseConfigRequestSecurityPolicy.DATA_ACCESS_CONTROL,
+        ),
+        (GetWorkspaceWarehouseConfigResponseSecurityPolicy.NONE, SetWorkspaceWarehouseConfigRequestSecurityPolicy.NONE),
+    ],
+)
+def test_create_global_spn_set_warehouse_config_security_policy(get_security_policy, set_security_policy):
+    w = create_autospec(WorkspaceClient)
+    w.cluster_policies.get.return_value = Policy(definition=json.dumps({"foo": "bar"}))
+    w.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        security_policy=get_security_policy
+    )
+    rows = {
+        "SELECT \\* FROM hive_metastore.ucx.external_locations": [
+            ["abfss://[email protected]/folder1", "1"]
+        ]
+    }
+    location = ExternalLocations(w, MockBackend(rows=rows), "ucx")
+    installation = MockInstallation(DEFAULT_CONFIG.copy())
+    api_client = azure_api_client()
+    azure_resources = AzureResources(api_client, api_client, include_subscriptions="002")
+    azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
+    azure_resource_permission.create_uber_principal(
+        MockPrompts({"Enter a name for the uber service principal to be created*": "UCXServicePrincipal"})
+    )
+
+    assert w.warehouses.set_workspace_warehouse_config.call_args.kwargs["security_policy"] == set_security_policy
+
+
 def test_create_access_connectors_for_storage_accounts_logs_no_storage_accounts(caplog):
     """A warning should be logged when no storage account is present."""
     w = create_autospec(WorkspaceClient)