From b7c7821c578693ca17bb30f99600211a8d05a191 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Tue, 22 Sep 2020 09:46:59 -0700 Subject: [PATCH 01/13] added aad token functionality but aad token is only accepted as env variable --- databricks_cli/configure/cli.py | 20 +++++++++++++++++++- databricks_cli/configure/config.py | 5 ++++- databricks_cli/configure/provider.py | 24 +++++++++++++++++++++--- databricks_cli/sdk/api_client.py | 6 +++++- databricks_cli/sdk/service.py | 11 +++++++++-- databricks_cli/secrets/api.py | 4 ++-- databricks_cli/secrets/cli.py | 19 ++++++++++++++++--- 7 files changed, 76 insertions(+), 13 deletions(-) diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index a933ddef..2a6640ad 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -22,6 +22,7 @@ # limitations under the License. import click +import os from click import ParamType @@ -31,6 +32,8 @@ from databricks_cli.configure.config import profile_option, get_profile_from_context, debug_option PROMPT_HOST = 'Databricks Host (should begin with https://)' +PROMPT_RESOURCE_ID = 'Resource/Workspace ID' +PROMPT_AZ_TOKEN = 'Azure Token' PROMPT_USERNAME = 'Username' PROMPT_PASSWORD = 'Password' # NOQA PROMPT_TOKEN = 'Token' # NOQA @@ -44,6 +47,18 @@ def _configure_cli_token(profile, insecure): update_and_persist_config(profile, new_config) +def _configure_cli_aad_token(profile, insecure): + config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() + host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) + #token = click.prompt(PROMPT_TOKEN, default=config.token) + #az_token = click.prompt(PROMPT_AZ_TOKEN, default=config.az_token) + token = os.environ.get('DATABRICKS_TOKEN') + az_token = os.environ.get('DATABRICKS_AZ_TOKEN') + resource_id = click.prompt(PROMPT_RESOURCE_ID, default=config.resource_id) + new_config = DatabricksConfig.from_aad_token(host, token, az_token, resource_id, insecure) + update_and_persist_config(profile, new_config) + + def _configure_cli_password(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() if config.password: @@ -63,10 +78,11 @@ def _configure_cli_password(profile, insecure): @click.command(context_settings=CONTEXT_SETTINGS, short_help='Configures host and authentication info for the CLI.') @click.option('--token', show_default=True, is_flag=True, default=False) +@click.option('--aad-token', show_default=True, is_flag=True, default=False) @click.option('--insecure', show_default=True, is_flag=True, default=None) @debug_option @profile_option -def configure_cli(token, insecure): +def configure_cli(token, aad_token, insecure): """ Configures host and authentication info for the CLI. """ @@ -74,6 +90,8 @@ def configure_cli(token, insecure): insecure_str = str(insecure) if insecure is not None else None if token: _configure_cli_token(profile, insecure_str) + elif aad_token: + _configure_cli_aad_token(profile, insecure_str) else: _configure_cli_password(profile, insecure_str) diff --git a/databricks_cli/configure/config.py b/databricks_cli/configure/config.py index f6ca04e4..65bd33c7 100644 --- a/databricks_cli/configure/config.py +++ b/databricks_cli/configure/config.py @@ -83,7 +83,10 @@ def callback(ctx, param, value): # NOQA def _get_api_client(config, command_name=""): verify = config.insecure is None - if config.is_valid_with_token: + if config.is_valid_with_az_token: + return ApiClient(host=config.host, token=config.token, az_token=config.az_token, + resource_id=config.resource_id, verify=verify, command_name=command_name) + elif config.is_valid_with_token: return ApiClient(host=config.host, token=config.token, verify=verify, command_name=command_name) return ApiClient(user=config.username, password=config.password, diff --git a/databricks_cli/configure/provider.py b/databricks_cli/configure/provider.py index d8b96aae..1bf6fc57 100644 --- a/databricks_cli/configure/provider.py +++ b/databricks_cli/configure/provider.py @@ -36,6 +36,8 @@ PASSWORD = 'password' # NOQA TOKEN = 'token' INSECURE = 'insecure' +AZ_TOKEN = 'az_token' +RESOURCE_ID = 'resource_id' DEFAULT_SECTION = 'DEFAULT' # User-provided override for the DatabricksConfigProvider @@ -97,6 +99,8 @@ def update_and_persist_config(profile, databricks_config): _set_option(raw_config, profile, PASSWORD, databricks_config.password) _set_option(raw_config, profile, TOKEN, databricks_config.token) _set_option(raw_config, profile, INSECURE, databricks_config.insecure) + _set_option(raw_config, profile, AZ_TOKEN, databricks_config.az_token) + _set_option(raw_config, profile, RESOURCE_ID, databricks_config.resource_id) _overwrite_config(raw_config) @@ -238,7 +242,9 @@ def get_config(self): password = os.environ.get('DATABRICKS_PASSWORD') token = os.environ.get('DATABRICKS_TOKEN') insecure = os.environ.get('DATABRICKS_INSECURE') - config = DatabricksConfig(host, username, password, token, insecure) + az_token = os.environ.get('DATABRICKS_AZ_TOKEN') + resource_id = os.environ.get('DATABRICKS_RESOURCE_ID') + config = DatabricksConfig(host, username, password, token, insecure, az_token, resource_id) if config.is_valid: return config return None @@ -256,24 +262,32 @@ def get_config(self): password = _get_option_if_exists(raw_config, self.profile, PASSWORD) token = _get_option_if_exists(raw_config, self.profile, TOKEN) insecure = _get_option_if_exists(raw_config, self.profile, INSECURE) - config = DatabricksConfig(host, username, password, token, insecure) + az_token = _get_option_if_exists(raw_config, self.profile, AZ_TOKEN) + resource_id = _get_option_if_exists(raw_config, self.profile, RESOURCE_ID) + config = DatabricksConfig(host, username, password, token, insecure, az_token, resource_id) if config.is_valid: return config return None class DatabricksConfig(object): - def __init__(self, host, username, password, token, insecure): # noqa + def __init__(self, host, username, password, token, insecure, az_token=None, resource_id=None): # noqa self.host = host self.username = username self.password = password self.token = token self.insecure = insecure + self.az_token = az_token + self.resource_id = resource_id @classmethod def from_token(cls, host, token, insecure=None): return DatabricksConfig(host, None, None, token, insecure) + @classmethod + def from_aad_token(cls, host, token, az_token, resource_id, insecure=None): + return DatabricksConfig(host, None, None, token, insecure, az_token, resource_id) + @classmethod def from_password(cls, host, username, password, insecure=None): return DatabricksConfig(host, username, password, None, insecure) @@ -286,6 +300,10 @@ def empty(cls): def is_valid_with_token(self): return self.host is not None and self.token is not None + @property + def is_valid_with_az_token(self): + return self.host is not None and self.token is not None and self.az_token is not None and self.resource_id is not None + @property def is_valid_with_password(self): return self.host is not None and self.username is not None and self.password is not None diff --git a/databricks_cli/sdk/api_client.py b/databricks_cli/sdk/api_client.py index d5870a6a..dcbd1aa9 100644 --- a/databricks_cli/sdk/api_client.py +++ b/databricks_cli/sdk/api_client.py @@ -65,7 +65,7 @@ class ApiClient(object): A partial Python implementation of dbc rest api to be used by different versions of the client. """ - def __init__(self, user=None, password=None, host=None, token=None, + def __init__(self, user=None, password=None, host=None, token=None, az_token=None, resource_id=None, apiVersion=version.API_VERSION, default_headers={}, verify=True, command_name=""): if host[-1] == "/": host = host[:-1] @@ -81,6 +81,10 @@ def __init__(self, user=None, password=None, host=None, token=None, encoded_auth = (user + ":" + password).encode() user_header_data = "Basic " + base64.standard_b64encode(encoded_auth).decode() auth = {'Authorization': user_header_data, 'Content-Type': 'text/json'} + elif token is not None and az_token is not None and resource_id is not None: + auth = {'Authorization': 'Bearer {}'.format(token), 'Content-Type': 'text/json', + 'X-Databricks-Azure-SP-Management-Token': '{}'.format(az_token), + 'X-Databricks-Azure-Workspace-Resource-Id': '{}'.format(resource_id)} elif token is not None: auth = {'Authorization': 'Bearer {}'.format(token), 'Content-Type': 'text/json'} else: diff --git a/databricks_cli/sdk/service.py b/databricks_cli/sdk/service.py index c0b4b55b..56c5423d 100644 --- a/databricks_cli/sdk/service.py +++ b/databricks_cli/sdk/service.py @@ -616,7 +616,7 @@ class SecretService(object): def __init__(self, client): self.client = client - def create_scope(self, scope, initial_manage_principal=None, scope_backend_type=None, + def create_scope(self, scope, initial_manage_principal=None, scope_backend_type=None, azure_keyvault=None, headers=None): _data = {} if scope is not None: @@ -624,7 +624,14 @@ def create_scope(self, scope, initial_manage_principal=None, scope_backend_type= if initial_manage_principal is not None: _data['initial_manage_principal'] = initial_manage_principal if scope_backend_type is not None: - _data['scope_backend_type'] = scope_backend_type + if scope_backend_type == 'databricks': + _data['scope_backend_type_str'] = 1 + elif azure_keyvault is not None and scope_backend_type == 'azure_keyvault': + _data['scope_backend_type'] = 2 + _data['backend_azure_keyvault'] = { + 'resource_id': str(azure_keyvault['resource_id']), + 'dns_name': str(azure_keyvault['dns_name']), + } return self.client.perform_query('POST', '/secrets/scopes/create', data=_data, headers=headers) def delete_scope(self, scope, headers=None): diff --git a/databricks_cli/secrets/api.py b/databricks_cli/secrets/api.py index 6ebf510c..fe5b0a01 100644 --- a/databricks_cli/secrets/api.py +++ b/databricks_cli/secrets/api.py @@ -28,8 +28,8 @@ class SecretApi(object): def __init__(self, api_client): self.client = SecretService(api_client) - def create_scope(self, scope, initial_manage_principal): - return self.client.create_scope(scope, initial_manage_principal) + def create_scope(self, scope, initial_manage_principal, scope_backend_type, azure_keyvault_info): + return self.client.create_scope(scope, initial_manage_principal, scope_backend_type, azure_keyvault_info) def delete_scope(self, scope): return self.client.delete_scope(scope) diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py index 485dfdb9..b1c6959a 100644 --- a/databricks_cli/secrets/cli.py +++ b/databricks_cli/secrets/cli.py @@ -44,21 +44,34 @@ short_help="Creates a secret scope.") @click.option('--scope', required=True, type=SecretScopeClickType(), help=SecretScopeClickType.help) @click.option('--initial-manage-principal', - help='The initial principal that can manage the created secret scope.' + help='Sushi The initial principal that can manage the created secret scope.' ' If specified, the initial ACL with MANAGE permission applied to the scope is' ' assigned to the supplied principal (user or group). Currently, the only supported' ' principal for this option is the group "users", which contains all users in the' ' workspace. If not specified, the initial ACL with MANAGE permission applied to the' ' scope is assigned to the request issuer\'s user identity.') +@click.option('--scope-backend-type', type=click.Choice(['azure_keyvault', 'databricks'], case_sensitive=True), + default='databricks', help='The backend that will be used for this secret scope. ' + 'Options are: 1) \'azure_keyvault\' and 2) \'databricks\' (default option)') +@click.option('--subscription-id', default=None, type=click.STRING, + help='The subscription ID associated with the azure keyvault to be used as the backend' + ' for the secret scope. NOTE: Only use with azure-keyvault as backend') +@click.option('--dns-name', default=None, type=click.STRING, + help='The dns name associated with the azure keyvault to be used as the backed for the' + ' secret scope. NOTE: Only use with azure-keyvault as backend') @debug_option @profile_option @eat_exceptions @provide_api_client -def create_scope(api_client, scope, initial_manage_principal): +def create_scope(api_client, scope, initial_manage_principal, scope_backend_type, subscription_id, dns_name): """ Creates a new secret scope with given name. """ - SecretApi(api_client).create_scope(scope, initial_manage_principal) + azure_keyvault_info = { + 'resource_id': subscription_id, + 'dns_name': dns_name + } + SecretApi(api_client).create_scope(scope, initial_manage_principal, scope_backend_type, azure_keyvault_info) def _scopes_to_table(scopes_json): From 192b8d563bc0aa8125d9841cf70a2d1561457879 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Tue, 22 Sep 2020 13:22:30 -0700 Subject: [PATCH 02/13] ready for review --- databricks_cli/configure/cli.py | 4 +--- databricks_cli/secrets/cli.py | 5 +++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index 2a6640ad..dc07b276 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -50,8 +50,6 @@ def _configure_cli_token(profile, insecure): def _configure_cli_aad_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) - #token = click.prompt(PROMPT_TOKEN, default=config.token) - #az_token = click.prompt(PROMPT_AZ_TOKEN, default=config.az_token) token = os.environ.get('DATABRICKS_TOKEN') az_token = os.environ.get('DATABRICKS_AZ_TOKEN') resource_id = click.prompt(PROMPT_RESOURCE_ID, default=config.resource_id) @@ -78,7 +76,7 @@ def _configure_cli_password(profile, insecure): @click.command(context_settings=CONTEXT_SETTINGS, short_help='Configures host and authentication info for the CLI.') @click.option('--token', show_default=True, is_flag=True, default=False) -@click.option('--aad-token', show_default=True, is_flag=True, default=False) +@click.option('--az-token', show_default=True, is_flag=True, default=False) @click.option('--insecure', show_default=True, is_flag=True, default=None) @debug_option @profile_option diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py index b1c6959a..6393a808 100644 --- a/databricks_cli/secrets/cli.py +++ b/databricks_cli/secrets/cli.py @@ -44,7 +44,7 @@ short_help="Creates a secret scope.") @click.option('--scope', required=True, type=SecretScopeClickType(), help=SecretScopeClickType.help) @click.option('--initial-manage-principal', - help='Sushi The initial principal that can manage the created secret scope.' + help='The initial principal that can manage the created secret scope.' ' If specified, the initial ACL with MANAGE permission applied to the scope is' ' assigned to the supplied principal (user or group). Currently, the only supported' ' principal for this option is the group "users", which contains all users in the' @@ -52,7 +52,8 @@ ' scope is assigned to the request issuer\'s user identity.') @click.option('--scope-backend-type', type=click.Choice(['azure_keyvault', 'databricks'], case_sensitive=True), default='databricks', help='The backend that will be used for this secret scope. ' - 'Options are: 1) \'azure_keyvault\' and 2) \'databricks\' (default option)') + 'Options are (case-sensitive): 1) \'azure_keyvault\' and 2) \'databricks\' ' + '(default option)') @click.option('--subscription-id', default=None, type=click.STRING, help='The subscription ID associated with the azure keyvault to be used as the backend' ' for the secret scope. NOTE: Only use with azure-keyvault as backend') From 246408ae47f311d4c3952c2d7992f6d32f3b74e2 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Tue, 22 Sep 2020 14:47:12 -0700 Subject: [PATCH 03/13] improved functionality for input prompts --- databricks_cli/cli.py | 2 +- databricks_cli/configure/cli.py | 24 ++++++++++++++++++++---- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index 5825efc4..0b0a8f1b 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -67,4 +67,4 @@ def cli(): cli.add_command(pipelines_group, name='pipelines') if __name__ == "__main__": - cli() \ No newline at end of file + cli() diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index dc07b276..cdf7d809 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -37,6 +37,11 @@ PROMPT_USERNAME = 'Username' PROMPT_PASSWORD = 'Password' # NOQA PROMPT_TOKEN = 'Token' # NOQA +ENV_TOKEN = 'DATABRICKS_TOKEN' +ENV_AZ_TOKEN = 'DATABRICKS_AZ_TOKEN' +PROMPT_ENV_TOKEN = 'Have you correctly set the environment variable (\'DATABRICKS_TOKEN\') for the Bearer Token? (y/n)' +PROMPT_ENV_AZ_TOKEN = 'Have you correctly set the environment variable (\'DATABRICKS_AZ_TOKEN\') ' \ + 'for the Azure Token? (y/n)' def _configure_cli_token(profile, insecure): @@ -50,8 +55,19 @@ def _configure_cli_token(profile, insecure): def _configure_cli_aad_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) - token = os.environ.get('DATABRICKS_TOKEN') - az_token = os.environ.get('DATABRICKS_AZ_TOKEN') + + is_token_env_set = click.prompt(PROMPT_ENV_TOKEN, type=bool) + if not is_token_env_set or not ENV_TOKEN in os.environ: + print('Set Environment Variable \'DATABRICKS_TOKEN\' with your Bearer Token and run again.') + return + token = os.environ.get(ENV_TOKEN) + + is_az_token_env_set = click.prompt(PROMPT_ENV_AZ_TOKEN, type=bool) + if not is_az_token_env_set or not ENV_AZ_TOKEN in os.environ: + print('Set Environment Variable \'DATABRICKS_AZ_TOKEN\' with your Azure Token and run again.') + return + az_token = os.environ.get(ENV_AZ_TOKEN) + resource_id = click.prompt(PROMPT_RESOURCE_ID, default=config.resource_id) new_config = DatabricksConfig.from_aad_token(host, token, az_token, resource_id, insecure) update_and_persist_config(profile, new_config) @@ -80,7 +96,7 @@ def _configure_cli_password(profile, insecure): @click.option('--insecure', show_default=True, is_flag=True, default=None) @debug_option @profile_option -def configure_cli(token, aad_token, insecure): +def configure_cli(token, az_token, insecure): """ Configures host and authentication info for the CLI. """ @@ -88,7 +104,7 @@ def configure_cli(token, aad_token, insecure): insecure_str = str(insecure) if insecure is not None else None if token: _configure_cli_token(profile, insecure_str) - elif aad_token: + elif az_token: _configure_cli_aad_token(profile, insecure_str) else: _configure_cli_password(profile, insecure_str) From 7547177efc52010a8424a447c4270d8db4dce9c0 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Wed, 23 Sep 2020 10:34:26 -0700 Subject: [PATCH 04/13] changed workflow of accepting tokens --- databricks_cli/configure/cli.py | 33 ++++++++++++++++++++------------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index cdf7d809..8eeae825 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -47,27 +47,34 @@ def _configure_cli_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) - token = click.prompt(PROMPT_TOKEN, default=config.token, hide_input=True) + token = click.prompt(PROMPT_TOKEN, default=config.token) new_config = DatabricksConfig.from_token(host, token, insecure) update_and_persist_config(profile, new_config) -def _configure_cli_aad_token(profile, insecure): +def _configure_cli_az_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() - host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) - is_token_env_set = click.prompt(PROMPT_ENV_TOKEN, type=bool) - if not is_token_env_set or not ENV_TOKEN in os.environ: - print('Set Environment Variable \'DATABRICKS_TOKEN\' with your Bearer Token and run again.') + if ENV_TOKEN not in os.environ: + print('[ERROR] Set Environment Variable \'DATABRICKS_TOKEN\' with your Bearer Token and run again.\n') + print('Commands to run to get your Bearer token:\n' + '\t az login\n' + '\t token_response=$(az account get-access-token --resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n' + '\t export DATABRICKS_TOKEN=$(jq .accessToken -r <<< "$token_response")\n' + ) return - token = os.environ.get(ENV_TOKEN) - - is_az_token_env_set = click.prompt(PROMPT_ENV_AZ_TOKEN, type=bool) - if not is_az_token_env_set or not ENV_AZ_TOKEN in os.environ: - print('Set Environment Variable \'DATABRICKS_AZ_TOKEN\' with your Azure Token and run again.') + if ENV_AZ_TOKEN not in os.environ: + print('[ERROR] Set Environment Variable \'DATABRICKS_AZ_TOKEN\' with your Azure Token and run again.\n') + print('Commands to run to get your Azure token:\n' + '\t az login\n' + '\t token_response=$(az account get-access-token --resource https://management.core.windows.net/)\n' + '\t export DATABRICKS_AZ_TOKEN=$(jq .accessToken -r <<< "$token_response")\n' + ) return - az_token = os.environ.get(ENV_AZ_TOKEN) + host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) + token = os.environ.get(ENV_TOKEN) + az_token = os.environ.get(ENV_AZ_TOKEN) resource_id = click.prompt(PROMPT_RESOURCE_ID, default=config.resource_id) new_config = DatabricksConfig.from_aad_token(host, token, az_token, resource_id, insecure) update_and_persist_config(profile, new_config) @@ -105,7 +112,7 @@ def configure_cli(token, az_token, insecure): if token: _configure_cli_token(profile, insecure_str) elif az_token: - _configure_cli_aad_token(profile, insecure_str) + _configure_cli_az_token(profile, insecure_str) else: _configure_cli_password(profile, insecure_str) From 085adbf1a2b959c1e727e01263cde87f7dd37b44 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Thu, 24 Sep 2020 16:47:24 -0700 Subject: [PATCH 05/13] autogenerated sdk and removed azToken and resourceID --- databricks_cli/configure/cli.py | 39 ++-- databricks_cli/configure/config.py | 5 +- databricks_cli/configure/provider.py | 24 +-- databricks_cli/sdk/__init__.py | 2 +- databricks_cli/sdk/api_client.py | 12 +- databricks_cli/sdk/service.py | 279 +++++++++++++++++---------- databricks_cli/sdk/version.py | 0 databricks_cli/secrets/api.py | 4 +- databricks_cli/secrets/cli.py | 16 +- 9 files changed, 205 insertions(+), 176 deletions(-) mode change 100644 => 100755 databricks_cli/sdk/__init__.py mode change 100644 => 100755 databricks_cli/sdk/api_client.py mode change 100644 => 100755 databricks_cli/sdk/service.py mode change 100644 => 100755 databricks_cli/sdk/version.py diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index 8eeae825..90c38ad7 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -33,15 +33,10 @@ PROMPT_HOST = 'Databricks Host (should begin with https://)' PROMPT_RESOURCE_ID = 'Resource/Workspace ID' -PROMPT_AZ_TOKEN = 'Azure Token' PROMPT_USERNAME = 'Username' PROMPT_PASSWORD = 'Password' # NOQA PROMPT_TOKEN = 'Token' # NOQA -ENV_TOKEN = 'DATABRICKS_TOKEN' -ENV_AZ_TOKEN = 'DATABRICKS_AZ_TOKEN' -PROMPT_ENV_TOKEN = 'Have you correctly set the environment variable (\'DATABRICKS_TOKEN\') for the Bearer Token? (y/n)' -PROMPT_ENV_AZ_TOKEN = 'Have you correctly set the environment variable (\'DATABRICKS_AZ_TOKEN\') ' \ - 'for the Azure Token? (y/n)' +ENV_AAD_TOKEN = 'DATABRICKS_TOKEN' def _configure_cli_token(profile, insecure): @@ -52,31 +47,21 @@ def _configure_cli_token(profile, insecure): update_and_persist_config(profile, new_config) -def _configure_cli_az_token(profile, insecure): +def _configure_cli_aad_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() - if ENV_TOKEN not in os.environ: - print('[ERROR] Set Environment Variable \'DATABRICKS_TOKEN\' with your Bearer Token and run again.\n') - print('Commands to run to get your Bearer token:\n' + if ENV_AAD_TOKEN not in os.environ: + print('[ERROR] Set Environment Variable \'%s\' with your AAD Token and run again.\n' % ENV_AAD_TOKEN) + print('Commands to run to get your AAD token:\n' '\t az login\n' '\t token_response=$(az account get-access-token --resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n' - '\t export DATABRICKS_TOKEN=$(jq .accessToken -r <<< "$token_response")\n' - ) - return - if ENV_AZ_TOKEN not in os.environ: - print('[ERROR] Set Environment Variable \'DATABRICKS_AZ_TOKEN\' with your Azure Token and run again.\n') - print('Commands to run to get your Azure token:\n' - '\t az login\n' - '\t token_response=$(az account get-access-token --resource https://management.core.windows.net/)\n' - '\t export DATABRICKS_AZ_TOKEN=$(jq .accessToken -r <<< "$token_response")\n' + '\t export %s=$(jq .accessToken -r <<< "$token_response")\n' % ENV_AAD_TOKEN ) return host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) - token = os.environ.get(ENV_TOKEN) - az_token = os.environ.get(ENV_AZ_TOKEN) - resource_id = click.prompt(PROMPT_RESOURCE_ID, default=config.resource_id) - new_config = DatabricksConfig.from_aad_token(host, token, az_token, resource_id, insecure) + aad_token = os.environ.get(ENV_AAD_TOKEN) + new_config = DatabricksConfig.from_token(host, aad_token, insecure) update_and_persist_config(profile, new_config) @@ -99,11 +84,11 @@ def _configure_cli_password(profile, insecure): @click.command(context_settings=CONTEXT_SETTINGS, short_help='Configures host and authentication info for the CLI.') @click.option('--token', show_default=True, is_flag=True, default=False) -@click.option('--az-token', show_default=True, is_flag=True, default=False) +@click.option('--aad-token', show_default=True, is_flag=True, default=False) @click.option('--insecure', show_default=True, is_flag=True, default=None) @debug_option @profile_option -def configure_cli(token, az_token, insecure): +def configure_cli(token, aad_token, insecure): """ Configures host and authentication info for the CLI. """ @@ -111,8 +96,8 @@ def configure_cli(token, az_token, insecure): insecure_str = str(insecure) if insecure is not None else None if token: _configure_cli_token(profile, insecure_str) - elif az_token: - _configure_cli_az_token(profile, insecure_str) + elif aad_token: + _configure_cli_aad_token(profile, insecure_str) else: _configure_cli_password(profile, insecure_str) diff --git a/databricks_cli/configure/config.py b/databricks_cli/configure/config.py index 65bd33c7..f6ca04e4 100644 --- a/databricks_cli/configure/config.py +++ b/databricks_cli/configure/config.py @@ -83,10 +83,7 @@ def callback(ctx, param, value): # NOQA def _get_api_client(config, command_name=""): verify = config.insecure is None - if config.is_valid_with_az_token: - return ApiClient(host=config.host, token=config.token, az_token=config.az_token, - resource_id=config.resource_id, verify=verify, command_name=command_name) - elif config.is_valid_with_token: + if config.is_valid_with_token: return ApiClient(host=config.host, token=config.token, verify=verify, command_name=command_name) return ApiClient(user=config.username, password=config.password, diff --git a/databricks_cli/configure/provider.py b/databricks_cli/configure/provider.py index 1bf6fc57..e1a1eeac 100644 --- a/databricks_cli/configure/provider.py +++ b/databricks_cli/configure/provider.py @@ -36,8 +36,6 @@ PASSWORD = 'password' # NOQA TOKEN = 'token' INSECURE = 'insecure' -AZ_TOKEN = 'az_token' -RESOURCE_ID = 'resource_id' DEFAULT_SECTION = 'DEFAULT' # User-provided override for the DatabricksConfigProvider @@ -99,8 +97,6 @@ def update_and_persist_config(profile, databricks_config): _set_option(raw_config, profile, PASSWORD, databricks_config.password) _set_option(raw_config, profile, TOKEN, databricks_config.token) _set_option(raw_config, profile, INSECURE, databricks_config.insecure) - _set_option(raw_config, profile, AZ_TOKEN, databricks_config.az_token) - _set_option(raw_config, profile, RESOURCE_ID, databricks_config.resource_id) _overwrite_config(raw_config) @@ -242,9 +238,7 @@ def get_config(self): password = os.environ.get('DATABRICKS_PASSWORD') token = os.environ.get('DATABRICKS_TOKEN') insecure = os.environ.get('DATABRICKS_INSECURE') - az_token = os.environ.get('DATABRICKS_AZ_TOKEN') - resource_id = os.environ.get('DATABRICKS_RESOURCE_ID') - config = DatabricksConfig(host, username, password, token, insecure, az_token, resource_id) + config = DatabricksConfig(host, username, password, token, insecure) if config.is_valid: return config return None @@ -262,32 +256,24 @@ def get_config(self): password = _get_option_if_exists(raw_config, self.profile, PASSWORD) token = _get_option_if_exists(raw_config, self.profile, TOKEN) insecure = _get_option_if_exists(raw_config, self.profile, INSECURE) - az_token = _get_option_if_exists(raw_config, self.profile, AZ_TOKEN) - resource_id = _get_option_if_exists(raw_config, self.profile, RESOURCE_ID) - config = DatabricksConfig(host, username, password, token, insecure, az_token, resource_id) + config = DatabricksConfig(host, username, password, token, insecure) if config.is_valid: return config return None class DatabricksConfig(object): - def __init__(self, host, username, password, token, insecure, az_token=None, resource_id=None): # noqa + def __init__(self, host, username, password, token, insecure): # noqa self.host = host self.username = username self.password = password self.token = token self.insecure = insecure - self.az_token = az_token - self.resource_id = resource_id @classmethod def from_token(cls, host, token, insecure=None): return DatabricksConfig(host, None, None, token, insecure) - @classmethod - def from_aad_token(cls, host, token, az_token, resource_id, insecure=None): - return DatabricksConfig(host, None, None, token, insecure, az_token, resource_id) - @classmethod def from_password(cls, host, username, password, insecure=None): return DatabricksConfig(host, username, password, None, insecure) @@ -300,10 +286,6 @@ def empty(cls): def is_valid_with_token(self): return self.host is not None and self.token is not None - @property - def is_valid_with_az_token(self): - return self.host is not None and self.token is not None and self.az_token is not None and self.resource_id is not None - @property def is_valid_with_password(self): return self.host is not None and self.username is not None and self.password is not None diff --git a/databricks_cli/sdk/__init__.py b/databricks_cli/sdk/__init__.py old mode 100644 new mode 100755 index 2e42a61c..67c76dc0 --- a/databricks_cli/sdk/__init__.py +++ b/databricks_cli/sdk/__init__.py @@ -24,7 +24,7 @@ """ Databricks Python REST Client 2.0 for interacting with various services. -Currently supports services including clusters, clusters policies and jobs. +Currently supports services including clusters and jobs. Requires Python 2.7.9 or above. diff --git a/databricks_cli/sdk/api_client.py b/databricks_cli/sdk/api_client.py old mode 100644 new mode 100755 index dcbd1aa9..fc3c0523 --- a/databricks_cli/sdk/api_client.py +++ b/databricks_cli/sdk/api_client.py @@ -65,7 +65,7 @@ class ApiClient(object): A partial Python implementation of dbc rest api to be used by different versions of the client. """ - def __init__(self, user=None, password=None, host=None, token=None, az_token=None, resource_id=None, + def __init__(self, user=None, password=None, host=None, token=None, apiVersion=version.API_VERSION, default_headers={}, verify=True, command_name=""): if host[-1] == "/": host = host[:-1] @@ -81,10 +81,6 @@ def __init__(self, user=None, password=None, host=None, token=None, az_token=Non encoded_auth = (user + ":" + password).encode() user_header_data = "Basic " + base64.standard_b64encode(encoded_auth).decode() auth = {'Authorization': user_header_data, 'Content-Type': 'text/json'} - elif token is not None and az_token is not None and resource_id is not None: - auth = {'Authorization': 'Bearer {}'.format(token), 'Content-Type': 'text/json', - 'X-Databricks-Azure-SP-Management-Token': '{}'.format(az_token), - 'X-Databricks-Azure-Workspace-Resource-Id': '{}'.format(resource_id)} elif token is not None: auth = {'Authorization': 'Bearer {}'.format(token), 'Content-Type': 'text/json'} else: @@ -117,10 +113,10 @@ def perform_query(self, method, path, data = {}, headers = None): if method == 'GET': translated_data = {k: _translate_boolean_to_query_param(data[k]) for k in data} resp = self.session.request(method, self.url + path, params = translated_data, - verify = self.verify, headers = headers) + verify = self.verify, headers = headers) else: resp = self.session.request(method, self.url + path, data = json.dumps(data), - verify = self.verify, headers = headers) + verify = self.verify, headers = headers) try: resp.raise_for_status() except requests.exceptions.HTTPError as e: @@ -141,4 +137,4 @@ def _translate_boolean_to_query_param(value): return 'true' else: return 'false' - return value + return value \ No newline at end of file diff --git a/databricks_cli/sdk/service.py b/databricks_cli/sdk/service.py old mode 100644 new mode 100755 index 56c5423d..0933888a --- a/databricks_cli/sdk/service.py +++ b/databricks_cli/sdk/service.py @@ -78,7 +78,7 @@ def create_job(self, name=None, existing_cluster_id=None, new_cluster=None, libr if max_concurrent_runs is not None: _data['max_concurrent_runs'] = max_concurrent_runs return self.client.perform_query('POST', '/jobs/create', data=_data, headers=headers) - + def submit_run(self, run_name=None, existing_cluster_id=None, new_cluster=None, libraries=None, notebook_task=None, spark_jar_task=None, spark_python_task=None, spark_submit_task=None, timeout_seconds=None, headers=None): @@ -112,7 +112,7 @@ def submit_run(self, run_name=None, existing_cluster_id=None, new_cluster=None, if timeout_seconds is not None: _data['timeout_seconds'] = timeout_seconds return self.client.perform_query('POST', '/jobs/runs/submit', data=_data, headers=headers) - + def reset_job(self, job_id, new_settings, headers=None): _data = {} if job_id is not None: @@ -122,24 +122,24 @@ def reset_job(self, job_id, new_settings, headers=None): if not isinstance(new_settings, dict): raise TypeError('Expected databricks.JobSettings() or dict for field new_settings') return self.client.perform_query('POST', '/jobs/reset', data=_data, headers=headers) - + def delete_job(self, job_id, headers=None): _data = {} if job_id is not None: _data['job_id'] = job_id return self.client.perform_query('POST', '/jobs/delete', data=_data, headers=headers) - + def get_job(self, job_id, headers=None): _data = {} if job_id is not None: _data['job_id'] = job_id return self.client.perform_query('GET', '/jobs/get', data=_data, headers=headers) - + def list_jobs(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/jobs/list', data=_data, headers=headers) - + def run_now(self, job_id=None, jar_params=None, notebook_params=None, python_params=None, spark_submit_params=None, headers=None): _data = {} @@ -154,7 +154,7 @@ def run_now(self, job_id=None, jar_params=None, notebook_params=None, python_par if spark_submit_params is not None: _data['spark_submit_params'] = spark_submit_params return self.client.perform_query('POST', '/jobs/run-now', data=_data, headers=headers) - + def list_runs(self, job_id=None, active_only=None, completed_only=None, offset=None, limit=None, headers=None): _data = {} @@ -169,31 +169,31 @@ def list_runs(self, job_id=None, active_only=None, completed_only=None, offset=N if limit is not None: _data['limit'] = limit return self.client.perform_query('GET', '/jobs/runs/list', data=_data, headers=headers) - + def get_run(self, run_id=None, headers=None): _data = {} if run_id is not None: _data['run_id'] = run_id return self.client.perform_query('GET', '/jobs/runs/get', data=_data, headers=headers) - + def delete_run(self, run_id=None, headers=None): _data = {} if run_id is not None: _data['run_id'] = run_id return self.client.perform_query('POST', '/jobs/runs/delete', data=_data, headers=headers) - + def cancel_run(self, run_id, headers=None): _data = {} if run_id is not None: _data['run_id'] = run_id return self.client.perform_query('POST', '/jobs/runs/cancel', data=_data, headers=headers) - + def get_run_output(self, run_id, headers=None): _data = {} if run_id is not None: _data['run_id'] = run_id return self.client.perform_query('GET', '/jobs/runs/get-output', data=_data, headers=headers) - + def export_run(self, run_id, views_to_export=None, headers=None): _data = {} if run_id is not None: @@ -201,7 +201,7 @@ def export_run(self, run_id, views_to_export=None, headers=None): if views_to_export is not None: _data['views_to_export'] = views_to_export return self.client.perform_query('GET', '/jobs/runs/export', data=_data, headers=headers) - + class ClusterService(object): def __init__(self, client): @@ -209,15 +209,15 @@ def __init__(self, client): def list_clusters(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/clusters/list', data=_data, headers=headers) - + def create_cluster(self, num_workers=None, autoscale=None, cluster_name=None, spark_version=None, spark_conf=None, aws_attributes=None, node_type_id=None, driver_node_type_id=None, ssh_public_keys=None, custom_tags=None, - cluster_log_conf=None, init_scripts=None, spark_env_vars=None, - autotermination_minutes=None, enable_elastic_disk=None, cluster_source=None, - instance_pool_id=None, headers=None): + cluster_log_conf=None, spark_env_vars=None, autotermination_minutes=None, + enable_elastic_disk=None, cluster_source=None, instance_pool_id=None, + headers=None): _data = {} if num_workers is not None: _data['num_workers'] = num_workers @@ -247,8 +247,6 @@ def create_cluster(self, num_workers=None, autoscale=None, cluster_name=None, sp _data['cluster_log_conf'] = cluster_log_conf if not isinstance(cluster_log_conf, dict): raise TypeError('Expected databricks.ClusterLogConf() or dict for field cluster_log_conf') - if init_scripts is not None: - _data['init_scripts'] = init_scripts if spark_env_vars is not None: _data['spark_env_vars'] = spark_env_vars if autotermination_minutes is not None: @@ -260,36 +258,36 @@ def create_cluster(self, num_workers=None, autoscale=None, cluster_name=None, sp if instance_pool_id is not None: _data['instance_pool_id'] = instance_pool_id return self.client.perform_query('POST', '/clusters/create', data=_data, headers=headers) - + def start_cluster(self, cluster_id, headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('POST', '/clusters/start', data=_data, headers=headers) - + def list_spark_versions(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/clusters/spark-versions', data=_data, headers=headers) - + def delete_cluster(self, cluster_id, headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('POST', '/clusters/delete', data=_data, headers=headers) - + def permanent_delete_cluster(self, cluster_id, headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('POST', '/clusters/permanent-delete', data=_data, headers=headers) - + def restart_cluster(self, cluster_id, headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('POST', '/clusters/restart', data=_data, headers=headers) - + def resize_cluster(self, cluster_id, num_workers=None, autoscale=None, headers=None): _data = {} if cluster_id is not None: @@ -301,13 +299,13 @@ def resize_cluster(self, cluster_id, num_workers=None, autoscale=None, headers=N if not isinstance(autoscale, dict): raise TypeError('Expected databricks.AutoScale() or dict for field autoscale') return self.client.perform_query('POST', '/clusters/resize', data=_data, headers=headers) - + def edit_cluster(self, cluster_id, num_workers=None, autoscale=None, cluster_name=None, spark_version=None, spark_conf=None, aws_attributes=None, node_type_id=None, driver_node_type_id=None, ssh_public_keys=None, custom_tags=None, - cluster_log_conf=None, init_scripts=None, spark_env_vars=None, - autotermination_minutes=None, enable_elastic_disk=None, cluster_source=None, - instance_pool_id=None, headers=None): + cluster_log_conf=None, spark_env_vars=None, autotermination_minutes=None, + enable_elastic_disk=None, cluster_source=None, instance_pool_id=None, + headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id @@ -339,8 +337,6 @@ def edit_cluster(self, cluster_id, num_workers=None, autoscale=None, cluster_nam _data['cluster_log_conf'] = cluster_log_conf if not isinstance(cluster_log_conf, dict): raise TypeError('Expected databricks.ClusterLogConf() or dict for field cluster_log_conf') - if init_scripts is not None: - _data['init_scripts'] = init_scripts if spark_env_vars is not None: _data['spark_env_vars'] = spark_env_vars if autotermination_minutes is not None: @@ -352,35 +348,35 @@ def edit_cluster(self, cluster_id, num_workers=None, autoscale=None, cluster_nam if instance_pool_id is not None: _data['instance_pool_id'] = instance_pool_id return self.client.perform_query('POST', '/clusters/edit', data=_data, headers=headers) - + def get_cluster(self, cluster_id, headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('GET', '/clusters/get', data=_data, headers=headers) - + def pin_cluster(self, cluster_id, headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('POST', '/clusters/pin', data=_data, headers=headers) - + def unpin_cluster(self, cluster_id, headers=None): _data = {} if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('POST', '/clusters/unpin', data=_data, headers=headers) - + def list_node_types(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/clusters/list-node-types', data=_data, headers=headers) - + def list_available_zones(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/clusters/list-zones', data=_data, headers=headers) - + def get_events(self, cluster_id, start_time=None, end_time=None, order=None, event_types=None, offset=None, limit=None, headers=None): _data = {} @@ -400,6 +396,7 @@ def get_events(self, cluster_id, start_time=None, end_time=None, order=None, eve _data['limit'] = limit return self.client.perform_query('POST', '/clusters/events', data=_data, headers=headers) + class PolicyService(object): def __init__(self, client): self.client = client @@ -452,12 +449,12 @@ def cluster_status(self, cluster_id, headers=None): if cluster_id is not None: _data['cluster_id'] = cluster_id return self.client.perform_query('GET', '/libraries/cluster-status', data=_data, headers=headers) - + def all_cluster_statuses(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/libraries/all-cluster-statuses', data=_data, headers=headers) - + def install_libraries(self, cluster_id, libraries=None, headers=None): _data = {} if cluster_id is not None: @@ -465,7 +462,7 @@ def install_libraries(self, cluster_id, libraries=None, headers=None): if libraries is not None: _data['libraries'] = libraries return self.client.perform_query('POST', '/libraries/install', data=_data, headers=headers) - + def uninstall_libraries(self, cluster_id, libraries=None, headers=None): _data = {} if cluster_id is not None: @@ -473,7 +470,7 @@ def uninstall_libraries(self, cluster_id, libraries=None, headers=None): if libraries is not None: _data['libraries'] = libraries return self.client.perform_query('POST', '/libraries/uninstall', data=_data, headers=headers) - + class DbfsService(object): def __init__(self, client): @@ -488,19 +485,41 @@ def read(self, path, offset=None, length=None, headers=None): if length is not None: _data['length'] = length return self.client.perform_query('GET', '/dbfs/read', data=_data, headers=headers) - + + def read_test(self, path, offset=None, length=None, headers=None): + _data = {} + if path is not None: + _data['path'] = path + if offset is not None: + _data['offset'] = offset + if length is not None: + _data['length'] = length + return self.client.perform_query('GET', '/dbfs-testing/read', data=_data, headers=headers) + def get_status(self, path, headers=None): _data = {} if path is not None: _data['path'] = path return self.client.perform_query('GET', '/dbfs/get-status', data=_data, headers=headers) - + + def get_status_test(self, path, headers=None): + _data = {} + if path is not None: + _data['path'] = path + return self.client.perform_query('GET', '/dbfs-testing/get-status', data=_data, headers=headers) + def list(self, path, headers=None): _data = {} if path is not None: _data['path'] = path return self.client.perform_query('GET', '/dbfs/list', data=_data, headers=headers) - + + def list_test(self, path, headers=None): + _data = {} + if path is not None: + _data['path'] = path + return self.client.perform_query('GET', '/dbfs-testing/list', data=_data, headers=headers) + def put(self, path, contents=None, overwrite=None, headers=None): _data = {} if path is not None: @@ -510,13 +529,29 @@ def put(self, path, contents=None, overwrite=None, headers=None): if overwrite is not None: _data['overwrite'] = overwrite return self.client.perform_query('POST', '/dbfs/put', data=_data, headers=headers) - + + def put_test(self, path, contents=None, overwrite=None, headers=None): + _data = {} + if path is not None: + _data['path'] = path + if contents is not None: + _data['contents'] = contents + if overwrite is not None: + _data['overwrite'] = overwrite + return self.client.perform_query('POST', '/dbfs-testing/put', data=_data, headers=headers) + def mkdirs(self, path, headers=None): _data = {} if path is not None: _data['path'] = path return self.client.perform_query('POST', '/dbfs/mkdirs', data=_data, headers=headers) - + + def mkdirs_test(self, path, headers=None): + _data = {} + if path is not None: + _data['path'] = path + return self.client.perform_query('POST', '/dbfs-testing/mkdirs', data=_data, headers=headers) + def move(self, source_path, destination_path, headers=None): _data = {} if source_path is not None: @@ -524,7 +559,15 @@ def move(self, source_path, destination_path, headers=None): if destination_path is not None: _data['destination_path'] = destination_path return self.client.perform_query('POST', '/dbfs/move', data=_data, headers=headers) - + + def move_test(self, source_path, destination_path, headers=None): + _data = {} + if source_path is not None: + _data['source_path'] = source_path + if destination_path is not None: + _data['destination_path'] = destination_path + return self.client.perform_query('POST', '/dbfs-testing/move', data=_data, headers=headers) + def delete(self, path, recursive=None, headers=None): _data = {} if path is not None: @@ -532,7 +575,15 @@ def delete(self, path, recursive=None, headers=None): if recursive is not None: _data['recursive'] = recursive return self.client.perform_query('POST', '/dbfs/delete', data=_data, headers=headers) - + + def delete_test(self, path, recursive=None, headers=None): + _data = {} + if path is not None: + _data['path'] = path + if recursive is not None: + _data['recursive'] = recursive + return self.client.perform_query('POST', '/dbfs-testing/delete', data=_data, headers=headers) + def create(self, path, overwrite=None, headers=None): _data = {} if path is not None: @@ -540,7 +591,15 @@ def create(self, path, overwrite=None, headers=None): if overwrite is not None: _data['overwrite'] = overwrite return self.client.perform_query('POST', '/dbfs/create', data=_data, headers=headers) - + + def create_test(self, path, overwrite=None, headers=None): + _data = {} + if path is not None: + _data['path'] = path + if overwrite is not None: + _data['overwrite'] = overwrite + return self.client.perform_query('POST', '/dbfs-testing/create', data=_data, headers=headers) + def add_block(self, handle, data, headers=None): _data = {} if handle is not None: @@ -548,13 +607,27 @@ def add_block(self, handle, data, headers=None): if data is not None: _data['data'] = data return self.client.perform_query('POST', '/dbfs/add-block', data=_data, headers=headers) - + + def add_block_test(self, handle, data, headers=None): + _data = {} + if handle is not None: + _data['handle'] = handle + if data is not None: + _data['data'] = data + return self.client.perform_query('POST', '/dbfs-testing/add-block', data=_data, headers=headers) + def close(self, handle, headers=None): _data = {} if handle is not None: _data['handle'] = handle return self.client.perform_query('POST', '/dbfs/close', data=_data, headers=headers) - + + def close_test(self, handle, headers=None): + _data = {} + if handle is not None: + _data['handle'] = handle + return self.client.perform_query('POST', '/dbfs-testing/close', data=_data, headers=headers) + class WorkspaceService(object): def __init__(self, client): @@ -565,13 +638,13 @@ def mkdirs(self, path, headers=None): if path is not None: _data['path'] = path return self.client.perform_query('POST', '/workspace/mkdirs', data=_data, headers=headers) - + def list(self, path, headers=None): _data = {} if path is not None: _data['path'] = path return self.client.perform_query('GET', '/workspace/list', data=_data, headers=headers) - + def import_workspace(self, path, format=None, language=None, content=None, overwrite=None, headers=None): _data = {} @@ -586,7 +659,7 @@ def import_workspace(self, path, format=None, language=None, content=None, overw if overwrite is not None: _data['overwrite'] = overwrite return self.client.perform_query('POST', '/workspace/import', data=_data, headers=headers) - + def export_workspace(self, path, format=None, direct_download=None, headers=None): _data = {} if path is not None: @@ -596,7 +669,7 @@ def export_workspace(self, path, format=None, direct_download=None, headers=None if direct_download is not None: _data['direct_download'] = direct_download return self.client.perform_query('GET', '/workspace/export', data=_data, headers=headers) - + def delete(self, path, recursive=None, headers=None): _data = {} if path is not None: @@ -604,47 +677,44 @@ def delete(self, path, recursive=None, headers=None): if recursive is not None: _data['recursive'] = recursive return self.client.perform_query('POST', '/workspace/delete', data=_data, headers=headers) - + def get_status(self, path, headers=None): _data = {} if path is not None: _data['path'] = path return self.client.perform_query('GET', '/workspace/get-status', data=_data, headers=headers) - + class SecretService(object): def __init__(self, client): self.client = client - def create_scope(self, scope, initial_manage_principal=None, scope_backend_type=None, azure_keyvault=None, - headers=None): + def create_scope(self, scope, initial_manage_principal=None, scope_backend_type=None, + backend_azure_keyvault=None, headers=None): _data = {} if scope is not None: _data['scope'] = scope if initial_manage_principal is not None: _data['initial_manage_principal'] = initial_manage_principal if scope_backend_type is not None: - if scope_backend_type == 'databricks': - _data['scope_backend_type_str'] = 1 - elif azure_keyvault is not None and scope_backend_type == 'azure_keyvault': - _data['scope_backend_type'] = 2 - _data['backend_azure_keyvault'] = { - 'resource_id': str(azure_keyvault['resource_id']), - 'dns_name': str(azure_keyvault['dns_name']), - } + _data['scope_backend_type'] = scope_backend_type + if backend_azure_keyvault is not None: + _data['backend_azure_keyvault'] = backend_azure_keyvault + if not isinstance(backend_azure_keyvault, dict): + raise TypeError('Expected databricks.AzureKeyVaultSecretScopeMetadata() or dict for field backend_azure_keyvault') return self.client.perform_query('POST', '/secrets/scopes/create', data=_data, headers=headers) - + def delete_scope(self, scope, headers=None): _data = {} if scope is not None: _data['scope'] = scope return self.client.perform_query('POST', '/secrets/scopes/delete', data=_data, headers=headers) - + def list_scopes(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/secrets/scopes/list', data=_data, headers=headers) - + def put_secret(self, scope, key, string_value=None, bytes_value=None, headers=None): _data = {} if scope is not None: @@ -656,7 +726,7 @@ def put_secret(self, scope, key, string_value=None, bytes_value=None, headers=No if bytes_value is not None: _data['bytes_value'] = bytes_value return self.client.perform_query('POST', '/secrets/put', data=_data, headers=headers) - + def delete_secret(self, scope, key, headers=None): _data = {} if scope is not None: @@ -664,13 +734,13 @@ def delete_secret(self, scope, key, headers=None): if key is not None: _data['key'] = key return self.client.perform_query('POST', '/secrets/delete', data=_data, headers=headers) - + def list_secrets(self, scope, headers=None): _data = {} if scope is not None: _data['scope'] = scope return self.client.perform_query('GET', '/secrets/list', data=_data, headers=headers) - + def put_acl(self, scope, principal, permission, headers=None): _data = {} if scope is not None: @@ -680,7 +750,7 @@ def put_acl(self, scope, principal, permission, headers=None): if permission is not None: _data['permission'] = permission return self.client.perform_query('POST', '/secrets/acls/put', data=_data, headers=headers) - + def delete_acl(self, scope, principal, headers=None): _data = {} if scope is not None: @@ -688,13 +758,13 @@ def delete_acl(self, scope, principal, headers=None): if principal is not None: _data['principal'] = principal return self.client.perform_query('POST', '/secrets/acls/delete', data=_data, headers=headers) - + def list_acls(self, scope, headers=None): _data = {} if scope is not None: _data['scope'] = scope return self.client.perform_query('GET', '/secrets/acls/list', data=_data, headers=headers) - + def get_acl(self, scope, principal, headers=None): _data = {} if scope is not None: @@ -702,7 +772,7 @@ def get_acl(self, scope, principal, headers=None): if principal is not None: _data['principal'] = principal return self.client.perform_query('GET', '/secrets/acls/get', data=_data, headers=headers) - + class GroupsService(object): def __init__(self, client): @@ -713,7 +783,7 @@ def create_group(self, group_name, headers=None): if group_name is not None: _data['group_name'] = group_name return self.client.perform_query('POST', '/groups/create', data=_data, headers=headers) - + def add_to_group(self, parent_name, user_name=None, group_name=None, headers=None): _data = {} if user_name is not None: @@ -723,7 +793,7 @@ def add_to_group(self, parent_name, user_name=None, group_name=None, headers=Non if parent_name is not None: _data['parent_name'] = parent_name return self.client.perform_query('POST', '/groups/add-member', data=_data, headers=headers) - + def remove_from_group(self, parent_name, user_name=None, group_name=None, headers=None): _data = {} if user_name is not None: @@ -733,24 +803,24 @@ def remove_from_group(self, parent_name, user_name=None, group_name=None, header if parent_name is not None: _data['parent_name'] = parent_name return self.client.perform_query('POST', '/groups/remove-member', data=_data, headers=headers) - + def get_groups(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/groups/list', data=_data, headers=headers) - + def get_group_members(self, group_name, headers=None): _data = {} if group_name is not None: _data['group_name'] = group_name return self.client.perform_query('GET', '/groups/list-members', data=_data, headers=headers) - + def remove_group(self, group_name, headers=None): _data = {} if group_name is not None: _data['group_name'] = group_name return self.client.perform_query('POST', '/groups/delete', data=_data, headers=headers) - + def get_groups_for_principal(self, user_name=None, group_name=None, headers=None): _data = {} if user_name is not None: @@ -758,7 +828,7 @@ def get_groups_for_principal(self, user_name=None, group_name=None, headers=None if group_name is not None: _data['group_name'] = group_name return self.client.perform_query('GET', '/groups/list-parents', data=_data, headers=headers) - + class TokenService(object): def __init__(self, client): @@ -771,12 +841,12 @@ def create_token(self, lifetime_seconds=None, comment=None, headers=None): if comment is not None: _data['comment'] = comment return self.client.perform_query('POST', '/token/create', data=_data, headers=headers) - + def list_tokens(self, headers=None): _data = {} - + return self.client.perform_query('GET', '/token/list', data=_data, headers=headers) - + def revoke_token(self, token_id, headers=None): _data = {} if token_id is not None: @@ -900,7 +970,7 @@ def create(self, id=None, name=None, storage=None, configuration=None, clusters= if allow_duplicate_names is not None: _data['allow_duplicate_names'] = allow_duplicate_names return self.client.perform_query('POST', '/pipelines', data=_data, headers=headers) - + def deploy(self, pipeline_id=None, id=None, name=None, storage=None, configuration=None, clusters=None, libraries=None, trigger=None, filters=None, allow_duplicate_names=None, headers=None): @@ -927,21 +997,20 @@ def deploy(self, pipeline_id=None, id=None, name=None, storage=None, configurati raise TypeError('Expected databricks.Filters() or dict for field filters') if allow_duplicate_names is not None: _data['allow_duplicate_names'] = allow_duplicate_names - return self.client.perform_query('PUT', - '/pipelines/{pipeline_id}'.format(pipeline_id=pipeline_id), - data=_data, headers=headers) - + return self.client.perform_query('PUT', '/pipelines/{pipeline_id}'.format(pipeline_id=pipeline_id), data=_data, headers=headers) + def delete(self, pipeline_id=None, headers=None): _data = {} - + return self.client.perform_query('DELETE', '/pipelines/{pipeline_id}'.format(pipeline_id=pipeline_id), data=_data, headers=headers) - + def get(self, pipeline_id=None, headers=None): _data = {} - + return self.client.perform_query('GET', '/pipelines/{pipeline_id}'.format(pipeline_id=pipeline_id), data=_data, headers=headers) - + def reset(self, pipeline_id=None, headers=None): _data = {} - + return self.client.perform_query('POST', '/pipelines/{pipeline_id}/reset'.format(pipeline_id=pipeline_id), data=_data, headers=headers) + \ No newline at end of file diff --git a/databricks_cli/sdk/version.py b/databricks_cli/sdk/version.py old mode 100644 new mode 100755 diff --git a/databricks_cli/secrets/api.py b/databricks_cli/secrets/api.py index fe5b0a01..4ae4658d 100644 --- a/databricks_cli/secrets/api.py +++ b/databricks_cli/secrets/api.py @@ -28,8 +28,8 @@ class SecretApi(object): def __init__(self, api_client): self.client = SecretService(api_client) - def create_scope(self, scope, initial_manage_principal, scope_backend_type, azure_keyvault_info): - return self.client.create_scope(scope, initial_manage_principal, scope_backend_type, azure_keyvault_info) + def create_scope(self, scope, initial_manage_principal, scope_backend_type, backend_azure_keyvault): + return self.client.create_scope(scope, initial_manage_principal, scope_backend_type, backend_azure_keyvault) def delete_scope(self, scope): return self.client.delete_scope(scope) diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py index 6393a808..6a1c27ba 100644 --- a/databricks_cli/secrets/cli.py +++ b/databricks_cli/secrets/cli.py @@ -50,12 +50,12 @@ ' principal for this option is the group "users", which contains all users in the' ' workspace. If not specified, the initial ACL with MANAGE permission applied to the' ' scope is assigned to the request issuer\'s user identity.') -@click.option('--scope-backend-type', type=click.Choice(['azure_keyvault', 'databricks'], case_sensitive=True), - default='databricks', help='The backend that will be used for this secret scope. ' +@click.option('--scope-backend-type', type=click.Choice(['AZURE_KEYVAULT', 'DATABRICKS'], case_sensitive=True), + default='DATABRICKS', help='The backend that will be used for this secret scope. ' 'Options are (case-sensitive): 1) \'azure_keyvault\' and 2) \'databricks\' ' '(default option)') -@click.option('--subscription-id', default=None, type=click.STRING, - help='The subscription ID associated with the azure keyvault to be used as the backend' +@click.option('--resource-id', default=None, type=click.STRING, + help='The resource ID associated with the azure keyvault to be used as the backend' ' for the secret scope. NOTE: Only use with azure-keyvault as backend') @click.option('--dns-name', default=None, type=click.STRING, help='The dns name associated with the azure keyvault to be used as the backed for the' @@ -64,15 +64,15 @@ @profile_option @eat_exceptions @provide_api_client -def create_scope(api_client, scope, initial_manage_principal, scope_backend_type, subscription_id, dns_name): +def create_scope(api_client, scope, initial_manage_principal, scope_backend_type, resource_id, dns_name): """ Creates a new secret scope with given name. """ - azure_keyvault_info = { - 'resource_id': subscription_id, + backend_azure_keyvault = { + 'resource_id': resource_id, 'dns_name': dns_name } - SecretApi(api_client).create_scope(scope, initial_manage_principal, scope_backend_type, azure_keyvault_info) + SecretApi(api_client).create_scope(scope, initial_manage_principal, scope_backend_type, backend_azure_keyvault) def _scopes_to_table(scopes_json): From e2cff965b65a0edfefd62fd5d24f1d22ba934033 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Thu, 24 Sep 2020 16:51:18 -0700 Subject: [PATCH 06/13] added more detailed description for create-scope command --- databricks_cli/secrets/cli.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py index 6a1c27ba..afaff2c7 100644 --- a/databricks_cli/secrets/cli.py +++ b/databricks_cli/secrets/cli.py @@ -53,7 +53,9 @@ @click.option('--scope-backend-type', type=click.Choice(['AZURE_KEYVAULT', 'DATABRICKS'], case_sensitive=True), default='DATABRICKS', help='The backend that will be used for this secret scope. ' 'Options are (case-sensitive): 1) \'azure_keyvault\' and 2) \'databricks\' ' - '(default option)') + '(default option)' + '\nNote: To create an Azure Keyvault, be sure to configure an AAD Token using' + '\'databricks-cli configure --aad-token\'') @click.option('--resource-id', default=None, type=click.STRING, help='The resource ID associated with the azure keyvault to be used as the backend' ' for the secret scope. NOTE: Only use with azure-keyvault as backend') From 888fc15ef6aa7f5c278912315db8a60ca958aff4 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Fri, 25 Sep 2020 11:34:05 -0700 Subject: [PATCH 07/13] ready for review --- databricks_cli/configure/cli.py | 1 - databricks_cli/sdk/__init__.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index 90c38ad7..b6ea8fdc 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -32,7 +32,6 @@ from databricks_cli.configure.config import profile_option, get_profile_from_context, debug_option PROMPT_HOST = 'Databricks Host (should begin with https://)' -PROMPT_RESOURCE_ID = 'Resource/Workspace ID' PROMPT_USERNAME = 'Username' PROMPT_PASSWORD = 'Password' # NOQA PROMPT_TOKEN = 'Token' # NOQA diff --git a/databricks_cli/sdk/__init__.py b/databricks_cli/sdk/__init__.py index 67c76dc0..2e42a61c 100755 --- a/databricks_cli/sdk/__init__.py +++ b/databricks_cli/sdk/__init__.py @@ -24,7 +24,7 @@ """ Databricks Python REST Client 2.0 for interacting with various services. -Currently supports services including clusters and jobs. +Currently supports services including clusters, clusters policies and jobs. Requires Python 2.7.9 or above. From b10d948e4bf003069a1d29c06392cd1fe875b02f Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Mon, 28 Sep 2020 13:27:44 -0700 Subject: [PATCH 08/13] ready for review - reset sdk files to og and changed variable name --- databricks_cli/configure/cli.py | 10 +++++----- databricks_cli/sdk/__init__.py | 0 databricks_cli/sdk/api_client.py | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) mode change 100755 => 100644 databricks_cli/sdk/__init__.py mode change 100755 => 100644 databricks_cli/sdk/api_client.py diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index b6ea8fdc..11a33f27 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -35,7 +35,7 @@ PROMPT_USERNAME = 'Username' PROMPT_PASSWORD = 'Password' # NOQA PROMPT_TOKEN = 'Token' # NOQA -ENV_AAD_TOKEN = 'DATABRICKS_TOKEN' +ENV_DATABRICKS_TOKEN = 'DATABRICKS_TOKEN' def _configure_cli_token(profile, insecure): @@ -49,17 +49,17 @@ def _configure_cli_token(profile, insecure): def _configure_cli_aad_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() - if ENV_AAD_TOKEN not in os.environ: - print('[ERROR] Set Environment Variable \'%s\' with your AAD Token and run again.\n' % ENV_AAD_TOKEN) + if ENV_DATABRICKS_TOKEN not in os.environ: + print('[ERROR] Set Environment Variable \'%s\' with your AAD Token and run again.\n' % ENV_DATABRICKS_TOKEN) print('Commands to run to get your AAD token:\n' '\t az login\n' '\t token_response=$(az account get-access-token --resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n' - '\t export %s=$(jq .accessToken -r <<< "$token_response")\n' % ENV_AAD_TOKEN + '\t export %s=$(jq .accessToken -r <<< "$token_response")\n' % ENV_DATABRICKS_TOKEN ) return host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) - aad_token = os.environ.get(ENV_AAD_TOKEN) + aad_token = os.environ.get(ENV_DATABRICKS_TOKEN) new_config = DatabricksConfig.from_token(host, aad_token, insecure) update_and_persist_config(profile, new_config) diff --git a/databricks_cli/sdk/__init__.py b/databricks_cli/sdk/__init__.py old mode 100755 new mode 100644 diff --git a/databricks_cli/sdk/api_client.py b/databricks_cli/sdk/api_client.py old mode 100755 new mode 100644 index fc3c0523..d5870a6a --- a/databricks_cli/sdk/api_client.py +++ b/databricks_cli/sdk/api_client.py @@ -113,10 +113,10 @@ def perform_query(self, method, path, data = {}, headers = None): if method == 'GET': translated_data = {k: _translate_boolean_to_query_param(data[k]) for k in data} resp = self.session.request(method, self.url + path, params = translated_data, - verify = self.verify, headers = headers) + verify = self.verify, headers = headers) else: resp = self.session.request(method, self.url + path, data = json.dumps(data), - verify = self.verify, headers = headers) + verify = self.verify, headers = headers) try: resp.raise_for_status() except requests.exceptions.HTTPError as e: @@ -137,4 +137,4 @@ def _translate_boolean_to_query_param(value): return 'true' else: return 'false' - return value \ No newline at end of file + return value From 49f83ce82b4a7171f8cbee21f4437590dcde66a2 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Mon, 28 Sep 2020 13:52:25 -0700 Subject: [PATCH 09/13] addressed review comments --- databricks_cli/configure/cli.py | 10 +++++----- databricks_cli/secrets/cli.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index 11a33f27..b06eae23 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -35,7 +35,7 @@ PROMPT_USERNAME = 'Username' PROMPT_PASSWORD = 'Password' # NOQA PROMPT_TOKEN = 'Token' # NOQA -ENV_DATABRICKS_TOKEN = 'DATABRICKS_TOKEN' +ENV_AAD_TOKEN = 'DATABRICKS_AAD_TOKEN' def _configure_cli_token(profile, insecure): @@ -49,17 +49,17 @@ def _configure_cli_token(profile, insecure): def _configure_cli_aad_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() - if ENV_DATABRICKS_TOKEN not in os.environ: - print('[ERROR] Set Environment Variable \'%s\' with your AAD Token and run again.\n' % ENV_DATABRICKS_TOKEN) + if ENV_AAD_TOKEN not in os.environ: + print('[ERROR] Set Environment Variable \'%s\' with your AAD Token and run again.\n' % ENV_AAD_TOKEN) print('Commands to run to get your AAD token:\n' '\t az login\n' '\t token_response=$(az account get-access-token --resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n' - '\t export %s=$(jq .accessToken -r <<< "$token_response")\n' % ENV_DATABRICKS_TOKEN + '\t export %s=$(jq .accessToken -r <<< "$token_response")\n' % ENV_AAD_TOKEN ) return host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) - aad_token = os.environ.get(ENV_DATABRICKS_TOKEN) + aad_token = os.environ.get(ENV_AAD_TOKEN) new_config = DatabricksConfig.from_token(host, aad_token, insecure) update_and_persist_config(profile, new_config) diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py index afaff2c7..07df7560 100644 --- a/databricks_cli/secrets/cli.py +++ b/databricks_cli/secrets/cli.py @@ -52,7 +52,7 @@ ' scope is assigned to the request issuer\'s user identity.') @click.option('--scope-backend-type', type=click.Choice(['AZURE_KEYVAULT', 'DATABRICKS'], case_sensitive=True), default='DATABRICKS', help='The backend that will be used for this secret scope. ' - 'Options are (case-sensitive): 1) \'azure_keyvault\' and 2) \'databricks\' ' + 'Options are (case-sensitive): 1) \'AZURE_KEYVAULT\' and 2) \'DATABRICKS\' ' '(default option)' '\nNote: To create an Azure Keyvault, be sure to configure an AAD Token using' '\'databricks-cli configure --aad-token\'') From 78ca2a4fb6dcb0748818d8afcb9354d76161f240 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Mon, 28 Sep 2020 13:58:10 -0700 Subject: [PATCH 10/13] added extra space --- databricks_cli/secrets/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py index 07df7560..564e5923 100644 --- a/databricks_cli/secrets/cli.py +++ b/databricks_cli/secrets/cli.py @@ -54,7 +54,7 @@ default='DATABRICKS', help='The backend that will be used for this secret scope. ' 'Options are (case-sensitive): 1) \'AZURE_KEYVAULT\' and 2) \'DATABRICKS\' ' '(default option)' - '\nNote: To create an Azure Keyvault, be sure to configure an AAD Token using' + '\nNote: To create an Azure Keyvault, be sure to configure an AAD Token using ' '\'databricks-cli configure --aad-token\'') @click.option('--resource-id', default=None, type=click.STRING, help='The resource ID associated with the azure keyvault to be used as the backend' From e980cc43abe86daa1e8259bea62dfdc6a01dbdae Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Mon, 28 Sep 2020 14:39:19 -0700 Subject: [PATCH 11/13] reverted version.py to og --- databricks_cli/sdk/version.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) mode change 100755 => 100644 databricks_cli/sdk/version.py diff --git a/databricks_cli/sdk/version.py b/databricks_cli/sdk/version.py old mode 100755 new mode 100644 From f0b9341e87592cd20b9a01255139dd065305f465 Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Mon, 28 Sep 2020 15:00:57 -0700 Subject: [PATCH 12/13] resolved merge conflict --- databricks_cli/sdk/service.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/databricks_cli/sdk/service.py b/databricks_cli/sdk/service.py index 0933888a..0d6a05f3 100755 --- a/databricks_cli/sdk/service.py +++ b/databricks_cli/sdk/service.py @@ -1013,4 +1013,13 @@ def reset(self, pipeline_id=None, headers=None): _data = {} return self.client.perform_query('POST', '/pipelines/{pipeline_id}/reset'.format(pipeline_id=pipeline_id), data=_data, headers=headers) - \ No newline at end of file + + def run(self, pipeline_id=None, headers=None): + _data = {} + + return self.client.perform_query('POST', '/pipelines/{pipeline_id}/run'.format(pipeline_id=pipeline_id), data=_data, headers=headers) + + def stop(self, pipeline_id=None, headers=None): + _data = {} + + return self.client.perform_query('POST', '/pipelines/{pipeline_id}/stop'.format(pipeline_id=pipeline_id), data=_data, headers=headers) From 86352cd01ca3b2a97cd3ef278a914a875f0184ab Mon Sep 17 00:00:00 2001 From: sushi1998 Date: Mon, 28 Sep 2020 15:13:26 -0700 Subject: [PATCH 13/13] fixed lint errors --- databricks_cli/configure/cli.py | 8 +++++--- databricks_cli/secrets/api.py | 6 ++++-- databricks_cli/secrets/cli.py | 22 +++++++++++++--------- 3 files changed, 22 insertions(+), 14 deletions(-) diff --git a/databricks_cli/configure/cli.py b/databricks_cli/configure/cli.py index b06eae23..3e8303da 100644 --- a/databricks_cli/configure/cli.py +++ b/databricks_cli/configure/cli.py @@ -21,8 +21,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -import click import os +import click from click import ParamType @@ -50,10 +50,12 @@ def _configure_cli_aad_token(profile, insecure): config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty() if ENV_AAD_TOKEN not in os.environ: - print('[ERROR] Set Environment Variable \'%s\' with your AAD Token and run again.\n' % ENV_AAD_TOKEN) + print('[ERROR] Set Environment Variable \'%s\' with your ' + 'AAD Token and run again.\n' % ENV_AAD_TOKEN) print('Commands to run to get your AAD token:\n' '\t az login\n' - '\t token_response=$(az account get-access-token --resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n' + '\t token_response=$(az account get-access-token ' + '--resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n' '\t export %s=$(jq .accessToken -r <<< "$token_response")\n' % ENV_AAD_TOKEN ) return diff --git a/databricks_cli/secrets/api.py b/databricks_cli/secrets/api.py index 4ae4658d..c3487705 100644 --- a/databricks_cli/secrets/api.py +++ b/databricks_cli/secrets/api.py @@ -28,8 +28,10 @@ class SecretApi(object): def __init__(self, api_client): self.client = SecretService(api_client) - def create_scope(self, scope, initial_manage_principal, scope_backend_type, backend_azure_keyvault): - return self.client.create_scope(scope, initial_manage_principal, scope_backend_type, backend_azure_keyvault) + def create_scope(self, scope, initial_manage_principal, scope_backend_type, + backend_azure_keyvault): + return self.client.create_scope(scope, initial_manage_principal, + scope_backend_type, backend_azure_keyvault) def delete_scope(self, scope): return self.client.delete_scope(scope) diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py index 564e5923..3695e9d4 100644 --- a/databricks_cli/secrets/cli.py +++ b/databricks_cli/secrets/cli.py @@ -50,23 +50,26 @@ ' principal for this option is the group "users", which contains all users in the' ' workspace. If not specified, the initial ACL with MANAGE permission applied to the' ' scope is assigned to the request issuer\'s user identity.') -@click.option('--scope-backend-type', type=click.Choice(['AZURE_KEYVAULT', 'DATABRICKS'], case_sensitive=True), +@click.option('--scope-backend-type', + type=click.Choice(['AZURE_KEYVAULT', 'DATABRICKS'], case_sensitive=True), default='DATABRICKS', help='The backend that will be used for this secret scope. ' - 'Options are (case-sensitive): 1) \'AZURE_KEYVAULT\' and 2) \'DATABRICKS\' ' - '(default option)' - '\nNote: To create an Azure Keyvault, be sure to configure an AAD Token using ' + 'Options are (case-sensitive): 1) \'AZURE_KEYVAULT\' and ' + '2) \'DATABRICKS\' (default option)' + '\nNote: To create an Azure Keyvault, be sure ' + 'to configure an AAD Token using ' '\'databricks-cli configure --aad-token\'') @click.option('--resource-id', default=None, type=click.STRING, help='The resource ID associated with the azure keyvault to be used as the backend' ' for the secret scope. NOTE: Only use with azure-keyvault as backend') @click.option('--dns-name', default=None, type=click.STRING, - help='The dns name associated with the azure keyvault to be used as the backed for the' - ' secret scope. NOTE: Only use with azure-keyvault as backend') + help='The dns name associated with the azure keyvault to be used as the' + ' backed for the secret scope. NOTE: Only use with azure-keyvault as backend') @debug_option @profile_option @eat_exceptions @provide_api_client -def create_scope(api_client, scope, initial_manage_principal, scope_backend_type, resource_id, dns_name): +def create_scope(api_client, scope, initial_manage_principal, + scope_backend_type, resource_id, dns_name): """ Creates a new secret scope with given name. """ @@ -74,7 +77,8 @@ def create_scope(api_client, scope, initial_manage_principal, scope_backend_type 'resource_id': resource_id, 'dns_name': dns_name } - SecretApi(api_client).create_scope(scope, initial_manage_principal, scope_backend_type, backend_azure_keyvault) + SecretApi(api_client).create_scope(scope, initial_manage_principal, + scope_backend_type, backend_azure_keyvault) def _scopes_to_table(scopes_json): @@ -332,7 +336,7 @@ def get_acl(api_client, scope, principal, output): @debug_option @profile_option @eat_exceptions -def secrets_group(): +def secrets_group(): # pragma: no cover """ Utility to interact with secret API. """