CLI support for setting Delta Sharing recipient properties (#584)

Co-authored-by: Xiaotong Sun <Xiaotong Sun>
Co-authored-by: Pieter Noordhuis <[email protected]>

Author: zhuansunxt

databricks_cli/unity_catalog/api.py

@@ -193,8 +193,10 @@ def update_share_permissions(self, name, perm_spec):
 
     # Recipient APIs
 
-    def create_recipient(self, name, comment, sharing_id, allowed_ip_addresses):
-        return self.client.create_recipient(name, comment, sharing_id, allowed_ip_addresses)
+    def create_recipient(self, name, comment, sharing_id,
+                         allowed_ip_addresses, custom_properties):
+        return self.client.create_recipient(name, comment, sharing_id,
+                                            allowed_ip_addresses, custom_properties)
 
     def list_recipients(self):
         return self.client.list_recipients()

databricks_cli/unity_catalog/delta_sharing_cli.py

@@ -383,6 +383,17 @@ def delete_share_cli(api_client, name):
 
 ##############  Recipient Commands  ##############
 
+def parse_recipient_custom_properties(custom_property_list):
+    custom_properties = []
+    for property_str in custom_property_list:
+        tokens = property_str.split('=', 2)
+        if len(tokens) != 2:
+            raise ValueError('Invalid format for property. '
+                             + 'The format should be <key>=<value>.')
+        custom_properties.append({"key": tokens[0], "value": tokens[1]})
+    return custom_properties    
+
+
 @click.command(context_settings=CONTEXT_SETTINGS,
                short_help='Create a new recipient.')
 @click.option('--name', required=True, help='Name of new recipient.')
@@ -394,16 +405,23 @@ def delete_share_cli(api_client, name):
               help=(
                   'IP address in CIDR notation that is allowed to use delta sharing. '
                   '(can be specified multiple times).'))
+@click.option('--property', 'custom_property', default=None, required=False, multiple=True,
+              help=(
+                  'Properties of the recipient. Key and value should be provided '
+                  'at the same time separated by an equal sign. '
+                  'Example: --property country=US.'))
 @debug_option
 @profile_option
 @eat_exceptions
 @provide_api_client
-def create_recipient_cli(api_client, name, comment, sharing_id, allowed_ip_address):
+def create_recipient_cli(api_client, name, comment, sharing_id,
+                         allowed_ip_address, custom_property):
     """
     Create a new recipient.
     """
     recipient_json = UnityCatalogApi(api_client).create_recipient(
-        name, comment, sharing_id, allowed_ip_address)
+        name, comment, sharing_id,
+        allowed_ip_address, parse_recipient_custom_properties(custom_property))
     click.echo(mc_pretty_format(recipient_json))
 
 
@@ -451,6 +469,12 @@ def get_recipient_cli(api_client, name):
                   'IP address in CIDR notation that is allowed to use delta sharing '
                   '(can be specified multiple times). Specify a single empty string to disable '
                   'IP allowlist.'))
+@click.option('--property', 'custom_property', default=None, required=False, multiple=True,
+              help=(
+                  'Properties of the recipient. Key and value should be provided '
+                  'at the same time separated by an equal sign. '
+                  'Example: --property country=US. '
+                  'Specify a single empty string to remove all properties.'))
 @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(),
@@ -460,21 +484,26 @@ def get_recipient_cli(api_client, name):
 @eat_exceptions
 @provide_api_client
 def update_recipient_cli(api_client, name, new_name, comment, owner,
-                         allowed_ip_address, json_file, json):
+                         allowed_ip_address, custom_property, json_file, json):
     """
     Update a recipient.
 
     The public specification for the JSON request is in development.
     """
     if ((new_name is not None) or (comment is not None) or (owner is not None) or
-        len(allowed_ip_address) > 0):
+        len(allowed_ip_address) or len(custom_property) > 0):
         if (json_file is not None) or (json is not None):
             raise ValueError('Cannot specify JSON if any other update flags are specified')
         data = {'name': new_name, 'comment': comment, 'owner': owner}
         if len(allowed_ip_address) > 0:
             data['ip_access_list'] = {}
             if len(allowed_ip_address) != 1 or allowed_ip_address[0] != '':
                 data['ip_access_list']['allowed_ip_addresses'] = allowed_ip_address
+        if len(custom_property) > 0:
+            data['properties_kvpairs'] = {}
+            if len(custom_property) != 1 or custom_property[0] != '':
+                data['properties_kvpairs']['properties'] = parse_recipient_custom_properties(
+                    custom_property)
         recipient_json = UnityCatalogApi(api_client).update_recipient(name, data)
         click.echo(mc_pretty_format(recipient_json))
     else:

databricks_cli/unity_catalog/uc_service.py

@@ -379,7 +379,7 @@ def update_share_permissions(self, name, perm_spec, headers=None):
     # Recipient Operations
 
     def create_recipient(self, name, comment=None, sharing_id=None,
-                         allowed_ip_addresses=None, headers=None):
+                         allowed_ip_addresses=None, custom_properties=None, headers=None):
         _data = {
             'name': name,
         }
@@ -394,6 +394,10 @@ def create_recipient(self, name, comment=None, sharing_id=None,
             _data['ip_access_list'] = {
                 'allowed_ip_addresses': allowed_ip_addresses,
             }
+        if custom_properties is not None:
+            _data['properties_kvpairs'] = {
+                'properties': custom_properties
+            }
 
         return self.client.perform_query('POST', '/unity-catalog/recipients', data=_data,
                                          headers=headers)

tests/unity_catalog/test_delta_sharing_cli.py

@@ -463,12 +463,33 @@ def test_create_recipient_cli(api_mock, echo_mock):
             '--comment', 'comment',
             '--allowed-ip-address', '8.8.8.8',
             '--allowed-ip-address', '8.8.4.4',
+            '--property', 'k1=v1',
+            '--property', 'k2=v2'
         ])
     api_mock.create_recipient.assert_called_once_with(
-        RECIPIENT_NAME, 'comment', None, ('8.8.8.8', '8.8.4.4'))
+        RECIPIENT_NAME,
+        'comment',
+        None,
+        ('8.8.8.8', '8.8.4.4'),
+        [{"key": "k1", "value": "v1"}, {"key": "k2", "value": "v2"}]
+    )
     echo_mock.assert_called_once_with(mc_pretty_format(RECIPIENT))
 
 
+@provide_conf
+def test_create_recipient_cli_invalid_custom_property(api_mock):
+    api_mock.create_recipient.return_value = RECIPIENT
+    runner = CliRunner()
+    runner.invoke(
+        delta_sharing_cli.create_recipient_cli,
+        args=[
+            '--name', RECIPIENT_NAME,
+            '--property', 'k1=v1=v2'
+        ])
+
+    assert not api_mock.create_recipient.called
+
+
 @provide_conf
 def test_create_recipient_cli_with_sharing_id(api_mock, echo_mock):
     api_mock.create_recipient.return_value = RECIPIENT
@@ -480,7 +501,7 @@ def test_create_recipient_cli_with_sharing_id(api_mock, echo_mock):
             '--sharing-id', '123e4567-e89b-12d3-a456-426614174000'
         ])
     api_mock.create_recipient.assert_called_once_with(
-        RECIPIENT_NAME, None, '123e4567-e89b-12d3-a456-426614174000', ())
+        RECIPIENT_NAME, None, '123e4567-e89b-12d3-a456-426614174000', (), [])
     echo_mock.assert_called_once_with(mc_pretty_format(RECIPIENT))
 
 
@@ -516,7 +537,9 @@ def test_update_recipient_cli(api_mock, echo_mock):
             '--comment', 'comment',
             '--owner', 'owner',
             '--allowed-ip-address', '8.8.8.8',
-            '--allowed-ip-address', '8.8.4.4'
+            '--allowed-ip-address', '8.8.4.4',
+            '--property', 'k1=v1',
+            '--property', 'k2=v2'
         ])
     expected_data = {
         'name': 'new_recipient_name',
@@ -527,6 +550,12 @@ def test_update_recipient_cli(api_mock, echo_mock):
                 '8.8.8.8',
                 '8.8.4.4'
             )
+        },
+        'properties_kvpairs': {
+            'properties': [
+                {"key": "k1", "value": "v1"}, 
+                {"key": "k2", "value": "v2"}
+            ]
         }
     }
     api_mock.update_recipient.assert_called_once_with(RECIPIENT_NAME, expected_data)