def bring_your_own_workgroup()

in migration/athena/athena_workgroup_migration.py [0:0]

• 25 lines of code
• 1 McCabe index (conditional complexity)

def bring_your_own_workgroup(workgroup_name, domain_id, project_id, account_id, region):
    print(f"Tagging Athena workgroup {workgroup_name} with DataZone project ID...")
    # Call Athena tag-resource API with the given workgroup_name
    athena = boto3.client('athena', region_name=region)
    athena.tag_resource(
        ResourceARN=f'arn:aws:athena:{region}:{account_id}:workgroup/{workgroup_name}',
        Tags=[{'Key': 'AmazonDataZoneProject', 'Value': project_id}]
    )
    print(f"Tagged Athena workgroup {workgroup_name} with DataZone project ID.")

    print(f"Updating default Athena connection with workgroup {workgroup_name}...")
    # Call Datazone list-connections API to find the default Athena connection
    datazone = boto3.client('datazone', region_name=region)
    default_athena_connection = datazone.list_connections(
        domainIdentifier=domain_id,
        projectIdentifier=project_id,
        type='ATHENA'
    )
    # Call DataZone update-connection API to update the default Athena connection with the given workgroup_name
    datazone.update_connection(
        domainIdentifier=domain_id,
        identifier=default_athena_connection['items'][0]['connectionId'],
        props={
            'athenaProperties': {
                'workgroupName': workgroup_name
            }
        }
    )
    print(f"Updated default Athena connection with workgroup {workgroup_name}.")