adalab-cli 1.4.0__py3-none-any.whl

adalab_cli/__init__.py

@@ -0,0 +1,11 @@
+"""The adalab-cli base module exposes the CLI application.
+"""
+
+import importlib.metadata
+
+_DISTRIBUTION_METADATA = importlib.metadata.metadata("adalab-cli")
+__project__ = _DISTRIBUTION_METADATA["name"]
+__version__ = _DISTRIBUTION_METADATA["version"]
+__description__ = _DISTRIBUTION_METADATA["description"]
+
+__all__ = ["cli"]

adalab_cli/cli/__init__.py

@@ -0,0 +1,5 @@
+from .app import app
+
+__all__ = ["app"]
+
+__title__ = "adalab-cli"

adalab_cli/cli/app.py

@@ -0,0 +1,191 @@
+"""Main module for the adalib CLI.
+
+This module initializes the Typer application and integrates various
+submodules that handle different command groups for the adalib CLI.
+
+Submodules:
+- oci_commands: Handles commands related to Open Container Initiative (OCI).
+- gallery_commands: Manages commands for gallery operations.
+- kernel_commands: Contains commands for kernel-related operations.
+- user_commands: Deals with user-related commands.
+"""
+
+import importlib.metadata
+import os
+
+import requests
+import typer
+from loguru import logger
+from prompt_toolkit.completion import NestedCompleter
+from rich import print as rich_print
+from rich.console import Console
+from rich.table import Table
+
+from . import authentication_flows as auth
+from .completion_tree import completion_tree
+from .databases_commands import db_app
+from .gallery_commands import gallery_app
+from .interactive import interactive_mode
+from .logger import Logger
+from .login_commands import load_config
+from .schedule_commands import schedules_app
+from .user_commands import user_app
+
+app = typer.Typer(no_args_is_help=False)
+
+
+@app.callback("adalab", invoke_without_command=True, no_args_is_help=True)
+def main(
+    token: str = typer.Option(None, help="User token for authentication"),
+    adalab_url: str = typer.Option(None, help="URL of the AdaLab server"),
+    adalab_secret: str = typer.Option(None, help="Client secret for AdaLab"),
+    interactive: bool = typer.Option(
+        False, "-i", "--interactive", help="Activate interactive mode"
+    ),
+    verbose: int = typer.Option(
+        0,
+        "-v",
+        count=True,
+        help="Verbosity level of the logs: -v (default): INFO | -vv: DEBUG | -vvv: TRACE",
+        min=0,
+        max=3,
+    ),
+    no_output: bool = typer.Option(False, "--no-output", help="Disable output from the CLI."),
+):
+    if no_output:
+        logger.remove()
+    else:
+        Logger(verbosity=verbose)
+    logger.info("Starting adalab CLI.")
+    logger.debug("Debugging mode activated.")
+    logger.trace("Tracing mode activated.")
+
+    # Attempt to configure adalib with user token.
+    # This works when token is either provided or stored in environment.
+    try:
+        auth.in_line_authentication(adalab_url, adalab_secret, token)
+    except AssertionError as e:
+        logger.info(f"{str(e)}")
+
+    # Start the interactive mode
+    if interactive:
+        rich_print("🏡 Welcome to AdaLab CLI interactive mode.")
+        rich_print("👋 Type 'exit' to quit.")
+        rich_print("ℹ️  Type 'help' for help.")
+        interactive_mode(
+            this_app=app,
+            title="adalab",
+            nested=True,
+            color="seagreen",
+            completion=NestedCompleter.from_nested_dict(completion_tree),
+        )
+
+
+@app.command("login")
+def login_commands(
+    auth_flow: str = typer.Option(
+        default="browser",
+        help="Authentication flow. Options: browser, device-code, user-token",
+    ),
+    token: str = typer.Option(None, help="Token for authentication"),
+):
+    """
+    Login to adalab
+    """
+    stored_configuration = load_config()
+    adalab_server_url = stored_configuration.get(
+        "adalab_url"
+    )  # Look for the URL in the config file
+    if adalab_server_url is None:
+        adalab_server_url = os.getenv("ADALAB_URL")  # Look for the URL in env vars
+        if adalab_server_url is None:
+            adalab_server_url = typer.prompt("What is the AdaLab url?")  # Ask for the URL
+    logger.info(f"AdaLab server URL: {adalab_server_url}")
+
+    adalab_client_secret = stored_configuration.get(
+        "adalab_secret"
+    )  # Look for the client secret in the config file
+    if adalab_client_secret is None:
+        adalab_client_secret = os.getenv("ADALAB_CLIENT_SECRET")
+        if "ADALAB_CLIENT_SECRET" not in os.environ:
+            adalab_client_secret = typer.prompt(
+                "What is the AdaLab client secret?", hide_input=True
+            )
+    logger.info("AdaLab client secret: ********")
+    # Check validity of URL
+    response = requests.get(adalab_server_url + "/adaboard/api", timeout=10)
+    if response.status_code != 200:
+        typer.echo("Invalid URL. Please try again.")
+        raise typer.Exit()
+
+    # Check which of the three kinds of authentication flows was set
+    if auth_flow == "browser":
+        auth.start_browser_auth(
+            adalab_server_url=adalab_server_url, adalab_client_secret=adalab_client_secret
+        )
+    elif auth_flow == "user-token":
+        auth.start_user_token_auth(
+            adalab_server_url=adalab_server_url,
+            token=token,
+            adalab_client_secret=adalab_client_secret,
+        )
+    elif auth_flow == "device-code":
+        auth.start_device_code_auth(
+            adalab_server_url=adalab_server_url, adalab_client_secret=adalab_client_secret
+        )
+
+    else:
+        typer.echo(
+            "Invalid authentication flow. Please try again."
+            "Options: browser, device-code, user-token"
+        )
+        raise typer.Exit()
+
+
+@app.command("version")
+def get_version():
+    """
+    Prints the version of adalib.
+    """
+
+    version_string_adalib = importlib.metadata.version("adalib")
+    version_string_adalib_auth = importlib.metadata.version("adalib-auth")
+    version_string_adalab_cli = importlib.metadata.version("adalab-cli")
+
+    table = Table("Package", "Version")
+    table.add_row("adalib", version_string_adalib)
+    table.add_row("adalib-auth", version_string_adalib_auth)
+    table.add_row("adalab-cli", version_string_adalab_cli)
+
+    console = Console()
+
+    console.print(table)
+
+
+app.add_typer(
+    db_app,
+    name="databases",
+    help="Commands for interacting with databases.",
+    no_args_is_help=True,
+)
+
+app.add_typer(
+    gallery_app,
+    name="cards",
+    help="Commands for interacting with the gallery.",
+    no_args_is_help=True,
+)
+
+app.add_typer(
+    user_app,
+    name="user",
+    help="Commands for interacting with user information.",
+    no_args_is_help=True,
+)
+
+app.add_typer(
+    schedules_app,
+    name="schedules",
+    help="Commands for interacting with schedules.",
+    no_args_is_help=True,
+)

adalab_cli/cli/authentication_flows.py

@@ -0,0 +1,147 @@
+import os
+from time import sleep
+
+import requests
+import typer
+from adalib_auth.config import get_config
+from keycloak import exceptions
+from loguru import logger
+from rich import print as rich_print
+
+from .login_commands import CLIENT_ID, KeycloakAuth, save_credentials
+
+
+def start_browser_auth(adalab_server_url: str = None, adalab_client_secret: str = None):
+    """Starts the browser-based authentication flow.
+
+    :param adalab_server_url: The URL of the Adalab server, defaults to None
+    :type adalab_server_url: str, optional
+    """
+    auth = KeycloakAuth(
+        adalab_server_url=adalab_server_url, adalab_client_secret=adalab_client_secret
+    )
+    auth.login()
+    logger.success("Authentication successful. Credentials saved.")
+    rich_print(
+        "🎉 [bold green]Authentication successful[/bold green]. You can now enjoy the AdaLab CLI."
+    )
+
+
+def start_user_token_auth(
+    token: str = None,
+    adalab_server_url: str = None,
+    adalab_client_secret: str = None,
+    store_credentials: bool = True,
+):
+    """Starts the user token authentication flow.
+
+    This function exchanges the user token for an access and refresh token. The token can be given
+    directly in the prompt or as an environment variable.
+    The access token and refresh token are saved to a file for future use.
+
+    :param token: The user token, it will be retrieve from ADALAB_USER_TOKEN environment variable
+        or prompted from the user.
+    :type token: str, optional
+    :param adalab_server_url: The URL of the ADALab server. If not provided, it will be retrieved
+        from the adalab_server_url argument.
+    :type adalab_server_url: str, optional
+    """
+
+    if token is None:
+        token = os.getenv("ADALAB_USER_TOKEN") or typer.prompt(
+            "What is your user token", hide_input=True
+        )
+
+    resp = requests.get(
+        adalab_server_url + "/adaboard/api/adalib/token",
+        headers={"Authorization": f"Token {token}"},
+    )
+    if resp.status_code != 200:
+        raise AssertionError(f"Failed to authenticate with user token. {resp.text}")
+
+    jh_token = resp.json()["access_token"]
+    new_token = KeycloakAuth(
+        adalab_server_url=adalab_server_url, adalab_client_secret=adalab_client_secret
+    ).client.exchange_token(token=jh_token, audience=CLIENT_ID)
+    if store_credentials:
+        save_credentials(
+            my_access_token=new_token["access_token"],
+            my_refresh_token=new_token["refresh_token"],
+        )
+    logger.success("Authentication successful. Credentials saved.")
+    rich_print(
+        "🎉 [bold green]Authentication successful[/bold green]. You can now enjoy the AdaLab CLI."
+    )
+
+
+def start_device_code_auth(adalab_server_url: str = None, adalab_client_secret: str = None):
+    """
+    Starts the device code authentication flow.
+
+    This function initiates the device code authentication flow using the AdaLab server URL.
+    It retrieves the device code information from the Keycloak authentication client,
+    prompts the user to visit the verification URI and enter the user code,
+    and waits for the user authentication to complete.
+    Once the authentication is successful, the access token and refresh token are saved.
+
+    :param adalab_server_url: The AdaLab server URL, defaults to None
+    :type adalab_server_url: str, optional
+    """
+    auth = KeycloakAuth(
+        adalab_server_url=adalab_server_url, adalab_client_secret=adalab_client_secret
+    )
+    device_code_info = auth.client.device()
+    device_code = device_code_info["device_code"]
+    user_code = device_code_info["user_code"]
+    verification_uri = device_code_info["verification_uri"]
+    rich_print("Please visit %s and enter the code: %s", verification_uri, user_code)
+    rich_print("This code expires in %s seconds", device_code_info["expires_in"])
+    token = None
+    while token is None:
+        try:
+            token = auth.client.token(
+                grant_type="urn:ietf:params:oauth:grant-type:device_code",
+                device_code=device_code,
+            )
+        except exceptions.KeycloakPostError as e:
+            if "authorization_pending" in str(e):
+                logger.info("Waiting for user authentication...")
+                sleep(device_code_info["interval"])
+            else:
+                logger.error("Error during device polling: %s", str(e))
+                exit(1)
+    save_credentials(
+        my_access_token=token["access_token"],
+        my_refresh_token=token["refresh_token"],
+    )
+    logger.success("Authentication successful. Credentials saved.")
+    rich_print(
+        "🎉 [bold green]Authentication successful[/bold green]. You can now enjoy the AdaLab CLI."
+    )
+
+
+def in_line_authentication(
+    adalab_server_url: str = None, adalab_secret: str = None, user_token: str = None
+) -> None:
+    """Auxiliary function for in-line authentication. This will not store credentials in the
+    configuration file. Instead it will be used for one-time authentication. Only valid for the
+    user token authentication flow.
+
+    :param adalab_server_url: The AdaLab server URL, defaults to None
+    :type adalab_server_url: str, optional
+    :param adalab_secret: The AdaLab JupyterHub client secret, defaults to None
+    :type adalab_secret: str, optional
+    :param user_token: The user token retrieved from AdaLab, defaults to None
+    :type user_token: str, optional
+    """
+    if adalab_server_url is None:
+        adalab_server_url = os.getenv("ADALAB_URL") or None
+    if adalab_secret is None:
+        adalab_secret = os.getenv("ADALAB_CLIENT_SECRET") or None
+    if user_token is None:
+        user_token = os.getenv("ADALAB_USER_TOKEN") or None
+
+    if None in (adalab_server_url, adalab_secret, user_token):
+        raise AssertionError("Missing required variables for in-line authentication.")
+
+    get_config(adaboard_api_url=adalab_server_url + "/adaboard/api", token=user_token)

adalab_cli/cli/completion_tree.py

@@ -0,0 +1,67 @@
+"""Map of the CLI commands and subcommands for the CLI completion tree."""
+
+completion_tree = {
+    "cards": {
+        "approve-card": {"--help": None},
+        "expose-card": {"--help": None},
+        "get-card": {"--help": None},
+        "hide-card": {"--help": None},
+        "list-cards": {
+            "all": None,
+            "notebook": None,
+            "voila": None,
+            "url": None,
+            "group": None,
+            "--help": None,
+        },
+        "help": None,
+        "exit": None,
+    },
+    "databases": {
+        "list-databases": {"--no-pretty": None, "--help": None},
+        "list-datasets": {"--no-pretty": None, "--return-sql": None, "--help": None},
+        "list-tables": {"--no-pretty": None, "--help": None},
+        "view-dataset": {"--head": None, "--no-pretty": None, "--help": None},
+        "help": None,
+        "exit": None,
+    },
+    "schedules": {
+        "create-schedule": {
+            "--schedule": None,
+            "--pool": None,
+            "--active": None,
+            "--concurrent": None,
+            "--cleanup": None,
+            "--timeout": None,
+            "--help": None,
+        },
+        "delete-run": {"--help": None},
+        "delete-schedule": {"--help": None},
+        "start-run": {
+            "--cleanup": None,
+            "--pool": None,
+            "--timeout": None,
+            "--help": None,
+        },
+        "stop-run": {"--help": None},
+        "update-schedule": {
+            "--schedule": None,
+            "--pool": None,
+            "--active": None,
+            "--concurrent": None,
+            "--cleanup": None,
+            "--timeout": None,
+            "--help": None,
+        },
+        "help": None,
+        "exit": None,
+    },
+    "user": {
+        "whoami": {"--full-information": None, "--with-notifications": None, "--help": None},
+        "help": None,
+        "exit": None,
+    },
+    "version": None,
+    "help": None,
+    "exit": None,
+}

adalab_cli/cli/databases_commands.py

@@ -0,0 +1,221 @@
+""" docstrings"""
+
+import sys
+
+import typer
+from loguru import logger
+from prompt_toolkit.completion import NestedCompleter
+from rich import print as rich_print
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+# from adalib.superset import databases, datasets
+from typing_extensions import Annotated
+
+from .completion_tree import completion_tree
+from .factory import check_authentication
+
+# from . import utils
+from .interactive import interactive_mode
+
+db_app = typer.Typer()
+
+
+@db_app.command("list-databases")
+def database_list_databases(
+    pretty: Annotated[
+        bool,
+        typer.Option(help="Set this flag to get the output in a nice tabular view."),
+    ] = True
+):
+    """Prints all databases available."""
+    from adalib.superset import databases  # noqa: E402
+
+    from . import utils
+
+    try:
+        with Progress(
+            SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True
+        ) as progress:
+            progress.add_task(description="Fetching data...", total=None)
+            my_databases = databases.all().as_df()
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+        sys.exit(1)
+    # my_databases = databases.all().as_df()
+    utils.print_dataframe(dataframe=my_databases, pretty=pretty, title="Databases")
+
+
+@db_app.command("list-datasets")
+def database_list_datasets(
+    database_id: Annotated[
+        int,
+        typer.Argument(help="Specific database from which datasets are listed."),
+    ] = None,
+    return_sql: Annotated[
+        bool,
+        typer.Option(help="Set this flag to get the SQL query in the output."),
+    ] = False,
+    pretty: Annotated[
+        bool,
+        typer.Option(help="Set this flag to get the output in a nice tabular view."),
+    ] = True,
+):
+    """
+    Prints all datasets available in the databases DATABASE_ID.
+    If DATABASE_ID is not specified then datasets in all available databases are returned.
+
+    Example usage:
+
+    adalab databases list-datasets --return-sql --pretty 1
+
+    will print the available datasets in the database identified with DATABASE_ID=1, including
+    the SQL query (if available) in a nice tabular view.
+    """
+    from . import utils
+
+    try:
+        with Progress(
+            SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True
+        ) as progress:
+            progress.add_task(description="Fetching data...", total=None)
+            my_datasets = utils.get_datasets()
+
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+        sys.exit(1)
+
+    filtered_datasets = utils.filter_datasets_by_database_id(
+        dataset_list=my_datasets, database_id=database_id
+    )
+
+    if filtered_datasets is False:
+        rich_print(f"The database with index {database_id} is not available.")
+
+    else:
+        format_datasets = utils.check_sql(dataset_list=filtered_datasets, include_sql=return_sql)
+        utils.print_dataframe(
+            dataframe=format_datasets,
+            pretty=pretty,
+            title=(
+                f"Datasets in database {database_id}"
+                if (database_id is not None)
+                else "All Datasets"
+            ),
+        )
+
+
+@db_app.command("view-dataset", no_args_is_help=True)
+def database_view_dataset_by_index(
+    ds_index: Annotated[int, typer.Argument(help="Index of the dataset to print.")] = None,
+    head: Annotated[
+        int,
+        typer.Option(help="Set --head N for printing the first N elements in the dataframe."),
+    ] = 10,
+    pretty: Annotated[
+        bool,
+        typer.Option(help="Set this flag to get the output in a nice tabular view."),
+    ] = True,
+):
+    """
+    Prints the first elements of the dataset indexed by DS_INDEX, by
+    default it prints the first 10 elements.
+
+    Example usage:
+
+    adalab databases view-dataset 15 --head 11 --pretty
+
+    will print the first 11 elements of the dataset with index 15
+    in a nice tabular view.
+    """
+    from adalib.superset import datasets
+
+    from . import utils
+
+    if ds_index not in datasets.all().as_df().index:
+        return rich_print(
+            f"Dataset with index {ds_index} [bold red]not available[/bold red].\n"
+            + "Run [bold blue]list-datasets[/bold blue] for available options"
+        )
+
+    try:
+        with Progress(
+            SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True
+        ) as progress:
+            progress.add_task(description="Fetching data...", total=None)
+            my_dataset = datasets.get(id=ds_index)
+
+            if my_dataset.kind == "virtual":
+                my_df = utils.retrieve_virtual_dataset(my_dataset=my_dataset, head=head)
+            elif my_dataset.kind == "physical":
+                my_df = utils.retrieve_physical_dataset(my_dataset=my_dataset, head=head)
+
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+        sys.exit(1)
+
+    utils.print_dataframe(dataframe=my_df, pretty=pretty, title=f"Dataset {ds_index}")
+
+
+@db_app.command("list-tables", no_args_is_help=True)
+def list_tables(
+    db_index: Annotated[int, typer.Argument(help="Database index.")] = None,
+    pretty: Annotated[
+        bool,
+        typer.Option(help="Set this flag to get the output in a nice tabular view."),
+    ] = True,
+):
+    """
+    List tables in the database with the specified index.
+
+    Example usage:
+
+    adalab databases list-tables 1 --pretty
+
+    will print the tables in the databes with database index 1 in a nice tabular view.
+    """
+    from adalib.superset import databases
+
+    from . import utils
+
+    try:
+        available_index = databases.all().as_df().index
+        assert db_index in available_index, (
+            f"Database with index {db_index} not available."
+            + " Run list-databases for available options"
+        )
+    except AssertionError as ae:
+        logger.error(f"Error: {str(ae)}")
+        sys.exit(1)
+
+    try:
+        with Progress(
+            SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True
+        ) as progress:
+            progress.add_task(description="Fetching data...", total=None)
+            db = databases.get(id=db_index)
+            tables = db.run(
+                query="SELECT table_name "
+                "FROM information_schema.tables  "
+                "WHERE table_schema = 'public';"
+            ).as_df()
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+        sys.exit(1)
+    utils.print_dataframe(dataframe=tables, pretty=pretty, title=f"Tables in database {db_index}")
+
+
+@db_app.callback("callback", invoke_without_command=True)
+def main(
+    interactive: bool = typer.Option(
+        False, "-i", "--interactive", help="Activate interactive mode"
+    )
+):
+    if interactive:
+        interactive_mode(
+            this_app=db_app,
+            title="databases",
+            color="maroon",
+            completion=NestedCompleter.from_nested_dict(completion_tree["databases"]),
+        )
+    else:
+        check_authentication()