[PLAT-21956] Add az-token config functionality to db cli

databricks_cli/cli.py

@@ -67,4 +67,4 @@ def cli():
 cli.add_command(pipelines_group, name='pipelines')
 
 if __name__ == "__main__":
-    cli()
+    cli()

databricks_cli/configure/cli.py

@@ -21,6 +21,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import click
 
 from click import ParamType
@@ -34,16 +35,37 @@
 PROMPT_USERNAME = 'Username'
 PROMPT_PASSWORD = 'Password' #  NOQA
 PROMPT_TOKEN = 'Token' #  NOQA
+ENV_AAD_TOKEN = 'DATABRICKS_AAD_TOKEN'
 
 
 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):
+    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('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
+              )
+        return
+
+    host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
+    aad_token = os.environ.get(ENV_AAD_TOKEN)
+    new_config = DatabricksConfig.from_token(host, aad_token, 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,17 +85,20 @@ 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.
     """
     profile = get_profile_from_context()
     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)
 

databricks_cli/configure/provider.py

@@ -263,7 +263,7 @@ def get_config(self):
 
 
 class DatabricksConfig(object):
-    def __init__(self, host, username, password, token, insecure): # noqa
+    def __init__(self, host, username, password, token, insecure):  # noqa
         self.host = host
         self.username = username
         self.password = password