diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index 0b4f6ecf..6960ee9a 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -41,6 +41,7 @@ from databricks_cli.instance_pools.cli import instance_pools_group from databricks_cli.pipelines.cli import pipelines_group from databricks_cli.repos.cli import repos_group +from databricks_cli.unity_catalog.cli import unity_catalog_group @click.group(context_settings=CONTEXT_SETTINGS) @@ -67,6 +68,7 @@ def cli(): cli.add_command(instance_pools_group, name="instance-pools") cli.add_command(pipelines_group, name='pipelines') cli.add_command(repos_group, name='repos') +cli.add_command(unity_catalog_group, name='unity-catalog') if __name__ == "__main__": cli() diff --git a/databricks_cli/click_types.py b/databricks_cli/click_types.py index 463efe82..4ef647ce 100644 --- a/databricks_cli/click_types.py +++ b/databricks_cli/click_types.py @@ -112,6 +112,16 @@ class PipelineIdClickType(ParamType): help = 'The pipeline ID.' +class MetastoreIdClickType(ParamType): + name = 'METASTORE_ID' + help = 'ID of the Metastore' + + +class WorkspaceIdClickType(ParamType): + name = 'WORKSPACE_ID' + help = 'ID of the Workspace' + + class OneOfOption(Option): def __init__(self, *args, **kwargs): self.one_of = kwargs.pop('one_of') diff --git a/databricks_cli/unity_catalog/__init__.py b/databricks_cli/unity_catalog/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/databricks_cli/unity_catalog/api.py b/databricks_cli/unity_catalog/api.py new file mode 100644 index 00000000..2ec2a1bf --- /dev/null +++ b/databricks_cli/unity_catalog/api.py @@ -0,0 +1,250 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from databricks_cli.unity_catalog.uc_service import UnityCatalogService + + +class UnityCatalogApi(object): + def __init__(self, api_client): + self.client = UnityCatalogService(api_client) + + # Metastore APIs + + def create_metastore(self, name, storage_root): + return self.client.create_metastore(name, storage_root) + + def list_metastores(self): + return self.client.list_metastores() + + def get_metastore(self, metastore_id): + return self.client.get_metastore(metastore_id) + + def update_metastore(self, metastore_id, metastore_spec): + return self.client.update_metastore(metastore_id, metastore_spec) + + def delete_metastore(self, metastore_id, force): + return self.client.delete_metastore(metastore_id, force) + + def get_metastore_summary(self): + return self.client.get_metastore_summary() + + def create_metastore_assignment(self, workspace_id, metastore_id, default_catalog_name): + return self.client.create_metastore_assignment(workspace_id, metastore_id, + default_catalog_name) + + def update_metastore_assignment(self, workspace_id, metastore_id, default_catalog_name): + return self.client.update_metastore_assignment(workspace_id, metastore_id, + default_catalog_name) + + def delete_metastore_assignment(self, workspace_id, metastore_id): + return self.client.delete_metastore_assignment(workspace_id, metastore_id) + + # External Location APIs + + def create_external_location(self, loc_spec, skip_validation): + return self.client.create_external_location(loc_spec, skip_validation) + + def list_external_locations(self): + return self.client.list_external_locations() + + def get_external_location(self, name): + return self.client.get_external_location(name) + + def update_external_location(self, name, loc_spec, force, skip_validation): + return self.client.update_external_location(name, loc_spec, force, skip_validation) + + def delete_external_location(self, name, force): + return self.client.delete_external_location(name, force) + + def validate_external_location(self, validation_spec): + return self.client.validate_external_location(validation_spec) + + # Data Access Configuration APIs + + def create_dac(self, metastore_id, dac_spec, skip_validation): + return self.client.create_dac(metastore_id, dac_spec, skip_validation) + + def list_dacs(self, metastore_id): + return self.client.list_dacs(metastore_id) + + def get_dac(self, metastore_id, dac_id): + return self.client.get_dac(metastore_id, dac_id) + + def delete_dac(self, metastore_id, dac_id): + return self.client.delete_dac(metastore_id, dac_id) + + # Storage Credentials + + def create_storage_credential(self, cred_spec, skip_validation): + return self.client.create_storage_credential(cred_spec, skip_validation) + + def list_storage_credentials(self, name_pattern): + return self.client.list_storage_credentials(name_pattern) + + def get_storage_credential(self, name): + return self.client.get_storage_credential(name) + + def update_storage_credential(self, name, cred_spec, skip_validation): + return self.client.update_storage_credential(name, cred_spec, skip_validation) + + def delete_storage_credential(self, name, force): + return self.client.delete_storage_credential(name, force) + + # Catalog APIs + + def create_catalog(self, catalog_name, comment, provider, share): + return self.client.create_catalog(catalog_name, comment, provider, share) + + def list_catalogs(self): + return self.client.list_catalogs() + + def get_catalog(self, name): + return self.client.get_catalog(name) + + def update_catalog(self, name, catalog_spec): + return self.client.update_catalog(name, catalog_spec) + + def delete_catalog(self, catalog_name): + return self.client.delete_catalog(catalog_name) + + # Schema APIs + + def create_schema(self, catalog_name, schema_name, comment): + return self.client.create_schema(catalog_name, schema_name, comment) + + def list_schemas(self, catalog_name, name_pattern): + return self.client.list_schemas(catalog_name, name_pattern) + + def get_schema(self, full_name): + return self.client.get_schema(full_name) + + def update_schema(self, full_name, schema_spec): + return self.client.update_schema(full_name, schema_spec) + + def delete_schema(self, schema_full_name): + return self.client.delete_schema(schema_full_name) + + # Table APIs + + def create_table(self, table_spec): + return self.client.create_table(table_spec) + + def list_tables(self, catalog_name, schema_name, name_pattern): + return self.client.list_tables(catalog_name, schema_name, name_pattern) + + def list_table_summaries(self, catalog_name): + return self.client.list_table_summaries(catalog_name) + + def get_table(self, full_name): + return self.client.get_table(full_name) + + def update_table(self, full_name, table_spec): + return self.client.update_table(full_name, table_spec) + + def delete_table(self, table_full_name): + return self.client.delete_table(table_full_name) + + # Share APIs + + def create_share(self, name): + return self.client.create_share(name) + + def list_shares(self): + return self.client.list_shares() + + def get_share(self, name, include_shared_data): + return self.client.get_share(name, include_shared_data) + + def update_share(self, name, share_spec): + return self.client.update_share(name, share_spec) + + def delete_share(self, name): + return self.client.delete_share(name) + + def list_share_permissions(self, name): + return self.client.list_share_permissions(name) + + def update_share_permissions(self, name, perm_spec): + return self.client.update_share_permissions(name, perm_spec) + + # Recipient APIs + + def create_recipient(self, name, comment, sharing_code, allowed_ip_addresses): + return self.client.create_recipient(name, comment, sharing_code, allowed_ip_addresses) + + def list_recipients(self): + return self.client.list_recipients() + + def get_recipient(self, name): + return self.client.get_recipient(name) + + def update_recipient(self, name, recipient_spec): + return self.client.update_recipient(name, recipient_spec) + + def rotate_recipient_token(self, name, existing_token_expire_in_seconds): + return self.client.rotate_recipient_token(name, existing_token_expire_in_seconds) + + def get_recipient_share_permissions(self, name): + return self.client.get_recipient_share_permissions(name) + + def delete_recipient(self, name): + return self.client.delete_recipient(name) + + # Provider APIs + + def create_provider(self, name, comment, recipient_profile): + return self.client.create_provider(name, comment, recipient_profile) + + def list_providers(self): + return self.client.list_providers() + + def get_provider(self, name): + return self.client.get_provider(name) + + def update_provider(self, name, new_name=None, comment=None, recipient_profile=None): + return self.client.update_provider(name, new_name, comment, recipient_profile) + + def delete_provider(self, name): + return self.client.delete_provider(name) + + def list_provider_shares(self, name): + return self.client.list_provider_shares(name) + + # Permissions APIs + + def get_permissions(self, sec_type, sec_name): + return self.client.get_permissions(sec_type, sec_name) + + def update_permissions(self, sec_type, sec_name, diff_spec): + return self.client.update_permissions(sec_type, sec_name, diff_spec) + + def replace_permissions(self, sec_type, sec_name, perm_spec): + return self.client.replace_permissions(sec_type, sec_name, perm_spec) + + # Lineage APIs + + def list_lineages_by_table(self, table_name): + return self.client.list_lineages_by_table(table_name) + + def list_lineages_by_column(self, table_name, column_name): + return self.client.list_lineages_by_column(table_name, column_name) diff --git a/databricks_cli/unity_catalog/catalog_cli.py b/databricks_cli/unity_catalog/catalog_cli.py new file mode 100644 index 00000000..57b756dc --- /dev/null +++ b/databricks_cli/unity_catalog/catalog_cli.py @@ -0,0 +1,155 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import JsonClickType +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create a new catalog.') +@click.option('--name', required=True, help='Name of new catalog.') +@click.option('--comment', default=None, required=False, + help='Free-form text description.') +@click.option('--provider', default=None, required=False, + help='Name of the Provider (for creating Delta Sharing Catalog).') +@click.option('--share', default=None, required=False, + help='Name of the Share under the Provider to create a Delta Sharing Catalog.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_catalog_cli(api_client, name, comment, provider, share): + """ + Create a new catalog. + """ + catalog_json = UnityCatalogApi(api_client).create_catalog(name, comment, + provider, share) + click.echo(mc_pretty_format(catalog_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List catalogs.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_catalogs_cli(api_client): + """ + List catalogs. + """ + catalogs_json = UnityCatalogApi(api_client).list_catalogs() + click.echo(mc_pretty_format(catalogs_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a catalog.') +@click.option('--name', required=True, + help='Name of the catalog to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_catalog_cli(api_client, name): + """ + Get a catalog. + """ + catalog_json = UnityCatalogApi(api_client).get_catalog(name) + click.echo(mc_pretty_format(catalog_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a catalog.') +@click.option('--name', required=True, + help='Name of the catalog to update.') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/catalogs/{name}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/catalogs/{name}')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_catalog_cli(api_client, name, json_file, json): + """ + Update a catalog. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_catalog(name, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a catalog.') +@click.option('--name', required=True, + help='Name of the catalog to delete.') +@click.option('--purge', '-p', is_flag=True, default=False, + help='Purge all child schemas and tables of catalog.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_catalog_cli(api_client, name, purge): + """ + Delete a catalog. + """ + if purge: + tables_response = UnityCatalogApi(api_client).list_table_summaries(name) + for t in tables_response.get('tables', []): + click.echo("Deleting table: %s" % (t['full_name'])) + UnityCatalogApi(api_client).delete_table(t['full_name']) + + schemas_response = UnityCatalogApi(api_client).list_schemas(name, None) + for s in schemas_response.get('schemas', []): + click.echo("Purging schema: %s" % (s['full_name'])) + UnityCatalogApi(api_client).delete_schema(s['full_name']) + + UnityCatalogApi(api_client).delete_catalog(name) + + +@click.group() +def catalogs_group(): # pragma: no cover + pass + + +def register_catalog_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_catalog_cli), name='create-catalog') + cmd_group.add_command(hide(list_catalogs_cli), name='list-catalogs') + cmd_group.add_command(hide(get_catalog_cli), name='get-catalog') + cmd_group.add_command(hide(update_catalog_cli), name='update-catalog') + cmd_group.add_command(hide(delete_catalog_cli), name='delete-catalog') + + # Register command group. + catalogs_group.add_command(create_catalog_cli, name='create') + catalogs_group.add_command(list_catalogs_cli, name='list') + catalogs_group.add_command(get_catalog_cli, name='get') + catalogs_group.add_command(update_catalog_cli, name='update') + catalogs_group.add_command(delete_catalog_cli, name='delete') + cmd_group.add_command(catalogs_group, name='catalogs') diff --git a/databricks_cli/unity_catalog/cli.py b/databricks_cli/unity_catalog/cli.py new file mode 100644 index 00000000..dab0f9bc --- /dev/null +++ b/databricks_cli/unity_catalog/cli.py @@ -0,0 +1,62 @@ +# Databricks CLI +# Copyright 2021 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.utils import CONTEXT_SETTINGS +from databricks_cli.version import print_version_callback, version + +from databricks_cli.unity_catalog.metastore_cli import register_metastore_commands +from databricks_cli.unity_catalog.catalog_cli import register_catalog_commands +from databricks_cli.unity_catalog.schema_cli import register_schema_commands +from databricks_cli.unity_catalog.table_cli import register_table_commands +from databricks_cli.unity_catalog.ext_loc_cli import register_ext_loc_commands +from databricks_cli.unity_catalog.cred_cli import register_cred_commands +from databricks_cli.unity_catalog.delta_sharing_cli import register_delta_sharing_commands +from databricks_cli.unity_catalog.perms_cli import register_perms_commands +from databricks_cli.unity_catalog.lineage_cli import register_lineage_commands + + +@click.group(context_settings=CONTEXT_SETTINGS, + help='Utility to interact with Databricks Unity Catalog.\n\n' + + '**********************************************************************\n' + + 'WARNING: these commands are EXPERIMENTAL and not officially supported.\n' + + '**********************************************************************') +@click.option('--version', '-v', is_flag=True, callback=print_version_callback, + expose_value=False, is_eager=True, help=version) +def unity_catalog_group(): # pragma: no cover + """ + Utility to interact with Databricks Unity Catalog. + """ + pass + + +register_metastore_commands(unity_catalog_group) +register_ext_loc_commands(unity_catalog_group) +register_cred_commands(unity_catalog_group) +register_catalog_commands(unity_catalog_group) +register_schema_commands(unity_catalog_group) +register_table_commands(unity_catalog_group) +register_delta_sharing_commands(unity_catalog_group) +register_perms_commands(unity_catalog_group) +register_lineage_commands(unity_catalog_group) diff --git a/databricks_cli/unity_catalog/cred_cli.py b/databricks_cli/unity_catalog/cred_cli.py new file mode 100644 index 00000000..5cb9544c --- /dev/null +++ b/databricks_cli/unity_catalog/cred_cli.py @@ -0,0 +1,161 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import JsonClickType +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +############# Storage Credential Commands ############ + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create storage credential.') +@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(), + help=json_file_help(method='POST', path='/storage-credentials')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='POST', path='/storage-credentials')) +@debug_option +@profile_option +# UC's createStorageCredential returns a 401 when validation fails; that translates to +# a misleading error when eat_exceptions is enabled: +# Your authentication information may be incorrect. Please reconfigure with ``dbfs configure`` +# 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): + """ + 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) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List storage credentials.') +@click.option('--name-pattern', default=None, + help='SQL LIKE pattern that the credential name must match to be in list.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_credentials_cli(api_client, name_pattern): + """ + List storage credentials. + """ + creds_json = UnityCatalogApi(api_client).list_storage_credentials(name_pattern) + click.echo(mc_pretty_format(creds_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a storage credential.') +@click.option('--name', required=True, + help='Name of the storage credential to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_credential_cli(api_client, name): + """ + Get a storage credential. + """ + cred_json = UnityCatalogApi(api_client).get_storage_credential(name) + click.echo(mc_pretty_format(cred_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a storage credential.') +@click.option('--name', required=True, + help='Name of the storage credential to update.') +@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(), + help=json_file_help(method='PATCH', path='/storage-credentials/{name}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/storage-credentials/{name}')) +@debug_option +@profile_option +# See comment for create-storage-credential +#@eat_exceptions +@provide_api_client +def update_credential_cli(api_client, name, 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) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a storage credential.') +@click.option('--name', required=True, + help='Name of the storage credential to delete.') +@click.option('--force', '-f', is_flag=True, default=False, + help='Force deletion even if credential has dependent tables/locations') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_credential_cli(api_client, name, force): + """ + Delete a storage credential. + """ + UnityCatalogApi(api_client).delete_storage_credential(name, force) + + +@click.group() +def storage_credentials_group(): # pragma: no cover + pass + + +def register_cred_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_credential_cli), name='create-storage-credential') + cmd_group.add_command(hide(list_credentials_cli), name='list-storage-credentials') + cmd_group.add_command(hide(get_credential_cli), name='get-storage-credential') + cmd_group.add_command(hide(update_credential_cli), name='update-storage-credential') + cmd_group.add_command(hide(delete_credential_cli), name='delete-storage-credential') + + # Register command group. + storage_credentials_group.add_command(create_credential_cli, name='create') + storage_credentials_group.add_command(list_credentials_cli, name='list') + storage_credentials_group.add_command(get_credential_cli, name='get') + storage_credentials_group.add_command(update_credential_cli, name='update') + storage_credentials_group.add_command(delete_credential_cli, name='delete') + cmd_group.add_command(storage_credentials_group, name='storage-credentials') diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py new file mode 100644 index 00000000..6672f3ac --- /dev/null +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -0,0 +1,511 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import JsonClickType +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +############## Share Commands ############## + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create a new share.') +@click.option('--name', required=True, help='Name of new share.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_share_cli(api_client, name): + """ + Create a new share. + """ + share_json = UnityCatalogApi(api_client).create_share(name) + click.echo(mc_pretty_format(share_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List shares.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_shares_cli(api_client): + """ + List shares. + """ + shares_json = UnityCatalogApi(api_client).list_shares() + click.echo(mc_pretty_format(shares_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a share.') +@click.option('--name', required=True, + help='Name of the share to get.') +@click.option('--include-shared-data', default=True, + help='Whether to include shared data in the response.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_share_cli(api_client, name, include_shared_data): + """ + Get a share. + """ + share_json = UnityCatalogApi(api_client).get_share(name, include_shared_data) + click.echo(mc_pretty_format(share_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List permissions on a share.') +@click.option('--name', required=True, + help='Name of the share to list permissions on.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_share_permissions_cli(api_client, name): + """ + List permissions on a share. + """ + perms_json = UnityCatalogApi(api_client).list_share_permissions(name) + click.echo(mc_pretty_format(perms_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update permissions on a share.') +@click.option('--name', required=True, + help='Name of the share whose permissions are updated.') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='POST', path='/shares/{name}/permissions')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='POST', path='/shares/{name}/permissions')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_share_permissions_cli(api_client, name, json_file, json): + """ + Update permissions on a share. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_share_permissions(name, json)) + + +def shared_data_object(name): + return {'name': name, 'data_object_type': 'TABLE'} + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a share.') +@click.option('--name', required=True, + help='Name of the share to update.') +@click.option('--add-table', default=None, multiple=True, + metavar='NAME', + help='Full name of table to add to share (can be specified multiple times).') +@click.option('--remove-table', default=None, multiple=True, + metavar='NAME', + help='Full name of table to remove from share (can be specified multiple times).') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/shares/{name}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/shares/{name}')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_share_cli(api_client, name, add_table, remove_table, json_file, json): + """ + Update a share. + + The public specification for the JSON request is in development. + """ + if len(add_table) > 0 or len(remove_table) > 0: + updates = [] + for a in add_table: + updates.append({'action': 'ADD', 'data_object': shared_data_object(a)}) + for r in remove_table: + updates.append({'action': 'REMOVE', 'data_object': shared_data_object(r)}) + d = {'updates': updates} + share_json = UnityCatalogApi(api_client).update_share(name, d) + click.echo(mc_pretty_format(share_json)) + else: + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_share(name, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a share.') +@click.option('--name', required=True, + help='Name of the share to delete.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_share_cli(api_client, name): + """ + Delete a share. + """ + UnityCatalogApi(api_client).delete_share(name) + + +############## Recipient Commands ############## + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create a new recipient.') +@click.option('--name', required=True, help='Name of new recipient.') +@click.option('--comment', default=None, required=False, + help='Free-form text description.') +@click.option('--sharing-code', default=None, required=False, + help='A one-time sharing code shared by the data recipient offline.') +@click.option('--allowed_ip_address', default=None, required=False, multiple=True, + help=( + 'IP address in CIDR notation that is allowed to use delta sharing. ' + 'Supports multiple options.')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_recipient_cli(api_client, name, comment, sharing_code, allowed_ip_address): + """ + Create a new recipient. + """ + recipient_json = UnityCatalogApi(api_client).create_recipient( + name, comment, sharing_code, allowed_ip_address) + click.echo(mc_pretty_format(recipient_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List recipients.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_recipients_cli(api_client): + """ + List recipients. + """ + recipients_json = UnityCatalogApi(api_client).list_recipients() + click.echo(mc_pretty_format(recipients_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a recipient.') +@click.option('--name', required=True, + help='Name of the recipient to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_recipient_cli(api_client, name): + """ + Get a recipient. + """ + recipient_json = UnityCatalogApi(api_client).get_recipient(name) + click.echo(mc_pretty_format(recipient_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a recipient.') +@click.option('--name', required=True, + help='Name of the recipient who needs to be updated.') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/recipients/{name}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/recipients/{name}')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_recipient_cli(api_client, name, json_file, json): + """ + Update a recipient. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_recipient(name, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Rotate token for the recipient.') +@click.option('--name', required=True, help='Name of new recipient.') +@click.option('--existing-token-expire-in-seconds', default=None, required=False, + help='Expire the existing token in number of seconds from now,' + + ' 0 to expire it immediately.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def rotate_recipient_token_cli(api_client, name, existing_token_expire_in_seconds): + """ + Rotate recipient token. + """ + recipient_json = \ + UnityCatalogApi(api_client).rotate_recipient_token(name, existing_token_expire_in_seconds) + click.echo(mc_pretty_format(recipient_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List share permissions of a recipient.') +@click.option('--name', required=True, + help='Name of the recipient.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_recipient_permissions_cli(api_client, name): + """ + List a recipient's share permissions. + """ + recipient_json = UnityCatalogApi(api_client).get_recipient_share_permissions(name) + click.echo(mc_pretty_format(recipient_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a recipient.') +@click.option('--name', required=True, + help='Name of the recipient to delete.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_recipient_cli(api_client, name): + """ + Delete a recipient. + """ + UnityCatalogApi(api_client).delete_recipient(name) + + +############## Provider Commands ############## + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create a provider.') +@click.option('--name', required=True, help='Name of the new provider.') +@click.option('--comment', default=None, required=False, + help='Free-form text description.') +@click.option('--recipient-profile-json-file', default=None, required=False, type=click.Path(), + help='File containing recipient profile in JSON format.') +@click.option('--recipient-profile-json', default=None, required=False, type=JsonClickType(), + help='JSON string containing recipient profile.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_provider_cli(api_client, name, comment, recipient_profile_json_file, + recipient_profile_json): + """ + Create a provider. + + The public specification for the JSON request is in development. + """ + if recipient_profile_json is None and recipient_profile_json_file is None: + created_provider = UnityCatalogApi(api_client).create_provider( + name, comment, recipient_profile=None) + click.echo(mc_pretty_format(created_provider)) + json_cli_base(recipient_profile_json_file, recipient_profile_json, + lambda json: UnityCatalogApi(api_client).create_provider(name, comment, json), + error_msg='Either --recipient-profile-json-file or ' + + '--recipient-profile-json should be provided') + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List providers.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_providers_cli(api_client): + """ + List providers. + """ + proviers_json = UnityCatalogApi(api_client).list_providers() + click.echo(mc_pretty_format(proviers_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a provider.') +@click.option('--name', required=True, + help='Name of the provider to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_provider_cli(api_client, name): + """ + Get a provider. + """ + provier_json = UnityCatalogApi(api_client).get_provider(name) + click.echo(mc_pretty_format(provier_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a provider.') +@click.option('--name', required=True, help='Name of the provider to update.') +@click.option('--new_name', default=None, help='New name of the provider.') +@click.option('--comment', default=None, required=False, + help='Free-form text description.') +@click.option('--recipient-profile-json-file', default=None, type=click.Path(), + help='File containing recipient profile in JSON format.') +@click.option('--recipient-profile-json', default=None, type=JsonClickType(), + help='JSON string containing recipient profile.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_provider_cli(api_client, name, new_name, comment, recipient_profile_json_file, + recipient_profile_json): + """ + Update a provider. + + The public specification for the JSON request is in development. + """ + if recipient_profile_json is None and recipient_profile_json_file is None: + updated_provider = UnityCatalogApi(api_client).update_provider(name, new_name, comment) + click.echo(mc_pretty_format(updated_provider)) + else: + json_cli_base(recipient_profile_json_file, recipient_profile_json, + lambda json: UnityCatalogApi(api_client).update_provider(name, new_name, + comment, json), + error_msg='Either --recipient-profile-json-file or ' + + '--recipient-profile-json should be provided') + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List shares of a provider.') +@click.option('--name', required=True, + help='Name of the provider.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_provider_shares_cli(api_client, name): + """ + List a provider's shares. + """ + shares_json = UnityCatalogApi(api_client).list_provider_shares(name) + click.echo(mc_pretty_format(shares_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a provider.') +@click.option('--name', required=True, + help='Name of the provider to delete.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_provider_cli(api_client, name): + """ + Delete a provider. + """ + UnityCatalogApi(api_client).delete_provider(name) + + +@click.group() +def shares_group(): # pragma: no cover + pass + + +def register_shares_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_share_cli), name='create-share') + cmd_group.add_command(hide(list_shares_cli), name='list-shares') + cmd_group.add_command(hide(get_share_cli), name='get-share') + cmd_group.add_command(hide(update_share_cli), name='update-share') + cmd_group.add_command(hide(delete_share_cli), name='delete-share') + cmd_group.add_command(hide(list_share_permissions_cli), name='list-share-permissions') + cmd_group.add_command(hide(update_share_permissions_cli), name='update-share-permissions') + + # Register command group. + shares_group.add_command(create_share_cli, name='create') + shares_group.add_command(list_shares_cli, name='list') + shares_group.add_command(get_share_cli, name='get') + shares_group.add_command(update_share_cli, name='update') + shares_group.add_command(delete_share_cli, name='delete') + shares_group.add_command(list_share_permissions_cli, name='list-permissions') + shares_group.add_command(update_share_permissions_cli, name='update-permissions') + cmd_group.add_command(shares_group, name='shares') + + +@click.group() +def recipients_group(): # pragma: no cover + pass + + +def register_recipients_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_recipient_cli), name='create-recipient') + cmd_group.add_command(hide(list_recipients_cli), name='list-recipients') + cmd_group.add_command(hide(get_recipient_cli), name='get-recipient') + cmd_group.add_command(hide(update_recipient_cli), name='update-recipient') + cmd_group.add_command(hide(rotate_recipient_token_cli), name='rotate-recipient-token') + cmd_group.add_command(hide(list_recipient_permissions_cli), name='list-recipient-permissions') + cmd_group.add_command(hide(delete_recipient_cli), name='delete-recipient') + + # Register command group. + recipients_group.add_command(create_recipient_cli, name='create') + recipients_group.add_command(list_recipients_cli, name='list') + recipients_group.add_command(get_recipient_cli, name='get') + recipients_group.add_command(update_recipient_cli, name='update') + recipients_group.add_command(rotate_recipient_token_cli, name='rotate-token') + recipients_group.add_command(list_recipient_permissions_cli, name='list-permissions') + recipients_group.add_command(delete_recipient_cli, name='delete') + cmd_group.add_command(recipients_group, name='recipients') + + +@click.group() +def providers_group(): # pragma: no cover + pass + + +def register_providers_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_provider_cli), name='create-provider') + cmd_group.add_command(hide(list_providers_cli), name='list-providers') + cmd_group.add_command(hide(get_provider_cli), name='get-provider') + cmd_group.add_command(hide(update_provider_cli), name='update-provider') + cmd_group.add_command(hide(delete_provider_cli), name='delete-provider') + cmd_group.add_command(hide(list_provider_shares_cli), name='list-provider-shares') + + # Register command group. + providers_group.add_command(create_provider_cli, name='create') + providers_group.add_command(list_providers_cli, name='list') + providers_group.add_command(get_provider_cli, name='get') + providers_group.add_command(update_provider_cli, name='update') + providers_group.add_command(delete_provider_cli, name='delete') + providers_group.add_command(list_provider_shares_cli, name='list-shares') + cmd_group.add_command(providers_group, name='providers') + + +def register_delta_sharing_commands(cmd_group): + register_shares_commands(cmd_group) + register_recipients_commands(cmd_group) + register_providers_commands(cmd_group) diff --git a/databricks_cli/unity_catalog/ext_loc_cli.py b/databricks_cli/unity_catalog/ext_loc_cli.py new file mode 100644 index 00000000..3221f570 --- /dev/null +++ b/databricks_cli/unity_catalog/ext_loc_cli.py @@ -0,0 +1,232 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import JsonClickType +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import del_none, hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create External Location.') +@click.option('--name', default=None, + help='Name of new external location') +@click.option('--url', default=None, + help='Path URL for the new external location') +@click.option('--storage-credential-name', default=None, + help='Name of storage credential to use with new external location') +@click.option('--skip-validation', '-s', 'skip_val', is_flag=True, default=False, + help='Skip the validation of location\'s storage credential before creation') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='POST', path='/external-locations')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='POST', path='/external-locations')) +@debug_option +@profile_option +# UC's createExternalLocation returns a 401 when the validation of the external location's +# storage credential fails; that translates to a misleading error when eat_exceptions is enabled: +# Your authentication information may be incorrect. Please reconfigure with ``dbfs configure`` +# Until that is fixed (should return a 400), show full error trace. +#@eat_exceptions +@provide_api_client +def create_location_cli(api_client, name, url, storage_credential_name, skip_val, json_file, json): + """ + Create new external location. + + The public specification for the JSON request is in development. + """ + if (name is not None) and (url is not None) and (storage_credential_name is not None): + if (json_file is not None) or (json is not None): + raise ValueError('Cannot specify JSON if both name and url are given') + data = {"name": name, "url": url, "credential_name": storage_credential_name} + loc_json = UnityCatalogApi(api_client).create_external_location(data, skip_val) + click.echo(mc_pretty_format(loc_json)) + elif (json is None) and (json_file is None): + raise ValueError('Must provide name, url and storage-credential-name' + + ' or use JSON specification') + else: + json_cli_base(json_file, json, + lambda json: + UnityCatalogApi(api_client).create_external_location(json, skip_val), + encode_utf8=True) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List external locations.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_locations_cli(api_client, ): + """ + List external locations. + """ + locs_json = UnityCatalogApi(api_client).list_external_locations() + click.echo(mc_pretty_format(locs_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get an external location.') +@click.option('--name', required=True, + help='Name of the external location to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_location_cli(api_client, name): + """ + Get an external location. + """ + loc_json = UnityCatalogApi(api_client).get_external_location(name) + click.echo(mc_pretty_format(loc_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update an external location.') +@click.option('--name', required=True, + help='Name of the external location to update.') +@click.option('--force', '-f', is_flag=True, default=False, + help='Force update even if location has dependent tables/mounts') +@click.option('--skip-validation', '-s', 'skip_val', is_flag=True, default=False, + help='Skip the validation of location\'s storage credential before creation') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/external-locations/{name}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/external-locations/{name}')) +@debug_option +@profile_option +# See comment for create_location_cli +#@eat_exceptions +@provide_api_client +def update_location_cli(api_client, name, force, skip_val, json_file, json): + """ + Update an external location. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_external_location(name, json, + force, + skip_val), + encode_utf8=True) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete an external location.') +@click.option('--name', required=True, + help='Name of the external location to delete.') +@click.option('--force', '-f', is_flag=True, default=False, + help='Force deletion even if location has dependent tables/mounts') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_location_cli(api_client, name, force): + """ + Delete an external location. + """ + UnityCatalogApi(api_client).delete_external_location(name, force) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Validate a external location/credential pair') +@click.option('--name', default=None, + help='Name of the external location to validate.') +@click.option('--url', default=None, + help='A storage URL to validate.') +@click.option('--cred-name', default=None, + help='Name of the storage credential to use for validation.') +@click.option('--cred-aws-iam-role', default=None, + help='An aws role to validate') +@click.option('--cred-az-directory-id', default=None, + help='An Azure directory id to validate') +@click.option('--cred-az-application-id', default=None, + help='An Azure application id to validate') +@click.option('--cred-az-client-secret', default=None, + help='An Azure directory id to validate') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def validate_location_cli(api_client, name, url, cred_name, cred_aws_iam_role, cred_az_directory_id, + cred_az_application_id, cred_az_client_secret): + """ + Validate an external location/credential combination. + + This call will attempt to read/list/write/delete with the given credentials and + external location. + + One of name/url must be provided. If both are specified, the given credential + name will be excluded from path overlap checks (used to validate a potential + update of that credential). + + One of cred-name, or cloud provider specific credential parameters must be + provided. + """ + validation_spec = { + "external_location_name": name, + "url": url, + "storage_credential_name": cred_name, + } + if cred_aws_iam_role is not None: + validation_spec["aws_iam_role"] = { + "role_arn": cred_aws_iam_role + } + + if cred_az_directory_id is not None: + validation_spec["azure_service_principal"] = { + "directory_id": cred_az_directory_id, + "application_id": cred_az_application_id, + "client_secret": cred_az_client_secret + } + del_none(validation_spec) + validation_json = UnityCatalogApi(api_client).validate_external_location(validation_spec) + click.echo(mc_pretty_format(validation_json)) + + +@click.group() +def external_locations_group(): # pragma: no cover + pass + + +def register_ext_loc_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_location_cli), name='create-external-location') + cmd_group.add_command(hide(list_locations_cli), name='list-external-locations') + cmd_group.add_command(hide(get_location_cli), name='get-external-location') + cmd_group.add_command(hide(update_location_cli), name='update-external-location') + cmd_group.add_command(hide(delete_location_cli), name='delete-external-location') + cmd_group.add_command(hide(validate_location_cli), name='validate-external-location') + + # Register command group. + external_locations_group.add_command(create_location_cli, name='create') + external_locations_group.add_command(list_locations_cli, name='lists') + external_locations_group.add_command(get_location_cli, name='get') + external_locations_group.add_command(update_location_cli, name='update') + external_locations_group.add_command(delete_location_cli, name='delete') + external_locations_group.add_command(validate_location_cli, name='validate') + cmd_group.add_command(external_locations_group, name='external-locations') diff --git a/databricks_cli/unity_catalog/lineage_cli.py b/databricks_cli/unity_catalog/lineage_cli.py new file mode 100644 index 00000000..d7da61a3 --- /dev/null +++ b/databricks_cli/unity_catalog/lineage_cli.py @@ -0,0 +1,214 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import mc_pretty_format, hide +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, to_graph + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List table lineage.') +@click.option('--table-name', required=True, + help='Name of the table with 3L namespace') +@click.option('--level', required=False, type=int, default=1, + help='level of lineage to retrieve') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_table_lineages_cli(api_client, table_name, level): + """ + List table lineage by table name. + :table_name str: name of the table with 3L format. E.g catalog.schema.table + + example response: + + digraph "lineage graph of main.lineage.user_account" { + "main.lineage.user_account" -> "main.lineage.user_transaction"; + "main.lineage.dinner_price" -> "main.lineage.price_entry","main.lineage.user_account"; + } + + Returns the specified levels of downstream/upstream + + """ + node_to_downstream = list_table_lineages_recursive_cli(api_client, table_name, level) + click.echo(to_graph(node_to_downstream, "lineage graph of {}".format(table_name))) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List column lineage.') +@click.option('--table-name', required=True, + help='Name of the table with 3L namespace') +@click.option('--column-name', required=True, + help='Name of the column for lineage analysis') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_column_lineages_cli(api_client, table_name, column_name): + """ + List column lineage by table name and column name. + :table_name str: name of the table with 3L format. E.g catalog.schema.table + :column_name str: name of the column + + example response: + { + "downstream_cols": [ + { + "workspace_id": 6051921418418893, + "table_type": "TABLE", + "catalog_name": "main", + "table_name": "dinner_price", + "schema_name": "lineage", + "name": "full_menu" + } + ], + "upstream_cols": [ + { + "workspace_id": 6051921418418893, + "table_type": "TABLE", + "catalog_name": "main", + "table_name": "menu", + "schema_name": "lineage", + "name": "app" + }, + { + "workspace_id": 6051921418418893, + "table_type": "TABLE", + "catalog_name": "main", + "table_name": "menu", + "schema_name": "lineage", + "name": "desert" + }, + { + "workspace_id": 6051921418418893, + "table_type": "TABLE", + "catalog_name": "main", + "table_name": "menu", + "schema_name": "lineage", + "name": "main" + } + ] + } + + Returns the downstream/upstream of a given column + """ + + schemas_json = UnityCatalogApi(api_client).list_lineages_by_column(table_name, column_name) + click.echo(mc_pretty_format(schemas_json)) + + +@click.group() +def lineage_group(): # pragma: no cover + pass + + +def register_lineage_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(list_table_lineages_cli), name='list-table-lineages') + cmd_group.add_command(hide(list_column_lineages_cli), name='list-column-lineages') + + # Register command group. + # Note: we deviate from the "noun-verb" pattern here because it would be awkward to have to + # spell out "list" or "list-for". Table and column lineage is read only by definition. + lineage_group.add_command(list_table_lineages_cli, name='table') + lineage_group.add_command(list_column_lineages_cli, name='column') + cmd_group.add_command(lineage_group, name='lineage') + + +def get_table_name(table_node): + return "{}.{}.{}".format( + table_node['catalog_name'], + table_node['schema_name'], + table_node['name'] + ) + + +def list_table_lineages_recursive_cli(api_client, table_name, level): + node_to_downstream = {} + level_count = 0 + current_level = [table_name] + next_level = [] + initial_upstream = [] + # go downstream + while level_count < level: + for current_table in current_level: + if current_table in node_to_downstream: + # skip if the table is visited before + continue + lineage_json = UnityCatalogApi(api_client).list_lineages_by_table(current_table) + if level == 0: + initial_upstream = [ + get_table_name( + table_node + ) for table_node in lineage_json['upstream_tables'] + ] if 'upstream_tables' in lineage_json else [] + cur_downstream = [ + get_table_name( + table_node + ) for table_node in lineage_json['downstream_tables'] + ] if 'downstream_tables' in lineage_json else [] + next_level.extend(cur_downstream) + if len(cur_downstream) > 0: + node_to_downstream[current_table] = cur_downstream + level_count = level_count + 1 + current_level = next_level + next_level = [] + # go upstream + level_count = 1 + current_level = initial_upstream + connect_upstream_tables(initial_upstream, table_name, node_to_downstream) + next_level = [] + while level_count <= level: + for current_table in current_level: + lineage_json = UnityCatalogApi(api_client).list_lineages_by_table(current_table) + upstream_of_current = [ + get_table_name( + table_node + ) for table_node in lineage_json['upstream_tables'] + ] if 'upstream_tables' in lineage_json else [] + next_level.extend(upstream_of_current) + cur_downstream = [ + get_table_name( + table_node + ) for table_node in lineage_json['downstream_tables'] + ] if 'downstream_tables' in lineage_json else [] + if len(cur_downstream) > 0: + if current_table in node_to_downstream: + node_to_downstream[current_table] = cur_downstream + level_count = level_count + 1 + current_level = next_level + next_level = [] + return node_to_downstream + + +def connect_upstream_tables(upstream_tables, current_table, node_to_downstream): + """ + fill node_to_downstream dict based with give upstreams and current table + """ + if len(upstream_tables) > 0: + if upstream_tables not in node_to_downstream: + node_to_downstream[upstream_tables] = [current_table] diff --git a/databricks_cli/unity_catalog/metastore_cli.py b/databricks_cli/unity_catalog/metastore_cli.py new file mode 100644 index 00000000..12dcda17 --- /dev/null +++ b/databricks_cli/unity_catalog/metastore_cli.py @@ -0,0 +1,206 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import MetastoreIdClickType, WorkspaceIdClickType, JsonClickType +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +################# Metastore Commands ##################### + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create a metastore.') +@click.option('--name', required=True, help='Name of the new metastore.') +@click.option('--storage-root', required=True, + help='Storage root URL for the new metastore.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_metastore_cli(api_client, name, storage_root): + """ + Create new metastore. + """ + metastore_json = UnityCatalogApi(api_client).create_metastore(name, storage_root) + click.echo(mc_pretty_format(metastore_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List metastores.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_metastores_cli(api_client): + """ + List metastores. + """ + metastores_json = UnityCatalogApi(api_client).list_metastores() + click.echo(mc_pretty_format(metastores_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a metastore.') +@click.option('--id', 'metastore_id', required=True, type=MetastoreIdClickType(), + help='Unique identifier of the metastore to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_metastore_cli(api_client, metastore_id): + """ + Get a metastore. + """ + metastore_json = UnityCatalogApi(api_client).get_metastore(metastore_id) + click.echo(mc_pretty_format(metastore_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a metastore.') +@click.option('--id', 'metastore_id', required=True, type=MetastoreIdClickType(), + help='Unique identifier of the metastore to update.') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/metastores/{id}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/metastores/{id}')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_metastore_cli(api_client, metastore_id, json_file, json): + """ + Update a metastore. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_metastore(metastore_id, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a metastore.') +@click.option('--id', 'metastore_id', required=True, type=MetastoreIdClickType(), + help='Unique identifier of the metastore to delete.') +@click.option('--force', '-f', is_flag=True, default=False) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_metastore_cli(api_client, metastore_id, force): + """ + Delete a metastore. + """ + UnityCatalogApi(api_client).delete_metastore(metastore_id, force) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get summary info of current metastore.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def metastore_summary_cli(api_client): + """ + Get metastore summary. + """ + summary_json = UnityCatalogApi(api_client).get_metastore_summary() + click.echo(mc_pretty_format(summary_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Assign a metastore to a workspace.') +@click.option('--workspace-id', 'workspace_id', required=True, type=WorkspaceIdClickType(), + help='Unique identifier of the workspace for the metastore assignment.') +@click.option('--metastore-id', 'metastore_id', required=True, type=MetastoreIdClickType(), + help='Unique identifier of the metastore to assign to the workspace.') +@click.option('--default-catalog-name', 'default_catalog_name', required=False, + default='hive_metastore', + help='Name of the default catalog to use with the metastore ' + + '(default: "hive_metastore").') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def assign_metastore_cli(api_client, workspace_id, metastore_id, default_catalog_name): + """ + Assign a metastore to a specified workspace. + + If the workspace already has a metastore assigned, it is updated. + """ + resp = UnityCatalogApi(api_client).create_metastore_assignment(workspace_id, metastore_id, + default_catalog_name) + # resp will just be an empty object ('{}') but it's good to print *something* + click.echo(mc_pretty_format(resp)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Unassigns a metastore from a workspace.') +@click.option('--workspace-id', 'workspace_id', required=True, type=WorkspaceIdClickType(), + help='Unique identifier of the workspace.') +@click.option('--metastore-id', 'metastore_id', required=True, type=MetastoreIdClickType(), + help='Unique identifier of the metastore to unassign from the workspace.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def unassign_metastore_cli(api_client, workspace_id, metastore_id): + """ + Unassign a metastore from a workspace. + """ + resp = UnityCatalogApi(api_client).delete_metastore_assignment(workspace_id, metastore_id) + # resp will just be an empty object ('{}') but it's good to print *something* + click.echo(mc_pretty_format(resp)) + + +@click.group() +def metastores_group(): # pragma: no cover + pass + + +def register_metastore_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_metastore_cli), name='create-metastore') + cmd_group.add_command(hide(list_metastores_cli), name='list-metastores') + cmd_group.add_command(hide(get_metastore_cli), name='get-metastore') + cmd_group.add_command(hide(update_metastore_cli), name='update-metastore') + cmd_group.add_command(hide(delete_metastore_cli), name='delete-metastore') + cmd_group.add_command(hide(metastore_summary_cli), name='metastore-summary') + cmd_group.add_command(hide(assign_metastore_cli), name='assign-metastore') + cmd_group.add_command(hide(unassign_metastore_cli), name='unassign-metastore') + + # Register command group. + metastores_group.add_command(create_metastore_cli, name='create') + metastores_group.add_command(list_metastores_cli, name='list') + metastores_group.add_command(get_metastore_cli, name='get') + metastores_group.add_command(update_metastore_cli, name='update') + metastores_group.add_command(delete_metastore_cli, name='delete') + metastores_group.add_command(metastore_summary_cli, name='get-summary') + metastores_group.add_command(assign_metastore_cli, name='assign') + metastores_group.add_command(unassign_metastore_cli, name='unassign') + cmd_group.add_command(metastores_group, name='metastores') diff --git a/databricks_cli/unity_catalog/perms_cli.py b/databricks_cli/unity_catalog/perms_cli.py new file mode 100644 index 00000000..ea6e19a3 --- /dev/null +++ b/databricks_cli/unity_catalog/perms_cli.py @@ -0,0 +1,140 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import JsonClickType, OneOfOption +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +PERMISSIONS_OBJ_TYPES = [ + 'catalog', 'schema', 'table', 'storage-credential', 'external-location' +] + + +def _get_perm_securable_name_and_type(catalog_name, schema_full_name, table_full_name, + credential_name, location_name): + if catalog_name: + return ('catalog', catalog_name) + elif schema_full_name: + return ('schema', schema_full_name) + elif table_full_name: + return ('table', table_full_name) + elif credential_name: + return ('storage-credential', credential_name) + else: + return ('external-location', location_name) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get permissions on a securable.') +@click.option('--catalog', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Name of catalog of interest') +@click.option('--schema', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Full name of schema of interest') +@click.option('--table', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Full name of table of interest') +@click.option('--storage-credential', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Name of the storage credential of interest') +@click.option('--external-location', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Name of the external location of interest') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_permissions_cli(api_client, catalog, schema, table, storage_credential, + external_location): + """ + Get permissions on a securable. + """ + sec_type, sec_name = _get_perm_securable_name_and_type(catalog, schema, table, + storage_credential, external_location) + + perm_json = UnityCatalogApi(api_client).get_permissions(sec_type, sec_name) + click.echo(mc_pretty_format(perm_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='update permissions on a securable.') +@click.option('--catalog', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Name of catalog of interest') +@click.option('--schema', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Full name of schema of interest') +@click.option('--table', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Full name of table of interest') +@click.option('--storage-credential', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Name of the storage credential of interest') +@click.option('--external-location', cls=OneOfOption, default=None, + one_of=PERMISSIONS_OBJ_TYPES, + help='Name of the external location of interest') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/permissions/{securable}/{id}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/permissions/{securable}/{id}')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_permissions_cli(api_client, catalog, schema, table, storage_credential, + external_location, json_file, json): + """ + Update permissions on a securable. + + The public specification for the JSON request is in development. + """ + sec_type, sec_name = _get_perm_securable_name_and_type(catalog, schema, table, + storage_credential, external_location) + + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_permissions(sec_type, sec_name, + json), + encode_utf8=True) + + +@click.group() +def permissions_group(): # pragma: no cover + pass + + +def register_perms_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(get_permissions_cli), name='get-permissions') + cmd_group.add_command(hide(update_permissions_cli), name='update-permissions') + + # Register command group. + permissions_group.add_command(get_permissions_cli, name='get') + permissions_group.add_command(update_permissions_cli, name='update') + cmd_group.add_command(permissions_group, name='permissions') diff --git a/databricks_cli/unity_catalog/schema_cli.py b/databricks_cli/unity_catalog/schema_cli.py new file mode 100644 index 00000000..f2b2745b --- /dev/null +++ b/databricks_cli/unity_catalog/schema_cli.py @@ -0,0 +1,154 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import JsonClickType +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create a new schema.') +@click.option('--catalog-name', required=True, help='Parent catalog of new schema.') +@click.option('--name', required=True, + help='Name of new schema, relative to parent catalog.') +@click.option('--comment', default=None, required=False, + help='Free-form text description.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_schema_cli(api_client, catalog_name, name, comment): + """ + Create a new schema in the specified catalog. + """ + schema_json = UnityCatalogApi(api_client).create_schema(catalog_name, name, comment) + click.echo(mc_pretty_format(schema_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List schemas.') +@click.option('--catalog-name', required=True, + help='Name of the parent catalog for schemas of interest.') +@click.option('--name-pattern', default=None, + help='SQL LIKE pattern that the schema name must match to be in list.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_schemas_cli(api_client, catalog_name, name_pattern): + """ + List schemas. + """ + schemas_json = UnityCatalogApi(api_client).list_schemas(catalog_name, name_pattern) + click.echo(mc_pretty_format(schemas_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a schema.') +@click.option('--full-name', required=True, + help='Full name (.) of the schema to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_schema_cli(api_client, full_name): + """ + Get a schema. + """ + schema_json = UnityCatalogApi(api_client).get_schema(full_name) + click.echo(mc_pretty_format(schema_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a schema.') +@click.option('--full-name', required=True, + help='Full name (.) of the schema to update.') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/schemas/{name}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/schemas/{name}')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_schema_cli(api_client, full_name, json_file, json): + """ + Update a schema. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_schema(full_name, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a schema.') +@click.option('--full-name', required=True, + help='Full name (.) of the schema to delete.') +@click.option('--purge', '-p', is_flag=True, default=False, + help='Purge all child schemas and tables of catalog.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_schema_cli(api_client, full_name, purge): + """ + Delete a schema. + """ + if purge: + (catalog_name, schema_name) = full_name.split('.') + click.echo("Purging all tables from schema %s in catalog %s" % (schema_name, catalog_name)) + + tables_response = UnityCatalogApi(api_client).list_tables(catalog_name, schema_name, None) + for t in tables_response['tables']: + click.echo("Deleting table: %s" % (t['full_name'])) + UnityCatalogApi(api_client).delete_table(t['full_name']) + + UnityCatalogApi(api_client).delete_schema(full_name) + + +@click.group() +def schemas_group(): # pragma: no cover + pass + + +def register_schema_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_schema_cli), name='create-schema') + cmd_group.add_command(hide(list_schemas_cli), name='list-schemas') + cmd_group.add_command(hide(get_schema_cli), name='get-schema') + cmd_group.add_command(hide(update_schema_cli), name='update-schema') + cmd_group.add_command(hide(delete_schema_cli), name='delete-schema') + + # Register command group. + schemas_group.add_command(create_schema_cli, name='create') + schemas_group.add_command(list_schemas_cli, name='list') + schemas_group.add_command(get_schema_cli, name='get') + schemas_group.add_command(update_schema_cli, name='update') + schemas_group.add_command(delete_schema_cli, name='delete') + cmd_group.add_command(schemas_group, name='schemas') diff --git a/databricks_cli/unity_catalog/table_cli.py b/databricks_cli/unity_catalog/table_cli.py new file mode 100644 index 00000000..8873c9e9 --- /dev/null +++ b/databricks_cli/unity_catalog/table_cli.py @@ -0,0 +1,178 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import JsonClickType +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.unity_catalog.api import UnityCatalogApi +from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \ + mc_pretty_format +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Create a table. [DO NOT USE]', + hidden=True) +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='POST', path='/tables')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='POST', path='/tables')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_table_cli(api_client, json_file, json): + """ + Create new table specified by the JSON input. + + WARNING: Creating table metadata via the UC API may create a table + that is unusable in DBR. Instead, use SQL commands (CREATE TABLE) in DBR. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).create_table(json), + encode_utf8=True) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List tables.') +@click.option('--catalog-name', required=True, + help='Name of the parent catalog for tables of interest.') +@click.option('--schema-name', required=True, + help='Name of the parent schema for tables of interest.') +@click.option('--name-pattern', default=None, + help='SQL LIKE pattern that the table name must match to be in list.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_tables_cli(api_client, catalog_name, schema_name, name_pattern): + """ + List tables. + """ + tables_json = UnityCatalogApi(api_client).list_tables(catalog_name, schema_name, name_pattern) + click.echo(mc_pretty_format(tables_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='List table summaries.') +@click.option('--catalog-name', required=True, + help='Name of the parent catalog for tables of interest.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_table_summaries_cli(api_client, catalog_name): + """ + List table summaries (in bulk). + """ + tables_json = UnityCatalogApi(api_client).list_table_summaries(catalog_name) + click.echo(mc_pretty_format(tables_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Get a table.') +@click.option('--full-name', required=True, + help='Full name (..) of the table to get.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_table_cli(api_client, full_name): + """ + Get a table. + """ + table_json = UnityCatalogApi(api_client).get_table(full_name) + click.echo(mc_pretty_format(table_json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a table. [DO NOT USE]', + hidden=True) +@click.option('--full-name', required=True, + help='Full name (..
) of the table to update.') +@click.option('--json-file', default=None, type=click.Path(), + help=json_file_help(method='PATCH', path='/tables/{name}')) +@click.option('--json', default=None, type=JsonClickType(), + help=json_string_help(method='PATCH', path='/tables/{name}')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_table_cli(api_client, full_name, json_file, json): + """ + Update a table. + + WARNING: Altering table metadata via the UC API may cause the table + to be unusable in DBR. Instead, use SQL commands (ALTER TABLE) in DBR. + + The public specification for the JSON request is in development. + """ + json_cli_base(json_file, json, + lambda json: UnityCatalogApi(api_client).update_table(full_name, json), + encode_utf8=True) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Delete a table.') +@click.option('--full-name', required=True, + help='Full name (..
) of the table to delete.') +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_table_cli(api_client, full_name): + """ + Delete a table. + """ + UnityCatalogApi(api_client).delete_table(full_name) + + +@click.group() +def tables_group(): # pragma: no cover + """ + Note: To create or update tables, please run the appropriate SQL commands + on a cluster or endpoint (CREATE TABLE and ALTER TABLE). + """ + pass + + +def register_table_commands(cmd_group): + # Register deprecated "verb-noun" commands for backward compatibility. + cmd_group.add_command(hide(create_table_cli), name='create-table') + cmd_group.add_command(hide(list_tables_cli), name='list-tables') + cmd_group.add_command(hide(list_table_summaries_cli), name='list-table-summaries') + cmd_group.add_command(hide(get_table_cli), name='get-table') + cmd_group.add_command(hide(update_table_cli), name='update-table') + cmd_group.add_command(hide(delete_table_cli), name='delete-table') + + # Register command group. + tables_group.add_command(create_table_cli, name='create') + tables_group.add_command(list_tables_cli, name='list') + tables_group.add_command(list_table_summaries_cli, name='list-summaries') + tables_group.add_command(get_table_cli, name='get') + tables_group.add_command(update_table_cli, name='update') + tables_group.add_command(delete_table_cli, name='delete') + cmd_group.add_command(tables_group, name='tables') diff --git a/databricks_cli/unity_catalog/uc_service.py b/databricks_cli/unity_catalog/uc_service.py new file mode 100644 index 00000000..f19d7cf8 --- /dev/null +++ b/databricks_cli/unity_catalog/uc_service.py @@ -0,0 +1,493 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from databricks_cli.unity_catalog.utils import mc_pretty_format + + +class UnityCatalogService(object): + def __init__(self, client): + self.client = client + + # Metastore Operations + + def create_metastore(self, name, storage_root, headers=None): + _data = { + 'name': name, + 'storage_root': storage_root, + } + return self.client.perform_query('POST', '/unity-catalog/metastores', data=_data, + headers=headers) + + def list_metastores(self, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/metastores', data=_data, + headers=headers) + + def get_metastore(self, metastore_id, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/metastores/%s' % (metastore_id), + data=_data, headers=headers) + + def get_metastore_summary(self, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/metastore_summary', data=_data, + headers=headers) + + def update_metastore(self, metastore_id, metastore_spec, headers=None): + return self.client.perform_query('PATCH', '/unity-catalog/metastores/%s' % (metastore_id), + data=metastore_spec, headers=headers) + + def delete_metastore(self, metastore_id, force=None, headers=None): + _data = {} + if force is not None: + _data['force'] = force + + return self.client.perform_query('DELETE', '/unity-catalog/metastores/%s' % (metastore_id), + data=_data, headers=headers) + + def create_metastore_assignment(self, workspace_id, metastore_id, default_catalog_name=None, + headers=None): + _data = { + 'metastore_id': metastore_id + } + if default_catalog_name is not None: + _data['default_catalog_name'] = default_catalog_name + url = '/unity-catalog/workspaces/%s/metastore' % (workspace_id) + return self.client.perform_query('PUT', url, data=_data, headers=headers) + + def update_metastore_assignment(self, workspace_id, metastore_id, default_catalog_name, + headers=None): + _data = { + 'metastore_id': metastore_id, + 'default_catalog_name': default_catalog_name + } + url = '/unity-catalog/workspaces/%s/metastore' % (workspace_id) + return self.client.perform_query('PATCH', url, data=_data, headers=headers) + + def delete_metastore_assignment(self, workspace_id, metastore_id, headers=None): + _data = { + 'metastore_id': metastore_id + } + url = '/unity-catalog/workspaces/%s/metastore' % (workspace_id) + return self.client.perform_query('DELETE', url, data=_data, headers=headers) + + # External Location Operations + + def create_external_location(self, loc_spec, skip_validation, headers=None): + # Merge the skip_validation arg, since it's not a query arg and the + # ExternalLocationInfo spec is 'inline' + if skip_validation: + loc_spec['skip_validation'] = skip_validation + url = '/unity-catalog/external-locations' + return self.client.perform_query('POST', url, data=loc_spec, headers=headers) + + def list_external_locations(self, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/external-locations', data=_data, + headers=headers) + + def get_external_location(self, name, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/external-locations/%s' % (name), + data=_data, headers=headers) + + def update_external_location(self, name, loc_spec, force, skip_validation, headers=None): + _data = loc_spec + # Merge the skip_validation arg, since it's not a query arg and the + # ExternalLocationInfo spec is 'inline' + if skip_validation: + _data['skip_validation'] = skip_validation + # Same for the 'force' field. + if force: + _data["force"] = True + + return self.client.perform_query('PATCH', '/unity-catalog/external-locations/%s' % (name), + data=_data, headers=headers) + + def delete_external_location(self, name, force, headers=None): + _data = { + "force": force + } + return self.client.perform_query('DELETE', '/unity-catalog/external-locations/%s' % (name), + data=_data, headers=headers) + + def validate_external_location(self, validation_spec, headers=None): + return self.client.perform_query('POST', '/unity-catalog/validate-storage-credentials', + data=validation_spec, headers=headers) + + # Data Access Configuration Operations + + def create_dac(self, metastore_id, dac_spec, skip_validation, headers=None): + if skip_validation: + dac_spec['skip_validation'] = skip_validation + url = '/unity-catalog/metastores/%s/data-access-configurations' % (metastore_id) + return self.client.perform_query('POST', url, data=dac_spec, headers=headers) + + def list_dacs(self, metastore_id, headers=None): + _data = {} + url = '/unity-catalog/metastores/%s/data-access-configurations' % (metastore_id) + return self.client.perform_query('GET', url, data=_data, headers=headers) + + def get_dac(self, metastore_id, dac_id, headers=None): + url = '/unity-catalog/metastores/%s/data-access-configurations/%s' % (metastore_id, dac_id) + return self.client.perform_query('GET', url, headers=headers) + + def delete_dac(self, metastore_id, dac_id, headers=None): + url = '/unity-catalog/metastores/%s/data-access-configurations/%s' % (metastore_id, dac_id) + return self.client.perform_query('DELETE', url, headers=headers) + + # Storage Credential Operations + + def create_storage_credential(self, cred_spec, skip_validation, headers=None): + # Merge the skip_validation arg, since it's not a query arg and the + # StorageCredentialInfo spec is 'inline' + if skip_validation: + cred_spec['skip_validation'] = skip_validation + url = '/unity-catalog/storage-credentials' + return self.client.perform_query('POST', url, data=cred_spec, headers=headers) + + def list_storage_credentials(self, name_pattern=None, headers=None): + _data = {} + if name_pattern is not None: + _data['name_pattern'] = name_pattern + + return self.client.perform_query('GET', '/unity-catalog/storage-credentials', + data=_data, headers=headers) + + def get_storage_credential(self, name, headers=None): + _data = {} + + return self.client.perform_query('GET', '/unity-catalog/storage-credentials/%s' % (name), + data=_data, headers=headers) + + def update_storage_credential(self, name, cred_spec, skip_validation, headers=None): + # Merge the skip_validation arg, since it's not a query arg and the + # StorageCredentialInfo spec is 'inline' + if skip_validation: + cred_spec['skip_validation'] = skip_validation + return self.client.perform_query('PATCH', '/unity-catalog/storage-credentials/%s' % (name), + data=cred_spec, headers=headers) + + def delete_storage_credential(self, name, force, headers=None): + _data = {} + if force: + _data["force"] = True + + return self.client.perform_query('DELETE', '/unity-catalog/storage-credentials/%s' % (name), + data=_data, headers=headers) + + # Catalog Operations + + def create_catalog(self, name, comment=None, provider=None, share=None, headers=None): + _data = { + 'name': name, + } + if comment is not None: + _data['comment'] = comment + if provider is not None: + _data['provider_name'] = provider + if share is not None: + _data['share_name'] = share + return self.client.perform_query('POST', '/unity-catalog/catalogs', data=_data, + headers=headers) + + def list_catalogs(self, headers=None): + _data = {} + + return self.client.perform_query('GET', '/unity-catalog/catalogs', data=_data, + headers=headers) + + def get_catalog(self, name, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/catalogs/%s' % (name), + data=_data, headers=headers) + + def update_catalog(self, name, catalog_spec, headers=None): + return self.client.perform_query('PATCH', '/unity-catalog/catalogs/%s' % (name), + data=catalog_spec, headers=headers) + + def delete_catalog(self, name, headers=None): + _data = {} + return self.client.perform_query('DELETE', '/unity-catalog/catalogs/%s' % (name), + data=_data, headers=headers) + + # Schema Operations + + def create_schema(self, catalog_name, new_schema_name, comment=None, headers=None): + _data = { + 'catalog_name': catalog_name, + 'name': new_schema_name, + } + if comment is not None: + _data['comment'] = comment + return self.client.perform_query('POST', '/unity-catalog/schemas', data=_data, + headers=headers) + + def list_schemas(self, catalog_name=None, name_regex=None, headers=None): + _data = {} + if catalog_name is not None: + _data['catalog_name'] = catalog_name + if name_regex is not None: + _data['schema_name_regex'] = name_regex + + return self.client.perform_query('GET', '/unity-catalog/schemas', data=_data, + headers=headers) + + def list_lineages_by_table(self, table_name=None, headers=None): + """ + List table lineage by table name + """ + _data = {} + if table_name is not None: + _data['table_name'] = table_name + + return self.client.perform_query('GET', '/lineage-tracking/table-lineage/get', data=_data, + headers=headers) + + def list_lineages_by_column(self, table_name=None, column_name=None, headers=None): + """ + List column lineage by table name and comlumn name + """ + _data = {} + if table_name is not None: + _data['table_name'] = table_name + if column_name is not None: + _data['column_name'] = column_name + + return self.client.perform_query('GET', '/lineage-tracking/column-lineage/get', data=_data, + headers=headers) + + def get_schema(self, full_name, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/schemas/%s' % (full_name), + data=_data, headers=headers) + + def update_schema(self, full_name, schema_spec, headers=None): + return self.client.perform_query('PATCH', '/unity-catalog/schemas/%s' % (full_name), + data=schema_spec, headers=headers) + + def delete_schema(self, full_name, headers=None): + _data = {} + return self.client.perform_query('DELETE', '/unity-catalog/schemas/%s' % (full_name), + data=_data, headers=headers) + + # Table Operations + + def create_table(self, table_spec, headers=None): + return self.client.perform_query('POST', '/unity-catalog/tables', data=table_spec, + headers=headers) + + def list_tables(self, catalog_name, schema_name=None, name_regex=None, headers=None): + _data = { + 'catalog_name': catalog_name + } + if schema_name is not None: + _data['schema_name'] = schema_name + if name_regex is not None: + _data['table_name_regex'] = name_regex + + return self.client.perform_query('GET', '/unity-catalog/tables', data=_data, + headers=headers) + + def list_table_summaries(self, catalog_name, headers=None): + _data = { + 'catalog_name': catalog_name + } + return self.client.perform_query('GET', '/unity-catalog/table-summaries', data=_data, + headers=headers) + + def get_table(self, full_name, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/tables/%s' % (full_name), + data=_data, headers=headers) + + def update_table(self, full_name, table_spec, headers=None): + return self.client.perform_query('PATCH', '/unity-catalog/tables/%s' % (full_name), + data=table_spec, headers=headers) + + def delete_table(self, full_name, headers=None): + _data = {} + return self.client.perform_query('DELETE', '/unity-catalog/tables/%s' % (full_name), + data=_data, headers=headers) + + # Share Operations + + def create_share(self, name, headers=None): + _data = { + 'name': name + } + return self.client.perform_query('POST', '/unity-catalog/shares', data=_data, + headers=headers) + + def list_shares(self, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/shares', data=_data, + headers=headers) + + def get_share(self, name, include_shared_data, headers=None): + _data = {'include_shared_data': include_shared_data} + + return self.client.perform_query('GET', '/unity-catalog/shares/%s' % (name), + data=_data, headers=headers) + + def update_share(self, name, share_spec, headers=None): + return self.client.perform_query('PATCH', '/unity-catalog/shares/%s' % (name), + data=share_spec, headers=headers) + + def delete_share(self, name, headers=None): + _data = {} + return self.client.perform_query('DELETE', '/unity-catalog/shares/%s' % (name), + data=_data, headers=headers) + + def list_share_permissions(self, name, headers=None): + _data = {} + return self.client.perform_query('GET', '/unity-catalog/shares/%s/permissions' % (name), + data=_data, headers=headers) + + def update_share_permissions(self, name, perm_spec, headers=None): + return self.client.perform_query('PATCH', '/unity-catalog/shares/%s/permissions' % (name), + data=perm_spec, headers=headers) + + # Recipient Operations + + def create_recipient(self, name, comment=None, sharing_code=None, + allowed_ip_addresses=None, headers=None): + _data = { + 'name': name, + } + if comment is not None: + _data['comment'] = comment + if sharing_code is not None: + _data['sharing_code'] = sharing_code + _data['authentication_type'] = 'DATABRICKS' + else: + _data['authentication_type'] = 'TOKEN' + if allowed_ip_addresses is not None: + _data['ip_access_list'] = { + 'allowed_ip_addresses': allowed_ip_addresses, + } + + return self.client.perform_query('POST', '/unity-catalog/recipients', data=_data, + headers=headers) + + def list_recipients(self, headers=None): + _data = {} + + return self.client.perform_query('GET', '/unity-catalog/recipients', data=_data, + headers=headers) + + def get_recipient(self, name, headers=None): + _data = {} + + return self.client.perform_query('GET', '/unity-catalog/recipients/%s' % (name), + data=_data, headers=headers) + + def update_recipient(self, name, recipient_spec, headers=None): + return self.client.perform_query('PATCH', '/unity-catalog/recipients/%s' % (name), + data=recipient_spec, headers=headers) + + def rotate_recipient_token(self, name, existing_token_expire_in_seconds=None, headers=None): + _data = { + 'name': name, + } + if existing_token_expire_in_seconds is not None: + _data['existing_token_expire_in_seconds'] = existing_token_expire_in_seconds + return self.client.perform_query('POST', '/unity-catalog/recipients/%s/rotate-token' % + (name), data=_data, headers=headers) + + def get_recipient_share_permissions(self, name, headers=None): + _data = {} + + return self.client.perform_query('GET', '/unity-catalog/recipients/%s/share-permissions' + % (name), data=_data, headers=headers) + + def delete_recipient(self, name, headers=None): + _data = {} + + return self.client.perform_query('DELETE', '/unity-catalog/recipients/%s' % (name), + data=_data, headers=headers) + + # Provider Operations + + def create_provider(self, name, comment, recipient_profile=None, headers=None): + _data = { + 'name': name, + } + if comment is not None: + _data['comment'] = comment + if recipient_profile is not None: + _data['recipient_profile_str'] = mc_pretty_format(recipient_profile) + _data['authentication_type'] = 'TOKEN' + else: + _data['authentication_type'] = 'DATABRICKS' + return self.client.perform_query('POST', '/unity-catalog/providers/', + data=_data, headers=headers) + + def list_providers(self, headers=None): + return self.client.perform_query('GET', '/unity-catalog/providers', data={}, + headers=headers) + + def get_provider(self, name, headers=None): + return self.client.perform_query('GET', '/unity-catalog/providers/%s' % (name), + data={}, headers=headers) + + def update_provider(self, name, new_name, comment, recipient_profile, headers=None): + _data = {} + if new_name is not None: + _data['name'] = new_name + if recipient_profile is not None: + _data['recipient_profile_str'] = mc_pretty_format(recipient_profile) + if comment is not None: + _data['comment'] = comment + + return self.client.perform_query('PATCH', '/unity-catalog/providers/%s' % (name), + data=_data, headers=headers) + + def delete_provider(self, name, headers=None): + return self.client.perform_query('DELETE', '/unity-catalog/providers/%s' % (name), + data={}, headers=headers) + + def list_provider_shares(self, name, headers=None): + return self.client.perform_query('GET', '/unity-catalog/providers/%s/shares' % (name), + data={}, headers=headers) + + # Permissions Operations + + def _permissions_url(self, sec_type, sec_name): + return '/unity-catalog/permissions/%s/%s' % (sec_type, sec_name) + + def get_permissions(self, sec_type, sec_name, headers=None): + _data = {} + return self.client.perform_query('GET', self._permissions_url(sec_type, sec_name), + data=_data, headers=headers) + + def update_permissions(self, sec_type, sec_name, perm_diff_spec, headers=None): + _data = perm_diff_spec + return self.client.perform_query('PATCH', self._permissions_url(sec_type, sec_name), + data=_data, headers=headers) + + def replace_permissions(self, sec_type, sec_name, perm_spec, headers=None): + _data = perm_spec + return self.client.perform_query('PUT', self._permissions_url(sec_type, sec_name), + data=_data, headers=headers) diff --git a/databricks_cli/unity_catalog/utils.py b/databricks_cli/unity_catalog/utils.py new file mode 100644 index 00000000..e57f3e41 --- /dev/null +++ b/databricks_cli/unity_catalog/utils.py @@ -0,0 +1,65 @@ +# Databricks CLI +# Copyright 2022 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy + +from databricks_cli.utils import pretty_format + + +# Encode UTF-8 strings in JSON blobs +def mc_pretty_format(json): + return pretty_format(json, encode_utf8=True) + + +def del_none(json): + to_delete = [] + to_recurse = [] + for key, value in json.items(): + if value is None: + to_delete.append(key) + if value is dict: + to_recurse.append(key) + for key in to_delete: + del json[key] + for key in to_recurse: + del_none(json[key]) + + +def hide(cmd): + """ + Return a copy of specified Click command instance with `hidden = True`. + This requires Click >= v7.0. + """ + cmd_copy = copy.copy(cmd) + cmd_copy.hidden = True + return cmd_copy + + +def json_file_help(method, path): + path = "/api/2.0/unity-catalog" + path + return "File containing JSON request to {} to {}.".format(method, path) + + +def json_string_help(method, path): + path = "/api/2.0/unity-catalog" + path + return "JSON string to {} to {}.".format(method, path) diff --git a/databricks_cli/utils.py b/databricks_cli/utils.py index 3e02911c..0d4d5fe9 100644 --- a/databricks_cli/utils.py +++ b/databricks_cli/utils.py @@ -98,23 +98,27 @@ def error_and_quit(message): sys.exit(1) -def pretty_format(json): +def pretty_format(json, encode_utf8=False): + if encode_utf8: + return json_dumps(json, indent=2, ensure_ascii=False) return json_dumps(json, indent=2) -def json_cli_base(json_file, json, api, print_response=True): +def json_cli_base(json_file, json, api, error_msg='', print_response=True, encode_utf8=False): """ Takes json_file or json string and calls an function "api" with the json deserialized """ if not (json_file is None) ^ (json is None): - raise RuntimeError('Either --json-file or --json should be provided') + if not error_msg: + error_msg = 'Either --json-file or --json should be provided' + raise RuntimeError(error_msg) if json_file: with open(json_file, 'r') as f: json = f.read() res = api(json_loads(json)) if print_response: - click.echo(pretty_format(res)) + click.echo(pretty_format(res, encode_utf8)) def truncate_string(s, length=100): @@ -123,6 +127,12 @@ def truncate_string(s, length=100): return s[:length] + '...' +def to_graph(dict_obj, graph_name="Graph"): + entries = ["\"{}\" -> {};".format(key, ','.join( + '"{}"'.format(item) for item in dict_obj.get(key))) for key in dict_obj] + return "digraph \"{}\" ".format(graph_name) + "{\n\t" + "\n\t".join(entries) + "\n}" + + class InvalidConfigurationError(RuntimeError): @staticmethod def for_profile(profile): diff --git a/examples/unity_catalog/create-cred1.json b/examples/unity_catalog/create-cred1.json new file mode 100644 index 00000000..6cb413bb --- /dev/null +++ b/examples/unity_catalog/create-cred1.json @@ -0,0 +1,6 @@ +{ + "name": "acain_test4", + "aws_iam_role": { + "role_arn": "arn:aws:iam::1234567890:role/MyRole-AJJHDSKSDF" + } +} diff --git a/examples/unity_catalog/create-loc1.json b/examples/unity_catalog/create-loc1.json new file mode 100644 index 00000000..532424ca --- /dev/null +++ b/examples/unity_catalog/create-loc1.json @@ -0,0 +1,6 @@ +{ + "name": "myloc1", + "credential_name": "mycred", + "comment": "hey hey", + "url": "s3://my.bucket" +} diff --git a/examples/unity_catalog/replace-permissions1.json b/examples/unity_catalog/replace-permissions1.json new file mode 100644 index 00000000..4268031f --- /dev/null +++ b/examples/unity_catalog/replace-permissions1.json @@ -0,0 +1,18 @@ +{ + "privilege_assignments": [ + { + "privileges": [ + "SELECT", + "USAGE" + ], + "principal": "foo@bar.com" + }, + { + "privileges": [ + "OWN", + "MODIFY" + ], + "principal": "some_group" + } + ] +} diff --git a/examples/unity_catalog/update-catalog.json b/examples/unity_catalog/update-catalog.json new file mode 100644 index 00000000..ede80465 --- /dev/null +++ b/examples/unity_catalog/update-catalog.json @@ -0,0 +1,4 @@ +{ + "name": "my_catalog", + "comment": "updated comment" +} diff --git a/examples/unity_catalog/update-metastore.json b/examples/unity_catalog/update-metastore.json new file mode 100644 index 00000000..01878525 --- /dev/null +++ b/examples/unity_catalog/update-metastore.json @@ -0,0 +1,4 @@ +{ + "name": "my_metastore", + "storage_root": "s3://my.other.bucket" +} diff --git a/examples/unity_catalog/update-permissions1.json b/examples/unity_catalog/update-permissions1.json new file mode 100644 index 00000000..35b04201 --- /dev/null +++ b/examples/unity_catalog/update-permissions1.json @@ -0,0 +1,17 @@ +{ + "changes": [ + { + "principal": "adam.cain@databricks.com", + "add": ["SELECT"], + "remove": ["MODIFY"] + }, + { + "principal": "eng-data-security", + "remove": ["CREATE"] + }, + { + "principal": "users", + "add": ["USAGE"] + } + ] +} diff --git a/examples/unity_catalog/update-schema.json b/examples/unity_catalog/update-schema.json new file mode 100644 index 00000000..f43cf565 --- /dev/null +++ b/examples/unity_catalog/update-schema.json @@ -0,0 +1,4 @@ +{ + "name": "my_schema", + "comment": "updated comment" +} diff --git a/examples/unity_catalog/update-table.json b/examples/unity_catalog/update-table.json new file mode 100644 index 00000000..caa2b82d --- /dev/null +++ b/examples/unity_catalog/update-table.json @@ -0,0 +1,4 @@ +{ + "name": "my_table", + "comment": "updated comment" +} diff --git a/requirements.txt b/requirements.txt index 0a869e46..31c84c27 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ # Note: please keep this in sync with `setup.py`. -click>=6.7 +click>=7.0 pyjwt>=1.7.0 oauthlib>=3.1.0 requests>=2.17.3 diff --git a/setup.py b/setup.py index f4e9fab1..16902db2 100644 --- a/setup.py +++ b/setup.py @@ -35,7 +35,7 @@ packages=find_packages(exclude=['tests', 'tests.*']), install_requires=[ # Note: please keep this in sync with `requirements.txt`. - 'click>=6.7', + 'click>=7.0', 'pyjwt>=1.7.0', 'oauthlib>=3.1.0', 'requests>=2.17.3',