dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/cli/common.py

@@ -0,0 +1,120 @@
+import logging
+import sys
+from pathlib import Path
+from typing import Any
+
+import rich_click as click
+
+from dsgrid.dsgrid_rc import DsgridRuntimeConfig
+from dsgrid.exceptions import DSGBaseException
+
+
+logger = logging.getLogger(__name__)
+
+
+def check_output_directory(path: Path, fs_interface, force: bool):
+    """Ensures that the parameter path is an empty directory.
+
+    Parameters
+    ----------
+    path : Path
+    fs_interface : FilesystemInterface
+    force : bool
+        If False and the directory exists and has content, exit.
+    """
+    if path.exists():
+        if not bool(path.iterdir()):
+            return
+        if force:
+            fs_interface.rm_tree(path)
+        else:
+            print(
+                f"{path} already exists. Choose a different name or pass --force to overwrite it.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+    path.mkdir()
+
+
+def get_log_level_from_str(level):
+    """Convert a log level string to logging type."""
+    match level:
+        case "debug":
+            return logging.DEBUG
+        case "info":
+            return logging.INFO
+        case "warning":
+            return logging.WARNING
+        case "error":
+            return logging.ERROR
+        case _:
+            msg = f"Unsupported level={level}"
+            raise Exception(msg)
+
+
+def get_value_from_context(ctx, field) -> Any:
+    """Get the field value from the root of a click context."""
+    return ctx.find_root().params[field]
+
+
+def handle_dsgrid_exception(ctx, func, *args, **kwargs) -> tuple[Any, int]:
+    """Handle any dsgrid exceptions as specified by the CLI parameters."""
+    res = None
+    try:
+        res = func(*args, **kwargs)
+        return res, 0
+    except DSGBaseException:
+        exc_type, exc_value, exc_tb = sys.exc_info()
+        filename = exc_tb.tb_frame.f_code.co_filename
+        line = exc_tb.tb_lineno
+        msg = f'{func.__name__} failed: exception={exc_type.__name__} message="{exc_value}" {filename=} {line=}'
+        logger.error(msg)
+        if ctx.find_root().params["reraise_exceptions"]:
+            raise
+        return res, 1
+
+
+def handle_scratch_dir(*args):
+    """Handle the user input for scratch_dir. If a path is passed, ensure it exists."""
+    val = args[2]
+    if val is None:
+        return val
+    path = Path(val)
+    if not path.exists:
+        msg = f"scratch-dir={path} does not exist"
+        raise ValueError(msg)
+    return path
+
+
+def path_callback(*args) -> Path | None:
+    """Ensure that a Path CLI option value is returned as a Path object."""
+    val = args[2]
+    if val is None:
+        return val
+    return Path(val)
+
+
+# Copied from
+# https://stackoverflow.com/questions/45868549/creating-a-click-option-with-prompt-that-shows-only-if-default-value-is-empty
+# and modified for our desired password behavior.
+
+
+class OptionPromptPassword(click.Option):
+    """Custom class that only prompts for the password if the user set a different username value
+    than what is in the runtime config file."""
+
+    def get_default(self, ctx, **kwargs):
+        config = DsgridRuntimeConfig.load()
+        username = ctx.find_root().params.get("username")
+        if username != config.database_user:
+            return None
+        return config.database_password
+
+    def prompt_for_value(self, ctx):
+        default = self.get_default(ctx)
+
+        if default is None:
+            return super().prompt_for_value(ctx)
+
+        return default

dsgrid/cli/config.py

@@ -0,0 +1,176 @@
+"""CLI commands to manage the dsgrid runtime configuration"""
+
+import logging
+import sys
+
+import rich_click as click
+
+from dsgrid.common import BackendEngine
+from dsgrid.cli.common import handle_scratch_dir
+from dsgrid.dsgrid_rc import (
+    DsgridRuntimeConfig,
+    DEFAULT_THRIFT_SERVER_URL,
+    DEFAULT_BACKEND,
+)
+from dsgrid.exceptions import DSGInvalidParameter
+from dsgrid.registry.common import DatabaseConnection
+
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+def config():
+    """Config commands"""
+
+
+_config_epilog = """
+Create a dsgrid configuration file to store registry connection settings and
+other dsgrid parameters.
+
+Examples:\n
+$ dsgrid config create sqlite:///./registry.db\n
+$ dsgrid config create sqlite:////projects/dsgrid/registries/standard-scenarios/registry.db\n
+"""
+
+
+@click.command(epilog=_config_epilog)
+@click.argument("url")
+@click.option(
+    "-b",
+    "--backend-engine",
+    type=click.Choice([x.value for x in BackendEngine]),
+    default=DEFAULT_BACKEND.value,
+    help="Backend engine for SQL processing",
+)
+@click.option(
+    "-t",
+    "--thrift-server-url",
+    type=str,
+    default=DEFAULT_THRIFT_SERVER_URL,
+    help="URL for the Apache Thrift Server to be used by chronify. "
+    "Only applies if Spark is the backend engine.",
+)
+@click.option(
+    "-m",
+    "--use-hive-metastore",
+    is_flag=True,
+    default=False,
+    help="Set this flag to use a Hive metastore when sharing data with chronify. "
+    "Only applies if Spark is the backend engine.",
+)
+@click.option(
+    "--timings/--no-timings",
+    default=False,
+    is_flag=True,
+    show_default=True,
+    help="Enable tracking of function timings.",
+)
+@click.option(
+    "--use-absolute-db-path/--no-use-absolute-db-path",
+    default=True,
+    is_flag=True,
+    show_default=True,
+    help="Convert the SQLite database file path to an absolute path.",
+)
+# @click.option(
+#    "-U",
+#    "--username",
+#    type=str,
+#    default=getpass.getuser(),
+#    help="Database username",
+# )
+# @click.option(
+#    "-P",
+#    "--password",
+#    prompt=True,
+#    hide_input=True,
+#    type=str,
+#    default=DEFAULT_DB_PASSWORD,
+#    help="Database username",
+# )
+# @click.option(
+#    "-o",
+#    is_flag=True,
+#    default=False,
+#    show_default=True,
+#    help="Run registry commands in offline mode. WARNING: any commands you perform in offline "
+#    "mode run the risk of being out-of-sync with the latest dsgrid registry, and any write "
+#    "commands will not be officially synced with the remote registry",
+# )
+@click.option(
+    "--console-level",
+    default="info",
+    show_default=True,
+    help="Console log level.",
+)
+@click.option(
+    "--file-level",
+    default="info",
+    show_default=True,
+    help="File log level.",
+)
+@click.option(
+    "-r",
+    "--reraise-exceptions",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    help="Re-raise any dsgrid exception. Default is to log the exception and exit.",
+)
+@click.option(
+    "-s",
+    "--scratch-dir",
+    default=None,
+    callback=handle_scratch_dir,
+    help="Base directory for dsgrid temporary directories. Must be accessible on all compute "
+    "nodes. Defaults to the current directory.",
+)
+def create(
+    url,
+    backend_engine,
+    thrift_server_url,
+    use_hive_metastore,
+    timings,
+    use_absolute_db_path,
+    # username,
+    # password,
+    # offline,
+    console_level,
+    file_level,
+    reraise_exceptions,
+    scratch_dir,
+):
+    """Create a local dsgrid runtime configuration file."""
+    conn = DatabaseConnection(url=url)
+    try:
+        db_filename = conn.get_filename()
+        if use_absolute_db_path and not db_filename.is_absolute():
+            conn.url = f"sqlite:///{db_filename.resolve()}"
+
+    except DSGInvalidParameter as exc:
+        print(str(exc), file=sys.stderr)
+        sys.exit(1)
+
+    if not db_filename.exists():
+        print(f"The registry database file {db_filename} does not exist.", file=sys.stderr)
+        sys.exit(1)
+
+    dsgrid_config = DsgridRuntimeConfig(
+        backend_engine=backend_engine,
+        thrift_server_url=thrift_server_url,
+        use_hive_metastore=use_hive_metastore,
+        timings=timings,
+        database_url=conn.url,
+        # database_user=username,
+        # database_password=password,
+        offline=True,
+        console_level=console_level,
+        file_level=file_level,
+        reraise_exceptions=reraise_exceptions,
+        scratch_dir=scratch_dir,
+    )
+    dsgrid_config.dump()
+
+
+config.add_command(create)

dsgrid/cli/download.py

@@ -0,0 +1,13 @@
+"""Download a dataset."""
+
+import sys
+
+import rich_click as click
+
+
+@click.command()
+@click.argument("dataset")
+def download(dataset):
+    """Download a dataset."""
+    print("not currently functional")
+    sys.exit(1)

dsgrid/cli/dsgrid.py

@@ -0,0 +1,157 @@
+"""Main CLI command for dsgrid."""
+
+import logging
+from pathlib import Path
+
+import rich_click as click
+
+import dsgrid
+from chronify.utils.path_utils import check_overwrite
+from dsgrid.common import LOCAL_REGISTRY
+from dsgrid.registry.common import DatabaseConnection, DataStoreType
+from dsgrid.registry.registry_manager import RegistryManager
+from dsgrid.utils.timing import timer_stats_collector
+from dsgrid.cli.common import get_log_level_from_str, handle_scratch_dir
+from dsgrid.cli.config import config
+from dsgrid.cli.download import download
+from dsgrid.cli.install_notebooks import install_notebooks
+from dsgrid.cli.query import query
+from dsgrid.cli.registry import registry
+from dsgrid.loggers import setup_logging, check_log_file_size, disable_console_logging
+
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+@click.option(
+    "-c",
+    "--console-level",
+    default=dsgrid.runtime_config.console_level,
+    show_default=True,
+    help="Console log level.",
+)
+@click.option(
+    "-f",
+    "--file-level",
+    default=dsgrid.runtime_config.file_level,
+    show_default=True,
+    help="File log level.",
+)
+@click.option("-l", "--log-file", type=Path, default="dsgrid.log", help="Log to this file.")
+@click.option(
+    "-n", "--no-prompts", default=False, is_flag=True, show_default=True, help="Do not prompt."
+)
+@click.option(
+    "--timings/--no-timings",
+    default=dsgrid.runtime_config.timings,
+    is_flag=True,
+    show_default=True,
+    help="Enable tracking of function timings.",
+)
+@click.option(
+    "-u",
+    "--url",
+    type=str,
+    default=dsgrid.runtime_config.database_url,
+    envvar="DSGRID_REGISTRY_DATABASE_URL",
+    help="Database URL. Ex: http://localhost:8529",
+)
+@click.option(
+    "-r",
+    "--reraise-exceptions",
+    is_flag=True,
+    default=dsgrid.runtime_config.reraise_exceptions,
+    show_default=True,
+    help="Re-raise any dsgrid exception. Default is to log the exception and exit.",
+)
+@click.option(
+    "-s",
+    "--scratch-dir",
+    default=dsgrid.runtime_config.scratch_dir,
+    callback=handle_scratch_dir,
+    help="Base directory for dsgrid temporary directories. Must be accessible on all compute "
+    "nodes. Defaults to the current directory.",
+)
+@click.pass_context
+def cli(
+    ctx,
+    console_level,
+    file_level,
+    log_file,
+    no_prompts,
+    timings,
+    url,
+    reraise_exceptions,
+    scratch_dir,
+):
+    """dsgrid commands"""
+    if timings:
+        timer_stats_collector.enable()
+    else:
+        timer_stats_collector.disable()
+    path = Path(log_file)
+    check_log_file_size(path, no_prompts=no_prompts)
+    ctx.params["console_level"] = get_log_level_from_str(console_level)
+    ctx.params["file_level"] = get_log_level_from_str(file_level)
+    setup_logging(
+        "dsgrid",
+        path,
+        console_level=ctx.params["console_level"],
+        file_level=ctx.params["file_level"],
+        mode="a",
+    )
+
+
+@cli.result_callback()
+def callback(*args, **kwargs):
+    with disable_console_logging(name="dsgrid"):
+        timer_stats_collector.log_stats()
+
+
+_create_registry_epilog = """
+Examples:\n
+$ dsgrid create-registry sqlite:////projects/dsgrid/my_project/registry.db -p /projects/dsgrid/my_project/registry-data\n
+"""
+
+
+@click.command(name="create-registry", epilog=_create_registry_epilog)
+@click.argument("url")
+@click.option(
+    "-p",
+    "--data-path",
+    default=LOCAL_REGISTRY,
+    show_default=True,
+    callback=lambda *x: Path(x[2]),
+    help="Local dsgrid registry data path. Must not contain the registry file listed in URL.",
+)
+@click.option(
+    "-f",
+    "--overwrite",
+    "--force",
+    is_flag=True,
+    default=False,
+    help="Delete registry_path and the database if they already exist.",
+)
+@click.option(
+    "-t",
+    "--data-store-type",
+    type=click.Choice([x.value for x in DataStoreType]),
+    default=DataStoreType.FILESYSTEM.value,
+    show_default=True,
+    help="Type of store to use for the registry data.",
+    callback=lambda *x: DataStoreType(x[2]),
+)
+def create_registry(url: str, data_path: Path, overwrite: bool, data_store_type: DataStoreType):
+    """Create a new registry."""
+    check_overwrite(data_path, overwrite)
+    conn = DatabaseConnection(url=url)
+    RegistryManager.create(conn, data_path, overwrite=overwrite, data_store_type=data_store_type)
+
+
+cli.add_command(config)
+cli.add_command(create_registry)
+cli.add_command(download)
+cli.add_command(install_notebooks)
+cli.add_command(query)
+cli.add_command(registry)

dsgrid/cli/dsgrid_admin.py

@@ -0,0 +1,92 @@
+"""CLI for dsgrid admin commands (testing purposes only)."""
+
+import logging
+from pathlib import Path
+
+import rich_click as click
+
+from dsgrid.config.simple_models import RegistrySimpleModel
+from dsgrid.dsgrid_rc import DsgridRuntimeConfig
+from dsgrid.loggers import setup_logging, check_log_file_size
+from dsgrid.registry.common import DatabaseConnection
+from dsgrid.registry.filter_registry_manager import FilterRegistryManager
+from dsgrid.registry.registry_manager import RegistryManager
+from dsgrid.utils.files import load_data
+
+
+logger = logging.getLogger(__name__)
+_config = DsgridRuntimeConfig.load()
+
+
+@click.group()
+@click.option("-l", "--log-file", default="dsgrid_admin.log", type=str, help="Log to this file.")
+@click.option(
+    "-n", "--no-prompts", default=False, is_flag=True, show_default=True, help="Do not prompt."
+)
+@click.option(
+    "--verbose", is_flag=True, default=False, show_default=True, help="Enable verbose log output."
+)
+def cli(log_file, no_prompts, verbose):
+    """dsgrid-admin commands (for testing purposes only)"""
+    path = Path(log_file)
+    level = logging.DEBUG if verbose else logging.INFO
+    check_log_file_size(path, no_prompts=no_prompts)
+    setup_logging("dsgrid", path, console_level=level, file_level=level, mode="a")
+
+
+@click.command()
+@click.option(
+    "--src-database-url",
+    required=True,
+    help="Source dsgrid registry database URL.",
+)
+@click.option(
+    "--dst-database-url",
+    default="dsgrid",
+    required=True,
+    help="Destination dsgrid registry database URL.",
+)
+@click.argument("dst_data_path", type=click.Path(exists=False), callback=lambda *x: Path(x[2]))
+@click.argument("config_file", type=click.Path(exists=True), callback=lambda *x: Path(x[2]))
+@click.option(
+    "-m",
+    "--mode",
+    default="data-symlinks",
+    type=click.Choice(["copy", "data-symlinks", "rsync"]),
+    show_default=True,
+    help="Controls whether to copy all data, make symlinks to data files, or sync data with the "
+    "rsync utility (not available on Windows).",
+)
+@click.option(
+    "-f",
+    "--overwrite",
+    "--force",
+    default=False,
+    is_flag=True,
+    show_default=True,
+    help="Overwrite dst_registry_path if it already exists. Does not apply if using rsync.",
+)
+def make_filtered_registry(
+    src_database_url,
+    dst_database_url,
+    dst_data_path: Path,
+    config_file: Path,
+    mode,
+    overwrite,
+):
+    """Make a filtered registry for testing purposes."""
+    simple_model = RegistrySimpleModel(**load_data(config_file))
+    src_conn = DatabaseConnection(url=src_database_url)
+    dst_conn = DatabaseConnection(url=dst_database_url)
+    RegistryManager.copy(
+        src_conn,
+        dst_conn,
+        dst_data_path,
+        mode=mode,
+        force=overwrite,
+    )
+    with FilterRegistryManager.load(dst_conn, offline_mode=True, use_remote_data=False) as mgr:
+        mgr.filter(simple_model=simple_model)
+
+
+cli.add_command(make_filtered_registry)

dsgrid/cli/install_notebooks.py

@@ -0,0 +1,62 @@
+import shutil
+import sys
+from pathlib import Path
+
+import rich_click as click
+
+import dsgrid
+
+NOTEBOOKS_DIRNAME = "dsgrid-notebooks"
+
+
+@click.command()
+@click.option(
+    "-p",
+    "--path",
+    default=Path.home(),
+    show_default=True,
+    type=click.Path(),
+    help="Path to install dsgrid notebooks.",
+    callback=lambda _, __, x: Path(x) / NOTEBOOKS_DIRNAME,
+)
+@click.option(
+    "-f",
+    "--overwrite",
+    "--force",
+    default=False,
+    show_default=True,
+    is_flag=True,
+    help="If true, overwrite existing files.",
+)
+def install_notebooks(path, overwrite):
+    """Install dsgrid notebooks to a local path."""
+    src_path = Path(dsgrid.__path__[0]) / "notebooks"
+    if not src_path.exists():
+        print(f"Unexpected error: dsgrid notebooks are not stored in {src_path}", file=sys.stderr)
+        sys.exit(1)
+
+    path.mkdir(exist_ok=True, parents=True)
+    to_copy = []
+    existing = []
+    for src_file in src_path.iterdir():
+        if src_file.suffix in (".ipynb", ".sh"):
+            dst = path / src_file.name
+            if dst.exists() and not overwrite:
+                existing.append(dst)
+            else:
+                to_copy.append((src_file, dst))
+    if existing:
+        print(
+            f"Existing files: {[str(x) for x in existing]}. "
+            "Choose a different location or set overwrite=true to overwrite.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    if not to_copy:
+        print("No notebook files found", file=sys.stderr)
+        sys.exit(1)
+
+    for src, dst in to_copy:
+        shutil.copyfile(src, dst)
+        print(f"Installed {dst}")