Integration for configuring workspace bindings (#628)

Author: willwwt

databricks_cli/unity_catalog/api.py

@@ -131,6 +131,12 @@ def update_catalog(self, name, catalog_spec):
     def delete_catalog(self, catalog_name):
         return self.client.delete_catalog(catalog_name)
 
+    def get_catalog_bindings(self, name):
+        return self.client.get_catalog_bindings(name)
+
+    def update_catalog_bindings(self, name, workspace_bindings_spec):
+        return self.client.update_catalog_bindings(name, workspace_bindings_spec)
+
     # Schema APIs
 
     def create_schema(self, catalog_name, schema_name, comment):

databricks_cli/unity_catalog/catalog_cli.py

@@ -137,6 +137,44 @@ def delete_catalog_cli(api_client, name, purge):
     UnityCatalogApi(api_client).delete_catalog(name)
 
 
+@click.command(context_settings=CONTEXT_SETTINGS,
+               short_help='Get workspace bindings of a catalog.')
+@click.option('--name', required=True,
+              help='Name of the catalog to get bindings for.')
+@debug_option
+@profile_option
+@eat_exceptions
+@provide_api_client
+def get_catalog_bindings_cli(api_client, name):
+    """
+    Get workspace bindings of a catalog.
+    """
+    catalog_json = UnityCatalogApi(api_client).get_catalog_bindings(name)
+    click.echo(mc_pretty_format(catalog_json))
+
+
+@click.command(context_settings=CONTEXT_SETTINGS,
+               short_help='Update workspace bindings of a catalog.')
+@click.option('--name', required=True,
+              help='Name of the catalog to update bindings for.')
+@click.option('--json-file', default=None, type=click.Path(),
+              help=json_file_help(method='PATCH', path='/workspace-bindings/catalogs/{name}'))
+@click.option('--json', default=None, type=JsonClickType(),
+              help=json_string_help(method='PATCH', path='/workspace-bindings/catalogs/{name}'))
+@debug_option
+@profile_option
+@eat_exceptions
+@provide_api_client
+def update_catalog_bindings_cli(api_client, name, json_file, json):
+    """
+    Update workspace bindings of 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_bindings(name, json))
+
+
 @click.group()
 def catalogs_group():  # pragma: no cover
     pass
@@ -149,11 +187,15 @@ def register_catalog_commands(cmd_group):
     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')
+    cmd_group.add_command(hide(get_catalog_bindings_cli), name='get-catalog-bindings')
+    cmd_group.add_command(hide(update_catalog_bindings_cli), name='update-catalog-bindings')
 
     # 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')
+    catalogs_group.add_command(get_catalog_bindings_cli, name='get-bindings')
+    catalogs_group.add_command(update_catalog_bindings_cli, name='update-bindings')
     cmd_group.add_command(catalogs_group, name='catalogs')

databricks_cli/unity_catalog/uc_service.py

@@ -239,6 +239,17 @@ def delete_catalog(self, name, headers=None):
         return self.client.perform_query('DELETE', '/unity-catalog/catalogs/%s' % (name),
                                          data=_data, headers=headers)
 
+    def get_catalog_bindings(self, name, headers=None):
+        _data = {}
+        return self.client.perform_query('GET',
+                                         '/unity-catalog/workspace-bindings/catalogs/%s' % (name),
+                                         data=_data, headers=headers)
+
+    def update_catalog_bindings(self, name, workspace_bindings_spec, headers=None):
+        return self.client.perform_query('PATCH',
+                                         '/unity-catalog/workspace-bindings/catalogs/%s' % (name),
+                                         data=workspace_bindings_spec, headers=headers)
+
     # Schema Operations
 
     def create_schema(self, catalog_name, new_schema_name, comment=None, headers=None):

tests/unity_catalog/test_catalog_cli.py

@@ -0,0 +1,82 @@
+# Databricks CLI
+# Copyright 2017 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.
+
+# pylint:disable=redefined-outer-name
+
+import mock
+import pytest
+from click.testing import CliRunner
+from databricks_cli.unity_catalog.utils import mc_pretty_format
+
+from databricks_cli.unity_catalog import catalog_cli
+from tests.utils import provide_conf
+
+CATALOG_NAME = 'test_catalog_name'
+
+WORKSPACE_BINDINGS = {
+    "workspaces": [6051921418418893]
+}
+
+EMPTY_WORKSPACE_BINDINGS = {}
+
+@pytest.fixture()
+def api_mock():
+    with mock.patch(
+            'databricks_cli.unity_catalog.catalog_cli.UnityCatalogApi') as uc_api_mock:
+        _cred_api_mock = mock.MagicMock()
+        uc_api_mock.return_value = _cred_api_mock
+        yield _cred_api_mock
+
+
+@pytest.fixture()
+def echo_mock():
+    with mock.patch('databricks_cli.unity_catalog.catalog_cli.click.echo') as echo_mock:
+        yield echo_mock
+
+
+@provide_conf
+def test_get_catalog_bindings_cli(api_mock, echo_mock):
+    api_mock.get_catalog_bindings.return_value = WORKSPACE_BINDINGS
+    runner = CliRunner()
+    runner.invoke(
+        catalog_cli.get_catalog_bindings_cli,
+        args=['--name', CATALOG_NAME])
+    api_mock.get_catalog_bindings.assert_called_once()
+    echo_mock.assert_called_once_with(mc_pretty_format(WORKSPACE_BINDINGS))
+
+@provide_conf
+def test_update_catalog_bindings_cli_with_json(api_mock, echo_mock):
+    api_mock.update_catalog_bindings.return_value = EMPTY_WORKSPACE_BINDINGS
+    runner = CliRunner()
+    runner.invoke(
+        catalog_cli.update_catalog_bindings_cli,
+        args=[
+            '--name', CATALOG_NAME,
+            '--json', '{ "unassign_workspaces": [6051921418418893] }'
+        ])
+    api_mock.update_catalog_bindings.assert_called_once_with(
+        CATALOG_NAME,
+        {
+            'unassign_workspaces': [6051921418418893]
+        })
+    echo_mock.assert_called_once_with(mc_pretty_format(EMPTY_WORKSPACE_BINDINGS))