Add uc_api_version to ApiClient, UC CLI

databricks_cli/configure/cli.py

@@ -32,7 +32,7 @@
 from databricks_cli.configure.config import profile_option, get_profile_from_context, debug_option
 from databricks_cli.configure.provider import DatabricksConfig, update_and_persist_config, \
     ProfileConfigProvider
-from databricks_cli.sdk.version import API_VERSION, API_VERSIONS
+from databricks_cli.sdk.version import API_VERSION, API_VERSIONS, DEFAULT_UC_API_VERSION
 from databricks_cli.utils import CONTEXT_SETTINGS
 
 PROMPT_HOST = 'Databricks Host (should begin with https://)'
@@ -50,7 +50,8 @@
 ]
 
 
-def _configure_cli_token_file(profile, token_file, host, insecure, jobs_api_version):
+def _configure_cli_token_file(profile, token_file, host, insecure, jobs_api_version,
+                              uc_api_version):
     if not path.exists(token_file):
         raise RuntimeError('Unable to read token from "{}"'.format(token_file))
 
@@ -65,11 +66,12 @@ def _configure_cli_token_file(profile, token_file, host, insecure, jobs_api_vers
                                              token=token,
                                              refresh_token=None,
                                              insecure=insecure,
-                                             jobs_api_version=jobs_api_version)
+                                             jobs_api_version=jobs_api_version,
+                                             uc_api_version=uc_api_version)
     update_and_persist_config(profile, new_config)
 
 
-def _configure_cli_token(profile, insecure, host, jobs_api_version):
+def _configure_cli_token(profile, insecure, host, jobs_api_version, uc_api_version):
     config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
 
     if not host:
@@ -80,7 +82,8 @@ def _configure_cli_token(profile, insecure, host, jobs_api_version):
                                              token=token,
                                              refresh_token=None,
                                              insecure=insecure,
-                                             jobs_api_version=jobs_api_version)
+                                             jobs_api_version=jobs_api_version,
+                                             uc_api_version=uc_api_version)
     update_and_persist_config(profile, new_config)
 
 
@@ -90,7 +93,7 @@ def scope_format(user_input):
     return list(map(lambda x: choice(x.strip()), user_inputs))
 
 
-def _configure_cli_oauth(profile, insecure, host, scope, jobs_api_version):
+def _configure_cli_oauth(profile, insecure, host, scope, jobs_api_version, uc_api_version):
     config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
 
     if not host:
@@ -105,11 +108,12 @@ def _configure_cli_oauth(profile, insecure, host, scope, jobs_api_version):
                                              token=access_token,
                                              refresh_token=refresh_token,
                                              insecure=insecure,
-                                             jobs_api_version=jobs_api_version)
+                                             jobs_api_version=jobs_api_version,
+                                             uc_api_version=uc_api_version)
     update_and_persist_config(profile, new_config)
 
 
-def _configure_cli_aad_token(profile, insecure, host, jobs_api_version):
+def _configure_cli_aad_token(profile, insecure, host, jobs_api_version, uc_api_version):
     config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
 
     if ENV_AAD_TOKEN not in os.environ:
@@ -131,11 +135,12 @@ def _configure_cli_aad_token(profile, insecure, host, jobs_api_version):
                                              token=aad_token,
                                              refresh_token=None,
                                              insecure=insecure,
-                                             jobs_api_version=jobs_api_version)
+                                             jobs_api_version=jobs_api_version,
+                                             uc_api_version=uc_api_version)
     update_and_persist_config(profile, new_config)
 
 
-def _configure_cli_password(profile, insecure, host, jobs_api_version):
+def _configure_cli_password(profile, insecure, host, jobs_api_version, uc_api_version):
     config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
     if config.password:
         default_password = '*' * len(config.password)
@@ -151,7 +156,7 @@ def _configure_cli_password(profile, insecure, host, jobs_api_version):
     if password == default_password:
         password = config.password
     new_config = DatabricksConfig.from_password(host, username, password, insecure,
-                                                jobs_api_version)
+                                                jobs_api_version, uc_api_version)
     update_and_persist_config(profile, new_config)
 
 
@@ -173,9 +178,12 @@ def _configure_cli_password(profile, insecure, host, jobs_api_version):
               help='DO NOT verify SSL Certificates')
 @click.option('--jobs-api-version', show_default=True, default=API_VERSION,
               type=click.Choice(API_VERSIONS), help='API version to use for jobs.')
+@click.option('--uc-api-version', show_default=True, default=DEFAULT_UC_API_VERSION,
+              type=click.Choice(API_VERSIONS), help='API version to use for unity-catalog.')
 @debug_option
 @profile_option
-def configure_cli(token, aad_token, insecure, host, token_file, jobs_api_version, oauth, scope):
+def configure_cli(token, aad_token, insecure, host, token_file, jobs_api_version, uc_api_version,
+                  oauth, scope):
     """
     Configures host, authentication, and jobs-api version for the CLI.
     """
@@ -184,19 +192,21 @@ def configure_cli(token, aad_token, insecure, host, token_file, jobs_api_version
 
     if token:
         _configure_cli_token(profile=profile, insecure=insecure_str, host=host,
-                             jobs_api_version=jobs_api_version)
+                             jobs_api_version=jobs_api_version, uc_api_version=uc_api_version)
     elif token_file:
         _configure_cli_token_file(profile=profile, insecure=insecure_str, host=host,
-                                  token_file=token_file, jobs_api_version=jobs_api_version)
+                                  token_file=token_file, jobs_api_version=jobs_api_version,
+                                  uc_api_version=uc_api_version)
     elif oauth:
         _configure_cli_oauth(profile=profile, insecure=insecure_str, host=host,
-                             scope=scope, jobs_api_version=jobs_api_version)
+                             scope=scope, jobs_api_version=jobs_api_version,
+                             uc_api_version=uc_api_version)
     elif aad_token:
         _configure_cli_aad_token(profile=profile, insecure=insecure_str, host=host,
-                                 jobs_api_version=jobs_api_version)
+                                 jobs_api_version=jobs_api_version, uc_api_version=uc_api_version)
     else:
         _configure_cli_password(profile=profile, insecure=insecure_str, host=host,
-                                jobs_api_version=jobs_api_version)
+                                jobs_api_version=jobs_api_version, uc_api_version=uc_api_version)
 
 
 class _DbfsHost(ParamType):

databricks_cli/configure/config.py

@@ -102,7 +102,9 @@ def _get_api_client(config, command_name=""):
     verify = config.insecure is None
     if config.is_valid_with_token:
         return ApiClient(host=config.host, token=config.token, verify=verify,
-                         command_name=command_name, jobs_api_version=config.jobs_api_version)
+                         command_name=command_name, jobs_api_version=config.jobs_api_version,
+                         uc_api_version=config.uc_api_version)
     return ApiClient(user=config.username, password=config.password,
                      host=config.host, verify=verify, command_name=command_name,
-                     jobs_api_version=config.jobs_api_version)
+                     jobs_api_version=config.jobs_api_version,
+                     uc_api_version=config.uc_api_version)

databricks_cli/configure/provider.py

@@ -38,6 +38,7 @@
 REFRESH_TOKEN = 'refresh_token'
 INSECURE = 'insecure'
 JOBS_API_VERSION = 'jobs-api-version'
+UC_API_VERSION = 'uc-api-version'
 DEFAULT_SECTION = 'DEFAULT'
 
 # User-provided override for the DatabricksConfigProvider
@@ -101,6 +102,7 @@ def update_and_persist_config(profile, databricks_config):
     _set_option(raw_config, profile, REFRESH_TOKEN, databricks_config.refresh_token)
     _set_option(raw_config, profile, INSECURE, databricks_config.insecure)
     _set_option(raw_config, profile, JOBS_API_VERSION, databricks_config.jobs_api_version)
+    _set_option(raw_config, profile, UC_API_VERSION, databricks_config.uc_api_version)
     _overwrite_config(raw_config)
 
 
@@ -248,8 +250,9 @@ def get_config(self):
         refresh_token = os.environ.get('DATABRICKS_REFRESH_TOKEN')
         insecure = os.environ.get('DATABRICKS_INSECURE')
         jobs_api_version = os.environ.get('DATABRICKS_JOBS_API_VERSION')
+        uc_api_version = os.environ.get('DATABRICKS_UC_API_VERSION')
         config = DatabricksConfig(host, username, password, token,
-                                  refresh_token, insecure, jobs_api_version)
+                                  refresh_token, insecure, jobs_api_version, uc_api_version)
         if config.is_valid:
             return config
         return None
@@ -269,43 +272,50 @@ def get_config(self):
         refresh_token = _get_option_if_exists(raw_config, self.profile, REFRESH_TOKEN)
         insecure = _get_option_if_exists(raw_config, self.profile, INSECURE)
         jobs_api_version = _get_option_if_exists(raw_config, self.profile, JOBS_API_VERSION)
+        uc_api_version = _get_option_if_exists(raw_config, self.profile, UC_API_VERSION)
         config = DatabricksConfig(host, username, password, token,
-                                  refresh_token, insecure, jobs_api_version)
+                                  refresh_token, insecure, jobs_api_version, uc_api_version)
         if config.is_valid:
             return config
         return None
 
 
 class DatabricksConfig(object):
     def __init__(self, host, username, password, token,
-                 refresh_token=None, insecure=None, jobs_api_version=None):  # noqa
+                 refresh_token=None, insecure=None, jobs_api_version=None,
+                 uc_api_version=None):  # noqa
         self.host = host
         self.username = username
         self.password = password
         self.token = token
         self.refresh_token = refresh_token
         self.insecure = insecure
         self.jobs_api_version = jobs_api_version
+        self.uc_api_version = uc_api_version
 
     @classmethod
-    def from_token(cls, host, token, refresh_token=None, insecure=None, jobs_api_version=None):
+    def from_token(cls, host, token, refresh_token=None, insecure=None, jobs_api_version=None,
+                   uc_api_version=None):
         return DatabricksConfig(host=host,
                                 username=None,
                                 password=None,
                                 token=token,
                                 refresh_token=refresh_token,
                                 insecure=insecure,
-                                jobs_api_version=jobs_api_version)
+                                jobs_api_version=jobs_api_version,
+                                uc_api_version=uc_api_version)
 
     @classmethod
-    def from_password(cls, host, username, password, insecure=None, jobs_api_version=None):
+    def from_password(cls, host, username, password, insecure=None, jobs_api_version=None,
+                      uc_api_version=None):
         return DatabricksConfig(host=host,
                                 username=username,
                                 password=password,
                                 token=None,
                                 refresh_token=None,
                                 insecure=insecure,
-                                jobs_api_version=jobs_api_version)
+                                jobs_api_version=jobs_api_version,
+                                uc_api_version=uc_api_version)
 
     @classmethod
     def empty(cls):
@@ -315,7 +325,8 @@ def empty(cls):
                                 token=None,
                                 refresh_token=None,
                                 insecure=None,
-                                jobs_api_version=None)
+                                jobs_api_version=None,
+                                uc_api_version=None)
 
     @property
     def is_valid_with_token(self):

databricks_cli/sdk/api_client.py

@@ -51,6 +51,9 @@
     from urllib3.util.retry import Retry
 
 from databricks_cli.version import version as databricks_cli_version
+from databricks_cli.sdk.version import DEFAULT_UC_API_VERSION
+from databricks_cli.unity_catalog.uc_service import is_uc_path
+
 
 class TlsV1HttpAdapter(HTTPAdapter):
     """
@@ -68,7 +71,8 @@ class ApiClient(object):
     to be used by different versions of the client.
     """
     def __init__(self, user=None, password=None, host=None, token=None,
-                 api_version=version.API_VERSION, default_headers={}, verify=True, command_name="", jobs_api_version=None):
+                 api_version=version.API_VERSION, default_headers={}, verify=True, command_name="",
+                 jobs_api_version=None, uc_api_version=None):
         if host[-1] == "/":
             host = host[:-1]
 
@@ -104,6 +108,8 @@ def __init__(self, user=None, password=None, host=None, token=None,
         self.verify = verify
         self.api_version = api_version
         self.jobs_api_version = jobs_api_version
+        # Default to UC API version 2.1 if it's not overridden by profile config
+        self.uc_api_version = uc_api_version if uc_api_version else DEFAULT_UC_API_VERSION
 
     def close(self):
         """Close the client"""
@@ -152,6 +158,8 @@ def get_url(self, path, version=None):
             return self.url + version + path
         elif self.jobs_api_version and path and path.startswith('/jobs'):
             return self.url + self.jobs_api_version + path
+        elif self.uc_api_version and path and is_uc_path(path):
+            return self.url + self.uc_api_version + path
         return self.url + self.api_version + path
 
 

databricks_cli/sdk/version.py

@@ -23,5 +23,7 @@
 
 API_VERSION = '2.0'
 
+DEFAULT_UC_API_VERSION = '2.1'
+
 # Available API versions
 API_VERSIONS = ['2.0', '2.1']

databricks_cli/unity_catalog/api.py

@@ -21,11 +21,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from databricks_cli.sdk.version import DEFAULT_UC_API_VERSION
 from databricks_cli.unity_catalog.uc_service import UnityCatalogService
 
 
 class UnityCatalogApi(object):
     def __init__(self, api_client):
+        # Default to UC API version 2.1 if it's not overridden by profile config
+        if api_client.uc_api_version is None:
+            api_client.uc_api_version = DEFAULT_UC_API_VERSION
         self.client = UnityCatalogService(api_client)
 
     # Metastore APIs

databricks_cli/unity_catalog/cli.py

@@ -26,6 +26,7 @@
 from databricks_cli.utils import CONTEXT_SETTINGS
 from databricks_cli.version import print_version_callback, version
 
+from databricks_cli.unity_catalog.configure_cli import register_configure_commands
 from databricks_cli.unity_catalog.metastore_cli import register_metastore_commands
 from databricks_cli.unity_catalog.catalog_cli import register_catalog_commands
 from databricks_cli.unity_catalog.schema_cli import register_schema_commands
@@ -51,6 +52,7 @@ def unity_catalog_group():  # pragma: no cover
     pass
 
 
+register_configure_commands(unity_catalog_group)
 register_metastore_commands(unity_catalog_group)
 register_ext_loc_commands(unity_catalog_group)
 register_cred_commands(unity_catalog_group)

databricks_cli/unity_catalog/configure_cli.py

@@ -0,0 +1,53 @@
+# Databricks CLI
+# Copyright 2022 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.
+
+import click
+
+from databricks_cli.configure.config import get_config, get_profile_from_context, profile_option
+from databricks_cli.configure.provider import DatabricksConfig, update_and_persist_config, \
+    ProfileConfigProvider
+from databricks_cli.utils import CONTEXT_SETTINGS
+from databricks_cli.sdk.version import API_VERSIONS
+
+#################  Configure Commands  #####################
+
+
+@click.command(context_settings=CONTEXT_SETTINGS)
+@click.option('--version', show_default=True, default=None, type=click.Choice(API_VERSIONS),
+              help='API version to use for unity-catalog.')
+@profile_option
+def configure_cli(version):
+    profile = get_profile_from_context()
+    config = ProfileConfigProvider(profile).get_config() if profile else get_config()
+    if config:
+        new_config = config
+    else:
+        click.echo("Using empty configuration.")
+        new_config = DatabricksConfig.empty()
+
+    new_config.uc_api_version = version
+    update_and_persist_config(profile, new_config)
+
+
+def register_configure_commands(cmd_group):
+    cmd_group.add_command(configure_cli, name='configure')

databricks_cli/unity_catalog/uc_service.py

@@ -25,6 +25,12 @@
 from databricks_cli.unity_catalog.utils import mc_pretty_format
 
 
+# 'path' is assumed to have the '/api/v2.x' prefix removed.
+#
+def is_uc_path(path):
+    return path.startswith('/unity-catalog') or path.startswith('/lineage-tracking')
+
+
 class UnityCatalogService(object):
     def __init__(self, client):
         self.client = client