dsgrid-toolkit 0.2.0__py3-none-any.whl

dsgrid/cli/config.py

@@ -0,0 +1,177 @@
+"""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,
+    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",
+#    "--offline",
+#    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,142 @@
+"""Main CLI command for dsgrid."""
+
+import logging
+from pathlib import Path
+
+import rich_click as click
+
+import dsgrid
+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(
+    "--offline/--online",
+    is_flag=True,
+    default=dsgrid.runtime_config.offline,
+    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(
+    "--timings/--no-timings",
+    default=dsgrid.runtime_config.timings,
+    is_flag=True,
+    show_default=True,
+    help="Enable tracking of function timings.",
+)
+# Server-related options are commented-out because the registry is currently only
+# supported in SQLite. If/when we add postgres support, these can be added back.
+# @click.option(
+#    "-U",
+#    "--username",
+#    type=str,
+#    default=dsgrid.runtime_config.database_user,
+#    help="Database username",
+# )
+# @click.option(
+#    "-P",
+#    "--password",
+#    prompt=True,
+#    hide_input=True,
+#    cls=OptionPromptPassword,
+#    help="dsgrid registry password. Will prompt unless it is passed or the username matches the "
+#    "runtime config file.",
+# )
+@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,
+    offline,
+    timings,
+    # username,
+    # password,
+    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()
+
+
+cli.add_command(config)
+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,349 @@
+"""Main CLI command for dsgrid."""
+
+import logging
+import sys
+from pathlib import Path
+
+import rich_click as click
+from chronify.utils.path_utils import check_overwrite
+
+from dsgrid.cli.common import get_value_from_context
+from dsgrid.common import LOCAL_REGISTRY, REMOTE_REGISTRY
+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 (
+    DataStoreType,
+    DatabaseConnection,
+    DatasetRegistryStatus,
+    VersionUpdateType,
+)
+from dsgrid.registry.registry_manager import RegistryManager
+from dsgrid.registry.filter_registry_manager import FilterRegistryManager
+from dsgrid.utils.files import load_data
+
+
+logger = logging.getLogger(__name__)
+_config = DsgridRuntimeConfig.load()
+
+
+"""
+Click Group Definitions
+"""
+
+
+@click.group()
+# Server-related options are commented-out because the registry is currently only
+# supported in SQLite. If/when we add postgres support, these can be added back.
+# @click.option(
+#    "-U",
+#    "--username",
+#    default=_config.database_user,
+#    show_default=True,
+#    help="dsgrid registry user name",
+# )
+# @click.option(
+#    "-P",
+#    "--password",
+#    prompt=True,
+#    hide_input=True,
+#    cls=OptionPromptPassword,
+#    help="dsgrid registry password. Will prompt unless it is passed or the username matches the "
+#    "runtime config file.",
+# )
+@click.option(
+    "--url",
+    default=_config.database_url,
+    show_default=True,
+    envvar="DSGRID_REGISTRY_DATABASE_URL",
+    help="dsgrid registry database URL. Override with the environment variable "
+    "DSGRID_REGISTRY_DATABASE_URL",
+)
+@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(
+    "--offline/--online",
+    is_flag=True,
+    default=_config.offline,
+    show_default=True,
+    help="run in 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(
+    "--verbose", is_flag=True, default=False, show_default=True, help="Enable verbose log output."
+)
+# def cli(url, username, password, log_file, no_prompts, offline, verbose):
+def cli(url, log_file, no_prompts, offline, verbose):
+    """dsgrid-admin commands"""
+    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.group()
+@click.option(
+    "--remote-path",
+    default=REMOTE_REGISTRY,
+    show_default=True,
+    help="path to dsgrid remote registry",
+)
+@click.pass_context
+def registry(ctx, remote_path):
+    """Manage a registry."""
+    no_prompts = ctx.parent.params["no_prompts"]
+    if "--help" in sys.argv:
+        ctx.obj = None
+    else:
+        conn = DatabaseConnection(
+            url=get_value_from_context(ctx, "url"),
+            # username=get_value_from_context(ctx, "username"),
+            # password=get_value_from_context(ctx, "password"),
+        )
+        ctx.obj = RegistryManager.load(
+            conn,
+            remote_path,
+            offline_mode=get_value_from_context(ctx, "offline"),
+            no_prompts=no_prompts,
+        )
+
+
+@click.group()
+@click.pass_obj
+def dimensions(registry_manager: RegistryManager):
+    """Dimension subcommands"""
+
+
+@click.group()
+@click.pass_obj
+def dimension_mappings(registry_manager: RegistryManager):
+    """Dimension mapping subcommands"""
+
+
+@click.group()
+@click.pass_obj
+def projects(registry_manager: RegistryManager):
+    """Project subcommands"""
+
+
+@click.group()
+@click.pass_obj
+def datasets(registry_manager: RegistryManager):
+    """Dataset subcommands"""
+
+
+"""
+Registry Commands
+"""
+
+_create_epilog = """
+Examples:\n
+$ dsgrid-admin create-registry sqlite:////projects/dsgrid/my_project/registry.db -p /projects/dsgrid/my_project/registry-data\n
+"""
+
+
+@click.command(epilog=_create_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]),
+)
+@click.pass_context
+def create_registry(
+    ctx, url: str, data_path: Path, overwrite: bool, data_store_type: DataStoreType
+):
+    """Create a new registry."""
+    check_overwrite(data_path, overwrite)
+    conn = DatabaseConnection(
+        url=url,  # This may change if/when we support a server database.
+        # url=get_value_from_context(ctx, "url"),
+        # username=get_value_from_context(ctx, "username"),
+        # password=get_value_from_context(ctx, "password"),
+    )
+    RegistryManager.create(conn, data_path, overwrite=overwrite, data_store_type=data_store_type)
+
+
+"""
+Dimension Commands
+"""
+
+
+@click.command(name="remove")
+@click.argument("dimension-id")
+@click.pass_obj
+def remove_dimension(registry_manager: RegistryManager, dimension_id: str):
+    """Remove a dimension from the dsgrid repository."""
+    registry_manager.dimension_manager.remove(dimension_id)
+
+
+"""
+Dimension Mapping Commands
+"""
+
+
+@click.command(name="remove")
+@click.argument("dimension-mapping-id")
+@click.pass_obj
+def remove_dimension_mapping(registry_manager: RegistryManager, dimension_mapping_id: str):
+    """Remove a dimension mapping from the dsgrid repository."""
+    registry_manager.dimension_mapping_manager.remove(dimension_mapping_id)
+
+
+"""
+Project Commands
+"""
+
+
+@click.command(name="remove")
+@click.argument("project-id")
+@click.pass_obj
+def remove_project(registry_manager: RegistryManager, project_id: str):
+    """Remove a project from the dsgrid repository."""
+    registry_manager.project_manager.remove(project_id)
+
+
+"""
+Dataset Commands
+"""
+
+
+@click.command(name="remove")
+@click.argument("dataset-ids", nargs=-1)
+@click.pass_obj
+def remove_datasets(registry_manager: RegistryManager, dataset_ids: list[str]):
+    """Remove one or more datasets from the dsgrid repository."""
+    dataset_mgr = registry_manager.dataset_manager
+    project_mgr = registry_manager.project_manager
+
+    # Ensure that all dataset IDs are valid before removing any of them.
+    for dataset_id in dataset_ids:
+        dataset_mgr.get_by_id(dataset_id)
+
+    for dataset_id in dataset_ids:
+        registry_manager.dataset_manager.remove(dataset_id)
+
+    dataset_ids_set = set(dataset_ids)
+    for project_id in project_mgr.list_ids():
+        config = project_mgr.get_by_id(project_id)
+        removed_dataset_ids = []
+        for dataset in config.iter_datasets():
+            if (
+                dataset.dataset_id in dataset_ids_set
+                and dataset.status == DatasetRegistryStatus.REGISTERED
+            ):
+                dataset.status = DatasetRegistryStatus.UNREGISTERED
+                dataset.mapping_references.clear()
+                removed_dataset_ids.append(dataset.dataset_id)
+        if removed_dataset_ids:
+            ids = ", ".join(removed_dataset_ids)
+            msg = (
+                f"Set status for datasets {ids} to unregistered in project {project_id} "
+                "after removal."
+            )
+            project_mgr.update(config, VersionUpdateType.MAJOR, msg)
+
+
+@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.",
+)
+@click.pass_context
+def make_filtered_registry(
+    ctx,
+    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))
+    # username = get_value_from_context(ctx, "username")
+    # password = get_value_from_context(ctx, "password")
+    src_conn = DatabaseConnection(
+        url=src_database_url,
+        # username=username,
+        # password=password,
+    )
+    dst_conn = DatabaseConnection(
+        url=dst_database_url,
+        # username=username,
+        # password=password,
+    )
+    RegistryManager.copy(
+        src_conn,
+        dst_conn,
+        dst_data_path,
+        mode=mode,
+        force=overwrite,
+    )
+    mgr = FilterRegistryManager.load(dst_conn, offline_mode=True, use_remote_data=False)
+    mgr.filter(simple_model=simple_model)
+
+
+cli.add_command(registry)
+cli.add_command(create_registry)
+cli.add_command(make_filtered_registry)
+
+registry.add_command(dimensions)
+registry.add_command(dimension_mappings)
+registry.add_command(projects)
+registry.add_command(datasets)
+
+dimensions.add_command(remove_dimension)
+dimension_mappings.add_command(remove_dimension_mapping)
+projects.add_command(remove_project)
+datasets.add_command(remove_datasets)