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

esgpull/cli/__init__.py

@@ -7,9 +7,9 @@ from esgpull import __version__
 from esgpull.cli.add import add
 from esgpull.cli.config import config
 from esgpull.cli.convert import convert
-from esgpull.cli.datasets import datasets
 from esgpull.cli.download import download
 from esgpull.cli.login import login
+from esgpull.cli.plugins import plugins
 from esgpull.cli.remove import remove
 from esgpull.cli.retry import retry
 from esgpull.cli.search import search
@@ -35,13 +35,13 @@ SUBCOMMANDS: list[click.Command] = [
     # autoremove,
     config,
     convert,
-    datasets,
     download,
     # facet,
     # get,
     self,
     # install,
     login,
+    plugins,
     remove,
     retry,
     search,

esgpull/cli/add.py

@@ -89,6 +89,7 @@ def add(
         esg.ui.print(subgraph)
         empty = Query()
         empty.compute_sha()
+        next_steps_queries = []
         for query in queries:
             query.compute_sha()
             esg.graph.resolve_require(query)
@@ -99,7 +100,8 @@ def add(
                 esg.ui.print(f"Skipping existing query: {query.rich_name}")
             else:
                 esg.graph.add(query)
-                esg.ui.print(f"New query added: {query.rich_name}")
+                if query.tracked:
+                    next_steps_queries.append(query)
         new_queries = esg.graph.merge()
         nb = len(new_queries)
         ies = "ies" if nb > 1 else "y"
@@ -107,4 +109,8 @@ def add(
             esg.ui.print(f":+1: {nb} new quer{ies} added.")
         else:
             esg.ui.print(":stop_sign: No new query was added.")
+        if next_steps_queries:
+            esg.ui.print("\nNext steps:\n")
+            for query in next_steps_queries:
+                esg.ui.print(f"\tesgpull update {query.sha[:6]}")
         esg.ui.raise_maybe_record(Exit(0))

esgpull/cli/config.py

@@ -4,26 +4,11 @@ import click
 from click.exceptions import Abort, BadOptionUsage, Exit
 
 from esgpull.cli.decorators import args, opts
-from esgpull.cli.utils import init_esgpull
+from esgpull.cli.utils import extract_subdict, init_esgpull
 from esgpull.config import ConfigKind
 from esgpull.tui import Verbosity
 
 
-def extract_command(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
-
-
 @click.command()
 @args.key
 @args.value
@@ -73,7 +58,7 @@ def config(
                 )
             kind = esg.config.kind
             old_value = esg.config.update_item(key, value, empty_ok=True)
-            info = extract_command(esg.config.dump(), key)
+            info = extract_subdict(esg.config.dump(), key)
             esg.config.write()
             esg.ui.print(info, toml=True)
             if kind == ConfigKind.NoFile:
@@ -86,12 +71,12 @@ def config(
         elif key is not None:
             if default:
                 old_value = esg.config.set_default(key)
-                info = extract_command(esg.config.dump(), key)
+                info = extract_subdict(esg.config.dump(), key)
                 esg.config.write()
                 esg.ui.print(info, toml=True)
                 esg.ui.print(f"Previous value: {old_value}")
             else:
-                info = extract_command(esg.config.dump(), key)
+                info = extract_subdict(esg.config.dump(), key)
                 esg.ui.print(info, toml=True)
         elif generate:
             overwrite = False
@@ -102,8 +87,7 @@ def config(
                 )
                 esg.ui.raise_maybe_record(Exit(0))
             elif esg.config.kind == ConfigKind.Partial and esg.ui.ask(
-                "A config file already exists,"
-                " fill it with missing defaults?",
+                "A config file already exists, fill it with missing defaults?",
                 default=False,
             ):
                 overwrite = True

esgpull/cli/plugins.py

@@ -0,0 +1,398 @@
+from collections import OrderedDict
+from textwrap import dedent
+
+import click
+from click.exceptions import Abort, BadOptionUsage
+
+from esgpull import Esgpull
+from esgpull.cli.decorators import args, opts
+from esgpull.cli.utils import extract_subdict, init_esgpull, totable
+from esgpull.models import Dataset, File, FileStatus
+from esgpull.plugin import Event, emit
+from esgpull.tui import Verbosity
+from esgpull.version import __version__
+
+
+@click.group()
+def plugins():
+    """Manage plugins."""
+    pass
+
+
+def check_enabled(esg: Esgpull):
+    if not esg.plugin_manager.enabled:
+        esg.ui.print("Plugin system is [red]disabled[/]")
+        esg.ui.print("To enable it, run:")
+        esg.ui.print("    esgpull config plugins.enabled true")
+        raise Abort
+
+
+@plugins.command("ls")
+@click.option(
+    "--json", "json_output", is_flag=True, help="Output in JSON format"
+)
+@opts.verbosity
+def list_plugins(verbosity: Verbosity, json_output: bool = False):
+    """List all available plugins and their status."""
+    esg = init_esgpull(verbosity=verbosity)
+
+    with esg.ui.logging("plugins", onraise=Abort):
+        check_enabled(esg)
+        esg.plugin_manager.discover_plugins(
+            esg.config.paths.plugins, load_all=True
+        )
+        if json_output:
+            # Format for JSON output
+            result = {}
+            for name, p in esg.plugin_manager.plugins.items():
+                # Get handlers by event type
+                handlers = {}
+                for h in p.handlers:
+                    event_type = h.event.value
+                    if event_type not in handlers:
+                        handlers[event_type] = []
+                    handlers[event_type].append(
+                        {
+                            "function": h.func.__name__,
+                            "priority": h.priority,
+                        }
+                    )
+
+                result[name] = {
+                    "enabled": esg.plugin_manager.is_plugin_enabled(name),
+                    "handlers": handlers,
+                    "min_version": p.min_version,
+                    "max_version": p.max_version,
+                }
+
+            esg.ui.print(result, json=True)
+        else:
+            # Format for human-readable output
+            if not esg.plugin_manager.plugins:
+                esg.ui.print("No plugins found.")
+                return
+
+            # Prepare table data
+            table_data = []
+            for name, p in esg.plugin_manager.plugins.items():
+                enabled = esg.plugin_manager.is_plugin_enabled(name)
+                status_circle = "🟢" if enabled else "🔴"
+                plugin_name = f"{status_circle} {name}"
+
+                # Collect event handlers with their details
+                handler_rows = []
+                events_by_type = {}
+                for h in p.handlers:
+                    event_type = h.event.value
+                    if event_type not in events_by_type:
+                        events_by_type[event_type] = []
+                    events_by_type[event_type].append(h.func.__name__)
+
+                # Create handler rows with event type and function names
+                for event_type, handler_names in events_by_type.items():
+                    for handler_name in handler_names:
+                        handler_rows.append(
+                            {"event": event_type, "function": handler_name}
+                        )
+
+                # First row with plugin info and first handler
+                if handler_rows:
+                    first_handler = handler_rows[0]
+                    first_row = OrderedDict(
+                        [
+                            ("plugin", plugin_name),
+                            ("event", first_handler["event"]),
+                            ("function", first_handler["function"]),
+                        ]
+                    )
+                    table_data.append(first_row)
+
+                    # Additional rows for remaining handlers
+                    for handler in handler_rows[1:]:
+                        additional_row = OrderedDict(
+                            [
+                                ("plugin", ""),
+                                ("event", handler["event"]),
+                                ("function", handler["function"]),
+                            ]
+                        )
+                        table_data.append(additional_row)
+                else:
+                    # Plugin with no handlers
+                    row = OrderedDict(
+                        [
+                            ("plugin", plugin_name),
+                            ("event", ""),
+                            ("function", ""),
+                        ]
+                    )
+                    table_data.append(row)
+
+            # Create and print table
+            table = totable(table_data)
+            esg.ui.print(table)
+
+
+@plugins.command("enable")
+@click.argument("plugin_name")
+@opts.verbosity
+def enable_plugin_cmd(plugin_name: str, verbosity: Verbosity):
+    """Enable a plugin."""
+    esg = init_esgpull(verbosity=verbosity)
+    plugin_path = esg.config.paths.plugins / f"{plugin_name}.py"
+    if not plugin_path.exists():
+        esg.ui.print(f"Plugin file '{plugin_name}.py' not found.")
+        return
+
+    result = esg.plugin_manager.enable_plugin(plugin_name)
+    if result:
+        esg.ui.print(f"Plugin '{plugin_name}' enabled.")
+    else:
+        esg.ui.print(f"Failed to enable plugin '{plugin_name}'.")
+
+
+@plugins.command("disable")
+@click.argument("plugin_name")
+@opts.verbosity
+def disable_plugin_cmd(plugin_name: str, verbosity: Verbosity):
+    """Disable a plugin."""
+    esg = init_esgpull(verbosity=verbosity)
+
+    # Check if the plugin file exists (even if not loaded)
+    plugin_path = esg.config.paths.plugins / f"{plugin_name}.py"
+    if not plugin_path.exists():
+        esg.ui.print(f"Plugin file '{plugin_name}.py' not found.")
+        return
+
+    result = esg.plugin_manager.disable_plugin(plugin_name)
+    if result:
+        esg.ui.print(f"Plugin '{plugin_name}' disabled.")
+    else:
+        esg.ui.print(f"Failed to disable plugin '{plugin_name}'.")
+
+
+@plugins.command("config")
+@args.key
+@args.value
+@opts.default
+@opts.generate
+@opts.verbosity
+def config_plugin(
+    verbosity: Verbosity,
+    key: str | None,
+    value: str | None,
+    default: bool,
+    generate: bool,
+    # config_set: str | None = None,
+    # config_unset: str | None = None,
+):
+    """View or modify plugin configuration."""
+    esg = init_esgpull(verbosity=verbosity)
+
+    with esg.ui.logging("plugins", onraise=Abort):
+        check_enabled(esg)
+        esg.plugin_manager.discover_plugins(
+            esg.config.paths.plugins, load_all=True
+        )
+        if key is not None:
+            plugin_name = key.split(".", 1)[0]
+            plugin_path = esg.config.paths.plugins / f"{plugin_name}.py"
+            if not plugin_path.exists():
+                esg.ui.print(f"Plugin file '{plugin_name}.py' not found.")
+                raise Abort
+        if key is not None and value is not None:
+            """update config"""
+            if default:
+                raise BadOptionUsage(
+                    "default",
+                    dedent(
+                        f"""
+                        --default/-d is invalid with a value.
+                        Instead use:
+
+                        $ esgpull plugins config {key} -d
+                        """
+                    ),
+                )
+            # Split key into plugin_name and remaining_path
+            plugin_name, *rest = key.split(".", 1)
+            remaining_path = rest[0] if rest else None
+
+            # Set the configuration value
+            if remaining_path:
+                old_value = esg.plugin_manager.set_plugin_config(
+                    plugin_name, remaining_path, value
+                )
+                info = extract_subdict(esg.plugin_manager.config.plugins, key)
+                esg.ui.print(info, toml=True)
+                esg.ui.print(f"Previous value: {old_value}")
+            else:
+                esg.ui.print(f"Cannot set plugin name directly: {plugin_name}")
+                raise Abort
+        elif key is not None:
+            if default:
+                # Split key into plugin_name and remaining_path
+                plugin_name, *rest = key.split(".", 1)
+                remaining_path = rest[0] if rest else None
+
+                if remaining_path:
+                    esg.plugin_manager.unset_plugin_config(
+                        plugin_name, remaining_path
+                    )
+                    msg = f":+1: Config reset to default for {key}"
+                    esg.ui.print(msg)
+                else:
+                    esg.ui.print(
+                        f"Cannot reset plugin directly: {plugin_name}"
+                    )
+                    raise Abort
+            else:
+                config = esg.plugin_manager.config.plugins
+                config = extract_subdict(config, key)
+                esg.ui.print(config, toml=True)
+        elif generate:
+            # Write the config file with all current settings
+            esg.plugin_manager.write_config(generate_full_config=True)
+            msg = f":+1: Plugin config generated at {esg.plugin_manager.config_path}"
+            esg.ui.print(msg)
+
+        else:
+            esg.ui.rule(str(esg.plugin_manager.config_path))
+            esg.ui.print(esg.plugin_manager.config.plugins, toml=True)
+
+
+@plugins.command("test")
+@click.argument(
+    "event_type",
+    type=click.Choice([e.value for e in Event]),
+)
+@opts.verbosity
+def test_plugin(
+    event_type: str,
+    verbosity: Verbosity,
+):
+    """Test plugin events."""
+    esg = init_esgpull(verbosity=verbosity)
+    with esg.ui.logging("plugins", onraise=Abort):
+        check_enabled(esg)
+        event = Event(event_type)
+        # Create sample test data
+        file1 = File(
+            file_id="file1",
+            dataset_id="dataset",
+            master_id="master",
+            url="file",
+            version="v0",
+            filename="file.nc",
+            local_path="project/folder",
+            data_node="data_node",
+            checksum="0",
+            checksum_type="0",
+            size=2**42,
+            status=FileStatus.Queued,
+        )
+        file2 = file1.clone()
+        file2.file_id = "file2"
+        file2.status = FileStatus.Done
+        error = ValueError("Placeholder example error")
+        dataset = Dataset(dataset_id="dataset", total_files=2)
+        dataset.files = [file1, file2]
+        emit(
+            event,
+            file=file1,
+            dataset=dataset,
+            error=error,
+        )
+
+
+@plugins.command("create")
+@click.argument(
+    "events",
+    nargs=-1,
+    type=click.Choice([e.value for e in Event]),
+)
+@click.option("-n", "--name", required=True, type=str)
+@opts.verbosity
+def create_plugin(
+    name: str, verbosity: Verbosity, events: list[str] | None = None
+):
+    """Create a new plugin template."""
+    esg = init_esgpull(verbosity=verbosity)
+
+    # Create plugin file path
+    plugin_path = esg.config.paths.plugins / f"{name}.py"
+    if plugin_path.exists():
+        esg.ui.print(f"Plugin file already exists: {plugin_path}")
+        return
+
+    # If no events specified, include them all
+    if not events:
+        events = [e.value for e in Event]
+
+    # Generate template
+    template = f"""# {name} plugin for ESGPull
+# 
+# This plugin was auto-generated. Edit as needed.
+from datetime import datetime
+from pathlib import Path
+from logging import Logger
+
+from esgpull.models import File, Query
+from esgpull.plugin import Event, on
+
+# Specify version compatibility (optional)
+MIN_ESGPULL_VERSION = "{__version__}"
+MAX_ESGPULL_VERSION = None
+
+# Configuration class for the plugin (optional)
+class Config:
+    \"""Configuration for {name} plugin\"""
+    # Add your configuration options here
+    timeout = 30  # Example config option
+    log_level = "INFO"  # Example config option
+
+"""
+
+    # Add handlers for requested events
+    handlers = {
+        "file_complete": """
+# File complete event handler
+@on(Event.file_complete, priority="normal")
+def handle_file_complete(
+    file: File, 
+    destination: Path, 
+    start_time: datetime,
+    end_time: datetime,
+    logger: Logger
+):
+    \"""Handle file complete event\"""
+    logger.info(f"File downloaded: {file.filename}")
+    # Add your custom logic here
+""",
+        "file_error": """
+# Download error event handler
+@on(Event.file_error, priority="normal")
+def handle_file_error(file: File, exception: Exception, logger: Logger):
+    \"""Handle file error event\"""
+    logger.error(f"Download failed for {file.filename}: {exception}")
+    # Add your custom logic here
+""",
+        "dataset_complete": """
+# Dataset complete event handler
+@on(Event.dataset_complete, priority="normal")
+def handle_dataset_complete(dataset: Dataset, logger: Logger):
+    \"""Handle dataset complete event\"""
+    logger.error(f"Dataset downloaded: {dataset.dataset_id}")
+    # Add your custom logic here
+""",
+    }
+
+    for event in events:
+        template += handlers[event]
+
+    # Write the plugin file
+    with open(plugin_path, "w") as f:
+        f.write(template)
+
+    esg.ui.print(f"Plugin template created at: {plugin_path}")
+    esg.ui.print("Edit the file to implement your custom plugin logic.")

esgpull/cli/update.py

@@ -10,7 +10,7 @@ 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.models import Dataset, File, FileStatus, Query
 from esgpull.tui import Verbosity, logger
 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,26 @@ 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)
+            new_files = [f for f in qf.files if f not in esg.db]
+            new_datasets = [d for d in qf.datasets if d not in esg.db]
+            nb_datasets = len(new_datasets)
             nb_files = len(new_files)
             if not qf.query.tracked:
                 esg.db.add(qf.query)
                 continue
-            elif nb_files == 0:
+            elif nb_datasets == 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])
+            if size > 0:
+                queue_msg = " and send new files to download queue"
+            else:
+                queue_msg = ""
             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}: {nb_files} new"
+                f" files, {nb_datasets} new datasets"
+                f" ({format_size(size)})."
+                f"\nUpdate the database{queue_msg}?"
             )
             if yes:
                 choice = "y"
@@ -175,9 +213,14 @@ def update(
                 legacy = esg.legacy_query
                 has_legacy = legacy.state.persistent
                 with esg.db.commit_context():
+                    for dataset in esg.ui.track(
+                        new_datasets,
+                        description=f"{qf.query.rich_name} (datasets)",
+                    ):
+                        esg.db.session.add(dataset)
                     for file in esg.ui.track(
                         new_files,
-                        description=qf.query.rich_name,
+                        description=f"{qf.query.rich_name} (files)",
                     ):
                         file_db = esg.db.get(File, file.sha)
                         if file_db is None:

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