datachain 0.8.3__py3-none-any.whl → 0.8.4__py3-none-any.whl

datachain/cli/__init__.py

@@ -0,0 +1,311 @@
+import logging
+import os
+import sys
+import traceback
+from multiprocessing import freeze_support
+from typing import Optional
+
+from datachain.cli.utils import get_logging_level
+from datachain.telemetry import telemetry
+
+from .commands import (
+    clear_cache,
+    completion,
+    dataset_stats,
+    du,
+    edit_dataset,
+    garbage_collect,
+    index,
+    list_datasets,
+    ls,
+    query,
+    rm_dataset,
+    show,
+)
+from .parser import get_parser
+
+logger = logging.getLogger("datachain")
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    from datachain.catalog import get_catalog
+
+    # Required for Windows multiprocessing support
+    freeze_support()
+
+    datachain_parser = get_parser()
+    args = datachain_parser.parse_args(argv)
+
+    if args.command in ("internal-run-udf", "internal-run-udf-worker"):
+        return handle_udf(args.command)
+
+    logger.addHandler(logging.StreamHandler())
+    logging_level = get_logging_level(args)
+    logger.setLevel(logging_level)
+
+    client_config = {
+        "aws_endpoint_url": args.aws_endpoint_url,
+        "anon": args.anon,
+    }
+
+    if args.debug_sql:
+        # This also sets this environment variable for any subprocesses
+        os.environ["DEBUG_SHOW_SQL_QUERIES"] = "True"
+
+    error = None
+
+    try:
+        catalog = get_catalog(client_config=client_config)
+        return handle_command(args, catalog, client_config)
+    except BrokenPipeError as exc:
+        error, return_code = handle_broken_pipe_error(exc)
+        return return_code
+    except (KeyboardInterrupt, Exception) as exc:
+        error, return_code = handle_general_exception(exc, args, logging_level)
+        return return_code
+    finally:
+        telemetry.send_cli_call(args.command, error=error)
+
+
+def handle_command(args, catalog, client_config) -> int:
+    """Handle the different CLI commands."""
+    from datachain.studio import process_jobs_args, process_studio_cli_args
+
+    command_handlers = {
+        "cp": lambda: handle_cp_command(args, catalog),
+        "clone": lambda: handle_clone_command(args, catalog),
+        "dataset": lambda: handle_dataset_command(args, catalog),
+        "ds": lambda: handle_dataset_command(args, catalog),
+        "ls": lambda: handle_ls_command(args, client_config),
+        "show": lambda: handle_show_command(args, catalog),
+        "du": lambda: handle_du_command(args, catalog, client_config),
+        "find": lambda: handle_find_command(args, catalog),
+        "index": lambda: handle_index_command(args, catalog),
+        "completion": lambda: handle_completion_command(args),
+        "query": lambda: handle_query_command(args, catalog),
+        "clear-cache": lambda: clear_cache(catalog),
+        "gc": lambda: garbage_collect(catalog),
+        "studio": lambda: process_studio_cli_args(args),
+        "job": lambda: process_jobs_args(args),
+    }
+
+    handler = command_handlers.get(args.command)
+    if handler:
+        handler()
+        return 0
+    print(f"invalid command: {args.command}", file=sys.stderr)
+    return 1
+
+
+def handle_cp_command(args, catalog):
+    catalog.cp(
+        args.sources,
+        args.output,
+        force=bool(args.force),
+        update=bool(args.update),
+        recursive=bool(args.recursive),
+        edatachain_file=None,
+        edatachain_only=False,
+        no_edatachain_file=True,
+        no_glob=args.no_glob,
+    )
+
+
+def handle_clone_command(args, catalog):
+    catalog.clone(
+        args.sources,
+        args.output,
+        force=bool(args.force),
+        update=bool(args.update),
+        recursive=bool(args.recursive),
+        no_glob=args.no_glob,
+        no_cp=args.no_cp,
+        edatachain=args.edatachain,
+        edatachain_file=args.edatachain_file,
+    )
+
+
+def handle_dataset_command(args, catalog):
+    dataset_commands = {
+        "pull": lambda: catalog.pull_dataset(
+            args.dataset,
+            args.output,
+            local_ds_name=args.local_name,
+            local_ds_version=args.local_version,
+            cp=args.cp,
+            force=bool(args.force),
+            edatachain=args.edatachain,
+            edatachain_file=args.edatachain_file,
+        ),
+        "edit": lambda: edit_dataset(
+            catalog,
+            args.name,
+            new_name=args.new_name,
+            description=args.description,
+            labels=args.labels,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "ls": lambda: list_datasets(
+            catalog=catalog,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "rm": lambda: rm_dataset(
+            catalog,
+            args.name,
+            version=args.version,
+            force=args.force,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "remove": lambda: rm_dataset(
+            catalog,
+            args.name,
+            version=args.version,
+            force=args.force,
+            studio=args.studio,
+            local=args.local,
+            all=args.all,
+            team=args.team,
+        ),
+        "stats": lambda: dataset_stats(
+            catalog,
+            args.name,
+            args.version,
+            show_bytes=args.bytes,
+            si=args.si,
+        ),
+    }
+
+    handler = dataset_commands.get(args.datasets_cmd)
+    if handler:
+        return handler()
+    raise Exception(f"Unexpected command {args.datasets_cmd}")
+
+
+def handle_ls_command(args, client_config):
+    ls(
+        args.sources,
+        long=bool(args.long),
+        studio=args.studio,
+        local=args.local,
+        all=args.all,
+        team=args.team,
+        update=bool(args.update),
+        client_config=client_config,
+    )
+
+
+def handle_show_command(args, catalog):
+    show(
+        catalog,
+        args.name,
+        args.version,
+        limit=args.limit,
+        offset=args.offset,
+        columns=args.columns,
+        no_collapse=args.no_collapse,
+        schema=args.schema,
+    )
+
+
+def handle_du_command(args, catalog, client_config):
+    du(
+        catalog,
+        args.sources,
+        show_bytes=args.bytes,
+        depth=args.depth,
+        si=args.si,
+        update=bool(args.update),
+        client_config=client_config,
+    )
+
+
+def handle_find_command(args, catalog):
+    results_found = False
+    for result in catalog.find(
+        args.sources,
+        update=bool(args.update),
+        names=args.name,
+        inames=args.iname,
+        paths=args.path,
+        ipaths=args.ipath,
+        size=args.size,
+        typ=args.type,
+        columns=args.columns,
+    ):
+        print(result)
+        results_found = True
+    if not results_found:
+        print("No results")
+
+
+def handle_index_command(args, catalog):
+    index(
+        catalog,
+        args.sources,
+        update=bool(args.update),
+    )
+
+
+def handle_completion_command(args):
+    print(completion(args.shell))
+
+
+def handle_query_command(args, catalog):
+    query(
+        catalog,
+        args.script,
+        parallel=args.parallel,
+        params=args.param,
+    )
+
+
+def handle_broken_pipe_error(exc):
+    # Python flushes standard streams on exit; redirect remaining output
+    # to devnull to avoid another BrokenPipeError at shutdown
+    # See: https://docs.python.org/3/library/signal.html#note-on-sigpipe
+    error = str(exc)
+    devnull = os.open(os.devnull, os.O_WRONLY)
+    os.dup2(devnull, sys.stdout.fileno())
+    return error, 141  # 128 + 13 (SIGPIPE)
+
+
+def handle_general_exception(exc, args, logging_level):
+    error = str(exc)
+    if isinstance(exc, KeyboardInterrupt):
+        msg = "Operation cancelled by the user"
+    else:
+        msg = str(exc)
+    print("Error:", msg, file=sys.stderr)
+    if logging_level <= logging.DEBUG:
+        traceback.print_exception(
+            type(exc),
+            exc,
+            exc.__traceback__,
+            file=sys.stderr,
+        )
+    if args.pdb:
+        import pdb  # noqa: T100
+
+        pdb.post_mortem()
+    return error, 1
+
+
+def handle_udf(command):
+    if command == "internal-run-udf":
+        from datachain.query.dispatch import udf_entrypoint
+
+        return udf_entrypoint()
+
+    if command == "internal-run-udf-worker":
+        from datachain.query.dispatch import udf_worker_entrypoint
+
+        return udf_worker_entrypoint()

datachain/cli/commands/__init__.py

@@ -0,0 +1,29 @@
+from .datasets import (
+    dataset_stats,
+    edit_dataset,
+    list_datasets,
+    list_datasets_local,
+    rm_dataset,
+)
+from .du import du
+from .index import index
+from .ls import ls
+from .misc import clear_cache, completion, garbage_collect
+from .query import query
+from .show import show
+
+__all__ = [
+    "clear_cache",
+    "completion",
+    "dataset_stats",
+    "du",
+    "edit_dataset",
+    "garbage_collect",
+    "index",
+    "list_datasets",
+    "list_datasets_local",
+    "ls",
+    "query",
+    "rm_dataset",
+    "show",
+]

datachain/cli/commands/datasets.py

@@ -0,0 +1,129 @@
+import sys
+from typing import TYPE_CHECKING, Optional
+
+from tabulate import tabulate
+
+from datachain import utils
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+from datachain.cli.utils import determine_flavors
+from datachain.config import Config
+from datachain.error import DatasetNotFoundError
+
+
+def list_datasets(
+    catalog: "Catalog",
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import list_datasets
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    local_datasets = set(list_datasets_local(catalog)) if all or local else set()
+    studio_datasets = (
+        set(list_datasets(team=team)) if (all or studio) and token else set()
+    )
+
+    rows = [
+        _datasets_tabulate_row(
+            name=name,
+            version=version,
+            both=(all or (local and studio)) and token,
+            local=(name, version) in local_datasets,
+            studio=(name, version) in studio_datasets,
+        )
+        for name, version in local_datasets.union(studio_datasets)
+    ]
+
+    print(tabulate(rows, headers="keys"))
+
+
+def list_datasets_local(catalog: "Catalog"):
+    for d in catalog.ls_datasets():
+        for v in d.versions:
+            yield (d.name, v.version)
+
+
+def _datasets_tabulate_row(name, version, both, local, studio):
+    row = {
+        "Name": name,
+        "Version": version,
+    }
+    if both:
+        row["Studio"] = "\u2714" if studio else "\u2716"
+        row["Local"] = "\u2714" if local else "\u2716"
+    return row
+
+
+def rm_dataset(
+    catalog: "Catalog",
+    name: str,
+    version: Optional[int] = None,
+    force: Optional[bool] = False,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import remove_studio_dataset
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        try:
+            catalog.remove_dataset(name, version=version, force=force)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
+
+    if (all or studio) and token:
+        remove_studio_dataset(team, name, version, force)
+
+
+def edit_dataset(
+    catalog: "Catalog",
+    name: str,
+    new_name: Optional[str] = None,
+    description: Optional[str] = None,
+    labels: Optional[list[str]] = None,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+):
+    from datachain.studio import edit_studio_dataset
+
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        try:
+            catalog.edit_dataset(name, new_name, description, labels)
+        except DatasetNotFoundError:
+            print("Dataset not found in local", file=sys.stderr)
+
+    if (all or studio) and token:
+        edit_studio_dataset(team, name, new_name, description, labels)
+
+
+def dataset_stats(
+    catalog: "Catalog",
+    name: str,
+    version: int,
+    show_bytes=False,
+    si=False,
+):
+    stats = catalog.dataset_stats(name, version)
+
+    if stats:
+        print(f"Number of objects: {stats.num_objects}")
+        if show_bytes:
+            print(f"Total objects size: {stats.size}")
+        else:
+            print(f"Total objects size: {utils.sizeof_fmt(stats.size, si=si): >7}")

datachain/cli/commands/du.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+from datachain import utils
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def du(catalog: "Catalog", sources, show_bytes=False, si=False, **kwargs):
+    for path, size in catalog.du(sources, **kwargs):
+        if show_bytes:
+            print(f"{size} {path}")
+        else:
+            print(f"{utils.sizeof_fmt(size, si=si): >7} {path}")

datachain/cli/commands/index.py

@@ -0,0 +1,12 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def index(
+    catalog: "Catalog",
+    sources,
+    **kwargs,
+):
+    catalog.index(sources, **kwargs)

datachain/cli/commands/ls.py

@@ -0,0 +1,169 @@
+import shlex
+from collections.abc import Iterable, Iterator
+from itertools import chain
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+from datachain.cli.utils import determine_flavors
+from datachain.config import Config
+
+
+def ls(
+    sources,
+    long: bool = False,
+    studio: bool = False,
+    local: bool = False,
+    all: bool = True,
+    team: Optional[str] = None,
+    **kwargs,
+):
+    token = Config().read().get("studio", {}).get("token")
+    all, local, studio = determine_flavors(studio, local, all, token)
+
+    if all or local:
+        ls_local(sources, long=long, **kwargs)
+
+    if (all or studio) and token:
+        ls_remote(sources, long=long, team=team)
+
+
+def ls_local(
+    sources,
+    long: bool = False,
+    catalog: Optional["Catalog"] = None,
+    client_config=None,
+    **kwargs,
+):
+    from datachain import DataChain
+
+    if catalog is None:
+        from datachain.catalog import get_catalog
+
+        catalog = get_catalog(client_config=client_config)
+    if sources:
+        actual_sources = list(ls_urls(sources, catalog=catalog, long=long, **kwargs))
+        if len(actual_sources) == 1:
+            for _, entries in actual_sources:
+                for entry in entries:
+                    print(format_ls_entry(entry))
+        else:
+            first = True
+            for source, entries in actual_sources:
+                # print a newline between directory listings
+                if first:
+                    first = False
+                else:
+                    print()
+                if source:
+                    print(f"{source}:")
+                for entry in entries:
+                    print(format_ls_entry(entry))
+    else:
+        chain = DataChain.listings()
+        for ls in chain.collect("listing"):
+            print(format_ls_entry(f"{ls.uri}@v{ls.version}"))  # type: ignore[union-attr]
+
+
+def format_ls_entry(entry: str) -> str:
+    if entry.endswith("/") or not entry:
+        entry = shlex.quote(entry[:-1])
+        return f"{entry}/"
+    return shlex.quote(entry)
+
+
+def ls_remote(
+    paths: Iterable[str],
+    long: bool = False,
+    team: Optional[str] = None,
+):
+    from datachain.node import long_line_str
+    from datachain.remote.studio import StudioClient
+
+    client = StudioClient(team=team)
+    first = True
+    for path, response in client.ls(paths):
+        if not first:
+            print()
+        if not response.ok or response.data is None:
+            print(f"{path}:\n  Error: {response.message}\n")
+            continue
+
+        print(f"{path}:")
+        if long:
+            for row in response.data:
+                entry = long_line_str(
+                    row["name"] + ("/" if row["dir_type"] else ""),
+                    row["last_modified"],
+                )
+                print(format_ls_entry(entry))
+        else:
+            for row in response.data:
+                entry = row["name"] + ("/" if row["dir_type"] else "")
+                print(format_ls_entry(entry))
+        first = False
+
+
+def ls_urls(
+    sources,
+    catalog: "Catalog",
+    long: bool = False,
+    **kwargs,
+) -> Iterator[tuple[str, Iterator[str]]]:
+    curr_dir = None
+    value_iterables = []
+    for next_dir, values in _ls_urls_flat(sources, long, catalog, **kwargs):
+        if curr_dir is None or next_dir == curr_dir:  # type: ignore[unreachable]
+            value_iterables.append(values)
+        else:
+            yield curr_dir, chain(*value_iterables)  # type: ignore[unreachable]
+            value_iterables = [values]
+        curr_dir = next_dir
+    if curr_dir is not None:
+        yield curr_dir, chain(*value_iterables)
+
+
+def _node_data_to_ls_values(row, long_format=False):
+    from datachain.node import DirType, long_line_str
+
+    name = row[0]
+    is_dir = row[1] == DirType.DIR
+    ending = "/" if is_dir else ""
+    value = name + ending
+    if long_format:
+        last_modified = row[2]
+        timestamp = last_modified if not is_dir else None
+        return long_line_str(value, timestamp)
+    return value
+
+
+def _ls_urls_flat(
+    sources,
+    long: bool,
+    catalog: "Catalog",
+    **kwargs,
+) -> Iterator[tuple[str, Iterator[str]]]:
+    from datachain.client import Client
+    from datachain.node import long_line_str
+
+    for source in sources:
+        client_cls = Client.get_implementation(source)
+        if client_cls.is_root_url(source):
+            buckets = client_cls.ls_buckets(**catalog.client_config)
+            if long:
+                values = (long_line_str(b.name, b.created) for b in buckets)
+            else:
+                values = (b.name for b in buckets)
+            yield source, values
+        else:
+            found = False
+            fields = ["name", "dir_type"]
+            if long:
+                fields.append("last_modified")
+            for data_source, results in catalog.ls([source], fields=fields, **kwargs):
+                values = (_node_data_to_ls_values(r, long) for r in results)
+                found = True
+                yield data_source.dirname(), values
+            if not found:
+                raise FileNotFoundError(f"No such file or directory: {source}")

datachain/cli/commands/misc.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING
+
+import shtab
+
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
+
+def clear_cache(catalog: "Catalog"):
+    catalog.cache.clear()
+
+
+def garbage_collect(catalog: "Catalog"):
+    temp_tables = catalog.get_temp_table_names()
+    if not temp_tables:
+        print("Nothing to clean up.")
+    else:
+        print(f"Garbage collecting {len(temp_tables)} tables.")
+        catalog.cleanup_tables(temp_tables)
+
+
+def completion(shell: str) -> str:
+    from datachain.cli import get_parser
+
+    return shtab.complete(
+        get_parser(),
+        shell=shell,
+    )