Added PR feedback, logging etc.

bin/install.py

@@ -1,6 +1,7 @@
 import os
 import sys
 from io import BytesIO
+import logging
 import shutil
 import argparse
 import subprocess
@@ -9,6 +10,9 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.workspace import ImportFormat
 
+from databricks.labs.ucx.logger import _install
+
+
 INSTALL_NOTEBOOK = """
 # Databricks notebook source
 # MAGIC %md
@@ -24,76 +28,78 @@
 
 """
 
+# install logging backend
+_install()
+logger = logging.getLogger(__name__)
+
+# parse command line parameters
 parser = argparse.ArgumentParser(prog="ucx",
                                  description="Builds and installs ucx.")
 parser.add_argument("--folder", "-f", default="ucx",
                     help="name of folder in workspace, default: ucx")
-parser.add_argument("--verbose", "-v", action="store_true",
-                    help="increase output verbosity")
+parser.add_argument("--quiet", action="store_true",
+                    help="suppress extraneous information")
 parser.add_argument("--debug", action="store_true",
                     help="enable debug mode")
 args = parser.parse_args()
 
+# adjust logging levels as needed
+if args.debug:
+    logging.getLogger("databricks").setLevel("DEBUG")
+
 
 def delete_local_dir(dir):
+    """Helper to delete a directory"""
     try:
         shutil.rmtree(dir)
     except OSError as e:
-        if args.verbose:
-            print(f"Error: {e.filename} - {e.strerror}.")
+        logger.error(f"Error: {e.filename} - {e.strerror}.")
 
 
-def main():
-    # build wheel in temp directory
-    tmp_dir = tempfile.TemporaryDirectory()
-    if args.verbose:
-        print(f"Created temporary directory: {tmp_dir.name}")
-    if args.verbose:
-        subprocess.run([
-            "python3", "-m", "pip",
-            "wheel", "--no-deps",
-            "--wheel-dir", tmp_dir.name,
-            ".."],
-            check=True)
-    else:
-        subprocess.run([
-            "python3", "-m", "pip",
-            "wheel", "--no-deps", "--quiet",
-            "--wheel-dir", tmp_dir.name,
-            ".."],
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-            check=True)
-    # get wheel name as first file in the temp directory
-    files = os.listdir(tmp_dir.name)
-    wheel_file_name = files[0]
-    local_wheel_file = tmp_dir.name + '/' + wheel_file_name
-    if args.verbose:
-        print(f"Wheel file: {wheel_file_name}")
-    # upload wheel and starer notebook to workspace
-    ws = WorkspaceClient()
-    folder_base = f"/Users/{ws.current_user.me().user_name}/{args.folder}"
-    remote_wheel_file = f"{folder_base}/{wheel_file_name}"
-    remote_notebook_file = f"{folder_base}/install_ucx.py"
-    if args.verbose:
-        print(f"Remote wheel file: {remote_wheel_file}")
-        print(f"Remote notebook file: {remote_notebook_file}")
-        print("Uploading...")
+def folder_exists(folder_base, ws):
+    """Helper to check if a workspace folder exists"""
+    folder_files = []
     try:
-        folder_files = []
         for f in ws.workspace.list(folder_base):
             folder_files.append(f.path)
-        print(f"ERROR: Remote folder '{folder_base}' already exists!")
-        print(f"Found: {folder_files} - ABORTING!")
-        sys.exit(-1)
+        logger.debug(f"Folder files: {folder_files}")
+        return True
     except:
-        pass
+        return False
+
+
+def build_wheel():
+    """Helper to build the wheel package"""
+    tmp_dir = tempfile.TemporaryDirectory()
+    logger.debug(f"Created temporary directory: {tmp_dir.name}")
+    streams = {}
+    if args.quiet:
+        streams = {
+            "stdout": subprocess.DEVNULL,
+            "stderr": subprocess.DEVNULL,
+        }
+    subprocess.run([
+        "python3", "-m", "pip",
+        "wheel", "--no-deps",
+        "--wheel-dir", tmp_dir.name,
+        ".."],
+        **streams,
+        check=True)
+    return tmp_dir.name
+
+
+def upload_artifacts(folder_base, local_wheel_file, wheel_file_name, ws):
+    """Helper to upload artifacts into a workspace folder"""
+    remote_wheel_file = f"{folder_base}/{wheel_file_name}"
+    remote_notebook_file = f"{folder_base}/install_ucx.py"
+    logger.info(f"Remote wheel file: {remote_wheel_file}")
+    logger.info(f"Remote notebook file: {remote_notebook_file}")
+    logger.info("Uploading...")
     ws.workspace.mkdirs(folder_base)
     with open(local_wheel_file, "rb") as fh:
-        buf = BytesIO(fh.read())
         ws.workspace.upload(
             path=remote_wheel_file,
-            content=buf,
+            content=fh.read(),
             format=ImportFormat.AUTO
         )
     buf = BytesIO(INSTALL_NOTEBOOK.format(
@@ -102,11 +108,29 @@ def main():
         path=remote_notebook_file,
         content=buf
     )
+
+
+def main():
+    # preflight check
+    ws = WorkspaceClient()
+    folder_base = f"/Users/{ws.current_user.me().user_name}/{args.folder}"
+    if folder_exists(folder_base, ws):
+        logger.error(
+            f"ERROR: Remote folder '{folder_base}' already exists, aborting!")
+        sys.exit(-1)
+    # build wheel in temp directory
+    tmp_dir = build_wheel()
+    # get wheel name as first file in the temp directory
+    files = os.listdir(tmp_dir)
+    wheel_file_name = files[0]
+    local_wheel_file = tmp_dir + '/' + wheel_file_name
+    logger.info(f"Wheel file: {wheel_file_name}")
+    # upload wheel and starer notebook to workspace
+    upload_artifacts(folder_base, local_wheel_file, wheel_file_name, ws)
     # cleanup
-    delete_local_dir(tmp_dir.name)
-    if args.verbose:
-        print("DONE.")
+    delete_local_dir(tmp_dir)
+    logger.info("DONE.")
 
 
 if __name__ == "__main__":
-    main()
+    main()