diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index be4b48b9..2aea6f4d 100644 --- a/databricks_cli/configure/cli.py +++ b/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): diff --git a/databricks_cli/configure/config.py b/databricks_cli/configure/config.py index a8ac32a6..625f03cb 100644 --- a/databricks_cli/configure/config.py +++ b/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) diff --git a/databricks_cli/configure/provider.py b/databricks_cli/configure/provider.py index 828a10b3..1325965f 100644 --- a/databricks_cli/configure/provider.py +++ b/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,8 +272,9 @@ 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 @@ -278,7 +282,8 @@ def get_config(self): 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 @@ -286,26 +291,31 @@ def __init__(self, host, username, password, 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): diff --git a/databricks_cli/sdk/api_client.py b/databricks_cli/sdk/api_client.py index eebd9581..08112f92 100644 --- a/databricks_cli/sdk/api_client.py +++ b/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 diff --git a/databricks_cli/sdk/version.py b/databricks_cli/sdk/version.py index b5eb50ec..b220a789 100644 --- a/databricks_cli/sdk/version.py +++ b/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'] diff --git a/databricks_cli/unity_catalog/api.py b/databricks_cli/unity_catalog/api.py index 2ec2a1bf..9f84d00f 100644 --- a/databricks_cli/unity_catalog/api.py +++ b/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 diff --git a/databricks_cli/unity_catalog/cli.py b/databricks_cli/unity_catalog/cli.py index dab0f9bc..03dbdcee 100644 --- a/databricks_cli/unity_catalog/cli.py +++ b/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) diff --git a/databricks_cli/unity_catalog/configure_cli.py b/databricks_cli/unity_catalog/configure_cli.py new file mode 100644 index 00000000..bcf4e1aa --- /dev/null +++ b/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') diff --git a/databricks_cli/unity_catalog/uc_service.py b/databricks_cli/unity_catalog/uc_service.py index f19d7cf8..582cf984 100644 --- a/databricks_cli/unity_catalog/uc_service.py +++ b/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 diff --git a/tests/configure/test_cli.py b/tests/configure/test_cli.py index 8fb3b2b7..b8681f46 100644 --- a/tests/configure/test_cli.py +++ b/tests/configure/test_cli.py @@ -84,6 +84,20 @@ def test_configure_cli_insecure(): assert get_config().insecure == 'True' +def test_configure_cli_uc_api_version_default(): + runner = CliRunner() + runner.invoke(cli.configure_cli, ['--token'], + input=(TEST_HOST + '\n' + TEST_TOKEN + '\n')) + assert get_config().uc_api_version == "2.1" + + +def test_configure_cli_uc_api_version(): + runner = CliRunner() + runner.invoke(cli.configure_cli, ['--uc-api-version', '2.0', '--token'], + input=(TEST_HOST + '\n' + TEST_TOKEN + '\n')) + assert get_config().uc_api_version == '2.0' + + def test_configure_cli_jobs_api_version(): runner = CliRunner() runner.invoke(cli.configure_cli, ['--jobs-api-version', '2.1', '--token'], diff --git a/tests/sdk/test_api_client.py b/tests/sdk/test_api_client.py index d65a3f16..d15c9e89 100644 --- a/tests/sdk/test_api_client.py +++ b/tests/sdk/test_api_client.py @@ -109,6 +109,16 @@ def test_get_url(): assert client.get_url('/jobs/list') == 'https://databricks.com/api/2.1/jobs/list' assert client.get_url('/jobs/list', '3.0') == 'https://databricks.com/api/3.0/jobs/list' + assert client.uc_api_version == '2.1' + assert client.get_url('/unity-catalog/catalogs') == \ + 'https://databricks.com/api/2.1/unity-catalog/catalogs' + assert client.get_url('/unity-catalog/catalogs', '2.0') == \ + 'https://databricks.com/api/2.0/unity-catalog/catalogs' + client = ApiClient(host='https://databricks.com', uc_api_version = '2.0') + assert client.uc_api_version == '2.0' + assert client.get_url('/unity-catalog/catalogs') == \ + 'https://databricks.com/api/2.0/unity-catalog/catalogs' + def test_api_client_url_parsing(): client = ApiClient(host='https://databricks.com') diff --git a/tests/unity_catalog/__init__.py b/tests/unity_catalog/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unity_catalog/test_configure_cli.py b/tests/unity_catalog/test_configure_cli.py new file mode 100644 index 00000000..a0a33718 --- /dev/null +++ b/tests/unity_catalog/test_configure_cli.py @@ -0,0 +1,43 @@ +# 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. + +# pylint:disable=redefined-outer-name + +from click.testing import CliRunner + +from databricks_cli.configure.config import get_config +import databricks_cli.unity_catalog.configure_cli as cli +from tests.utils import provide_conf + + +@provide_conf +def test_configure(): + runner = CliRunner() + runner.invoke(cli.configure_cli, []) + assert get_config().uc_api_version is None + + runner.invoke(cli.configure_cli, ['--version=2.0']) + assert get_config().uc_api_version == '2.0' + + runner.invoke(cli.configure_cli, ['--version=2.1']) + assert get_config().uc_api_version == '2.1'