Skip to content

Commit 22a6f38

Browse files
author
Sushanth Varma Kanumuri
authored
[PLAT-21956] Add az-token config functionality to db cli (#325)
* added aad token functionality but aad token is only accepted as env variable * ready for review * improved functionality for input prompts * changed workflow of accepting tokens * autogenerated sdk and removed azToken and resourceID * added more detailed description for create-scope command * ready for review * ready for review - reset sdk files to og and changed variable name * addressed review comments * added extra space * reverted version.py to og * resolved merge conflict * fixed lint errors
1 parent f8e2f92 commit 22a6f38

6 files changed

Lines changed: 226 additions & 104 deletions

File tree

databricks_cli/cli.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,4 +67,4 @@ def cli():
6767
cli.add_command(pipelines_group, name='pipelines')
6868

6969
if __name__ == "__main__":
70-
cli()
70+
cli()

databricks_cli/configure/cli.py

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
# See the License for the specific language governing permissions and
2222
# limitations under the License.
2323

24+
import os
2425
import click
2526

2627
from click import ParamType
@@ -34,16 +35,37 @@
3435
PROMPT_USERNAME = 'Username'
3536
PROMPT_PASSWORD = 'Password' # NOQA
3637
PROMPT_TOKEN = 'Token' # NOQA
38+
ENV_AAD_TOKEN = 'DATABRICKS_AAD_TOKEN'
3739

3840

3941
def _configure_cli_token(profile, insecure):
4042
config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
4143
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
42-
token = click.prompt(PROMPT_TOKEN, default=config.token, hide_input=True)
44+
token = click.prompt(PROMPT_TOKEN, default=config.token)
4345
new_config = DatabricksConfig.from_token(host, token, insecure)
4446
update_and_persist_config(profile, new_config)
4547

4648

49+
def _configure_cli_aad_token(profile, insecure):
50+
config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
51+
52+
if ENV_AAD_TOKEN not in os.environ:
53+
print('[ERROR] Set Environment Variable \'%s\' with your '
54+
'AAD Token and run again.\n' % ENV_AAD_TOKEN)
55+
print('Commands to run to get your AAD token:\n'
56+
'\t az login\n'
57+
'\t token_response=$(az account get-access-token '
58+
'--resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n'
59+
'\t export %s=$(jq .accessToken -r <<< "$token_response")\n' % ENV_AAD_TOKEN
60+
)
61+
return
62+
63+
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
64+
aad_token = os.environ.get(ENV_AAD_TOKEN)
65+
new_config = DatabricksConfig.from_token(host, aad_token, insecure)
66+
update_and_persist_config(profile, new_config)
67+
68+
4769
def _configure_cli_password(profile, insecure):
4870
config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
4971
if config.password:
@@ -63,17 +85,20 @@ def _configure_cli_password(profile, insecure):
6385
@click.command(context_settings=CONTEXT_SETTINGS,
6486
short_help='Configures host and authentication info for the CLI.')
6587
@click.option('--token', show_default=True, is_flag=True, default=False)
88+
@click.option('--aad-token', show_default=True, is_flag=True, default=False)
6689
@click.option('--insecure', show_default=True, is_flag=True, default=None)
6790
@debug_option
6891
@profile_option
69-
def configure_cli(token, insecure):
92+
def configure_cli(token, aad_token, insecure):
7093
"""
7194
Configures host and authentication info for the CLI.
7295
"""
7396
profile = get_profile_from_context()
7497
insecure_str = str(insecure) if insecure is not None else None
7598
if token:
7699
_configure_cli_token(profile, insecure_str)
100+
elif aad_token:
101+
_configure_cli_aad_token(profile, insecure_str)
77102
else:
78103
_configure_cli_password(profile, insecure_str)
79104

databricks_cli/configure/provider.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,7 @@ def get_config(self):
263263

264264

265265
class DatabricksConfig(object):
266-
def __init__(self, host, username, password, token, insecure): # noqa
266+
def __init__(self, host, username, password, token, insecure): # noqa
267267
self.host = host
268268
self.username = username
269269
self.password = password

0 commit comments

Comments
 (0)