Add a --token-file flag to the databricks configure command (#330)

* Add a --token-file flag to the databricks configure command
This flag causes the databricks configure command to read the token from the specified file instead of from stdin.
This can be used for places where the token is exposed as a secret in a file, but a complete configuration is not.

* Fix whitespace formatting after merge to please lint again

* Fix merge bug where configure_cli_token was being called right after _configure_cli_token_file

Co-authored-by: Allen Reese <areese999@apple.com>

Author: areese

databricks_cli/configure/cli.py

@@ -20,59 +20,83 @@
 # 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 io
 import os
-import click
+from os import path
 
+import click
 from click import ParamType
 
+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.utils import CONTEXT_SETTINGS
-from databricks_cli.configure.config import profile_option, get_profile_from_context, debug_option
 
 PROMPT_HOST = 'Databricks Host (should begin with https://)'
 PROMPT_USERNAME = 'Username'
-PROMPT_PASSWORD = 'Password' #  NOQA
-PROMPT_TOKEN = 'Token' #  NOQA
+PROMPT_PASSWORD = 'Password'  # NOQA
+PROMPT_TOKEN = 'Token'  # NOQA
 ENV_AAD_TOKEN = 'DATABRICKS_AAD_TOKEN'
 
 
-def _configure_cli_token(profile, insecure):
+def _configure_cli_token_file(profile, token_file, host, insecure):
+    if not path.exists(token_file):
+        raise RuntimeError('Unable to read token from "{}"'.format(token_file))
+
+    with io.open(token_file, encoding='utf-8') as f:
+        token = f.readline().strip()
+
+    config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
+    if not host:
+        host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
+
+    new_config = DatabricksConfig.from_token(host, token, insecure)
+    update_and_persist_config(profile, new_config)
+
+
+def _configure_cli_token(profile, insecure, host=None):
     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)
+
+    if not host:
+        host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
+
+    token = click.prompt(PROMPT_TOKEN, default=config.token, hide_input=True)
     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_aad_token(profile, insecure, host=None):
     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
-              )
+        click.echo('[ERROR] Set Environment Variable \'%s\' with your '
+                   'AAD Token and run again.\n' % ENV_AAD_TOKEN)
+        click.echo('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())
+    if not host:
+        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):
+def _configure_cli_password(profile, insecure, host):
     config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
     if config.password:
         default_password = '*' * len(config.password)
     else:
         default_password = None
-    host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
+
+    if not host:
+        host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
+
     username = click.prompt(PROMPT_USERNAME, default=config.username)
     password = click.prompt(PROMPT_PASSWORD, default=default_password, hide_input=True,
                             confirmation_prompt=True)
@@ -84,29 +108,40 @@ 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('--token', '-t', 'token', show_default=True, is_flag=True, default=False)
+@click.option('--token-file', '-f', 'token_file', default=None,
+              help='Instead of reading the token from stdin, ' +
+                   'read the token from a file provided by a secret store.')
+@click.option('--host', show_default=True, default=None,
+              help='Host to connect to.')
 @click.option('--aad-token', show_default=True, is_flag=True, default=False)
-@click.option('--insecure', show_default=True, is_flag=True, default=None)
+@click.option('--insecure', show_default=True, is_flag=True, default=None,
+              help='DO NOT verify SSL Certificates')
 @debug_option
 @profile_option
-def configure_cli(token, aad_token, insecure):
+def configure_cli(token, aad_token, insecure, host, token_file):
     """
     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)
+        _configure_cli_token(profile=profile, insecure=insecure_str, host=host)
+    elif token_file:
+        _configure_cli_token_file(profile=profile, insecure=insecure_str, host=host,
+                                  token_file=token_file)
     elif aad_token:
-        _configure_cli_aad_token(profile, insecure_str)
+        _configure_cli_aad_token(profile=profile, insecure=insecure_str, host=host)
     else:
-        _configure_cli_password(profile, insecure_str)
+        _configure_cli_password(profile=profile, insecure=insecure_str, host=host)
 
 
 class _DbfsHost(ParamType):
     """
     Used to validate the configured host
     """
+
     def convert(self, value, param, ctx):
         if value.startswith('https://'):
             return value