feedback

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

@@ -215,6 +215,20 @@ def _update_cluster_policy_with_instance_profile(self, policy: Policy, iam_insta
 
         self._ws.cluster_policies.edit(str(policy.policy_id), str(policy.name), definition=json.dumps(definition_dict))
 
+    def _update_sql_dac_with_instance_profile(self, iam_instance_profile: AWSInstanceProfile, prompts: Prompts):
+        warehouse_config = self._ws.warehouses.get_workspace_warehouse_config()
+        if warehouse_config.instance_profile_arn is not None:
+            if not prompts.confirm(
+                f"There is an existing instance profile {warehouse_config.instance_profile_arn} specified in the "
+                f"workspace warehouse config. Do you want UCX to to update it with the uber instance profile?"
+            ):
+                return
+        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,
+        )
+
     def get_instance_profile(self, instance_profile_name: str) -> AWSInstanceProfile | None:
         instance_profile_arn = self._aws_resources.get_instance_profile(instance_profile_name)
 
@@ -283,6 +297,7 @@ def create_uber_principal(self, prompts: Prompts):
             config.uber_instance_profile = iam_instance_profile.instance_profile_arn
             self._installation.save(config)
             self._update_cluster_policy_with_instance_profile(cluster_policy, iam_instance_profile)
+            self._update_sql_dac_with_instance_profile(iam_instance_profile, prompts)
             logger.info(f"Cluster policy \"{cluster_policy.name}\" updated successfully")
         except PermissionError:
             self._aws_resources.delete_instance_profile(iam_role_name, iam_role_name)

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

@@ -186,7 +186,7 @@ def _update_sql_dac_with_spn(
                     ),
                     EndpointConfPair(
                         f"spark_conf.fs.azure.account.oauth2.client.secret.{storage.name}.dfs.core.windows.net",
-                        f"{{secrets/{inventory_database}/uber_principal_secret}}",
+                        "{{secrets/" + inventory_database + "/uber_principal_secret}}",
                     ),
                 ]
             )

tests/unit/aws/test_access.py

@@ -16,6 +16,7 @@
     StorageCredentialInfo,
 )
 from databricks.sdk.service.compute import InstanceProfile, Policy
+from databricks.sdk.service.sql import GetWorkspaceWarehouseConfigResponse, EndpointConfPair
 
 from databricks.labs.ucx.assessment.aws import (
     AWSPolicyAction,
@@ -218,11 +219,20 @@ def test_create_uber_principal_existing_role(mock_ws, mock_installation, backend
         policy_id="foo", name="Unity Catalog Migration (ucx) (me@example.com)", definition=json.dumps({"foo": "bar"})
     )
     mock_ws.cluster_policies.get.return_value = cluster_policy
+    mock_ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        instance_profile_arn="arn:aws:iam::12345:instance-profile/existing-role"
+    )
     instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
     aws = create_autospec(AWSResources)
     aws.get_instance_profile.return_value = instance_profile_arn
     locations = ExternalLocations(mock_ws, backend, "ucx")
-    prompts = MockPrompts({"We have identified existing UCX migration role *": "yes"})
+    prompts = MockPrompts(
+        {
+            "Do you want to create new migration role *": "yes",
+            "There is an existing instance profile *": "yes",
+            "We have identified existing UCX migration role *": "yes",
+        }
+    )
     aws_resource_permissions = AWSResourcePermissions(
         mock_installation,
         mock_ws,
@@ -234,21 +244,33 @@ def test_create_uber_principal_existing_role(mock_ws, mock_installation, backend
     mock_ws.cluster_policies.edit.assert_called_with(
         'foo', 'Unity Catalog Migration (ucx) (me@example.com)', definition=json.dumps(definition)
     )
+    mock_ws.warehouses.set_workspace_warehouse_config.assert_called_with(
+        data_access_config=None,
+        instance_profile_arn='arn:aws:iam::12345:instance-profile/role1',
+        sql_configuration_parameters=None,
+    )
 
 
 def test_create_uber_principal_no_existing_role(mock_ws, mock_installation, backend, locations):
     cluster_policy = Policy(
         policy_id="foo", name="Unity Catalog Migration (ucx) (me@example.com)", definition=json.dumps({"foo": "bar"})
     )
     mock_ws.cluster_policies.get.return_value = cluster_policy
+    mock_ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
+        data_access_config=[EndpointConfPair("jdbc", "jdbc:sqlserver://localhost:1433;databaseName=master")]
+    )
     aws = create_autospec(AWSResources)
     aws.role_exists.return_value = False
     instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
     aws.create_migration_role.return_value = instance_profile_arn
     aws.create_instance_profile.return_value = instance_profile_arn
     aws.get_instance_profile.return_value = instance_profile_arn
     locations = ExternalLocations(mock_ws, backend, "ucx")
-    prompts = MockPrompts({"Do you want to create new migration role *": "yes"})
+    prompts = MockPrompts(
+        {
+            "Do you want to create new migration role *": "yes",
+        }
+    )
     aws_resource_permissions = AWSResourcePermissions(
         mock_installation,
         mock_ws,
@@ -261,6 +283,11 @@ def test_create_uber_principal_no_existing_role(mock_ws, mock_installation, back
     mock_ws.cluster_policies.edit.assert_called_with(
         'foo', 'Unity Catalog Migration (ucx) (me@example.com)', definition=json.dumps(definition)
     )
+    mock_ws.warehouses.set_workspace_warehouse_config.assert_called_with(
+        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,
+    )
 
 
 def test_create_uber_principal_no_storage(mock_ws, mock_installation, locations):

tests/unit/azure/test_access.py

@@ -9,6 +9,7 @@
 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.workspace import GetSecretResponse
 
 from databricks.labs.ucx.azure.access import AzureResourcePermissions
@@ -300,6 +301,7 @@ def test_create_global_spn():
     )
     w.cluster_policies.get.return_value = cluster_policy
     w.secrets.get_secret.return_value = GetSecretResponse("uber_principal_secret", "mypwd")
+    w.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse
     rows = {"SELECT \\* FROM ucx.external_locations": [["abfss://container1@sto2.dfs.core.windows.net/folder1", "1"]]}
     backend = MockBackend(rows=rows)
     location = ExternalLocations(w, backend, "ucx")
@@ -355,7 +357,27 @@ def test_create_global_spn():
     w.secrets.create_scope.assert_called_with("ucx")
     w.secrets.put_secret.assert_called_with("ucx", "uber_principal_secret", string_value="mypwd")
     w.warehouses.get_workspace_warehouse_config.assert_called_once()
-    w.warehouses.set_workspace_warehouse_config.assert_called_once()
+    w.warehouses.set_workspace_warehouse_config.assert_called_with(
+        data_access_config=[
+            EndpointConfPair(
+                key='spark_conf.fs.azure.account.oauth2.client.id.sto2.dfs.core.windows.net', value='appIduser1'
+            ),
+            EndpointConfPair(
+                key='spark_conf.fs.azure.account.oauth.provider.type.sto2.dfs.core.windows.net',
+                value='org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider',
+            ),
+            EndpointConfPair(
+                key='spark_conf.fs.azure.account.oauth2.client.endpoint.sto2.dfs.core.windows.net',
+                value='https://login.microsoftonline.com/bar/oauth2/token',
+            ),
+            EndpointConfPair(key='spark_conf.fs.azure.account.auth.type.sto2.dfs.core.windows.net', value='OAuth'),
+            EndpointConfPair(
+                key='spark_conf.fs.azure.account.oauth2.client.secret.sto2.dfs.core.windows.net',
+                value='{{secrets/ucx/uber_principal_secret}}',
+            ),
+        ],
+        sql_configuration_parameters=None,
+    )
 
 
 def test_create_access_connectors_for_storage_accounts_logs_no_storage_accounts(caplog):