From 2640c53e962afd13e55abb857de73afd2d3f419b Mon Sep 17 00:00:00 2001 From: Andy Date: Fri, 3 Mar 2023 20:14:03 -0800 Subject: [PATCH 1/8] Add add-schema, update-schema, and remove-schema to Unity Catalog CLI --- .../unity_catalog/delta_sharing_cli.py | 217 +++++++++++++++--- tests/unity_catalog/test_delta_sharing_cli.py | 78 +++++++ 2 files changed, 265 insertions(+), 30 deletions(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index 988c5fbb..ad994d86 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -122,27 +122,6 @@ def update_share_permissions_cli(api_client, name, json_file, json): json_cli_base(json_file, json, lambda json: UnityCatalogApi(api_client).update_share_permissions(name, json)) - -def shared_data_object(name=None, comment=None, shared_as=None, - cdf_enabled=None, partitions=None, start_version=None): - val = { - 'data_object_type': 'TABLE' - } - if name is not None: - val['name'] = name - if comment is not None: - val['comment'] = comment - if shared_as is not None: - val['shared_as'] = shared_as - if cdf_enabled is not None: - val['cdf_enabled'] = cdf_enabled - if partitions is not None: - val['partitions'] = partitions - if start_version is not None: - val['start_version'] = start_version - return val - - @click.command(context_settings=CONTEXT_SETTINGS, short_help='Update a share.') @click.option('--name', required=True, @@ -180,9 +159,9 @@ def update_share_cli(api_client, name, new_name, comment, owner, raise ValueError('Cannot specify JSON if any other update flags are specified') updates = [] for a in add_table: - updates.append({'action': 'ADD', 'data_object': shared_data_object(a)}) + updates.append({'action': 'ADD', 'data_object': shared_table_object(a)}) for r in remove_table: - updates.append({'action': 'REMOVE', 'data_object': shared_data_object(r)}) + updates.append({'action': 'REMOVE', 'data_object': shared_table_object(r)}) data = {'name': new_name, 'comment': comment, 'owner': owner, 'updates': updates} share_json = UnityCatalogApi(api_client).update_share(name, data) click.echo(mc_pretty_format(share_json)) @@ -190,8 +169,183 @@ def update_share_cli(api_client, name, new_name, comment, owner, json_cli_base(json_file, json, lambda json: UnityCatalogApi(api_client).update_share(name, json)) +def shared_schema_object(name=None, comment=None): + val = { + 'data_object_type': 'SCHEMA' + } + if name is not None: + val['name'] = name + if comment is not None: + val['comment'] = comment + return val + +def create_common_shared_schema_options(f): + @click.option('--schema', required=True, + help='Full name of the shared schema.') + @click.option('--comment', default=None, + help='New comment of the shared schema.') + @functools.wraps(f) + def wrapper(*args, **kwargs): + f(*args, **kwargs) + return wrapper + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Add a shared schema.') +@click.option('--share', required=True, + help='Name of the share to update.') +@create_common_shared_schema_options +@click.option('--json-file', default=None, type=click.Path(), + help="Adds a shared table based on shared data object represented in JSON file.") +@click.option('--json', default=None, type=JsonClickType(), + help="Adds a shared table based on shared data object represented in JSON.") +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def add_share_schema_cli(api_client, share, schema, comment, json_file, json): + """ + Adds a shared schema. + + The public specification for the JSON request is in development. + """ + if comment is not None: + if (json_file is not None) or (json is not None): + raise ValueError('Cannot specify JSON if any other flags are specified') + data = { + 'updates': [ + { + 'action': 'ADD', + 'data_object': shared_schema_object( + name=schema, + comment=comment + ) + } + ] + } + share_json = UnityCatalogApi(api_client).update_share(share, data) + click.echo(mc_pretty_format(share_json)) + else: + json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'ADD', + 'data_object': d, + } + ] + })) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a shared table.') +@click.option('--share', required=True, + help='Name of the share to update.') +@create_common_shared_schema_options +@click.option('--json-file', default=None, type=click.Path(), + help="Updates the shared table to shared data object represented in JSON file.") +@click.option('--json', default=None, type=JsonClickType(), + help="Updates the shared table to shared data object represented in JSON.") +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def update_share_schema_cli(api_client, share, schema, comment, json_file, json): + """ + Updates a shared schema. + + The public specification for the JSON request is in development. + """ + if comment is not None: + if (json_file is not None) or (json is not None): + raise ValueError('Cannot specify JSON if any other flags are specified') + data = { + 'updates': [ + { + 'action': 'UPDATE', + 'data_object': shared_schema_object( + name=schema, + comment=comment + ) + } + ] + } + share_json = UnityCatalogApi(api_client).update_share(share, data) + click.echo(mc_pretty_format(share_json)) + else: + json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'UPDATE', + 'data_object': d, + } + ] + })) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help='Update a shared table.') +@click.option('--share', required=True, + help='Name of the share to update.') +@click.option('--schema', default=None, + help='Full name of the schema to remove from share.') +@click.option('--json-file', default=None, type=click.Path(), + help="Removes the shared table based on shared data object represented in JSON file.") +@click.option('--json', default=None, type=JsonClickType(), + help="Removes the shared table based on shared data object represented in JSON.") +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def remove_share_schema_cli(api_client, share, schema, json_file, json): + """ + Removes a shared schema either by full schema name. + + The public specification for the JSON request is in development. + """ + if schema is not None: + if (json_file is not None) or (json is not None): + raise ValueError('Cannot specify JSON if any other flags are specified') + data = { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': shared_schema_object( + name=schema, + ) + } + ] + } + share_json = UnityCatalogApi(api_client).update_share(share, data) + click.echo(mc_pretty_format(share_json)) + else: + json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': d, + } + ] + })) + +def shared_table_object(name=None, comment=None, shared_as=None, + cdf_enabled=None, partitions=None, start_version=None): + val = { + 'data_object_type': 'TABLE' + } + if name is not None: + val['name'] = name + if comment is not None: + val['comment'] = comment + if shared_as is not None: + val['shared_as'] = shared_as + if cdf_enabled is not None: + val['cdf_enabled'] = cdf_enabled + if partitions is not None: + val['partitions'] = partitions + if start_version is not None: + val['start_version'] = start_version + return val -def create_common_shared_data_object_options(f): +def create_common_shared_table_options(f): @click.option('--table', required=True, help='Full name of the shared table.') @click.option('--shared-as', default=None, @@ -214,7 +368,7 @@ def wrapper(*args, **kwargs): short_help='Add a shared table.') @click.option('--share', required=True, help='Name of the share to update.') -@create_common_shared_data_object_options +@create_common_shared_table_options @click.option('--json-file', default=None, type=click.Path(), help="Adds a shared table based on shared data object represented in JSON file.") @click.option('--json', default=None, type=JsonClickType(), @@ -238,7 +392,7 @@ def add_share_table_cli(api_client, share, table, shared_as, comment, 'updates': [ { 'action': 'ADD', - 'data_object': shared_data_object( + 'data_object': shared_table_object( name=table, shared_as=shared_as, comment=comment, @@ -266,7 +420,7 @@ def add_share_table_cli(api_client, share, table, shared_as, comment, short_help='Update a shared table.') @click.option('--share', required=True, help='Name of the share to update.') -@create_common_shared_data_object_options +@create_common_shared_table_options @click.option('--json-file', default=None, type=click.Path(), help="Updates the shared table to shared data object represented in JSON file.") @click.option('--json', default=None, type=JsonClickType(), @@ -290,7 +444,7 @@ def update_share_table_cli(api_client, share, table, shared_as, comment, 'updates': [ { 'action': 'UPDATE', - 'data_object': shared_data_object( + 'data_object': shared_table_object( name=table, shared_as=shared_as, comment=comment, @@ -319,7 +473,7 @@ def update_share_table_cli(api_client, share, table, shared_as, comment, @click.option('--share', required=True, help='Name of the share to update.') @click.option('--table', default=None, - help='Full name of the table to update from share.') + help='Full name of the table to remove from share.') @click.option('--shared-as', default=None, help='New name of the table inside the share.') @click.option('--json-file', default=None, type=click.Path(), @@ -346,7 +500,7 @@ def remove_share_table_cli(api_client, share, table, shared_as, json_file, json) 'updates': [ { 'action': 'REMOVE', - 'data_object': shared_data_object( + 'data_object': shared_table_object( name=table, shared_as=shared_as, ) @@ -722,6 +876,9 @@ def register_shares_commands(cmd_group): 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(add_share_schema_cli, name='add-schema') + shares_group.add_command(update_share_schema_cli, name='update-schema') + shares_group.add_command(remove_share_schema_cli, name='remove-schema') shares_group.add_command(add_share_table_cli, name='add-table') shares_group.add_command(update_share_table_cli, name='update-table') shares_group.add_command(remove_share_table_cli, name='remove-table') diff --git a/tests/unity_catalog/test_delta_sharing_cli.py b/tests/unity_catalog/test_delta_sharing_cli.py index 7e0dbb5b..4ebc7d3b 100644 --- a/tests/unity_catalog/test_delta_sharing_cli.py +++ b/tests/unity_catalog/test_delta_sharing_cli.py @@ -214,6 +214,84 @@ def test_update_share_cli(api_mock, echo_mock): api_mock.update_share.assert_called_once_with(SHARE_NAME, expected_data) echo_mock.assert_called_once_with(mc_pretty_format(SHARE)) +@provide_conf +def test_add_share_schema_cli(api_mock, echo_mock): + api_mock.update_share.return_value = SHARE + runner = CliRunner() + runner.invoke( + delta_sharing_cli.add_share_schema_cli, + args=[ + '--share', SHARE_NAME, + '--schema', 'catalog.schema', + '--comment', 'add.comment', + ]) + expected_data = { + 'updates': [ + { + 'action': 'ADD', + 'data_object': { + 'data_object_type': 'SCHEMA', + 'name': 'catalog.schema', + 'comment': 'add.comment', + } + } + ] + } + api_mock.update_share.assert_called_once_with(SHARE_NAME, expected_data) + echo_mock.assert_called_once_with(mc_pretty_format(SHARE)) + + +@provide_conf +def test_update_share_schema_cli(api_mock, echo_mock): + api_mock.update_share.return_value = SHARE + runner = CliRunner() + runner.invoke( + delta_sharing_cli.update_share_schema_cli, + args=[ + '--share', SHARE_NAME, + '--schema', 'catalog.schema', + '--comment', 'update.comment', + ]) + expected_data = { + 'updates': [ + { + 'action': 'UPDATE', + 'data_object': { + 'data_object_type': 'SCHEMA', + 'name': 'catalog.schema', + 'comment': 'update.comment', + } + } + ] + } + api_mock.update_share.assert_called_once_with(SHARE_NAME, expected_data) + echo_mock.assert_called_once_with(mc_pretty_format(SHARE)) + + +@provide_conf +def test_remove_share_schema_cli(api_mock, echo_mock): + api_mock.update_share.return_value = SHARE + runner = CliRunner() + runner.invoke( + delta_sharing_cli.remove_share_schema_cli, + args=[ + '--share', SHARE_NAME, + '--schema', 'catalog.schema', + ]) + expected_data = { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': { + 'data_object_type': 'SCHEMA', + 'name': 'catalog.schema', + } + } + ] + } + api_mock.update_share.assert_called_once_with(SHARE_NAME, expected_data) + echo_mock.assert_called_once_with(mc_pretty_format(SHARE)) + @provide_conf def test_add_share_table_cli(api_mock, echo_mock): From e944d25094a131f0788971b8ae0ee7de473edecd Mon Sep 17 00:00:00 2001 From: Andy Date: Tue, 7 Mar 2023 10:47:57 -0800 Subject: [PATCH 2/8] Fix parameter description + add validation for --json/--json-file for CRUD operations for schema/table --- .../unity_catalog/delta_sharing_cli.py | 135 +++++++++++------- 1 file changed, 81 insertions(+), 54 deletions(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index ad994d86..f1f58d83 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -195,9 +195,9 @@ def wrapper(*args, **kwargs): help='Name of the share to update.') @create_common_shared_schema_options @click.option('--json-file', default=None, type=click.Path(), - help="Adds a shared table based on shared data object represented in JSON file.") + help="Adds a shared schema based on shared data object represented in JSON file.") @click.option('--json', default=None, type=JsonClickType(), - help="Adds a shared table based on shared data object represented in JSON.") + help="Adds a shared schema based on shared data object represented in JSON.") @debug_option @profile_option @eat_exceptions @@ -225,14 +225,18 @@ def add_share_schema_cli(api_client, share, schema, comment, json_file, json): share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) else: - json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'ADD', - 'data_object': d, - } - ] - })) + def api_call(d): + if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": + raise ValueError('Must specify data_object_type as "SCHEMA"') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'ADD', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -241,9 +245,9 @@ def add_share_schema_cli(api_client, share, schema, comment, json_file, json): help='Name of the share to update.') @create_common_shared_schema_options @click.option('--json-file', default=None, type=click.Path(), - help="Updates the shared table to shared data object represented in JSON file.") + help="Updates the shared schema to shared data object represented in JSON file.") @click.option('--json', default=None, type=JsonClickType(), - help="Updates the shared table to shared data object represented in JSON.") + help="Updates the shared schema to shared data object represented in JSON.") @debug_option @profile_option @eat_exceptions @@ -271,14 +275,18 @@ def update_share_schema_cli(api_client, share, schema, comment, json_file, json) share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) else: - json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'UPDATE', - 'data_object': d, - } - ] - })) + def api_call(d): + if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": + raise ValueError('Must specify data_object_type as "SCHEMA"') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'UPDATE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -288,9 +296,9 @@ def update_share_schema_cli(api_client, share, schema, comment, json_file, json) @click.option('--schema', default=None, help='Full name of the schema to remove from share.') @click.option('--json-file', default=None, type=click.Path(), - help="Removes the shared table based on shared data object represented in JSON file.") + help="Removes the shared schema based on shared data object represented in JSON file.") @click.option('--json', default=None, type=JsonClickType(), - help="Removes the shared table based on shared data object represented in JSON.") + help="Removes the shared schema based on shared data object represented in JSON.") @debug_option @profile_option @eat_exceptions @@ -317,14 +325,18 @@ def remove_share_schema_cli(api_client, share, schema, json_file, json): share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) else: - json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'REMOVE', - 'data_object': d, - } - ] - })) + def api_call(d): + if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": + raise ValueError('Must specify data_object_type as "SCHEMA"') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) def shared_table_object(name=None, comment=None, shared_as=None, cdf_enabled=None, partitions=None, start_version=None): @@ -406,14 +418,19 @@ def add_share_table_cli(api_client, share, table, shared_as, comment, share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) else: - json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'ADD', - 'data_object': d, - } - ] - })) + def api_call(d): + if 'data_object_type' in d and d['data_object_type'] != "TABLE": + raise ValueError('Must specify data_object_type a "TABLE" ' + 'or not specify data_object_type at all') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -458,14 +475,19 @@ def update_share_table_cli(api_client, share, table, shared_as, comment, share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) else: - json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'UPDATE', - 'data_object': d, - } - ] - })) + def api_call(d): + if 'data_object_type' in d and d['data_object_type'] != "TABLE": + raise ValueError('Must specify data_object_type a "TABLE" ' + 'or not specify data_object_type at all') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'UPDATE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -510,14 +532,19 @@ def remove_share_table_cli(api_client, share, table, shared_as, json_file, json) share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) else: - json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'REMOVE', - 'data_object': d, - } - ] - })) + def api_call(d): + if 'data_object_type' in d and d['data_object_type'] != "TABLE": + raise ValueError('Must specify data_object_type a "TABLE" ' + 'or not specify data_object_type at all') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, From dcf390c92bd37d698dbfeb0467b466987a13bf40 Mon Sep 17 00:00:00 2001 From: Andy Date: Tue, 7 Mar 2023 11:36:05 -0800 Subject: [PATCH 3/8] Fix error message --- databricks_cli/unity_catalog/delta_sharing_cli.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index f1f58d83..0a4e72a4 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -420,7 +420,7 @@ def add_share_table_cli(api_client, share, table, shared_as, comment, else: def api_call(d): if 'data_object_type' in d and d['data_object_type'] != "TABLE": - raise ValueError('Must specify data_object_type a "TABLE" ' + raise ValueError('Must specify data_object_type as "TABLE" ' 'or not specify data_object_type at all') UnityCatalogApi(api_client).update_share(share, { 'updates': [ @@ -477,7 +477,7 @@ def update_share_table_cli(api_client, share, table, shared_as, comment, else: def api_call(d): if 'data_object_type' in d and d['data_object_type'] != "TABLE": - raise ValueError('Must specify data_object_type a "TABLE" ' + raise ValueError('Must specify data_object_type as "TABLE" ' 'or not specify data_object_type at all') UnityCatalogApi(api_client).update_share(share, { 'updates': [ @@ -534,7 +534,7 @@ def remove_share_table_cli(api_client, share, table, shared_as, json_file, json) else: def api_call(d): if 'data_object_type' in d and d['data_object_type'] != "TABLE": - raise ValueError('Must specify data_object_type a "TABLE" ' + raise ValueError('Must specify data_object_type as "TABLE" ' 'or not specify data_object_type at all') UnityCatalogApi(api_client).update_share(share, { 'updates': [ From f11b382c77c6906b2c25624532cb79560b1f06a0 Mon Sep 17 00:00:00 2001 From: Andy Date: Fri, 31 Mar 2023 08:22:09 -0700 Subject: [PATCH 4/8] Update comment --- databricks_cli/unity_catalog/delta_sharing_cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index 0a4e72a4..7f710f3d 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -240,7 +240,7 @@ def api_call(d): @click.command(context_settings=CONTEXT_SETTINGS, - short_help='Update a shared table.') + short_help='Update a shared schema.') @click.option('--share', required=True, help='Name of the share to update.') @create_common_shared_schema_options @@ -290,7 +290,7 @@ def api_call(d): @click.command(context_settings=CONTEXT_SETTINGS, - short_help='Update a shared table.') + short_help='Update a shared schema.') @click.option('--share', required=True, help='Name of the share to update.') @click.option('--schema', default=None, From 10a8625658896b9ae51f34f7fe128fb1447a9401 Mon Sep 17 00:00:00 2001 From: Andy Date: Fri, 31 Mar 2023 08:24:22 -0700 Subject: [PATCH 5/8] Update comment x2 --- databricks_cli/unity_catalog/delta_sharing_cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index 7f710f3d..e29bd7d3 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -290,7 +290,7 @@ def api_call(d): @click.command(context_settings=CONTEXT_SETTINGS, - short_help='Update a shared schema.') + short_help='Remove a shared schema.') @click.option('--share', required=True, help='Name of the share to update.') @click.option('--schema', default=None, @@ -491,7 +491,7 @@ def api_call(d): @click.command(context_settings=CONTEXT_SETTINGS, - short_help='Update a shared table.') + short_help='Remove a shared table.') @click.option('--share', required=True, help='Name of the share to update.') @click.option('--table', default=None, From e1a9b4791aa002858011af54ce9417af987aff27 Mon Sep 17 00:00:00 2001 From: Andy Date: Fri, 31 Mar 2023 15:13:07 -0700 Subject: [PATCH 6/8] Fix lint --- databricks_cli/unity_catalog/delta_sharing_cli.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index e29bd7d3..d6c0fed9 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -296,7 +296,8 @@ def api_call(d): @click.option('--schema', default=None, help='Full name of the schema to remove from share.') @click.option('--json-file', default=None, type=click.Path(), - help="Removes the shared schema based on shared data object represented in JSON file.") + help="Removes the shared schema based on shared data object" + + "represented in JSON file.") @click.option('--json', default=None, type=JsonClickType(), help="Removes the shared schema based on shared data object represented in JSON.") @debug_option From 2cbb1154cd2a92a583021eeb4a1fd0df74f89bca Mon Sep 17 00:00:00 2001 From: Andy Date: Mon, 3 Apr 2023 21:53:37 -0700 Subject: [PATCH 7/8] CLI fixes --- .../unity_catalog/delta_sharing_cli.py | 205 +++++++++--------- tests/unity_catalog/test_delta_sharing_cli.py | 27 ++- 2 files changed, 128 insertions(+), 104 deletions(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index d6c0fed9..f1e5b239 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -180,7 +180,7 @@ def shared_schema_object(name=None, comment=None): return val def create_common_shared_schema_options(f): - @click.option('--schema', required=True, + @click.option('--schema', default=None, help='Full name of the shared schema.') @click.option('--comment', default=None, help='New comment of the shared schema.') @@ -208,9 +208,22 @@ def add_share_schema_cli(api_client, share, schema, comment, json_file, json): The public specification for the JSON request is in development. """ - if comment is not None: - if (json_file is not None) or (json is not None): - raise ValueError('Cannot specify JSON if any other flags are specified') + if (json_file is not None) or (json is not None): + def api_call(d): + if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": + raise ValueError('Must specify data_object_type as "SCHEMA"') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'ADD', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) + else: + if schema is None: + raise ValueError("Must specify full schema name when adding shared schema") data = { 'updates': [ { @@ -224,19 +237,6 @@ def add_share_schema_cli(api_client, share, schema, comment, json_file, json): } share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) - else: - def api_call(d): - if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": - raise ValueError('Must specify data_object_type as "SCHEMA"') - UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'ADD', - 'data_object': d, - } - ] - }) - json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -258,9 +258,22 @@ def update_share_schema_cli(api_client, share, schema, comment, json_file, json) The public specification for the JSON request is in development. """ - if comment is not None: - if (json_file is not None) or (json is not None): - raise ValueError('Cannot specify JSON if any other flags are specified') + if (json_file is not None) or (json is not None): + def api_call(d): + if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": + raise ValueError('Must specify data_object_type as "SCHEMA"') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'UPDATE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) + else: + if schema is None: + raise ValueError("Must specify full schema name when updating shared schema") data = { 'updates': [ { @@ -274,19 +287,6 @@ def update_share_schema_cli(api_client, share, schema, comment, json_file, json) } share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) - else: - def api_call(d): - if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": - raise ValueError('Must specify data_object_type as "SCHEMA"') - UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'UPDATE', - 'data_object': d, - } - ] - }) - json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -310,22 +310,7 @@ def remove_share_schema_cli(api_client, share, schema, json_file, json): The public specification for the JSON request is in development. """ - if schema is not None: - if (json_file is not None) or (json is not None): - raise ValueError('Cannot specify JSON if any other flags are specified') - data = { - 'updates': [ - { - 'action': 'REMOVE', - 'data_object': shared_schema_object( - name=schema, - ) - } - ] - } - share_json = UnityCatalogApi(api_client).update_share(share, data) - click.echo(mc_pretty_format(share_json)) - else: + if (json_file is not None) or (json is not None): def api_call(d): if 'data_object_type' not in d or d['data_object_type'] != "SCHEMA": raise ValueError('Must specify data_object_type as "SCHEMA"') @@ -338,6 +323,22 @@ def api_call(d): ] }) json_cli_base(json_file, json, api_call) + else: + if schema is None: + raise ValueError("Must specify full schema name when removing shared schema") + data = { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': shared_schema_object( + name=schema, + ) + } + ] + } + share_json = UnityCatalogApi(api_client).update_share(share, data) + click.echo(mc_pretty_format(share_json)) + def shared_table_object(name=None, comment=None, shared_as=None, cdf_enabled=None, partitions=None, start_version=None): @@ -359,7 +360,7 @@ def shared_table_object(name=None, comment=None, shared_as=None, return val def create_common_shared_table_options(f): - @click.option('--table', required=True, + @click.option('--table', default=None, help='Full name of the shared table.') @click.option('--shared-as', default=None, help='New name of the table to be shared as.') @@ -397,10 +398,23 @@ def add_share_table_cli(api_client, share, table, shared_as, comment, The public specification for the JSON request is in development. """ - if ((shared_as is not None) or (comment is not None) or (partitions is not None) or - (cdf is not None) or (start_version is not None)): - if (json_file is not None) or (json is not None): - raise ValueError('Cannot specify JSON if any other flags are specified') + if (json_file is not None) or (json is not None): + def api_call(d): + if 'data_object_type' in d and d['data_object_type'] != "TABLE": + raise ValueError('Must specify data_object_type as "TABLE" ' + 'or not specify data_object_type at all') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) + else: + if table is None: + raise ValueError('Must specify table name when adding shared table') data = { 'updates': [ { @@ -418,20 +432,6 @@ def add_share_table_cli(api_client, share, table, shared_as, comment, } share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) - else: - def api_call(d): - if 'data_object_type' in d and d['data_object_type'] != "TABLE": - raise ValueError('Must specify data_object_type as "TABLE" ' - 'or not specify data_object_type at all') - UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'REMOVE', - 'data_object': d, - } - ] - }) - json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -454,10 +454,23 @@ def update_share_table_cli(api_client, share, table, shared_as, comment, The public specification for the JSON request is in development. """ - if ((shared_as is not None) or (comment is not None) or (partitions is not None) or - (cdf is not None) or (start_version is not None)): - if (json_file is not None) or (json is not None): - raise ValueError('Cannot specify JSON if any other flags are specified') + if (json_file is not None) or (json is not None): + def api_call(d): + if 'data_object_type' in d and d['data_object_type'] != "TABLE": + raise ValueError('Must specify data_object_type as "TABLE" ' + 'or not specify data_object_type at all') + UnityCatalogApi(api_client).update_share(share, { + 'updates': [ + { + 'action': 'UPDATE', + 'data_object': d, + } + ] + }) + json_cli_base(json_file, json, api_call) + else: + if table is None: + raise ValueError('Must specify table name when updating shared table') data = { 'updates': [ { @@ -475,20 +488,6 @@ def update_share_table_cli(api_client, share, table, shared_as, comment, } share_json = UnityCatalogApi(api_client).update_share(share, data) click.echo(mc_pretty_format(share_json)) - else: - def api_call(d): - if 'data_object_type' in d and d['data_object_type'] != "TABLE": - raise ValueError('Must specify data_object_type as "TABLE" ' - 'or not specify data_object_type at all') - UnityCatalogApi(api_client).update_share(share, { - 'updates': [ - { - 'action': 'UPDATE', - 'data_object': d, - } - ] - }) - json_cli_base(json_file, json, api_call) @click.command(context_settings=CONTEXT_SETTINGS, @@ -516,23 +515,7 @@ def remove_share_table_cli(api_client, share, table, shared_as, json_file, json) if table is not None and shared_as is not None: raise ValueError("You can only pass in either --table or --shared_as and not both.") - if table is not None or shared_as is not None: - if (json_file is not None) or (json is not None): - raise ValueError('Cannot specify JSON if any other flags are specified') - data = { - 'updates': [ - { - 'action': 'REMOVE', - 'data_object': shared_table_object( - name=table, - shared_as=shared_as, - ) - } - ] - } - share_json = UnityCatalogApi(api_client).update_share(share, data) - click.echo(mc_pretty_format(share_json)) - else: + if (json_file is not None) or (json is not None): def api_call(d): if 'data_object_type' in d and d['data_object_type'] != "TABLE": raise ValueError('Must specify data_object_type as "TABLE" ' @@ -546,6 +529,22 @@ def api_call(d): ] }) json_cli_base(json_file, json, api_call) + else: + if table is None or shared_as is None: + raise ValueError('Must specify full or shared as table name when removing shared table') + data = { + 'updates': [ + { + 'action': 'REMOVE', + 'data_object': shared_table_object( + name=table, + shared_as=shared_as, + ) + } + ] + } + share_json = UnityCatalogApi(api_client).update_share(share, data) + click.echo(mc_pretty_format(share_json)) @click.command(context_settings=CONTEXT_SETTINGS, diff --git a/tests/unity_catalog/test_delta_sharing_cli.py b/tests/unity_catalog/test_delta_sharing_cli.py index 4ebc7d3b..b2814bfc 100644 --- a/tests/unity_catalog/test_delta_sharing_cli.py +++ b/tests/unity_catalog/test_delta_sharing_cli.py @@ -215,7 +215,7 @@ def test_update_share_cli(api_mock, echo_mock): echo_mock.assert_called_once_with(mc_pretty_format(SHARE)) @provide_conf -def test_add_share_schema_cli(api_mock, echo_mock): +def test_add_share_schema_with_comment_cli(api_mock, echo_mock): api_mock.update_share.return_value = SHARE runner = CliRunner() runner.invoke( @@ -241,6 +241,31 @@ def test_add_share_schema_cli(api_mock, echo_mock): echo_mock.assert_called_once_with(mc_pretty_format(SHARE)) +@provide_conf +def test_add_share_schema_cli(api_mock, echo_mock): + api_mock.update_share.return_value = SHARE + runner = CliRunner() + runner.invoke( + delta_sharing_cli.add_share_schema_cli, + args=[ + '--share', SHARE_NAME, + '--schema', 'catalog.schema', + ]) + expected_data = { + 'updates': [ + { + 'action': 'ADD', + 'data_object': { + 'data_object_type': 'SCHEMA', + 'name': 'catalog.schema', + } + } + ] + } + api_mock.update_share.assert_called_once_with(SHARE_NAME, expected_data) + echo_mock.assert_called_once_with(mc_pretty_format(SHARE)) + + @provide_conf def test_update_share_schema_cli(api_mock, echo_mock): api_mock.update_share.return_value = SHARE From edbefe75a54785eb27aae2238402abd3d91f5aef Mon Sep 17 00:00:00 2001 From: Andy Date: Mon, 3 Apr 2023 21:57:01 -0700 Subject: [PATCH 8/8] Fix remove sschema description and fix remove table assertion --- databricks_cli/unity_catalog/delta_sharing_cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/databricks_cli/unity_catalog/delta_sharing_cli.py b/databricks_cli/unity_catalog/delta_sharing_cli.py index f1e5b239..a673d630 100644 --- a/databricks_cli/unity_catalog/delta_sharing_cli.py +++ b/databricks_cli/unity_catalog/delta_sharing_cli.py @@ -306,7 +306,7 @@ def api_call(d): @provide_api_client def remove_share_schema_cli(api_client, share, schema, json_file, json): """ - Removes a shared schema either by full schema name. + Removes a shared schema by full schema name. The public specification for the JSON request is in development. """ @@ -530,7 +530,7 @@ def api_call(d): }) json_cli_base(json_file, json, api_call) else: - if table is None or shared_as is None: + if table is None and shared_as is None: raise ValueError('Must specify full or shared as table name when removing shared table') data = { 'updates': [