From 003f93e9c310e8ce52398ec81a516082b2ad3dcd Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 14:11:43 +0200 Subject: [PATCH 01/11] Add trampoline into newer CLI versions if found in $PATH --- databricks_cli/cli.py | 115 +++++++++++++++++++++++++++++++++++++++++- setup.py | 2 +- 2 files changed, 115 insertions(+), 2 deletions(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index 6960ee9a..40ebd914 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -21,6 +21,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +import os +import sys + import click from databricks_cli.configure.config import profile_option, debug_option @@ -70,5 +74,114 @@ def cli(): cli.add_command(repos_group, name='repos') cli.add_command(unity_catalog_group, name='unity-catalog') + +def _trampoline_into_new_cli(): + trampoline_disable_env_var = 'DATABRICKS_CLI_DO_NOT_EXECUTE_NEWER_VERSION' + if os.environ.get(trampoline_disable_env_var) is not None: + return + + # Try to trampoline only if we're running the CLI as 'databricks'. + if os.path.basename(sys.argv[0]) != 'databricks': + return + + # Check to see if the new version of the CLI is in $PATH. + paths = os.environ['PATH'].split(os.pathsep) + self = None + self_version = version + candidate = None + for path in paths: + exec_path = os.path.join(path, 'databricks') + if os.name == 'nt': + exec_path = os.path.join(path, 'databricks.exe') + + if not os.path.exists(exec_path): + continue + + # Keep a pointer to the first 'databricks' in $PATH. + if not self: + self = exec_path + + # The new CLI is a single binary larger than 1MB. + # We use this as heuristic to tell if the new CLI is installed. + # Because this version of the CLI is much smaller, we do not + # need to dedup our own path to avoid an exec loop. + stat = os.stat(exec_path) + if stat.st_size < (1024 * 1024): + continue + + candidate = exec_path + break + + # Return if we cannot find the new CLI in $PATH. + if candidate is None: + return + + # Determine the version of the new CLI. + candidate_version_json = os.popen(f'{candidate} version --output json').read().strip() + candidate_version_obj = json.loads(candidate_version_json) + candidate_version = candidate_version_obj['Version'] + + def e(message, highlight=False, nl=False): + style = dict() + if not highlight: + style["fg"] = 'yellow' + + click.echo(click.style(message, **style), err=True, nl=nl) + + e(f"Databricks CLI ") + e(f"v{candidate_version}", highlight=True) + e(f" found at ") + e(candidate, highlight=True, nl=True) + + e(f"Your current $PATH prefers running CLI ") + e(f"v{self_version}", highlight=True) + e(f" at ") + e(self, highlight=True, nl=True) + + e(f"", nl=True) + + e(f"Because both are installed and available in $PATH, I assume you are trying to run the newer version.", nl=True) + e(f"If you want to disable this behavior you can set DATABRICKS_CLI_DO_NOT_EXECUTE_NEWER_VERSION=1.", nl=True) + + e(f"", nl=True) + + e(f"Executing CLI ") + e(f"v{candidate_version}", highlight=True) + e(f"...", nl=True) + e(f"-" * (len("Executing CLI v{candidate_version}...")), nl=True) + + # The new CLI is at least 1MB in size. This is a good heuristic to + # determine if the new CLI is installed. + os.execv(exec_path, sys.argv) + + +def main(): + _trampoline_into_new_cli() + + try: + rv = cli(standalone_mode=False) + if isinstance(rv, int): + sys.exit(rv) + sys.exit(0) + except click.ClickException as e: + e.show() + message = f""" +It seems like you might have entered an invalid command or used invalid flags. + +The version of the CLI you're using is v{version}. + +If you're trying to invoke the new version of our CLI, please note that the commands and +flags might be different. To help you transition, we've created a migration guide that +explains these changes and how to adapt your command line arguments accordingly. + +You can find the migration guide at: https://docs.databricks.com/dev-tools/cli/migrate.html +""" + click.echo(click.style(message, fg='yellow'), err=True) + sys.exit(e.exit_code) + except click.Abort as e: + click.utils.echo("Aborted!", file=sys.stderr) + sys.exit(1) + + if __name__ == "__main__": - cli() + main() diff --git a/setup.py b/setup.py index 741f309b..7a73a212 100644 --- a/setup.py +++ b/setup.py @@ -46,7 +46,7 @@ ], entry_points=''' [console_scripts] - databricks=databricks_cli.cli:cli + databricks=databricks_cli.cli:main dbfs=databricks_cli.dbfs.cli:dbfs_group ''', zip_safe=False, From 1a0cf366a997a025854aeb1251b24a406b178d15 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 14:23:39 +0200 Subject: [PATCH 02/11] Lint --- databricks_cli/cli.py | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index 40ebd914..d94b67b0 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -53,7 +53,7 @@ expose_value=False, is_eager=True, help=version) @debug_option @profile_option -def cli(): +def cli(**kwargs): pass @@ -122,33 +122,35 @@ def _trampoline_into_new_cli(): candidate_version = candidate_version_obj['Version'] def e(message, highlight=False, nl=False): - style = dict() + style = {} if not highlight: style["fg"] = 'yellow' click.echo(click.style(message, **style), err=True, nl=nl) - e(f"Databricks CLI ") + e("Databricks CLI ") e(f"v{candidate_version}", highlight=True) - e(f" found at ") + e(" found at ") e(candidate, highlight=True, nl=True) - e(f"Your current $PATH prefers running CLI ") + e("Your current $PATH prefers running CLI ") e(f"v{self_version}", highlight=True) - e(f" at ") + e(" at ") e(self, highlight=True, nl=True) - e(f"", nl=True) + e("", nl=True) - e(f"Because both are installed and available in $PATH, I assume you are trying to run the newer version.", nl=True) - e(f"If you want to disable this behavior you can set DATABRICKS_CLI_DO_NOT_EXECUTE_NEWER_VERSION=1.", nl=True) + e("Because both are installed and available in $PATH, " + + "I assume you are trying to run the newer version.", nl=True) + e("If you want to disable this behavior you can set" + + "DATABRICKS_CLI_DO_NOT_EXECUTE_NEWER_VERSION=1.", nl=True) - e(f"", nl=True) + e("", nl=True) - e(f"Executing CLI ") + e("Executing CLI ") e(f"v{candidate_version}", highlight=True) - e(f"...", nl=True) - e(f"-" * (len("Executing CLI v{candidate_version}...")), nl=True) + e("...", nl=True) + e("-" * (len(f"Executing CLI v{candidate_version}...")), nl=True) # The new CLI is at least 1MB in size. This is a good heuristic to # determine if the new CLI is installed. @@ -178,7 +180,7 @@ def main(): """ click.echo(click.style(message, fg='yellow'), err=True) sys.exit(e.exit_code) - except click.Abort as e: + except click.Abort: click.utils.echo("Aborted!", file=sys.stderr) sys.exit(1) From d448387612a945c8492cdefab6dbb7e97647a1a6 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 14:25:27 +0200 Subject: [PATCH 03/11] Lint? --- databricks_cli/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index d94b67b0..179023e0 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -53,7 +53,7 @@ expose_value=False, is_eager=True, help=version) @debug_option @profile_option -def cli(**kwargs): +def cli(**_): pass From 69341c5486c002d93952af203fa2838a0fcc82b0 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 14:27:49 +0200 Subject: [PATCH 04/11] Refer to env var by variable --- databricks_cli/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index 179023e0..1475863b 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -143,7 +143,7 @@ def e(message, highlight=False, nl=False): e("Because both are installed and available in $PATH, " + "I assume you are trying to run the newer version.", nl=True) e("If you want to disable this behavior you can set" + - "DATABRICKS_CLI_DO_NOT_EXECUTE_NEWER_VERSION=1.", nl=True) + f"{trampoline_disable_env_var}=1.", nl=True) e("", nl=True) From d7f59759144d7e9a0bea82dfa19014fe7cb7d6dc Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 14:42:18 +0200 Subject: [PATCH 05/11] Py2 compat --- databricks_cli/cli.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index 1475863b..c9b87472 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -117,7 +117,7 @@ def _trampoline_into_new_cli(): return # Determine the version of the new CLI. - candidate_version_json = os.popen(f'{candidate} version --output json').read().strip() + candidate_version_json = os.popen(candidate + ' version --output json').read().strip() candidate_version_obj = json.loads(candidate_version_json) candidate_version = candidate_version_obj['Version'] @@ -129,12 +129,12 @@ def e(message, highlight=False, nl=False): click.echo(click.style(message, **style), err=True, nl=nl) e("Databricks CLI ") - e(f"v{candidate_version}", highlight=True) + e("v{}".format(candidate_version), highlight=True) e(" found at ") e(candidate, highlight=True, nl=True) e("Your current $PATH prefers running CLI ") - e(f"v{self_version}", highlight=True) + e("v{}".format(self_version), highlight=True) e(" at ") e(self, highlight=True, nl=True) @@ -143,14 +143,14 @@ def e(message, highlight=False, nl=False): e("Because both are installed and available in $PATH, " + "I assume you are trying to run the newer version.", nl=True) e("If you want to disable this behavior you can set" + - f"{trampoline_disable_env_var}=1.", nl=True) + "{}=1.".format(trampoline_disable_env_var), nl=True) e("", nl=True) e("Executing CLI ") - e(f"v{candidate_version}", highlight=True) + e("v{}".format(candidate_version), highlight=True) e("...", nl=True) - e("-" * (len(f"Executing CLI v{candidate_version}...")), nl=True) + e("-" * (len("Executing CLI v{}...".format(candidate_version))), nl=True) # The new CLI is at least 1MB in size. This is a good heuristic to # determine if the new CLI is installed. @@ -167,17 +167,17 @@ def main(): sys.exit(0) except click.ClickException as e: e.show() - message = f""" + message = """ It seems like you might have entered an invalid command or used invalid flags. -The version of the CLI you're using is v{version}. +The version of the CLI you're using is v{}. If you're trying to invoke the new version of our CLI, please note that the commands and flags might be different. To help you transition, we've created a migration guide that explains these changes and how to adapt your command line arguments accordingly. You can find the migration guide at: https://docs.databricks.com/dev-tools/cli/migrate.html -""" +""".format(version) click.echo(click.style(message, fg='yellow'), err=True) sys.exit(e.exit_code) except click.Abort: From 6218807a15caf8a24f1738788794acd0ce0efe00 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 15:13:16 +0200 Subject: [PATCH 06/11] Try/except on JSON decode of candidate version --- databricks_cli/cli.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index c9b87472..a786bdaf 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -118,8 +118,11 @@ def _trampoline_into_new_cli(): # Determine the version of the new CLI. candidate_version_json = os.popen(candidate + ' version --output json').read().strip() - candidate_version_obj = json.loads(candidate_version_json) - candidate_version = candidate_version_obj['Version'] + try: + candidate_version_obj = json.loads(candidate_version_json) + candidate_version = candidate_version_obj['Version'] + except RuntimeError: + candidate_version = '' def e(message, highlight=False, nl=False): style = {} From 6207ff2cf6b51c74747410de2f44745e81ac56d9 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 15:16:39 +0200 Subject: [PATCH 07/11] Explicit errors --- databricks_cli/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index a786bdaf..a6ba8e81 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -121,7 +121,7 @@ def _trampoline_into_new_cli(): try: candidate_version_obj = json.loads(candidate_version_json) candidate_version = candidate_version_obj['Version'] - except RuntimeError: + except (RuntimeError, json.JSONDecodeError): candidate_version = '' def e(message, highlight=False, nl=False): From 4310443d7384b2aa66e401cc2a4aeb51e65b490c Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 15:17:51 +0200 Subject: [PATCH 08/11] Whitespace --- databricks_cli/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index a6ba8e81..db65d923 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -145,7 +145,7 @@ def e(message, highlight=False, nl=False): e("Because both are installed and available in $PATH, " + "I assume you are trying to run the newer version.", nl=True) - e("If you want to disable this behavior you can set" + + e("If you want to disable this behavior you can set " + "{}=1.".format(trampoline_disable_env_var), nl=True) e("", nl=True) From 7117725621dac5b7c3ad5e1952d3f759fb1d6532 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Fri, 16 Jun 2023 16:37:32 +0200 Subject: [PATCH 09/11] exec_path -> candidate --- databricks_cli/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index db65d923..c8c95bc4 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -157,7 +157,7 @@ def e(message, highlight=False, nl=False): # The new CLI is at least 1MB in size. This is a good heuristic to # determine if the new CLI is installed. - os.execv(exec_path, sys.argv) + os.execv(candidate, sys.argv) def main(): From e18144dd12e256aa89ccb17cb1da9eaf3dba39f5 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Tue, 20 Jun 2023 07:21:22 +0200 Subject: [PATCH 10/11] Catch errors and proceed if trampoline function fails --- databricks_cli/cli.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index c8c95bc4..59ccd1ff 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -161,7 +161,11 @@ def e(message, highlight=False, nl=False): def main(): - _trampoline_into_new_cli() + try: + _trampoline_into_new_cli() + except Exception as e: + # Log the error and continue; perhaps a permissions issue? + click.echo("Failed to look for newer version of CLI: {}".format(e), err=True) try: rv = cli(standalone_mode=False) From a037d8e8ba65eb9184e1f9a351f29de6d8e1c5f9 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Tue, 20 Jun 2023 08:29:57 +0200 Subject: [PATCH 11/11] Lint --- databricks_cli/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index 59ccd1ff..b7f39d10 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -163,7 +163,7 @@ def e(message, highlight=False, nl=False): def main(): try: _trampoline_into_new_cli() - except Exception as e: + except Exception as e: # noqa # Log the error and continue; perhaps a permissions issue? click.echo("Failed to look for newer version of CLI: {}".format(e), err=True)