databrickslabs-jupyterlab

#!/usr/bin/env python

import argparse
import os
import re
import sys
from enum import Enum
from version_parser import Version

import databrickslabs_jupyterlab
from databrickslabs_jupyterlab.remote import (
    configure_ssh,
    download_notebook,
    get_cluster,
    get_python_path,
    install_libs,
    is_reachable,
    version_check,
)
from databrickslabs_jupyterlab.local import (
    conda_version,
    create_kernelspec,
    get_db_config,
    prepare_ssh_config,
    print_error,
    print_ok,
    remove_kernelspecs,
    show_profiles,
    write_config,
)
from databrickslabs_jupyterlab.install import install
from databrickslabs_jupyterlab.utils import bye


class VCheck(Enum):
    a = "all"
    d = "diff"
    s = "same"

    def __str__(self):
        return self.value


parser = argparse.ArgumentParser(description="Configure remote Databricks access with JupyterLab")

parser.add_argument("profile", nargs="?", type=str, default="", help="A databricks-cli profile")
# parser.add_argument(
#     "-b",
#     "--bootstrap",
#     dest="bootstrap",
#     action="store_true",
#     help="Bootstrap the local databrickslabs-jupyterlab environment",
# )
parser.add_argument(
    "-d", "--delete", dest="delete", action="store_true", help="Delete a jupyter kernelspec"
)
parser.add_argument(
    "-m",
    "--mirror",
    dest="mirror",
    action="store_true",
    help="Mirror a a remote Databricks environment",
)
# parser.add_argument(
#     "-c",
#     "--clipboard",
#     dest="clipboard",
#     action="store_true",
#     help="Copy the personal access token to the clipboard",
# )
parser.add_argument(
    "-D", "--debug", dest="debug", action="store_true", help="Debug calls to REST API"
)
parser.add_argument(
    "-e",
    "--env",
    dest="extra_env",
    nargs="+",
    help="Environment variables to add to the remote notebook env",
)
parser.add_argument(
    "-i", "--id", dest="cluster_id", help="The cluster_id to avoid manual selection"
)
parser.add_argument(
    "-k",
    "--kernelspec",
    dest="kernelspec",
    action="store_true",
    help="Create a kernel specification",
)
parser.add_argument("-l", "--lab", dest="lab", action="store_true", help="Safely start JupyterLab")
parser.add_argument("-n", "--notebook_url", dest="notebook_url", help="Download demo notebook")
parser.add_argument(
    "-o", "--organisation", dest="organisation", help="The organisation for Azure Databricks"
)
parser.add_argument(
    "-p",
    "--profiles",
    dest="profiles",
    action="store_true",
    help="Show all databricks cli profiles and check SSH key",
)
parser.add_argument(
    "-r",
    "--reconfigure",
    dest="reconfigure",
    action="store_true",
    help="Reconfigure cluster with id cluster_id",
)
parser.add_argument(
    "-s",
    "--ssh-config",
    dest="sshconfig",
    action="store_true",
    help="Configure SSH access for a cluster",
)
parser.add_argument(
    "-v",
    "--version",
    dest="version",
    action="store_true",
    help="Check version of databrickslabs-jupyterlab",
)
parser.add_argument(
    "-V",
    "--versioncheck",
    dest="versioncheck",
    type=VCheck,
    choices=list(VCheck),
    help="Check version of local env with remote env",
)
parser.add_argument(
    "-w",
    "--whitelist",
    dest="use_whitelist",
    action="store_true",
    help="Use a whitelist (include pkg) of packages to install instead of blacklist (exclude pkg)",
)
parser.add_argument(
    "-W",
    "--print-whitelist",
    dest="whitelist",
    action="store_true",
    help="Print whitelist (include pkg) of packages to install",
)
parser.add_argument(
    "-B",
    "--print-blacklist",
    dest="blacklist",
    action="store_true",
    help="Print blacklist (exclude pkg) of packages to install",
)
parser.add_argument(
    "-N",
    "--no-spark",
    dest="no_spark",
    action="store_true",
    help="Do not create a Spark Session",
)

args = parser.parse_args()

if args.debug:
    import http.client as http_client

    http_client.HTTPConnection.debuglevel = 1

#
# Print Version and exit
#
if args.version:
    import databrickslabs_jupyterlab._version

    print(databrickslabs_jupyterlab._version.__version__)  # pylint: disable=protected-access
    bye()

#
# Print White list and exit
#
if args.whitelist:
    import databrickslabs_jupyterlab.install

    print("   => If package exists in remote cluster an this list, it will be installed:")
    print(databrickslabs_jupyterlab.install.WHITELIST)
    bye()

#
# Print Black list and exit
#
if args.blacklist:
    import databrickslabs_jupyterlab.install

    print("   => These packages will note be installed:")
    print(databrickslabs_jupyterlab.install.BLACKLIST)
    bye()

#
# Print White list and exit
#
if args.delete:
    remove_kernelspecs()
    bye()

#
# Bootstrap the initial environment and exit
#
# if args.bootstrap:
#     update_local()
#     bye()

#
# Show all configure profiles
#
if args.profiles:
    show_profiles()
    bye()

#
# Download Demo notebook
#
if args.notebook_url:
    download_notebook(args.notebook_url)
    bye()

profile = args.profile

#
# Check that at least conda 4.7.5 is installed. If not, exit
#
version = conda_version()
if Version("4.7.5") < Version(version):
    print("Valid version of conda detected: %s" % version)
else:
    print("Too old conda version:")
    print("Please update conda to at least 4.7.5")
    sys.exit(1)

#
# Ensure profile argument is given
#
if args.profile == "":
    parser.print_help()
    print_error(
        "\ndatabrickslabs-jupyterlab: error: the following arguments are required: profile\n"
    )
    bye(1)


print("\n* Getting host and token from .databrickscfg")
host, token = get_db_config(profile)
if (
    ("azuredatabricks.net" in host)
    and (args.kernelspec or args.reconfigure)
    and not args.organisation
):
    if host.startswith("https://adb"):
        r = re.compile(r"https://adb-(\d+)\.\d+.azuredatabricks.net")
        result = r.match(host)
        args.organisation = result.group(1)
    else:
        print_error(
            "\n    Error: To configure an Azure Databricks cluster using the old CNAME,\n"
            + "(<region>.azuredatabricks.net) the organization id (-o id) is required"
        )
        bye(1)

#
# Configure SSH for cluster
#
if args.sshconfig:
    if not args.cluster_id:
        print_error("   => Error: Provide '-i cluster_id' with '-s' ")
        bye(1)
    configure_ssh(profile, args.cluster_id)
    bye()

#
# Check the ssh key for the given profile and retrieve host and token from ~/.databrickscfg
#
if not os.path.exists(os.path.expanduser("~/.ssh/id_%s.pub" % profile)):
    print_error("\n    Error: ssh key for profile '%s' is missing:" % profile)
    print("    Use 'databrickslabs-jupyterlab %s -s' to configure ssh access" % profile)
    bye(1)


#
# Copy PAT to clipboard. Can be combined with every other flag
#
# if args.clipboard:
#     import pyperclip

#     pyperclip.copy(token)
#     print_ok("   => Personal access token copied to clipboard")

#
# Select remote cluster if not given by either -i or being in the respective mirror environment
#
if args.kernelspec or args.reconfigure or args.lab or args.versioncheck or args.mirror:
    print("\n* Select remote cluster\n")
    cluster_id, endpoint, cluster_name, conda_env = get_cluster(profile, args.cluster_id)
    if cluster_name is None:
        bye(1)

    print("\n* Configuring ssh config for remote cluster")
    prepare_ssh_config(cluster_id, profile, endpoint)
    print("   => Testing whether cluster can be reached")
    if not is_reachable(endpoint):
        print_error("Cannot connect to remote cluster. Please check:")
        print_error(
            "- whether the address is reachable and the port is open of %s in cloud security group"
            % endpoint
        )
        print_error("- whether VPN is enabled if you need one to connect to the cluster")
        bye(1)
    print_ok("   => OK")


#
# Compare local and remote libray versions
#
if args.versioncheck:
    version_check(cluster_id, host, token, args.versioncheck)
    bye()


#
# Create a conda environment that mirrors the selected remote cluster and exit
#
if args.mirror:
    install(profile, host, token, cluster_id, cluster_name, args.use_whitelist)
    bye()

#
# Install databrickslabs_jupyterlab libraries on the driver
#
if args.kernelspec or args.reconfigure or args.lab:
    print("\n* Installing driver libraries")
    result = install_libs(cluster_id, host, token)
    if result[0] == 0:
        print_ok("   => OK")
    else:
        print_error(result[1])
        bye(1)

#
# Start Jupyter lab in a safely way with all prechecks from above
#
if args.lab:
    import jupyterlab.labapp

    sys.argv = sys.argv[:1]
    jupyterlab.labapp.main()

#
# Create a jupyter kernelspecification
#

if args.kernelspec:
    print("\n* Creating remote kernel spec")
    print("args.extra_env", args.extra_env)
    if args.extra_env is None:
        extra_env = ""
    else:
        extra_env = " ".join(args.extra_env)

    python_path = get_python_path(cluster_id)
    # Using local conda environment name avoids overwriting of kernelspecs from different
    # environments
    create_kernelspec(
        profile,
        args.organisation,
        host,
        cluster_id,
        cluster_name,
        os.environ["CONDA_DEFAULT_ENV"],
        python_path,
        args.no_spark,
        extra_env,
    )
    print_ok("   => OK")

    print("\n* Setting global config of jupyter lab (autorestart, timeout)")
    write_config()
    print_ok("   => OK")

print("")