Skip to content

Commit 102c78c

Browse files
nkvuongnkvuong
committed
feedback
1 parent 148212f commit 102c78c

File tree

4 files changed

+68
-4
lines changed

4 files changed

+68
-4
lines changed

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,20 @@ def _update_cluster_policy_with_instance_profile(self, policy: Policy, iam_insta
215215

216216
self._ws.cluster_policies.edit(str(policy.policy_id), str(policy.name), definition=json.dumps(definition_dict))
217217

218+
def _update_sql_dac_with_instance_profile(self, iam_instance_profile: AWSInstanceProfile, prompts: Prompts):
219+
warehouse_config = self._ws.warehouses.get_workspace_warehouse_config()
220+
if warehouse_config.instance_profile_arn is not None:
221+
if not prompts.confirm(
222+
f"There is an existing instance profile {warehouse_config.instance_profile_arn} specified in the "
223+
f"workspace warehouse config. Do you want UCX to to update it with the uber instance profile?"
224+
):
225+
return
226+
self._ws.warehouses.set_workspace_warehouse_config(
227+
data_access_config=warehouse_config.data_access_config,
228+
sql_configuration_parameters=warehouse_config.sql_configuration_parameters,
229+
instance_profile_arn=iam_instance_profile.instance_profile_arn,
230+
)
231+
218232
def get_instance_profile(self, instance_profile_name: str) -> AWSInstanceProfile | None:
219233
instance_profile_arn = self._aws_resources.get_instance_profile(instance_profile_name)
220234

@@ -283,6 +297,7 @@ def create_uber_principal(self, prompts: Prompts):
283297
config.uber_instance_profile = iam_instance_profile.instance_profile_arn
284298
self._installation.save(config)
285299
self._update_cluster_policy_with_instance_profile(cluster_policy, iam_instance_profile)
300+
self._update_sql_dac_with_instance_profile(iam_instance_profile, prompts)
286301
logger.info(f"Cluster policy \"{cluster_policy.name}\" updated successfully")
287302
except PermissionError:
288303
self._aws_resources.delete_instance_profile(iam_role_name, iam_role_name)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ def _update_sql_dac_with_spn(
236236
),
237237
EndpointConfPair(
238238
f"spark_conf.fs.azure.account.oauth2.client.secret.{storage.name}.dfs.core.windows.net",
239-
f"{{secrets/{inventory_database}/uber_principal_secret}}",
239+
"{{secrets/" + inventory_database + "/uber_principal_secret}}",
240240
),
241241
]
242242
)

‎tests/unit/aws/test_access.py‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
StorageCredentialInfo,
1717
)
1818
from databricks.sdk.service.compute import InstanceProfile, Policy
19+
from databricks.sdk.service.sql import GetWorkspaceWarehouseConfigResponse, EndpointConfPair
1920

2021
from databricks.labs.ucx.assessment.aws import (
2122
AWSPolicyAction,
@@ -218,11 +219,20 @@ def test_create_uber_principal_existing_role(mock_ws, mock_installation, backend
218219
policy_id="foo", name="Unity Catalog Migration (ucx) ([email protected])", definition=json.dumps({"foo": "bar"})
219220
)
220221
mock_ws.cluster_policies.get.return_value = cluster_policy
222+
mock_ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
223+
instance_profile_arn="arn:aws:iam::12345:instance-profile/existing-role"
224+
)
221225
instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
222226
aws = create_autospec(AWSResources)
223227
aws.get_instance_profile.return_value = instance_profile_arn
224228
locations = ExternalLocations(mock_ws, backend, "ucx")
225-
prompts = MockPrompts({"We have identified existing UCX migration role *": "yes"})
229+
prompts = MockPrompts(
230+
{
231+
"Do you want to create new migration role *": "yes",
232+
"There is an existing instance profile *": "yes",
233+
"We have identified existing UCX migration role *": "yes",
234+
}
235+
)
226236
aws_resource_permissions = AWSResourcePermissions(
227237
mock_installation,
228238
mock_ws,
@@ -234,21 +244,33 @@ def test_create_uber_principal_existing_role(mock_ws, mock_installation, backend
234244
mock_ws.cluster_policies.edit.assert_called_with(
235245
'foo', 'Unity Catalog Migration (ucx) ([email protected])', definition=json.dumps(definition)
236246
)
247+
mock_ws.warehouses.set_workspace_warehouse_config.assert_called_with(
248+
data_access_config=None,
249+
instance_profile_arn='arn:aws:iam::12345:instance-profile/role1',
250+
sql_configuration_parameters=None,
251+
)
237252

238253

239254
def test_create_uber_principal_no_existing_role(mock_ws, mock_installation, backend, locations):
240255
cluster_policy = Policy(
241256
policy_id="foo", name="Unity Catalog Migration (ucx) ([email protected])", definition=json.dumps({"foo": "bar"})
242257
)
243258
mock_ws.cluster_policies.get.return_value = cluster_policy
259+
mock_ws.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse(
260+
data_access_config=[EndpointConfPair("jdbc", "jdbc:sqlserver://localhost:1433;databaseName=master")]
261+
)
244262
aws = create_autospec(AWSResources)
245263
aws.role_exists.return_value = False
246264
instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
247265
aws.create_migration_role.return_value = instance_profile_arn
248266
aws.create_instance_profile.return_value = instance_profile_arn
249267
aws.get_instance_profile.return_value = instance_profile_arn
250268
locations = ExternalLocations(mock_ws, backend, "ucx")
251-
prompts = MockPrompts({"Do you want to create new migration role *": "yes"})
269+
prompts = MockPrompts(
270+
{
271+
"Do you want to create new migration role *": "yes",
272+
}
273+
)
252274
aws_resource_permissions = AWSResourcePermissions(
253275
mock_installation,
254276
mock_ws,
@@ -261,6 +283,11 @@ def test_create_uber_principal_no_existing_role(mock_ws, mock_installation, back
261283
mock_ws.cluster_policies.edit.assert_called_with(
262284
'foo', 'Unity Catalog Migration (ucx) ([email protected])', definition=json.dumps(definition)
263285
)
286+
mock_ws.warehouses.set_workspace_warehouse_config.assert_called_with(
287+
data_access_config=[EndpointConfPair("jdbc", "jdbc:sqlserver://localhost:1433;databaseName=master")],
288+
instance_profile_arn='arn:aws:iam::12345:instance-profile/role1',
289+
sql_configuration_parameters=None,
290+
)
264291

265292

266293
def test_create_uber_principal_no_storage(mock_ws, mock_installation, locations):

‎tests/unit/azure/test_access.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from databricks.sdk import WorkspaceClient
1010
from databricks.sdk.errors import NotFound
1111
from databricks.sdk.service.compute import Policy
12+
from databricks.sdk.service.sql import GetWorkspaceWarehouseConfigResponse, EndpointConfPair
1213
from databricks.sdk.service.workspace import GetSecretResponse
1314

1415
from databricks.labs.ucx.azure.access import AzureResourcePermissions
@@ -657,6 +658,7 @@ def test_create_global_spn():
657658
)
658659
w.cluster_policies.get.return_value = cluster_policy
659660
w.secrets.get_secret.return_value = GetSecretResponse("uber_principal_secret", "mypwd")
661+
w.warehouses.get_workspace_warehouse_config.return_value = GetWorkspaceWarehouseConfigResponse
660662
rows = {"SELECT \\* FROM ucx.external_locations": [["abfss://[email protected]/folder1", "1"]]}
661663
backend = MockBackend(rows=rows)
662664
location = ExternalLocations(w, backend, "ucx")
@@ -712,7 +714,27 @@ def test_create_global_spn():
712714
w.secrets.create_scope.assert_called_with("ucx")
713715
w.secrets.put_secret.assert_called_with("ucx", "uber_principal_secret", string_value="mypwd")
714716
w.warehouses.get_workspace_warehouse_config.assert_called_once()
715-
w.warehouses.set_workspace_warehouse_config.assert_called_once()
717+
w.warehouses.set_workspace_warehouse_config.assert_called_with(
718+
data_access_config=[
719+
EndpointConfPair(
720+
key='spark_conf.fs.azure.account.oauth2.client.id.sto2.dfs.core.windows.net', value='appIduser1'
721+
),
722+
EndpointConfPair(
723+
key='spark_conf.fs.azure.account.oauth.provider.type.sto2.dfs.core.windows.net',
724+
value='org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider',
725+
),
726+
EndpointConfPair(
727+
key='spark_conf.fs.azure.account.oauth2.client.endpoint.sto2.dfs.core.windows.net',
728+
value='https://login.microsoftonline.com/bar/oauth2/token',
729+
),
730+
EndpointConfPair(key='spark_conf.fs.azure.account.auth.type.sto2.dfs.core.windows.net', value='OAuth'),
731+
EndpointConfPair(
732+
key='spark_conf.fs.azure.account.oauth2.client.secret.sto2.dfs.core.windows.net',
733+
value='{{secrets/ucx/uber_principal_secret}}',
734+
),
735+
],
736+
sql_configuration_parameters=None,
737+
)
716738

717739

718740
def test_create_access_connectors_for_storage_accounts_logs_no_storage_accounts(caplog):

0 commit comments

Comments
 (0)