Skip to content

Commit 682fd3c

Browse files
andyl-dbandyl-db
authored
Add update-table command to Delta Sharing CLI (#568)
1 parent 1c444a7 commit 682fd3c

File tree

2 files changed

+363
-12
lines changed

2 files changed

+363
-12
lines changed

‎databricks_cli/unity_catalog/delta_sharing_cli.py‎

Lines changed: 202 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@
2121
# See the License for the specific language governing permissions and
2222
# limitations under the License.
2323

24+
import functools
25+
from json import loads as json_loads
26+
2427
import click
2528

2629
from databricks_cli.click_types import JsonClickType
@@ -120,8 +123,24 @@ def update_share_permissions_cli(api_client, name, json_file, json):
120123
lambda json: UnityCatalogApi(api_client).update_share_permissions(name, json))
121124

122125

123-
def shared_data_object(name):
124-
return {'name': name, 'data_object_type': 'TABLE'}
126+
def shared_data_object(name=None, comment=None, shared_as=None,
127+
cdf_enabled=None, partitions=None, start_version=None):
128+
val = {
129+
'data_object_type': 'TABLE'
130+
}
131+
if name is not None:
132+
val['name'] = name
133+
if comment is not None:
134+
val['comment'] = comment
135+
if shared_as is not None:
136+
val['shared_as'] = shared_as
137+
if cdf_enabled is not None:
138+
val['cdf_enabled'] = cdf_enabled
139+
if partitions is not None:
140+
val['partitions'] = partitions
141+
if start_version is not None:
142+
val['start_version'] = start_version
143+
return val
125144

126145

127146
@click.command(context_settings=CONTEXT_SETTINGS,
@@ -133,11 +152,12 @@ def shared_data_object(name):
133152
help='Free-form text description.')
134153
@click.option('--owner', default=None, required=False,
135154
help='Owner of the share.')
155+
# These options are hidden to encourage the usage of the new commands: add-table/update-table
136156
@click.option('--add-table', default=None, multiple=True,
137-
metavar='NAME',
157+
metavar='NAME', hidden=True,
138158
help='Full name of table to add to share (can be specified multiple times).')
139159
@click.option('--remove-table', default=None, multiple=True,
140-
metavar='NAME',
160+
metavar='NAME', hidden=True,
141161
help='Full name of table to remove from share (can be specified multiple times).')
142162
@click.option('--json-file', default=None, type=click.Path(),
143163
help=json_file_help(method='PATCH', path='/shares/{name}'))
@@ -171,6 +191,181 @@ def update_share_cli(api_client, name, new_name, comment, owner,
171191
lambda json: UnityCatalogApi(api_client).update_share(name, json))
172192

173193

194+
def create_common_shared_data_object_options(f):
195+
@click.option('--table', required=True,
196+
help='Full name of the shared table.')
197+
@click.option('--shared-as', default=None,
198+
help='New name of the table to be shared as.')
199+
@click.option('--comment', default=None,
200+
help='New comment of the shared table.')
201+
@click.option('--partitions', default=None, type=JsonClickType(),
202+
help='New partition specification of the shared table represented in JSON.')
203+
@click.option('--cdf/--no-cdf', is_flag=True, default=None,
204+
help='Toggles the change data feed for the shared table.')
205+
@click.option('--start-version', default=None, type=int,
206+
help='Specifies the current version of the shared table.')
207+
@functools.wraps(f)
208+
def wrapper(*args, **kwargs):
209+
f(*args, **kwargs)
210+
return wrapper
211+
212+
213+
@click.command(context_settings=CONTEXT_SETTINGS,
214+
short_help='Add a shared table.')
215+
@click.option('--share', required=True,
216+
help='Name of the share to update.')
217+
@create_common_shared_data_object_options
218+
@click.option('--json-file', default=None, type=click.Path(),
219+
help="Adds a shared table based on shared data object represented in JSON file.")
220+
@click.option('--json', default=None, type=JsonClickType(),
221+
help="Adds a shared table based on shared data object represented in JSON.")
222+
@debug_option
223+
@profile_option
224+
@eat_exceptions
225+
@provide_api_client
226+
def add_share_table_cli(api_client, share, table, shared_as, comment,
227+
partitions, cdf, start_version, json_file, json):
228+
"""
229+
Adds a shared table.
230+
231+
The public specification for the JSON request is in development.
232+
"""
233+
if ((shared_as is not None) or (comment is not None) or (partitions is not None) or
234+
(cdf is not None) or (start_version is not None)):
235+
if (json_file is not None) or (json is not None):
236+
raise ValueError('Cannot specify JSON if any other flags are specified')
237+
data = {
238+
'updates': [
239+
{
240+
'action': 'ADD',
241+
'data_object': shared_data_object(
242+
name=table,
243+
shared_as=shared_as,
244+
comment=comment,
245+
cdf_enabled=cdf,
246+
partitions=json_loads(partitions) if partitions is not None else None,
247+
start_version=start_version,
248+
)
249+
}
250+
]
251+
}
252+
share_json = UnityCatalogApi(api_client).update_share(share, data)
253+
click.echo(mc_pretty_format(share_json))
254+
else:
255+
json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, {
256+
'updates': [
257+
{
258+
'action': 'ADD',
259+
'data_object': d,
260+
}
261+
]
262+
}))
263+
264+
265+
@click.command(context_settings=CONTEXT_SETTINGS,
266+
short_help='Update a shared table.')
267+
@click.option('--share', required=True,
268+
help='Name of the share to update.')
269+
@create_common_shared_data_object_options
270+
@click.option('--json-file', default=None, type=click.Path(),
271+
help="Updates the shared table to shared data object represented in JSON file.")
272+
@click.option('--json', default=None, type=JsonClickType(),
273+
help="Updates the shared table to shared data object represented in JSON.")
274+
@debug_option
275+
@profile_option
276+
@eat_exceptions
277+
@provide_api_client
278+
def update_share_table_cli(api_client, share, table, shared_as, comment,
279+
partitions, cdf, start_version, json_file, json):
280+
"""
281+
Updates a shared table.
282+
283+
The public specification for the JSON request is in development.
284+
"""
285+
if ((shared_as is not None) or (comment is not None) or (partitions is not None) or
286+
(cdf is not None) or (start_version is not None)):
287+
if (json_file is not None) or (json is not None):
288+
raise ValueError('Cannot specify JSON if any other flags are specified')
289+
data = {
290+
'updates': [
291+
{
292+
'action': 'UPDATE',
293+
'data_object': shared_data_object(
294+
name=table,
295+
shared_as=shared_as,
296+
comment=comment,
297+
cdf_enabled=cdf,
298+
partitions=json_loads(partitions) if partitions is not None else None,
299+
start_version=start_version,
300+
)
301+
}
302+
]
303+
}
304+
share_json = UnityCatalogApi(api_client).update_share(share, data)
305+
click.echo(mc_pretty_format(share_json))
306+
else:
307+
json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, {
308+
'updates': [
309+
{
310+
'action': 'UPDATE',
311+
'data_object': d,
312+
}
313+
]
314+
}))
315+
316+
317+
@click.command(context_settings=CONTEXT_SETTINGS,
318+
short_help='Update a shared table.')
319+
@click.option('--share', required=True,
320+
help='Name of the share to update.')
321+
@click.option('--table', default=None,
322+
help='Full name of the table to update from share.')
323+
@click.option('--shared-as', default=None,
324+
help='New name of the table inside the share.')
325+
@click.option('--json-file', default=None, type=click.Path(),
326+
help="Removes the shared table based on shared data object represented in JSON file.")
327+
@click.option('--json', default=None, type=JsonClickType(),
328+
help="Removes the shared table based on shared data object represented in JSON.")
329+
@debug_option
330+
@profile_option
331+
@eat_exceptions
332+
@provide_api_client
333+
def remove_share_table_cli(api_client, share, table, shared_as, json_file, json):
334+
"""
335+
Removes a shared table either by table name or the shared-as table name.
336+
337+
The public specification for the JSON request is in development.
338+
"""
339+
if table is not None and shared_as is not None:
340+
raise ValueError("You can only pass in either --table or --shared_as and not both.")
341+
342+
if table is not None or shared_as is not None:
343+
if (json_file is not None) or (json is not None):
344+
raise ValueError('Cannot specify JSON if any other flags are specified')
345+
data = {
346+
'updates': [
347+
{
348+
'action': 'REMOVE',
349+
'data_object': shared_data_object(
350+
name=table,
351+
shared_as=shared_as,
352+
)
353+
}
354+
]
355+
}
356+
share_json = UnityCatalogApi(api_client).update_share(share, data)
357+
click.echo(mc_pretty_format(share_json))
358+
else:
359+
json_cli_base(json_file, json, lambda d: UnityCatalogApi(api_client).update_share(share, {
360+
'updates': [
361+
{
362+
'action': 'REMOVE',
363+
'data_object': d,
364+
}
365+
]
366+
}))
367+
368+
174369
@click.command(context_settings=CONTEXT_SETTINGS,
175370
short_help='Delete a share.')
176371
@click.option('--name', required=True,
@@ -498,6 +693,9 @@ def register_shares_commands(cmd_group):
498693
shares_group.add_command(list_shares_cli, name='list')
499694
shares_group.add_command(get_share_cli, name='get')
500695
shares_group.add_command(update_share_cli, name='update')
696+
shares_group.add_command(add_share_table_cli, name='add-table')
697+
shares_group.add_command(update_share_table_cli, name='update-table')
698+
shares_group.add_command(remove_share_table_cli, name='remove-table')
501699
shares_group.add_command(delete_share_cli, name='delete')
502700
shares_group.add_command(list_share_permissions_cli, name='list-permissions')
503701
shares_group.add_command(update_share_permissions_cli, name='update-permissions')

0 commit comments

Comments
 (0)