Skip to content

Add more configurable flags for UC metastore, external locations, and storage credentials. #562

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pietern merged 4 commits into from
Oct 3, 2022
Merged

Add more configurable flags for UC metastore, external locations, and storage credentials. #562

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
77a9a42
initial
wchau Sep 23, 2022
0551a72
.
wchau Sep 23, 2022
64902b6
Common options
wchau Sep 28, 2022
70e3ca9
newline
wchau Sep 28, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions .vscode/settings.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
{
"python.testing.pytestArgs": [
"tests"
],
"python.testing.unittestEnabled": false,
"python.testing.pytestEnabled": true,
"python.linting.pylintEnabled": false,
"python.linting.prospectorEnabled": true,
"python.linting.prospectorArgs": [
"-t", "dodgy",
"-t", "mccabe",
"-t", "profile-validator",
"-t", "pyflakes",
"-t", "pylint"
],
"python.linting.enabled": true
}
164 changes: 153 additions & 11 deletions databricks_cli/unity_catalog/cred_cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import functools

import click

from databricks_cli.click_types import JsonClickType
Expand All @@ -33,9 +35,91 @@

############# Storage Credential Commands ############

def fill_credential(
data, aws_iam_role_arn, az_sp_directory_id, az_sp_application_id,
az_sp_client_secret, az_mi_access_connector_id, az_mi_id, gcp_sak_email,
gcp_sak_private_key_id, gcp_sak_private_key):
if aws_iam_role_arn is not None:
data['aws_iam_role'] = {
'role_arn': aws_iam_role_arn
}

if ((az_sp_directory_id is not None) or (az_sp_application_id is not None) or
(az_sp_client_secret is not None)):
data['azure_service_principal'] = {
'directory_id': az_sp_directory_id,
'application_id': az_sp_application_id,
'client_secret': az_sp_client_secret
}

if (az_mi_access_connector_id is not None) or (az_mi_id is not None):
data['azure_managed_identity'] = {
'access_connector_id': az_mi_access_connector_id,
'managed_identity_id': az_mi_id
}

if ((gcp_sak_email is not None) or (gcp_sak_private_key_id is not None) or
(gcp_sak_private_key is not None)):
data['gcp_service_account_key'] = {
'email': gcp_sak_email,
'private_key_id': gcp_sak_private_key_id,
'private_key': gcp_sak_private_key
}


def create_update_common_options(f):
@click.option('--aws-iam-role-arn', default=None,
help='The Amazon Resource Name (ARN) of the AWS IAM role for S3 data access.')
@click.option('--az-sp-directory-id', default=None,
help=(
'The directory ID corresponding to the Azure Active Directory (AAD) '
'tenant of the application.'))
@click.option('--az-sp-application-id', default=None,
help=(
'The application ID of the application registration within the referenced '
'AAD tenant.'))
@click.option('--az-sp-client-secret', default=None,
help='The client secret generated for the above app ID in AAD.')
@click.option('--az-mi-access-connector-id', default=None,
help=(
'The Azure resource ID of the Azure Databricks Access Connector. '
'Use the format, '
'/subscriptions/{guid}/resourceGroups/{rg-name}/providers/Microsoft.Databricks'
'/accessConnectors/{connector-name} .'))
@click.option('--az-mi-id', default=None,
help=(
'The Azure resource ID of the managed identity. Use the format, '
'/subscriptions/{guid}/resourceGroups/{rg-name}/providers'
'/Microsoft.ManagedIdentity/userAssignedIdentities/{identity-name} .'
'This is only available for user-assigned identities. '
'For system-assigned identities, access-connector-id is used to identify '
'the identity. If this flag is not provided, '
'then we assume that it is using the system-assigned identity.'))
@click.option('--gcp-sak-email', default=None,
help=(
'Credential for GCP Service Account Key. '
'The email of the service account.'))
@click.option('--gcp-sak-private-key-id', default=None,
help=(
'Credential for GCP Service Account Key. '
'The ID of the service account\'s private key.'))
@click.option('--gcp-sak-private-key', default=None,
help=(
'Credential for GCP Service Account Key. '
'The service account\'s RSA private key.'))
@click.option('--comment', default=None,
help='Free-form text description.')
@functools.wraps(f)
def wrapper(*args, **kwargs):
f(*args, **kwargs)
return wrapper


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create storage credential.')
@click.option('--name', default=None,
help='Name of new storage credential')
@create_update_common_options
@click.option('--skip-validation', '-s', 'skip_val', is_flag=True, default=False,
help='Skip the validation of new credential info before creation')
@click.option('--json-file', default=None, type=click.Path(),
Expand All @@ -50,16 +134,42 @@
# Until that is fixed (should return a 400), show full error trace.
#@eat_exceptions
@provide_api_client
def create_credential_cli(api_client, skip_val, json_file, json):
def create_credential_cli(api_client, name, aws_iam_role_arn,
az_sp_directory_id, az_sp_application_id, az_sp_client_secret,
az_mi_access_connector_id, az_mi_id, gcp_sak_email,
gcp_sak_private_key_id, gcp_sak_private_key, comment,
skip_val, json_file, json):
"""
Create new storage credential.

The public specification for the JSON request is in development.
"""
json_cli_base(json_file, json,
lambda json: UnityCatalogApi(api_client).create_storage_credential(json,
skip_val),
encode_utf8=True)
has_credential_flag = (
(aws_iam_role_arn is not None) or
(az_sp_directory_id is not None) or (az_sp_application_id is not None) or
(az_sp_client_secret is not None) or (az_mi_access_connector_id is not None) or
(az_mi_id is not None) or (gcp_sak_email is not None) or
(gcp_sak_private_key_id is not None) or (gcp_sak_private_key is not None))
if ((name is not None) or has_credential_flag or (comment is not None)):
if (json_file is not None) or (json is not None):
raise ValueError('Cannot specify JSON if any other creation flags are specified')
data = {
'name': name,
'comment': comment
}
Copy link
Contributor

@adamcain-db adamcain-db Sep 28, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like we are leaving it completely to the server to validate the input -- e.g., we are not checking that data has required fields like name. I suppose that's reasonable (as we also do not validate JSON inputs), but I wanted to confirm that this is our approach for input validation at the CLI.

Copy link
Contributor Author

@wchau wchau Sep 28, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that's the idea. Instead of replicating the full validation logic, just letting the server do it and then return the error message would make things easier to maintain.

Copy link
Contributor

@adamcain-db adamcain-db Sep 28, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense.


fill_credential(
data, aws_iam_role_arn, az_sp_directory_id, az_sp_application_id,
az_sp_client_secret, az_mi_access_connector_id, az_mi_id, gcp_sak_email,
gcp_sak_private_key_id, gcp_sak_private_key)

cred_json = UnityCatalogApi(api_client).create_storage_credential(data, skip_val)
click.echo(mc_pretty_format(cred_json))
else:
json_cli_base(json_file, json,
lambda json: UnityCatalogApi(api_client).create_storage_credential(json,
skip_val),
encode_utf8=True)


@click.command(context_settings=CONTEXT_SETTINGS,
Expand Down Expand Up @@ -98,6 +208,10 @@ def get_credential_cli(api_client, name):
short_help='Update a storage credential.')
@click.option('--name', required=True,
help='Name of the storage credential to update.')
@click.option('--new-name', default=None, help='New name of the storage credential.')
@create_update_common_options
Copy link
Contributor

@adamcain-db adamcain-db Sep 28, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Slick!

@click.option('--owner', default=None,
help='Owner of the storage credential.')
@click.option('--skip-validation', '-s', 'skip_val', is_flag=True, default=False,
help='Skip the validation of new credential info before update')
@click.option('--json-file', default=None, type=click.Path(),
Expand All @@ -109,17 +223,45 @@ def get_credential_cli(api_client, name):
# See comment for create-storage-credential
#@eat_exceptions
@provide_api_client
def update_credential_cli(api_client, name, skip_val, json_file, json):
def update_credential_cli(api_client, name, new_name, aws_iam_role_arn,
az_sp_directory_id, az_sp_application_id, az_sp_client_secret,
az_mi_access_connector_id, az_mi_id, gcp_sak_email,
gcp_sak_private_key_id, gcp_sak_private_key, comment, owner,
skip_val, json_file, json):
"""
Update a storage credential.

The public specification for the JSON request is in development.
"""
json_cli_base(json_file, json,
lambda json: UnityCatalogApi(api_client).update_storage_credential(name,
json,
skip_val),
encode_utf8=True)
has_credential_flag = (
(aws_iam_role_arn is not None) or
(az_sp_directory_id is not None) or (az_sp_application_id is not None) or
(az_sp_client_secret is not None) or (az_mi_access_connector_id is not None) or
(az_mi_id is not None) or (gcp_sak_email is not None) or
(gcp_sak_private_key_id is not None) or (gcp_sak_private_key is not None))
if ((new_name is not None) or has_credential_flag or
(comment is not None) or (owner is not None)):
if (json_file is not None) or (json is not None):
raise ValueError('Cannot specify JSON if any other update flags are specified')
data = {
'name': new_name,
'comment': comment,
'owner': owner
}

fill_credential(
data, aws_iam_role_arn, az_sp_directory_id, az_sp_application_id,
az_sp_client_secret, az_mi_access_connector_id, az_mi_id, gcp_sak_email,
gcp_sak_private_key_id, gcp_sak_private_key)

cred_json = UnityCatalogApi(api_client).update_storage_credential(name, data, skip_val)
click.echo(mc_pretty_format(cred_json))
else:
json_cli_base(json_file, json,
lambda json: UnityCatalogApi(api_client).update_storage_credential(name,
json,
skip_val),
encode_utf8=True)


@click.command(context_settings=CONTEXT_SETTINGS,
Expand Down
Loading