Skip to content
New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Add add-schema, update-schema, and remove-schema to Delta Sharing CLI #608

Merged
merged 8 commits into from
Apr 6, 2023
Merged

Add add-schema, update-schema, and remove-schema to Delta Sharing CLI #608

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
2640c53
Add add-schema, update-schema, and remove-schema to Unity Catalog CLI
andyl-db Mar 4, 2023
e944d25
Fix parameter description + add validation for --json/--json-file for…
andyl-db Mar 7, 2023
dcf390c
Fix error message
andyl-db Mar 7, 2023
f11b382
Update comment
andyl-db Mar 31, 2023
10a8625
Update comment x2
andyl-db Mar 31, 2023
e1a9b47
Fix lint
andyl-db Mar 31, 2023
2cbb115
CLI fixes
andyl-db Apr 4, 2023
edbefe7
Fix remove sschema description and fix remove table assertion
andyl-db Apr 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
292 changes: 238 additions & 54 deletions databricks_cli/unity_catalog/delta_sharing_cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -180,18 +159,205 @@ 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))
else:
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 schema based on shared data object represented in JSON file.")
@click.option('--json', default=None, type=JsonClickType(),
help="Adds a shared schema 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:
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,
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 schema to shared data object represented in JSON file.")
@click.option('--json', default=None, type=JsonClickType(),
help="Updates the shared schema 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:
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,
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 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
@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:
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):
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,
Expand All @@ -214,7 +380,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(),
Expand All @@ -238,7 +404,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,
Expand All @@ -252,21 +418,26 @@ 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,
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(),
Expand All @@ -290,7 +461,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,
Expand All @@ -304,22 +475,27 @@ 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,
short_help='Update a shared table.')
@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(),
Expand All @@ -346,7 +522,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,
)
Expand All @@ -356,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,
Expand Down Expand Up @@ -722,6 +903,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')
Expand Down
Loading