esgpull 0.8.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

esgpull/cli/update.py

@@ -10,8 +10,8 @@ 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.models import Dataset, File, FileStatus, Query
+from esgpull.tui import Verbosity
 from esgpull.utils import format_size
 
 
@@ -20,10 +20,13 @@ class QueryFiles:
     query: Query
     expanded: Query
     skip: bool = False
+    datasets: list[Dataset] = field(default_factory=list)
     files: list[File] = field(default_factory=list)
+    dataset_hits: int = field(init=False)
     hits: int = field(init=False)
     hints: HintsDict = field(init=False)
     results: list[ResultSearch] = field(init=False)
+    dataset_results: list[ResultSearch] = field(init=False)
 
 
 @click.command()
@@ -80,20 +83,27 @@ def update(
             file=True,
             facets=["index_node"],
         )
-        for qf, qf_hints in zip(qfs, hints):
+        dataset_hits = esg.context.hits(
+            *[qf.expanded for qf in qfs],
+            file=False,
+        )
+        for qf, qf_hints, qf_dataset_hits in zip(qfs, hints, dataset_hits):
             qf.hits = sum(esg.context.hits_from_hints(qf_hints))
             if qf_hints:
                 qf.hints = qf_hints
+                qf.dataset_hits = qf_dataset_hits
             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}.")
+            esg.ui.print(
+                f"{qf.query.rich_name} -> {qf.hits} file{s} (before replica de-duplication)."
+            )
         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:
+        elif len(qfs) > 1:
             esg.ui.print(f"{total_hits} files found.")
         qfs = [qf for qf in qfs if not qf.skip]
         # Prepare optimally distributed requests to ESGF
@@ -101,6 +111,12 @@ def update(
         #   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_dataset_results = esg.context.prepare_search(
+                qf.expanded,
+                file=False,
+                hits=[qf.dataset_hits],
+                max_hits=None,
+            )
             if esg.config.api.use_custom_distribution_algorithm:
                 qf_results = esg.context.prepare_search_distributed(
                     qf.expanded,
@@ -115,7 +131,7 @@ def update(
                     hits=[qf.hits],
                     max_hits=None,
                 )
-            nb_req = len(qf_results)
+            nb_req = len(qf_dataset_results) + len(qf_results)
             if nb_req > 50:
                 msg = (
                     f"{nb_req} requests will be sent to ESGF to"
@@ -126,13 +142,32 @@ def update(
                         esg.ui.print(f"{qf.query.rich_name} is now untracked.")
                         qf.query.tracked = False
                         qf_results = []
+                        qf_dataset_results = []
                     case "n":
                         qf_results = []
+                        qf_dataset_results = []
                     case _:
                         ...
             qf.results = qf_results
+            qf.dataset_results = qf_dataset_results
         # Fetch files and update db
         # [?] TODO: dry_run to print urls here
+        with esg.ui.spinner("Fetching datasets"):
+            coros = []
+            for qf in qfs:
+                coro = esg.context._datasets(
+                    *qf.dataset_results, keep_duplicates=False
+                )
+                coros.append(coro)
+            datasets = esg.context.sync_gather(*coros)
+            for qf, qf_datasets in zip(qfs, datasets):
+                qf.datasets = [
+                    Dataset(
+                        dataset_id=record.dataset_id,
+                        total_files=record.number_of_files,
+                    )
+                    for record in qf_datasets
+                ]
         with esg.ui.spinner("Fetching files"):
             coros = []
             for qf in qfs:
@@ -142,23 +177,37 @@ def update(
             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])
+            with esg.db.commit_context():
+                unregistered_datasets = [
+                    f for f in qf.datasets if f not in esg.db
+                ]
+                if len(unregistered_datasets) > 0:
+                    esg.ui.print(
+                        f"Adding {len(unregistered_datasets)} new datasets to database."
+                    )
+                    esg.db.session.add_all(unregistered_datasets)
+                files_from_db = [
+                    esg.db.get(File, f.sha) for f in qf.files if f in esg.db
+                ]
+                registered_files = [f for f in files_from_db if f is not None]
+                unregistered_files = [f for f in qf.files if f not in esg.db]
+                if len(unregistered_files) > 0:
+                    esg.ui.print(
+                        f"Adding {len(unregistered_files)} new files to database."
+                    )
+                    esg.db.session.add_all(unregistered_files)
+                files = registered_files + unregistered_files
+            not_done_files = [
+                file for file in files if file.status != FileStatus.Done
+            ]
+            download_size = sum(file.size for file in not_done_files)
             msg = (
-                f"\nUpdating {qf.query.rich_name} with {nb_files}"
-                f" new files ({format_size(size)})."
-                "\nSend to download queue?"
+                f"\n{qf.query.rich_name}: {len(not_done_files)} "
+                f" files ({format_size(download_size)}) to download."
+                f"\nLink to query and send to download queue?"
             )
             if yes:
                 choice = "y"
@@ -174,25 +223,19 @@ def update(
             if choice == "y":
                 legacy = esg.legacy_query
                 has_legacy = legacy.state.persistent
+                applied_changes = False
                 with esg.db.commit_context():
                     for file in esg.ui.track(
-                        new_files,
-                        description=qf.query.rich_name,
+                        files,
+                        description=f"{qf.query.rich_name}",
                     ):
-                        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
+                        if file.status != FileStatus.Done:
                             file.status = FileStatus.Queued
-                            esg.db.session.add(file)
-                        elif has_legacy and legacy in file_db.queries:
-                            esg.db.unlink(query=legacy, file=file_db)
-                        esg.db.link(query=qf.query, file=file)
-                    qf.query.updated_at = datetime.now(timezone.utc)
+                        if has_legacy and legacy in file.queries:
+                            _ = esg.db.unlink(query=legacy, file=file)
+                        changed = esg.db.link(query=qf.query, file=file)
+                        applied_changes = applied_changes or changed
+                    if applied_changes:
+                        qf.query.updated_at = datetime.now(timezone.utc)
                     esg.db.session.add(qf.query)
         esg.ui.raise_maybe_record(Exit(0))

esgpull/cli/utils.py

@@ -103,7 +103,7 @@ def totable(docs: list[OrderedDict[str, Any]]) -> Table:
     table = Table(box=MINIMAL_DOUBLE_HEAD, show_edge=False)
     for key in docs[0].keys():
         justify: Literal["left", "right", "center"]
-        if key in ["file", "dataset"]:
+        if key in ["file", "dataset", "plugin"]:
             justify = "left"
         else:
             justify = "right"
@@ -243,3 +243,18 @@ def get_queries(
             kids = graph.get_all_children(query.sha)
             queries.extend(kids)
     return queries
+
+
+def extract_subdict(doc: dict, key: str | None) -> dict:
+    if key is None:
+        return doc
+    for part in key.split("."):
+        if not part:
+            raise KeyError(key)
+        elif part in doc:
+            doc = doc[part]
+        else:
+            raise KeyError(part)
+    for part in key.split(".")[::-1]:
+        doc = {part: doc}
+    return doc

esgpull/config.py

@@ -4,10 +4,10 @@ import logging
 from collections.abc import Iterator, Mapping
 from enum import Enum, auto
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, TypeVar, Union, cast, overload
 
 import tomlkit
-from attrs import Factory, define, field, fields
+from attrs import Factory, define, field
 from attrs import has as attrs_has
 from cattrs import Converter
 from cattrs.gen import make_dict_unstructure_fn, override
@@ -20,6 +20,66 @@ from esgpull.models.options import Options
 
 logger = logging.getLogger("esgpull")
 
+T = TypeVar("T")
+
+
+@overload
+def cast_value(
+    target: str, value: Union[str, int, bool, float], key: str
+) -> str: ...
+
+
+@overload
+def cast_value(
+    target: bool, value: Union[str, int, bool, float], key: str
+) -> bool: ...
+
+
+@overload
+def cast_value(
+    target: int, value: Union[str, int, bool, float], key: str
+) -> int: ...
+
+
+@overload
+def cast_value(
+    target: float, value: Union[str, int, bool, float], key: str
+) -> float: ...
+
+
+def cast_value(
+    target: Any, value: Union[str, int, bool, float], key: str
+) -> Any:
+    if isinstance(value, type(target)):
+        return value
+    elif attrs_has(type(target)):
+        raise KeyError(key)
+    elif isinstance(target, str):
+        return str(value)
+    elif isinstance(target, float):
+        try:
+            return float(value)
+        except Exception:
+            raise ValueError(value)
+    elif isinstance(target, bool):
+        if isinstance(value, str):
+            if value.lower() in ["on", "true"]:
+                return True
+            elif value.lower() in ["off", "false"]:
+                return False
+            else:
+                raise ValueError(value)
+        else:
+            raise TypeError(value)
+    elif isinstance(target, int):
+        # int must be after bool, because isinstance(True, int) == True
+        try:
+            return int(value)
+        except Exception:
+            raise ValueError(value)
+    else:
+        raise TypeError(value)
+
 
 @define
 class Paths:
@@ -28,6 +88,7 @@ class Paths:
     db: Path = field(converter=Path)
     log: Path = field(converter=Path)
     tmp: Path = field(converter=Path)
+    plugins: Path = field(converter=Path)
 
     @auth.default
     def _auth_factory(self) -> Path:
@@ -69,12 +130,21 @@ class Paths:
             root = InstallConfig.default
         return root / "tmp"
 
-    def __iter__(self) -> Iterator[Path]:
+    @plugins.default
+    def _plugins_factory(self) -> Path:
+        if InstallConfig.current is not None:
+            root = InstallConfig.current.path
+        else:
+            root = InstallConfig.default
+        return root / "plugins"
+
+    def values(self) -> Iterator[Path]:
         yield self.auth
         yield self.data
         yield self.db
         yield self.log
         yield self.tmp
+        yield self.plugins
 
 
 @define
@@ -213,6 +283,13 @@ def iter_keys(
             yield local_path
 
 
+@define
+class Plugins:
+    """Configuration for the plugin system"""
+
+    enabled: bool = False
+
+
 @define
 class Config:
     paths: Paths = Factory(Paths)
@@ -221,6 +298,7 @@ class Config:
     db: Db = Factory(Db)
     download: Download = Factory(Download)
     api: API = Factory(API)
+    plugins: Plugins = Factory(Plugins)
     _raw: TOMLDocument | None = field(init=False, default=None)
     _config_file: Path | None = field(init=False, default=None)
 
@@ -287,7 +365,7 @@ class Config:
     def update_item(
         self,
         key: str,
-        value: int | str,
+        value: str | int | bool,
         empty_ok: bool = False,
     ) -> int | str | None:
         if self._raw is None and empty_ok:
@@ -302,29 +380,8 @@ class Config:
             doc.setdefault(part, {})
             doc = doc[part]
             obj = getattr(obj, part)
-        value_type = getattr(fields(type(obj)), last).type
         old_value = getattr(obj, last)
-        if attrs_has(value_type):
-            raise KeyError(key)
-        elif value_type is str:
-            ...
-        elif value_type is int:
-            try:
-                value = value_type(value)
-            except Exception:
-                ...
-        elif value_type is bool:
-            if isinstance(value, bool):
-                ...
-            elif isinstance(value, str):
-                if value.lower() in ["on", "true"]:
-                    value = True
-                elif value.lower() in ["off", "false"]:
-                    value = False
-                else:
-                    raise ValueError(value)
-            else:
-                raise TypeError(value)
+        value = cast_value(old_value, value, key)
         setattr(obj, last, value)
         doc[last] = value
         return old_value

esgpull/constants.py

@@ -1,6 +1,9 @@
+import os
+
 CONFIG_FILENAME = "config.toml"
 INSTALLS_PATH_ENV = "ESGPULL_INSTALLS_PATH"
 ROOT_ENV = "ESGPULL_CURRENT"
+ESGPULL_DEBUG = os.environ.get("ESGPULL_DEBUG", "0") == "1"
 
 IDP = "/esgf-idp/openid/"
 CEDA_IDP = "/OpenID/Provider/server/"

esgpull/context.py

@@ -16,7 +16,7 @@ from rich.pretty import pretty_repr
 
 from esgpull.config import Config
 from esgpull.exceptions import SolrUnstableQueryError
-from esgpull.models import Dataset, File, Query
+from esgpull.models import DatasetRecord, File, Query
 from esgpull.tui import logger
 from esgpull.utils import format_date_iso, index2url, sync
 
@@ -151,7 +151,7 @@ class ResultHints(Result):
 
 @dataclass
 class ResultSearch(Result):
-    data: Sequence[File | Dataset] = field(init=False, repr=False)
+    data: Sequence[File | DatasetRecord] = field(init=False, repr=False)
 
     def process(self) -> None:
         raise NotImplementedError
@@ -159,14 +159,14 @@ class ResultSearch(Result):
 
 @dataclass
 class ResultDatasets(Result):
-    data: Sequence[Dataset] = field(init=False, repr=False)
+    data: Sequence[DatasetRecord] = field(init=False, repr=False)
 
     def process(self) -> None:
         self.data = []
         if self.success:
             for doc in self.json["response"]["docs"]:
                 try:
-                    dataset = Dataset.serialize(doc)
+                    dataset = DatasetRecord.serialize(doc)
                     self.data.append(dataset)
                 except KeyError as exc:
                     logger.exception(exc)
@@ -492,8 +492,8 @@ class Context:
         self,
         *results: ResultSearch,
         keep_duplicates: bool,
-    ) -> list[Dataset]:
-        datasets: list[Dataset] = []
+    ) -> list[DatasetRecord]:
+        datasets: list[DatasetRecord] = []
         ids: set[str] = set()
         async for result in self._fetch(*results):
             dataset_result = result.to(ResultDatasets)
@@ -501,7 +501,7 @@ class Context:
             if dataset_result.processed:
                 for d in dataset_result.data:
                     if not keep_duplicates and d.dataset_id in ids:
-                        logger.warning(f"Duplicate dataset {d.dataset_id}")
+                        logger.debug(f"Duplicate dataset {d.dataset_id}")
                     else:
                         datasets.append(d)
                         ids.add(d.dataset_id)
@@ -520,7 +520,7 @@ class Context:
             if files_result.processed:
                 for file in files_result.data:
                     if not keep_duplicates and file.sha in shas:
-                        logger.warning(f"Duplicate file {file.file_id}")
+                        logger.debug(f"Duplicate file {file.file_id}")
                     else:
                         files.append(file)
                         shas.add(file.sha)
@@ -627,7 +627,7 @@ class Context:
         date_from: datetime | None = None,
         date_to: datetime | None = None,
         keep_duplicates: bool = True,
-    ) -> list[Dataset]:
+    ) -> list[DatasetRecord]:
         if hits is None:
             hits = self.hits(*queries, file=False)
         results = self.prepare_search(

esgpull/database.py

@@ -12,11 +12,13 @@ import sqlalchemy.orm
 from alembic.config import Config as AlembicConfig
 from alembic.migration import MigrationContext
 from alembic.script import ScriptDirectory
+from sqlalchemy.inspection import inspect
 from sqlalchemy.orm import Session, joinedload, make_transient
 
 from esgpull import __file__
 from esgpull.config import Config
 from esgpull.models import File, Query, Table, sql
+from esgpull.models.base import Base, BaseNoSHA
 from esgpull.version import __version__
 
 # from esgpull.exceptions import NoClauseError
@@ -145,14 +147,26 @@ class Database:
         for item in items:
             make_transient(item)
 
-    def link(self, query: Query, file: File):
-        self.session.execute(sql.query_file.link(query, file))
-
-    def unlink(self, query: Query, file: File):
-        self.session.execute(sql.query_file.unlink(query, file))
+    def link(self, query: Query, file: File) -> bool:
+        if not self.session.scalar(sql.query_file.is_linked(query, file)):
+            self.session.execute(sql.query_file.link(query, file))
+            return True
+        else:
+            return False
 
-    def __contains__(self, item: Table) -> bool:
-        return self.scalars(sql.count(item))[0] > 0
+    def unlink(self, query: Query, file: File) -> bool:
+        if self.session.scalar(sql.query_file.is_linked(query, file)):
+            self.session.execute(sql.query_file.unlink(query, file))
+            return True
+        else:
+            return False
+
+    def __contains__(self, item: Base | BaseNoSHA) -> bool:
+        mapper = inspect(item.__class__)
+        pk_col = mapper.primary_key[0]
+        pk_value = getattr(item, pk_col.name)
+        stmt = sa.exists().where(pk_col == pk_value)
+        return self.scalars(sa.select(stmt))[0]
 
     def has_file_id(self, file: File) -> bool:
         return len(self.scalars(sql.file.with_file_id(file.file_id))) == 1

esgpull/download.py

@@ -1,6 +1,7 @@
 # from math import ceil
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass
+from datetime import datetime
 
 from httpx import AsyncClient
 
@@ -19,6 +20,7 @@ class DownloadCtx:
     completed: int = 0
     chunk: bytes | None = None
     digest: Digest | None = None
+    start_time: datetime | None = None
 
     @property
     def finished(self) -> bool:
@@ -54,6 +56,7 @@ class Simple(BaseDownloader):
         ctx: DownloadCtx,
         chunk_size: int,
     ) -> AsyncGenerator[DownloadCtx, None]:
+        ctx.start_time = datetime.now()
         async with client.stream("GET", ctx.file.url) as resp:
             resp.raise_for_status()
             async for chunk in resp.aiter_bytes(chunk_size=chunk_size):

esgpull/esgpull.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import logging
 from collections.abc import AsyncIterator
 from dataclasses import dataclass
+from datetime import datetime
 from functools import cached_property, partial
 from pathlib import Path
 from warnings import warn
@@ -25,6 +26,7 @@ from esgpull.auth import Auth, Credentials
 from esgpull.config import Config
 from esgpull.context import Context
 from esgpull.database import Database
+from esgpull.download import DownloadCtx
 from esgpull.exceptions import (
     DownloadCancelled,
     InvalidInstallPath,
@@ -44,6 +46,13 @@ from esgpull.models import (
     sql,
 )
 from esgpull.models.utils import short_sha
+from esgpull.plugin import (
+    Event,
+    PluginManager,
+    emit,
+    get_plugin_manager,
+    set_plugin_manager,
+)
 from esgpull.processor import Processor
 from esgpull.result import Err, Ok, Result
 from esgpull.tui import UI, DummyLive, Verbosity, logger
@@ -117,6 +126,18 @@ class Esgpull:
         if load_db:
             self.db = Database.from_config(self.config)
             self.graph = Graph(self.db)
+        # Initialize plugin system
+        plugin_config_path = self.config.paths.plugins / "plugins.toml"
+        try:
+            self.plugin_manager = get_plugin_manager()
+            self.plugin_manager.__init__(config_path=plugin_config_path)
+        except ValueError:
+            self.plugin_manager = PluginManager(config_path=plugin_config_path)
+            set_plugin_manager(self.plugin_manager)
+        if self.config.plugins.enabled:
+            self.plugin_manager.enabled = True
+            self.config.paths.plugins.mkdir(exist_ok=True, parents=True)
+            self.plugin_manager.discover_plugins(self.config.paths.plugins)
 
     def fetch_index_nodes(self) -> list[str]:
         """
@@ -309,7 +330,7 @@ class Esgpull:
         progress: Progress,
         task_ids: dict[str, TaskID],
         live: Live | DummyLive,
-    ) -> AsyncIterator[Result]:
+    ) -> AsyncIterator[Result[DownloadCtx]]:
         async for result in processor.process():
             task_idx = progress.task_ids.index(task_ids[result.data.file.sha])
             task = progress.tasks[task_idx]
@@ -348,7 +369,7 @@ class Esgpull:
                                 yield result
                             case Err(_, err):
                                 progress.remove_task(task.id)
-                                yield Err(result.data, err)
+                                yield Err(result.data, err=err)
                 case Err():
                     progress.remove_task(task.id)
                     yield result
@@ -438,15 +459,38 @@ class Esgpull:
                     match result:
                         case Ok():
                             main_progress.update(main_task_id, advance=1)
-                            result.data.file.status = FileStatus.Done
-                            files.append(result.data.file)
-                        case Err():
+                            file = result.data.file
+                            file.status = FileStatus.Done
+                            files.append(file)
+                            emit(
+                                Event.file_complete,
+                                file=file,
+                                destination=self.fs[file].drs,
+                                start_time=result.data.start_time,
+                                end_time=datetime.now(),
+                            )
+                            if file.dataset is not None:
+                                is_dataset_complete = self.db.scalars(
+                                    sql.dataset.is_complete(file.dataset)
+                                )[0]
+                                if is_dataset_complete:
+                                    emit(
+                                        Event.dataset_complete,
+                                        dataset=file.dataset,
+                                    )
+                        case Err(_, err):
                             queue_size -= 1
                             main_progress.update(
                                 main_task_id, total=queue_size
                             )
                             result.data.file.status = FileStatus.Error
                             errors.append(result)
+                            emit(
+                                Event.file_error,
+                                file=result.data.file,
+                                exception=err,
+                            )
+
                     if use_db:
                         self.db.add(result.data.file)
                     remaining_dict.pop(result.data.file.sha, None)