esgpull 0.6.3__py3-none-any.whl

esgpull/cli/self.py

@@ -0,0 +1,238 @@
+import os
+from configparser import ConfigParser
+from pathlib import Path
+
+import click
+from click.exceptions import Abort, Exit
+from rich.table import Table
+
+from esgpull import Esgpull
+from esgpull.cli.decorators import args, opts
+from esgpull.cli.utils import init_esgpull
+from esgpull.config import Config
+from esgpull.exceptions import (
+    InvalidInstallPath,
+    UnknownInstallName,
+    UnregisteredInstallPath,
+)
+from esgpull.install_config import InstallConfig
+from esgpull.tui import TempUI, Verbosity
+
+
+def get_synda_db_path(sdt_home: str | None = None) -> Path | None:
+    if sdt_home is None:
+        sdt_home = os.getenv("SDT_HOME")
+    if sdt_home is not None:
+        sdt_path = Path(sdt_home).expanduser().resolve()
+        conf_path = sdt_path / "conf" / "sdt.conf"
+        conf = ConfigParser()
+        conf.read(conf_path)
+        path = Path(conf["core"]["db_path"]) / "sdt.db"
+        if path.is_file():
+            db_path = path
+        else:
+            db_path = None
+    else:
+        db_path = None
+    return db_path
+
+
+@click.group()
+def self():
+    """
+    Manage esgpull installations / import synda database
+    """
+    ...
+
+
+@self.command()
+@args.path
+@opts.name
+@opts.verbosity
+def install(
+    path: Path | None,
+    name: str | None,
+    verbosity: Verbosity,
+):
+    with TempUI.logging():
+        TempUI.rule("[b green]esgpull[/] installation")
+        if path is None:
+            default = str(InstallConfig.default)
+            path = Path(TempUI.prompt("Install location", default=default))
+            idx = InstallConfig.index(path=path)
+            if idx > -1:
+                TempUI.print(
+                    (
+                        f"\n:stop_sign: {path} is already installed:\n"
+                        f"  {InstallConfig.installs[idx]}\n\n"
+                        f"{InstallConfig.activate_msg(idx)}"
+                    ),
+                    err=True,
+                )
+                raise Exit(1)
+            if name is None:
+                name = TempUI.prompt("Name (optional)") or None
+        idx = InstallConfig.add(path, name)
+        InstallConfig.choose(idx=idx)
+        path = InstallConfig.installs[idx].path
+        if path.is_dir():
+            TempUI.print(f"Using existing install at {path}")
+        else:
+            TempUI.print(f"Creating install directory and files at {path}")
+        InstallConfig.write()
+        TempUI.print(f"Install config added to {InstallConfig.path}")
+        esg = Esgpull(verbosity=verbosity, install=True)
+    with esg.ui.logging("init", onraise=Abort):
+        # with esg.ui.spinner("Fetching facets"):
+        #     if esg.fetch_facets(update=False):
+        #         esg.ui.print(":+1: Facets are initialised.")
+        #     else:
+        #         esg.ui.print(":+1: Facets were already initialised.")
+        sdt_home = os.getenv("SDT_HOME")
+        if sdt_home is not None:
+            db_name = get_synda_db_path(sdt_home) or ""
+            msg = (
+                f"Found existing synda installation at {sdt_home}\n"
+                "You can import its database by running:\n"
+                f"$ esgpull self import_synda {db_name}"
+            )
+            esg.ui.print(msg)
+
+
+@self.command()
+@args.path
+@opts.name
+def activate(
+    path: Path | None,
+    name: str | None,
+):
+    with TempUI.logging():
+        idx = InstallConfig.index(path=path, name=name)
+        if idx < 0:
+            if path is not None and InstallConfig.index(name=path.name) > -1:
+                raise ValueError(
+                    f"{InstallConfig.fullpath(path)} is not installed, "
+                    "did you mean this?\n\n"
+                    f"$ esgpull self activate --name {path}"
+                )
+            raise InvalidInstallPath(path=path)
+        TempUI.print(InstallConfig.activate_needs_eval(idx))
+
+
+@self.command()
+@args.path
+@opts.name
+def choose(
+    path: Path | None,
+    name: str | None,
+):
+    with TempUI.logging():
+        if path is None and name is None:
+            table = Table(
+                title="Install locations",
+                title_style="bold",
+                title_justify="left",
+                show_header=False,
+                show_footer=False,
+                show_edge=False,
+                box=None,
+                padding=(0, 1),
+            )
+            table.add_column(style="b yellow")
+            table.add_column(style="magenta")
+            table.add_column(style="b green")
+            for i, inst in enumerate(InstallConfig.installs):
+                table.add_row(
+                    "*" if i == InstallConfig.current_idx else "",
+                    str(inst.path),
+                    inst.name or "",
+                )
+            TempUI.print(table)
+            raise Exit(0)
+        InstallConfig.choose(path=path, name=name)
+        if InstallConfig.current is None:
+            if name is not None:
+                raise UnknownInstallName(name)
+            elif path is not None:
+                raise UnregisteredInstallPath(path)
+        else:
+            InstallConfig.write()
+
+
+@self.command()
+def reset():
+    with TempUI.logging():
+        if InstallConfig.current is not None:
+            idx = InstallConfig.current_idx
+            path = InstallConfig.current.path
+            InstallConfig.choose()
+            InstallConfig.write()
+            TempUI.print(
+                f"Install location is not set anymore as {path}\n"
+                # "To set it back to its previous location, run:\n"
+                # f"$ esgpull self choose {name}"
+            )
+            TempUI.print(InstallConfig.activate_msg(idx))
+            raise Exit(0)
+        else:
+            TempUI.print(":stop_sign: No install found.")
+            TempUI.print("To install esgpull, run:\n")
+            TempUI.print("$ esgpull self install")
+            raise Exit(1)
+
+
+@self.command()
+def delete():
+    with TempUI.logging():
+        if InstallConfig.current is None:
+            msg = (
+                "None\n"
+                "Please choose an existing install before trying to delete it."
+                "\n\n$ esgpull self choose ..."
+            )
+            raise UnregisteredInstallPath(msg)
+        else:
+            path = InstallConfig.current.path
+            TempUI.print(f"You are going to delete: {path}")
+            choice = TempUI.prompt(f"Please enter {path.name!r} to continue")
+            if choice == path.name:
+                TempUI.print(f"Deleting {path} from config...")
+                TempUI.print("To remove all files from this install, run:\n")
+                config = Config.load(path=path)
+                for p in config.paths:
+                    if not p.is_relative_to(path):
+                        TempUI.print(f"$ rm -rf {p}")
+                TempUI.print(f"$ rm -rf {path}")
+                InstallConfig.remove_current()
+                InstallConfig.write()
+            else:
+                raise Abort
+
+
+@self.command()
+@args.path
+@opts.verbosity
+def import_synda(
+    path: Path | None,
+    verbosity: Verbosity,
+):
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("import_synda", onraise=Abort):
+        if path is None:
+            sdt_home = os.getenv("SDT_HOME")
+            prompt_title = "Enter synda database location"
+            if sdt_home is not None:
+                esg.ui.print(
+                    "Found existing synda installation at"
+                    f" SDT_HOME={sdt_home}"
+                )
+                default = str(get_synda_db_path(sdt_home))
+                path = Path(esg.ui.prompt(prompt_title, default=str(default)))
+            else:
+                path = Path(esg.ui.prompt(prompt_title))
+        else:
+            path = path.expanduser().resolve()
+        if not path.is_file():
+            raise FileNotFoundError(path)
+        nb_imported = esg.import_synda(url=path, track=True, ask=True)
+        esg.ui.print(f"Imported {nb_imported} new files from {path}")

esgpull/cli/show.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import click
+from click.exceptions import Abort, BadArgumentUsage, Exit
+
+from esgpull.cli.decorators import args, groups, opts
+from esgpull.cli.utils import get_queries, init_esgpull, valid_name_tag
+from esgpull.tui import Verbosity
+
+
+@click.command()
+@args.query_id
+@opts.tag
+@groups.show
+@groups.json_yaml
+@opts.files
+@opts.shas
+@opts.verbosity
+def show(
+    query_id: str | None,
+    tag: str | None,
+    children: bool,
+    parents: bool,
+    expand: bool,
+    files: bool,
+    json: bool,
+    yaml: bool,
+    shas: bool,
+    verbosity: Verbosity,
+) -> None:
+    """
+    View query tree
+    """
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("show", onraise=Abort):
+        if not valid_name_tag(esg.graph, esg.ui, query_id, tag):
+            raise Exit(1)
+        if expand and query_id is not None:
+            esg.ui.print(esg.graph.expand(query_id))
+            raise Exit(0)
+        if query_id is None and tag is None:
+            esg.graph.load_db()
+            graph = esg.graph
+        else:
+            queries = get_queries(esg.graph, query_id, tag)
+            graph = esg.graph.subgraph(
+                *queries,
+                children=children,
+                parents=parents,
+                keep_db=True,
+            )
+        if tag is not None:
+            tag_db = esg.graph.get_tag(tag)
+            if tag_db is not None and tag_db.description is not None:
+                esg.ui.print(tag_db.description)
+        if shas:
+            esg.ui.print(list(graph.queries.keys()), json=True)
+        if json:
+            esg.ui.print(graph.asdict(files=files), json=True)
+        elif yaml:
+            esg.ui.print(graph.asdict(files=files), yaml=True)
+        elif files:
+            msg = "--files can only be used with --json or --yaml"
+            raise BadArgumentUsage(msg)
+        else:
+            esg.ui.print(graph)

esgpull/cli/status.py

@@ -0,0 +1,67 @@
+import click
+from click.exceptions import Abort, Exit
+from rich.box import MINIMAL_DOUBLE_HEAD
+from rich.table import Table
+from rich.text import Text
+
+from esgpull.cli.decorators import opts
+from esgpull.cli.utils import init_esgpull
+from esgpull.models import sql
+from esgpull.tui import Verbosity
+from esgpull.utils import format_size
+
+
+@click.command()
+@opts.simple
+@opts.all
+@opts.verbosity
+def status(
+    simple: bool,
+    all_: bool,
+    verbosity: Verbosity,
+):
+    """
+    View file queue status
+
+    Use the `--all` flag to include already downloaded files (`done` status).
+    """
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("status", onraise=Abort):
+        status_count_size = list(esg.db.rows(sql.file.status_count_size(all_)))
+        table = Table(box=MINIMAL_DOUBLE_HEAD, show_edge=False)
+        table.add_column("status", justify="right", style="bold blue")
+        table.add_column("files", justify="center")
+        table.add_column("size", justify="right", style="magenta")
+        if not status_count_size:
+            esg.ui.print("Queue is empty.")
+            raise Exit(0)
+        if simple:
+            for status, count, size in status_count_size:
+                table.add_row(status.value, str(count), format_size(size))
+        else:
+            esg.graph.load_db()
+            first_line = True
+            for status, count, total_size in status_count_size:
+                first_row = True
+                for query in esg.graph.queries.values():
+                    files = [f for f in query.files if f.status == status]
+                    if files:
+                        if first_row:
+                            if first_line:
+                                first_line = False
+                            else:
+                                table.add_row(end_section=True)
+                            table.add_row(
+                                Text(status.value, justify="left"),
+                                str(count),
+                                format_size(total_size),
+                                style="bold",
+                                end_section=True,
+                            )
+                            first_row = False
+                        table.add_row(
+                            query.rich_name,
+                            str(len(files)),
+                            format_size(sum([f.size for f in files])),
+                        )
+        esg.ui.print(table)

esgpull/cli/track.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import click
+from click.exceptions import Abort, Exit
+from rich.box import MINIMAL_DOUBLE_HEAD
+from rich.table import Table
+from rich.text import Text
+
+from esgpull.cli.decorators import args, opts
+from esgpull.cli.utils import init_esgpull, valid_name_tag
+from esgpull.tui import Verbosity
+
+
+@click.command()
+@args.query_ids
+@opts.record
+@opts.verbosity
+def track(
+    query_ids: tuple[str],
+    record: bool,
+    verbosity: Verbosity,
+) -> None:
+    """
+    Track queries
+
+    As a side effect, tracking a query applies all default options to the query,
+    so that modifications of the config's default options have no impact on
+    previouly tracked queries.
+    """
+    esg = init_esgpull(verbosity, record=record)
+    with esg.ui.logging("track", onraise=Abort):
+        for sha in query_ids:
+            if not valid_name_tag(esg.graph, esg.ui, sha, None):
+                esg.ui.raise_maybe_record(Exit(1))
+            query = esg.graph.get(sha)
+            if query.tracked:
+                esg.ui.print(f"{query.rich_name} is already tracked.")
+                esg.ui.raise_maybe_record(Exit(0))
+            if esg.graph.get_children(query.sha):
+                msg = f"{query.rich_name} has children, track anyway?"
+                if not esg.ui.ask(msg, default=False):
+                    esg.ui.raise_maybe_record(Abort)
+            expanded = esg.graph.expand(query.sha)
+            tracked_query = query.clone(compute_sha=False)
+            tracked_query.track(expanded.options)
+            if query.sha != tracked_query.sha:
+                msg = f"For {query.rich_name} to become tracked, options must be set."
+                esg.ui.print(msg)
+                table = Table(
+                    box=MINIMAL_DOUBLE_HEAD,
+                    show_edge=False,
+                    show_lines=True,
+                )
+                table.add_column(Text("before", justify="center"))
+                table.add_column(Text("after", justify="center"))
+                table.add_row(query, tracked_query)
+                esg.ui.print(table)
+                if not esg.ui.ask("Apply changes?"):
+                    esg.ui.raise_maybe_record(Abort)
+            esg.graph.replace(query, tracked_query)
+            esg.graph.merge()
+            esg.ui.print(f":+1: {tracked_query.rich_name} is now tracked.")
+        esg.ui.raise_maybe_record(Exit(0))
+
+
+@click.command()
+@args.query_ids
+@opts.verbosity
+def untrack(
+    query_ids: tuple[str],
+    verbosity: Verbosity,
+) -> None:
+    """
+    Untrack queries
+    """
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("untrack", onraise=Abort):
+        for sha in query_ids:
+            if not valid_name_tag(esg.graph, esg.ui, sha, None):
+                raise Exit(1)
+            query = esg.graph.get(sha)
+            if not query.tracked:
+                esg.ui.print(f"Query {query.rich_name} is already untracked.")
+                raise Exit(0)
+            query.untrack()
+            esg.graph.merge()
+            esg.ui.print(f":+1: Query {query.rich_name} is no longer tracked.")

esgpull/cli/update.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import click
+from click.exceptions import Abort, Exit
+
+from esgpull.cli.decorators import args, opts
+from esgpull.cli.utils import get_queries, init_esgpull, valid_name_tag
+from esgpull.context import HintsDict, ResultSearch
+from esgpull.exceptions import UnsetOptionsError
+from esgpull.models import File, FileStatus, Query
+from esgpull.tui import Verbosity, logger
+from esgpull.utils import format_size
+
+
+@dataclass
+class QueryFiles:
+    query: Query
+    expanded: Query
+    skip: bool = False
+    files: list[File] = field(default_factory=list)
+    hits: int = field(init=False)
+    hints: HintsDict = field(init=False)
+    results: list[ResultSearch] = field(init=False)
+
+
+@click.command()
+@args.query_id
+@opts.tag
+@opts.children
+@opts.yes
+@opts.record
+@opts.verbosity
+def update(
+    query_id: str | None,
+    tag: str | None,
+    children: bool,
+    yes: bool,
+    record: bool,
+    verbosity: Verbosity,
+) -> None:
+    """
+    Fetch files, link files <-> queries, send files to download queue
+    """
+    esg = init_esgpull(verbosity, record=record)
+    with esg.ui.logging("update", onraise=Abort):
+        # Select which queries to update + setup
+        if query_id is None and tag is None:
+            esg.graph.load_db()
+            queries = list(esg.graph.queries.values())
+        else:
+            if not valid_name_tag(esg.graph, esg.ui, query_id, tag):
+                esg.ui.raise_maybe_record(Exit(1))
+            queries = get_queries(
+                esg.graph,
+                query_id,
+                tag,
+                children=children,
+            )
+        qfs: list[QueryFiles] = []
+        for query in queries:
+            expanded = esg.graph.expand(query.sha)
+            if query.tracked and not expanded.trackable():
+                esg.ui.print(query)
+                raise UnsetOptionsError(query.name)
+            elif query.tracked:
+                qfs.append(QueryFiles(query, expanded))
+        queries = [
+            query
+            for query in queries
+            if query.tracked and query.sha != "LEGACY"
+        ]
+        if not qfs:
+            esg.ui.print(":stop_sign: Trying to update untracked queries.")
+            esg.ui.raise_maybe_record(Exit(0))
+        hints = esg.context.hints(
+            *[qf.expanded for qf in qfs],
+            file=True,
+            facets=["index_node"],
+        )
+        for qf, qf_hints in zip(qfs, hints):
+            qf.hits = sum(esg.context.hits_from_hints(qf_hints))
+            if qf_hints:
+                qf.hints = qf_hints
+            else:
+                qf.skip = True
+        for qf in qfs:
+            s = "s" if qf.hits > 1 else ""
+            esg.ui.print(f"{qf.query.rich_name} -> {qf.hits} file{s}.")
+        total_hits = sum([qf.hits for qf in qfs])
+        if total_hits == 0:
+            esg.ui.print("No files found.")
+            esg.ui.raise_maybe_record(Exit(0))
+        else:
+            esg.ui.print(f"{total_hits} files found.")
+        qfs = [qf for qf in qfs if not qf.skip]
+        # Prepare optimally distributed requests to ESGF
+        # [?] TODO: fetch FastFile first to determine what to fetch in detail later
+        #   It might be interesting for the special case where all files already
+        #   exist in db, then the detailed fetch could be skipped.
+        for qf in qfs:
+            qf_results = esg.context.prepare_search_distributed(
+                qf.expanded,
+                file=True,
+                hints=[qf.hints],
+                max_hits=None,
+            )
+            nb_req = len(qf_results)
+            if nb_req > 50:
+                msg = (
+                    f"{nb_req} requests will be sent to ESGF to"
+                    f" update {qf.query.rich_name}. Send anyway?"
+                )
+                match esg.ui.choice(msg, ["y", "n", "u"], default="n"):
+                    case "u":
+                        esg.ui.print(f"{qf.query.rich_name} is now untracked.")
+                        qf.query.tracked = False
+                        qf_results = []
+                    case "n":
+                        qf_results = []
+                    case _:
+                        ...
+            qf.results = qf_results
+        # Fetch files and update db
+        # [?] TODO: dry_run to print urls here
+        with esg.ui.spinner("Fetching files"):
+            coros = []
+            for qf in qfs:
+                coro = esg.context._files(*qf.results, keep_duplicates=False)
+                coros.append(coro)
+            files = esg.context.sync_gather(*coros)
+            for qf, qf_files in zip(qfs, files):
+                qf.files = qf_files
+        for qf in qfs:
+            shas = {f.sha for f in qf.query.files}
+            new_files: list[File] = []
+            for file in qf.files:
+                if file.sha not in shas:
+                    new_files.append(file)
+            nb_files = len(new_files)
+            if not qf.query.tracked:
+                esg.db.add(qf.query)
+                continue
+            elif nb_files == 0:
+                esg.ui.print(f"{qf.query.rich_name} is already up-to-date.")
+                continue
+            size = sum([file.size for file in new_files])
+            msg = (
+                f"\nUpdating {qf.query.rich_name} with {nb_files}"
+                f" new files ({format_size(size)})."
+                "\nSend to download queue?"
+            )
+            if yes:
+                choice = "y"
+            else:
+                while (
+                    choice := esg.ui.choice(
+                        msg,
+                        choices=["y", "n", "show"],
+                        show_choices=True,
+                    )
+                ) and choice == "show":
+                    esg.ui.print(esg.graph.subgraph(qf.query, parents=True))
+            if choice == "y":
+                legacy = esg.legacy_query
+                has_legacy = legacy.state.persistent
+                for file in new_files:
+                    file_db = esg.db.get(File, file.sha)
+                    if file_db is None:
+                        if esg.db.has_file_id(file):
+                            logger.error(
+                                "File id already exists in database, "
+                                "there might be an error with its checksum"
+                                f"\n{file}"
+                            )
+                            continue
+                        file.status = FileStatus.Queued
+                        file_db = esg.db.merge(file)
+                    elif has_legacy and legacy in file_db.queries:
+                        file_db.queries.remove(legacy)
+                    file_db.queries.append(qf.query)
+                    esg.db.add(file_db)
+        esg.ui.raise_maybe_record(Exit(0))