Skip to content

Commit 6a4b042

Browse files
nkvuongnkvuong
committed
make fmt
1 parent 212f65b commit 6a4b042

File tree

2 files changed

+73
-53
lines changed

2 files changed

+73
-53
lines changed

‎src/databricks/labs/ucx/cli.py‎

Lines changed: 72 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -203,12 +203,12 @@ def revert_migrated_tables(w: WorkspaceClient, schema: str, table: str, *, delet
203203

204204
@ucx.command
205205
def move(
206-
w: WorkspaceClient,
207-
from_catalog: str,
208-
from_schema: str,
209-
from_table: str,
210-
to_catalog: str,
211-
to_schema: str,
206+
w: WorkspaceClient,
207+
from_catalog: str,
208+
from_schema: str,
209+
from_table: str,
210+
to_catalog: str,
211+
to_schema: str,
212212
):
213213
"""move a uc table/tables from one schema to another schema in same or different catalog"""
214214
logger.info("Running move command")
@@ -230,12 +230,12 @@ def move(
230230

231231
@ucx.command
232232
def alias(
233-
w: WorkspaceClient,
234-
from_catalog: str,
235-
from_schema: str,
236-
from_table: str,
237-
to_catalog: str,
238-
to_schema: str,
233+
w: WorkspaceClient,
234+
from_catalog: str,
235+
from_schema: str,
236+
from_table: str,
237+
to_catalog: str,
238+
to_schema: str,
239239
):
240240
"""move a uc table/tables from one schema to another schema in same or different catalog"""
241241
if from_catalog == "" or to_catalog == "":
@@ -253,13 +253,13 @@ def alias(
253253

254254

255255
def _execute_for_cloud(
256-
w: WorkspaceClient,
257-
func_azure: Callable,
258-
func_aws: Callable,
259-
azure_resource_permissions: AzureResourcePermissions | None = None,
260-
subscription_id: str | None = None,
261-
aws_permissions: AWSResourcePermissions = None,
262-
aws_profile: str | None = None,
256+
w: WorkspaceClient,
257+
func_azure: Callable,
258+
func_aws: Callable,
259+
azure_resource_permissions: AzureResourcePermissions | None = None,
260+
subscription_id: str | None = None,
261+
aws_permissions: AWSResourcePermissions | None = None,
262+
aws_profile: str | None = None,
263263
):
264264
if w.config.is_azure:
265265
if w.config.auth_type != "azure-cli":
@@ -287,66 +287,85 @@ def _execute_for_cloud(
287287

288288

289289
@ucx.command
290-
def create_uber_principal(w: WorkspaceClient, subscription_id: str | None = None,
291-
azure_resource_permissions: AzureResourcePermissions | None = None,
292-
aws_profile: str | None = None,
293-
aws_resource_permissions: AWSResourcePermissions | None = None):
290+
def create_uber_principal(
291+
w: WorkspaceClient,
292+
subscription_id: str | None = None,
293+
azure_resource_permissions: AzureResourcePermissions | None = None,
294+
aws_profile: str | None = None,
295+
aws_resource_permissions: AWSResourcePermissions | None = None,
296+
):
294297
"""For azure cloud, creates a service principal and gives STORAGE BLOB READER access on all the storage account
295298
used by tables in the workspace and stores the spn info in the UCX cluster policy. For aws,
296299
it identifies all s3 buckets used by the Instance Profiles configured in the workspace.
297300
Pass subscription_id for azure and aws_profile for aws."""
298-
return _execute_for_cloud(w, _azure_setup_uber_principal, _aws_setup_uber_principal,
299-
azure_resource_permissions, subscription_id,
300-
aws_resource_permissions, aws_profile)
301+
return _execute_for_cloud(
302+
w,
303+
_azure_setup_uber_principal,
304+
_aws_setup_uber_principal,
305+
azure_resource_permissions,
306+
subscription_id,
307+
aws_resource_permissions,
308+
aws_profile,
309+
)
301310

302311

303-
def _azure_setup_uber_principal(w: WorkspaceClient, subscription_id: str,
304-
azure_resource_permissions: AzureResourcePermissions = None):
312+
def _azure_setup_uber_principal(
313+
w: WorkspaceClient, subscription_id: str, azure_resource_permissions: AzureResourcePermissions | None = None
314+
):
305315
prompts = Prompts()
306316
include_subscriptions = [subscription_id] if subscription_id else None
307-
if not azure_resource_permissions:
317+
if azure_resource_permissions is None:
308318
azure_resource_permissions = AzureResourcePermissions.for_cli(w, include_subscriptions=include_subscriptions)
309319
azure_resource_permissions.create_uber_principal(prompts)
310320

311321

312-
def _aws_setup_uber_principal(w: WorkspaceClient, aws_profile: str,
313-
aws_resource_permissions: AWSResourcePermissions = None):
322+
def _aws_setup_uber_principal(
323+
w: WorkspaceClient, aws_profile: str, aws_resource_permissions: AWSResourcePermissions | None = None
324+
):
314325
prompts = Prompts()
315326
installation = Installation.current(w, 'ucx')
316327
config = installation.load(WorkspaceConfig)
317328
sql_backend = StatementExecutionBackend(w, config.warehouse_id)
318329
aws = AWSResources(aws_profile)
319-
aws_resource_permissions = AWSResourcePermissions.for_cli(
320-
w, installation, sql_backend, aws, config.inventory_database
321-
)
330+
if aws_resource_permissions is None:
331+
aws_resource_permissions = AWSResourcePermissions.for_cli(
332+
w, installation, sql_backend, aws, config.inventory_database
333+
)
322334
aws_resource_permissions.create_uber_principal(prompts)
323335

324336

325337
@ucx.command
326-
def principal_prefix_access(w: WorkspaceClient,
327-
subscription_id: str | None = None,
328-
azure_resource_permissions: AzureResourcePermissions | None = None,
329-
aws_profile: str | None = None,
330-
aws_resource_permissions: AWSResourcePermissions | None = None
331-
):
338+
def principal_prefix_access(
339+
w: WorkspaceClient,
340+
subscription_id: str | None = None,
341+
azure_resource_permissions: AzureResourcePermissions | None = None,
342+
aws_profile: str | None = None,
343+
aws_resource_permissions: AWSResourcePermissions | None = None,
344+
):
332345
"""For azure cloud, identifies all storage accounts used by tables in the workspace, identify spn and its
333346
permission on each storage accounts. For aws, identifies all the Instance Profiles configured in the workspace and
334347
its access to all the S3 buckets, along with AWS roles that are set with UC access and its access to S3 buckets.
335348
The output is stored in the workspace install folder.
336349
Pass subscription_id for azure and aws_profile for aws."""
337350
return _execute_for_cloud(
338-
w, _azure_principal_prefix_access, _aws_principal_prefix_access, azure_resource_permissions,
339-
subscription_id, aws_resource_permissions, aws_profile
351+
w,
352+
_azure_principal_prefix_access,
353+
_aws_principal_prefix_access,
354+
azure_resource_permissions,
355+
subscription_id,
356+
aws_resource_permissions,
357+
aws_profile,
340358
)
341359

342360

343-
def _azure_principal_prefix_access(w: WorkspaceClient, subscription_id: str,
344-
azure_resource_permissions: AzureResourcePermissions = None):
361+
def _azure_principal_prefix_access(
362+
w: WorkspaceClient, subscription_id: str, azure_resource_permissions: AzureResourcePermissions | None = None
363+
):
345364
if w.config.auth_type != "azure-cli":
346365
logger.error("In order to obtain AAD token, Please run azure cli to authenticate.")
347366
return
348367
include_subscriptions = [subscription_id] if subscription_id else None
349-
if not azure_resource_permissions:
368+
if azure_resource_permissions is None:
350369
azure_resource_permissions = AzureResourcePermissions.for_cli(w, include_subscriptions=include_subscriptions)
351370
logger.info("Generating azure storage accounts and service principal permission info")
352371
path = azure_resource_permissions.save_spn_permissions()
@@ -355,8 +374,9 @@ def _azure_principal_prefix_access(w: WorkspaceClient, subscription_id: str,
355374
return
356375

357376

358-
def _aws_principal_prefix_access(w: WorkspaceClient, aws_profile: str,
359-
aws_resource_permissions: AWSResourcePermissions = None):
377+
def _aws_principal_prefix_access(
378+
w: WorkspaceClient, aws_profile: str, aws_resource_permissions: AWSResourcePermissions | None = None
379+
):
360380
if not shutil.which("aws"):
361381
logger.error("Couldn't find AWS CLI in path. Please install the CLI from https://aws.amazon.com/cli/")
362382
return
@@ -365,9 +385,10 @@ def _aws_principal_prefix_access(w: WorkspaceClient, aws_profile: str,
365385
config = installation.load(WorkspaceConfig)
366386
sql_backend = StatementExecutionBackend(w, config.warehouse_id)
367387
aws = AWSResources(aws_profile)
368-
if not aws_resource_permissions:
369-
aws_resource_permissions = AWSResourcePermissions.for_cli(w, installation, sql_backend, aws,
370-
config.inventory_database)
388+
if aws_resource_permissions is None:
389+
aws_resource_permissions = AWSResourcePermissions.for_cli(
390+
w, installation, sql_backend, aws, config.inventory_database
391+
)
371392
instance_role_path = aws_resource_permissions.save_instance_profile_permissions()
372393
logger.info(f"Instance profile and bucket info saved {instance_role_path}")
373394
logger.info("Generating UC roles and bucket permission info")
@@ -376,8 +397,7 @@ def _aws_principal_prefix_access(w: WorkspaceClient, aws_profile: str,
376397

377398

378399
@ucx.command
379-
def migrate_credentials(w: WorkspaceClient, aws_profile: str | None = None,
380-
aws_resources: AWSResources | None = None):
400+
def migrate_credentials(w: WorkspaceClient, aws_profile: str | None = None, aws_resources: AWSResources | None = None):
381401
"""For Azure, this command migrates Azure Service Principals, which have Storage Blob Data Contributor,
382402
Storage Blob Data Reader, Storage Blob Data Owner roles on ADLS Gen2 locations that are being used in
383403
Databricks, to UC storage credentials.

‎tests/unit/test_cli.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -345,7 +345,7 @@ def test_migrate_credentials_aws(ws, mocker):
345345
ws.config.is_aws = True
346346
ws.config.is_gcp = False
347347
aws_resources = create_autospec(AWSResources)
348-
aws_resources.validate_connection.return_value={"Account": "123456789012"}
348+
aws_resources.validate_connection.return_value = {"Account": "123456789012"}
349349
with patch("databricks.labs.blueprint.tui.Prompts.confirm", return_value=True):
350350
migrate_credentials(ws, aws_profile="profile", aws_resources=aws_resources)
351351
ws.storage_credentials.list.assert_called()

0 commit comments

Comments
 (0)