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

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"
 
+    @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 __iter__(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
@@ -151,8 +153,12 @@ class Database:
     def unlink(self, query: Query, file: File):
         self.session.execute(sql.query_file.unlink(query, file))
 
-    def __contains__(self, item: Table) -> bool:
-        return self.scalars(sql.count(item))[0] > 0
+    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)

esgpull/graph.py

@@ -418,7 +418,7 @@ class Graph:
                 if keep_require:
                     query_tree = query._rich_tree()
                 else:
-                    query_tree = query.no_require()._rich_tree()
+                    query_tree = query._rich_tree(hide_require=True)
             if query_tree is not None:
                 tree.add(query_tree)
                 self.fill_tree(query, query_tree)

esgpull/migrations/versions/0.9.0_update_tables.py

@@ -0,0 +1,28 @@
+"""update tables
+
+Revision ID: 0.9.0
+Revises: d14f179e553c
+Create Date: 2025-07-07 14:54:58.433022
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '0.9.0'
+down_revision = 'd14f179e553c'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###

esgpull/migrations/versions/d14f179e553c_file_add_composite_index_dataset_id_.py

@@ -0,0 +1,32 @@
+"""file_add_composite_index_dataset_id_status
+
+Revision ID: d14f179e553c
+Revises: e7edab5d4e4b
+Create Date: 2025-06-18 16:05:35.721085
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'd14f179e553c'
+down_revision = 'e7edab5d4e4b'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('file', schema=None) as batch_op:
+        batch_op.create_index('ix_file_dataset_status', ['dataset_id', 'status'], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('file', schema=None) as batch_op:
+        batch_op.drop_index('ix_file_dataset_status')
+
+    # ### end Alembic commands ###

esgpull/migrations/versions/e7edab5d4e4b_add_dataset_tracking.py

@@ -0,0 +1,39 @@
+"""add_dataset_tracking
+
+Revision ID: e7edab5d4e4b
+Revises: 0.8.0
+Create Date: 2025-05-23 17:38:22.066153
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'e7edab5d4e4b'
+down_revision = '0.8.0'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('dataset',
+    sa.Column('dataset_id', sa.String(length=255), nullable=False),
+    sa.Column('total_files', sa.Integer(), nullable=False),
+    sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False),
+    sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False),
+    sa.PrimaryKeyConstraint('dataset_id')
+    )
+    with op.batch_alter_table('file', schema=None) as batch_op:
+        batch_op.create_foreign_key('fk_file_dataset', 'dataset', ['dataset_id'], ['dataset_id'])
+
+    # ### end Alembic commands ###
+    
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('file', schema=None) as batch_op:
+        batch_op.drop_constraint('fk_file_dataset', type_='foreignkey')
+
+    op.drop_table('dataset')
+    # ### end Alembic commands ###

esgpull/models/__init__.py

@@ -1,7 +1,7 @@
 from typing import TypeVar
 
 from esgpull.models.base import Base
-from esgpull.models.dataset import Dataset
+from esgpull.models.dataset import Dataset, DatasetRecord
 from esgpull.models.facet import Facet
 from esgpull.models.file import FastFile, FileStatus
 from esgpull.models.options import Option, Options
@@ -15,6 +15,7 @@ Table = TypeVar("Table", bound=Base)
 __all__ = [
     "Base",
     "Dataset",
+    "DatasetRecord",
     "Facet",
     "FastFile",
     "File",

esgpull/models/base.py

@@ -16,16 +16,10 @@ T = TypeVar("T")
 Sha = sa.String(40)
 
 
-class Base(MappedAsDataclass, DeclarativeBase):
+# Base class for all models - provides core SQLAlchemy functionality
+class _BaseModel(MappedAsDataclass, DeclarativeBase):
     __dataclass_fields__: ClassVar[dict[str, Field]]
-    __sql_attrs__ = ("id", "sha", "_sa_instance_state", "__dataclass_fields__")
-
-    sha: Mapped[str] = mapped_column(
-        Sha,
-        init=False,
-        repr=False,
-        primary_key=True,
-    )
+    __sql_attrs__ = ("id", "_sa_instance_state", "__dataclass_fields__")
 
     @property
     def _names(self) -> tuple[str, ...]:
@@ -36,15 +30,38 @@ class Base(MappedAsDataclass, DeclarativeBase):
             result += (name,)
         return result
 
+    @property
+    def state(self) -> InstanceState:
+        return cast(InstanceState, sa.inspect(self))
+
+    def asdict(self) -> Mapping[str, Any]:
+        raise NotImplementedError
+
+
+# Base class for models that use SHA as primary key
+class Base(_BaseModel):
+    __abstract__ = True
+    __sql_attrs__ = ("id", "sha", "_sa_instance_state", "__dataclass_fields__")
+
+    sha: Mapped[str] = mapped_column(
+        Sha,
+        init=False,
+        repr=False,
+        primary_key=True,
+    )
+
     def _as_bytes(self) -> bytes:
         raise NotImplementedError
 
     def compute_sha(self) -> None:
         self.sha = sha1(self._as_bytes()).hexdigest()
 
-    @property
-    def state(self) -> InstanceState:
-        return cast(InstanceState, sa.inspect(self))
 
-    def asdict(self) -> Mapping[str, Any]:
-        raise NotImplementedError
+# Base class for models that don't use SHA (e.g., Dataset)
+class BaseNoSHA(_BaseModel):
+    __abstract__ = True
+    __sql_attrs__ = ("id", "_sa_instance_state", "__dataclass_fields__")
+
+
+# Keep SHAKeyMixin for backward compatibility if needed
+SHAKeyMixin = Base

esgpull/models/dataset.py

@@ -1,12 +1,22 @@
 from __future__ import annotations
 
-from dataclasses import asdict, dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
 
+import sqlalchemy as sa
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from esgpull.models.base import BaseNoSHA
 from esgpull.models.utils import find_int, find_str
 
+if TYPE_CHECKING:
+    from esgpull.models.query import File
+
 
 @dataclass
-class Dataset:
+class DatasetRecord:
     dataset_id: str
     master_id: str
     version: str
@@ -15,7 +25,7 @@ class Dataset:
     number_of_files: int
 
     @classmethod
-    def serialize(cls, source: dict) -> Dataset:
+    def serialize(cls, source: dict) -> DatasetRecord:
         dataset_id = find_str(source["instance_id"]).partition("|")[0]
         master_id, version = dataset_id.rsplit(".", 1)
         data_node = find_str(source["data_node"])
@@ -30,5 +40,38 @@ class Dataset:
             number_of_files=number_of_files,
         )
 
-    def asdict(self) -> dict:
-        return asdict(self)
+
+class Dataset(BaseNoSHA):
+    __tablename__ = "dataset"
+
+    dataset_id: Mapped[str] = mapped_column(sa.String(255), primary_key=True)
+    total_files: Mapped[int] = mapped_column(sa.Integer)
+    created_at: Mapped[datetime] = mapped_column(
+        server_default=sa.func.now(),
+        default_factory=lambda: datetime.now(timezone.utc),
+        init=False,
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        server_default=sa.func.now(),
+        default_factory=lambda: datetime.now(timezone.utc),
+        init=False,
+    )
+    files: Mapped[list[File]] = relationship(
+        back_populates="dataset",
+        foreign_keys="[File.dataset_id]",
+        primaryjoin="Dataset.dataset_id==File.dataset_id",
+        default_factory=list,
+        init=False,
+        repr=False,
+    )
+
+    def asdict(self) -> Mapping[str, Any]:
+        return {
+            "dataset_id": self.dataset_id,
+            "total_files": self.total_files,
+            "created_at": self.created_at.isoformat(),
+            "updated_at": self.updated_at.isoformat(),
+        }
+
+    def __hash__(self) -> int:
+        return hash(self.dataset_id)