esgpull 0.6.3__py3-none-any.whl

esgpull/cli/facet.py

@@ -0,0 +1,23 @@
+import click
+from click.exceptions import Abort
+
+from esgpull.cli.decorators import args, opts
+from esgpull.cli.utils import init_esgpull
+from esgpull.models import sql
+from esgpull.tui import Verbosity
+
+
+@click.command()
+@args.key
+@opts.verbosity
+def facet(
+    key: str | None,
+    verbosity: Verbosity,
+):
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("facet", onraise=Abort):
+        if key is None:
+            results = esg.db.scalars(sql.facet.names())
+        else:
+            results = esg.db.scalars(sql.facet.values(key))
+        esg.ui.print(sorted(results))

esgpull/cli/get.py

@@ -0,0 +1,28 @@
+import click
+import rich
+from click_params import PUBLIC_URL
+
+from esgpull.context import Context
+
+
+@click.command()
+@click.option("--metadata", "-m", is_flag=True, default=False)
+@click.argument("url", nargs=1, type=PUBLIC_URL)
+def get(
+    url: str,
+    metadata: bool,  # display file metadata, no download
+):
+    """
+    Steps:
+        1. validate url (version+filename)
+        2. create Context
+        3. fetch File from Context
+        4. search database (if exist -> [Y/N])
+        5. check file exists in output directory (if exist -> [Y/N]) -> rename
+        6. download
+    """
+    rich.print(url)
+    ctx = Context()
+    print(ctx)
+
+    raise NotImplementedError()

esgpull/cli/install.py

@@ -0,0 +1,85 @@
+# from __future__ import annotations
+
+# import click
+# from click.exceptions import Abort, Exit
+
+# from esgpull import Esgpull
+# from esgpull.cli.decorators import args, opts
+# from esgpull.cli.utils import load_facets
+# from esgpull.db.models import File
+# from esgpull.tui import Verbosity
+# from esgpull.utils import format_size
+
+
+# @click.command()
+# @opts.distrib
+# @opts.dry_run
+# @opts.force
+# @opts.replica
+# @opts.since
+# @opts.selection_file
+# @opts.verbosity
+# @args.facets
+# def install(
+#     facets: list[str],
+#     distrib: bool,
+#     dry_run: bool,
+#     force: bool,
+#     replica: bool | None,
+#     selection_file: str | None,
+#     since: str | None,
+#     verbosity: Verbosity,
+# ) -> None:
+#     esg = Esgpull(verbosity=verbosity)
+#     with (
+#         esg.context(
+#             distrib=distrib,
+#             latest=True,
+#             since=since,
+#             replica=replica,
+#         ) as ctx,
+#         esg.ui.logging("install", onraise=Abort),
+#     ):
+#         load_facets(ctx.query, facets, selection_file)
+#         if not ctx.query.dump():
+#             raise click.UsageError("No search terms provided.")
+#         hits = ctx.file_hits
+#         nb_files = sum(hits)
+#         esg.ui.print(f"Found {nb_files} files.")
+#         if nb_files > 500 and distrib:
+#             # Enable better distrib
+#             ctx.index_nodes = esg.fetch_index_nodes()
+#         if dry_run:
+#             queries = ctx._build_queries_search(
+#                 hits, file=True, max_results=nb_files, offset=0
+#             )
+#             esg.ui.print(queries)
+#             raise Exit(0)
+#         if not force and nb_files > 5000:
+#             nb_req = nb_files // esg.config.api.page_limit
+#             message = f"{nb_req} requests will be send to ESGF. Continue?"
+#             if not esg.ui.ask(message, default=True):
+#                 raise Abort
+#         results = ctx.search(
+#             file=True,
+#             max_results=None,
+#             offset=0,
+#             hits=hits,
+#         )
+#         files = [File.from_dict(result) for result in results]
+#         total_size = sum([file.size for file in files])
+#         esg.ui.print(f"Total size: {format_size(total_size)}")
+#         if not force:
+#             if not esg.ui.ask("Continue?", default=True):
+#                 raise Abort
+#         to_download, already_on_disk = esg.install(*files)
+#         if to_download:
+#             nb = len(to_download)
+#             s = "s" if nb > 1 else ""
+#             esg.ui.print(f"Installed {nb} new file{s} ready for download.")
+#         if already_on_disk:
+#             nb = len(already_on_disk)
+#             s = "s" if nb > 1 else ""
+#             esg.ui.print(f"Tracking {nb} file{s} already downloaded.")
+#         if not to_download and not already_on_disk:
+#             esg.ui.print("All files are already in the database.")

esgpull/cli/link.py

@@ -0,0 +1,105 @@
+# from __future__ import annotations
+
+# import click
+# from click.exceptions import Abort, Exit
+
+# from esgpull import Esgpull
+# from esgpull.cli.decorators import args, opts
+# from esgpull.cli.utils import get_queries, valid_name_tag
+# from esgpull.graph import Graph
+# from esgpull.tui import Verbosity
+
+
+# @click.command()
+# @args.query_id
+# @opts.tag
+# @opts.children
+# @opts.verbosity
+# def link(
+#     query_id: str | None,
+#     tag: str | None,
+#     children: bool,
+#     verbosity: Verbosity,
+# ) -> None:
+#     """
+#     Link known files to query/tag.
+#     Offline mode to avoid searching esgf ?
+#     """
+#     esg = Esgpull(verbosity=verbosity)
+#     with esg.ui.logging("remove", onraise=Abort):
+#         if query_id is None and tag is None:
+#             raise click.UsageError("No query or tag provided.")
+#         if not valid_name_tag(esg.graph, esg.ui, query_id, tag):
+#             raise Exit(1)
+#         queries = get_queries(
+#             esg.graph,
+#             query_id,
+#             tag,
+#             children=children,
+#         )
+#         nb = len(queries)
+#         ies = "ies" if nb > 1 else "y"
+#         graph = Graph(None)
+#         graph.add(*queries)
+#         esg.ui.print(graph)
+#         msg = f"Remove {nb} quer{ies}?"
+#         if not esg.ui.ask(msg, default=True):
+#             raise Abort
+#         for query in queries:
+#             if not children and esg.graph.get_children(query.sha):
+#                 esg.ui.print(
+#                     ":stop_sign: Some queries block "
+#                     f"removal of {query.rich_name}."
+#                 )
+#                 if esg.ui.ask("Show blocking queries?", default=False):
+#                     esg.ui.print(esg.graph.subgraph(query, children=True))
+#                 raise Exit(1)
+#         esg.db.delete(*queries)
+#         esg.ui.print(":+1:")
+
+
+# @click.command()
+# @args.query_id
+# @opts.tag
+# @opts.children
+# @opts.verbosity
+# def unlink(
+#     query_id: str | None,
+#     tag: str | None,
+#     children: bool,
+#     verbosity: Verbosity,
+# ) -> None:
+#     """
+#     Unlink files that were associated to query/tag
+#     """
+#     esg = Esgpull(verbosity=verbosity)
+#     with esg.ui.logging("remove", onraise=Abort):
+#         if query_id is None and tag is None:
+#             raise click.UsageError("No query or tag provided.")
+#         if not valid_name_tag(esg.graph, esg.ui, query_id, tag):
+#             raise Exit(1)
+#         queries = get_queries(
+#             esg.graph,
+#             query_id,
+#             tag,
+#             children=children,
+#         )
+#         nb = len(queries)
+#         ies = "ies" if nb > 1 else "y"
+#         graph = Graph(None)
+#         graph.add(*queries)
+#         esg.ui.print(graph)
+#         msg = f"Remove {nb} quer{ies}?"
+#         if not esg.ui.ask(msg, default=True):
+#             raise Abort
+#         for query in queries:
+#             if not children and esg.graph.get_children(query.sha):
+#                 esg.ui.print(
+#                     ":stop_sign: Some queries block "
+#                     f"removal of {query.rich_name}."
+#                 )
+#                 if esg.ui.ask("Show blocking queries?", default=False):
+#                     esg.ui.print(esg.graph.subgraph(query, children=True))
+#                 raise Exit(1)
+#         esg.db.delete(*queries)
+#         esg.ui.print(":+1:")

esgpull/cli/login.py

@@ -0,0 +1,56 @@
+import click
+from click.exceptions import Abort
+
+from esgpull.auth import Auth, AuthStatus, Credentials
+from esgpull.cli.decorators import opts
+from esgpull.cli.utils import init_esgpull
+from esgpull.constants import PROVIDERS
+from esgpull.tui import Verbosity
+
+
+@click.command()
+@opts.verbosity
+@opts.force
+def login(verbosity: Verbosity, force: bool):
+    """
+    OpenID authentication and certificates renewal
+
+    The first call to `login` is a prompt asking for provider/username/password.
+
+    Subsequent calls check whether the login certificates are valid, renewing them if needed.
+    Renewal can be forced using the `--force` flag.
+    """
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("login", onraise=Abort):
+        cred_file = esg.config.paths.auth / esg.config.credentials.filename
+        if not cred_file.is_file():
+            esg.ui.print("No credentials found.")
+            choices = []
+            providers = list(PROVIDERS)
+            for i, provider in enumerate(providers):
+                choices.append(str(i))
+                esg.ui.print(f"  [{i}] [i green]{provider}[/]")
+            provider_idx = esg.ui.choice(
+                "Select a provider",
+                choices=choices,
+                show_choices=False,
+            )
+            provider = providers[int(provider_idx)]
+            user = esg.ui.prompt("User")
+            password = esg.ui.prompt("Password", password=True)
+            credentials = Credentials(provider, user, password)
+            credentials.write(cred_file)
+            esg.auth = Auth.from_config(esg.config, credentials)
+        renew = force
+        status = esg.auth.status
+        status_name = status.value[0]
+        status_color = status.value[1]
+        esg.ui.print(f"Certificates are [{status_color}]{status_name}[/].")
+        if esg.auth.status == AuthStatus.Expired:
+            renew = renew or esg.ui.ask("Renew?")
+        elif esg.auth.status == AuthStatus.Missing:
+            renew = True
+        if renew:
+            with esg.ui.spinner("Renewing certificates"):
+                esg.auth.renew()
+            esg.ui.print(":+1: Renewed successfully")

esgpull/cli/remove.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+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.graph import Graph
+from esgpull.models import FileStatus
+from esgpull.tui import Verbosity
+from esgpull.utils import format_size
+
+
+@click.command()
+@args.query_id
+@opts.tag
+@opts.children
+@opts.verbosity
+def remove(
+    query_id: str | None,
+    tag: str | None,
+    children: bool,
+    verbosity: Verbosity,
+) -> None:
+    """
+    Remove queries from the database
+
+    A query required by other queries cannot be removed.
+
+    Upon removal, no files are deleted from the database nor from the filesystem, only links to that query are deleted.
+
+    There is currently no mechanism to delete "orphaned" files, as this functionality should be designed carefully and rooted in practical needs.
+    """
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("remove", onraise=Abort):
+        if query_id is None and tag is None:
+            raise click.UsageError("No query or tag provided.")
+        if not valid_name_tag(esg.graph, esg.ui, query_id, tag):
+            raise Exit(1)
+        queries = get_queries(
+            esg.graph,
+            query_id,
+            tag,
+            children=children,
+        )
+        nb = len(queries)
+        ies = "ies" if nb > 1 else "y"
+        graph = Graph(None)
+        graph.add(*queries)
+        esg.ui.print(graph)
+        msg = f"Remove {nb} quer{ies}?"
+        if not esg.ui.ask(msg, default=True):
+            raise Abort
+        for query in queries:
+            if query.has_files:
+                nb, size = query.files_count_size(FileStatus.Done)
+                if nb:
+                    esg.ui.print(
+                        f":stop_sign: {query.rich_name} is linked"
+                        f" to {nb} downloaded files ({format_size(size)})."
+                    )
+                    if not esg.ui.ask("Delete anyway?", default=False):
+                        raise Abort
+            if not children and esg.graph.get_children(query.sha):
+                esg.ui.print(
+                    ":stop_sign: Some queries block"
+                    f" removal of {query.rich_name}."
+                )
+                if esg.ui.ask("Show blocking queries?", default=False):
+                    esg.ui.print(esg.graph.subgraph(query, children=True))
+                raise Exit(1)
+        esg.db.delete(*queries)
+        esg.ui.print(":+1:")

esgpull/cli/retry.py

@@ -0,0 +1,43 @@
+from collections import Counter
+from collections.abc import Sequence
+
+import click
+from click.exceptions import Abort, Exit
+
+from esgpull.cli.decorators import args, opts
+from esgpull.cli.utils import init_esgpull
+from esgpull.models import FileStatus, sql
+from esgpull.tui import Verbosity
+
+
+@click.command()
+@args.status
+@opts.verbosity
+def retry(
+    status: Sequence[FileStatus],
+    verbosity: Verbosity,
+):
+    """
+    Re-queue failed and cancelled downloads
+    """
+    if not status:
+        status = FileStatus.retryable()
+    esg = init_esgpull(verbosity)
+    with esg.ui.logging("retry", onraise=Abort):
+        assert FileStatus.Done not in status
+        assert FileStatus.Queued not in status
+        files = list(esg.db.scalars(sql.file.with_status(*status)))
+        status_str = "/".join(f"[bold red]{s.value}[/]" for s in status)
+        if not files:
+            esg.ui.print(f"No {status_str} files found.")
+            raise Exit(0)
+        counts = Counter(file.status for file in files)
+        for file in files:
+            file.status = FileStatus.Queued
+        esg.db.add(*files)
+        msg = "Sent back to the queue: "
+        msg += ", ".join(
+            f"{count} [bold red]{status.value}[/]"
+            for status, count in counts.items()
+        )
+        esg.ui.print(msg)

esgpull/cli/search.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+import click
+from click.exceptions import Abort, Exit
+
+from esgpull.cli.decorators import args, groups, opts
+from esgpull.cli.utils import filter_keys, init_esgpull, parse_query, totable
+from esgpull.exceptions import PageIndexError
+from esgpull.graph import Graph
+from esgpull.models import Query
+from esgpull.tui import Verbosity
+
+
+@click.command()
+@args.facets
+@groups.query_def
+@groups.query_date
+@groups.display
+@opts.dry_run
+@opts.file
+@opts.facets_hints
+@opts.hints
+@groups.json_yaml
+@opts.detail
+@opts.show
+@opts.yes
+@opts.record
+@opts.verbosity
+# @opts.selection_file
+def search(
+    facets: list[str],
+    ## query_def
+    tags: list[str],
+    require: str | None,
+    distrib: str | None,
+    latest: str | None,
+    replica: str | None,
+    retracted: str | None,
+    date_from: datetime | None,
+    date_to: datetime | None,
+    ## display
+    _all: bool,
+    zero: bool,
+    page: int,
+    ## json_yaml
+    json: bool,
+    yaml: bool,
+    ## ungrouped
+    show: bool,
+    detail: int | None,
+    dry_run: bool,
+    file: bool,
+    facets_hints: bool,
+    hints: list[str] | None,
+    yes: bool,
+    # selection_file: str | None,
+    record: bool,
+    verbosity: Verbosity,
+) -> None:
+    """
+    Search datasets and files on ESGF
+
+    More info
+    """
+    esg = init_esgpull(verbosity, safe=False, record=record)
+    with esg.ui.logging("search", onraise=Abort):
+        query = parse_query(
+            facets=facets,
+            tags=tags,
+            require=require,
+            distrib=distrib,
+            latest=latest,
+            replica=replica,
+            retracted=retracted,
+        )
+        query.compute_sha()
+        esg.graph.resolve_require(query)
+        if show:
+            if json:
+                esg.ui.print(query.asdict(), json=True)
+            elif yaml:
+                esg.ui.print(query.asdict(), yaml=True)
+            else:
+                try:
+                    graph = esg.graph.subgraph(query, parents=True)
+                    esg.ui.print(graph)
+                except KeyError:
+                    esg.ui.print(query)
+            esg.ui.raise_maybe_record(Exit(0))
+        esg.graph.add(query, force=True)
+        query = esg.graph.expand(query.sha)
+        hits = esg.context.hits(
+            query,
+            file=file,
+            date_from=date_from,
+            date_to=date_to,
+        )
+        nb = sum(hits)
+        page_size = esg.config.cli.page_size
+        if detail is not None:
+            page_size = 1
+            page = detail
+        nb_pages = (nb // page_size) or 1
+        offset = page * page_size
+        max_hits = min(page_size, nb - offset)
+        if page > nb_pages:
+            raise PageIndexError(page, nb_pages)
+        elif zero:
+            max_hits = 0
+        elif _all:
+            offset = 0
+            max_hits = nb
+        ids = range(offset, offset + max_hits)
+        if dry_run:
+            search_results = esg.context.prepare_search(
+                query,
+                file=file,
+                hits=hits,
+                offset=offset,
+                max_hits=max_hits,
+                date_from=date_from,
+                date_to=date_to,
+            )
+            for result in search_results:
+                esg.ui.print(result.request.url)
+            esg.ui.raise_maybe_record(Exit(0))
+        if facets_hints:
+            not_distrib_query = query << Query(options=dict(distrib=False))
+            facet_counts = esg.context.hints(
+                not_distrib_query,
+                file=file,
+                facets=["*"],
+                date_from=date_from,
+                date_to=date_to,
+            )
+            esg.ui.print(list(facet_counts[0]), json=True)
+            esg.ui.raise_maybe_record(Exit(0))
+        if hints is not None:
+            facet_counts = esg.context.hints(
+                query,
+                file=file,
+                facets=hints,
+                date_from=date_from,
+                date_to=date_to,
+            )
+            esg.ui.print(facet_counts, json=True)
+            esg.ui.raise_maybe_record(Exit(0))
+        if max_hits > 200 and not yes:
+            nb_req = max_hits // esg.config.api.page_limit
+            message = f"{nb_req} requests will be sent to ESGF. Send anyway?"
+            if not esg.ui.ask(message, default=True):
+                esg.ui.raise_maybe_record(Abort)
+        if detail is not None:
+            queries = esg.context.search_as_queries(
+                query,
+                file=file,
+                hits=hits,
+                offset=offset,
+                max_hits=max_hits,
+                keep_duplicates=True,
+                date_from=date_from,
+                date_to=date_to,
+            )
+            if json:
+                selections = [q.selection.asdict() for q in queries]
+                esg.ui.print(selections, json=True)
+            elif yaml:
+                selections = [q.selection.asdict() for q in queries]
+                esg.ui.print(selections, yaml=True)
+            else:
+                graph = Graph(None)
+                graph.add(*queries, clone=False)
+                esg.ui.print(graph)
+            esg.ui.raise_maybe_record(Exit(0))
+        results = esg.context.search(
+            query,
+            file=file,
+            hits=hits,
+            offset=offset,
+            max_hits=max_hits,
+            keep_duplicates=True,
+            date_from=date_from,
+            date_to=date_to,
+        )
+        if json:
+            esg.ui.print([f.asdict() for f in results], json=True)
+        elif yaml:
+            esg.ui.print([f.asdict() for f in results], yaml=True)
+        else:
+            f_or_d = "file" if file else "dataset"
+            s = "s" if nb != 1 else ""
+            esg.ui.print(f"Found {nb} {f_or_d}{s}.")
+            if results:
+                unique_ids = {r.master_id for r in results}
+                unique_nodes = {(r.master_id, r.data_node) for r in results}
+                needs_data_node = len(unique_nodes) > len(unique_ids)
+                docs = filter_keys(results, ids=ids, data_node=needs_data_node)
+                esg.ui.print(totable(docs))
+        esg.ui.raise_maybe_record(Exit(0))