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

esgpull/migrations/versions/0.8.0_update_tables.py

@@ -0,0 +1,28 @@
+"""update tables
+
+Revision ID: 0.8.0
+Revises: 14c72daea083
+Create Date: 2025-05-15 11:28:10.755003
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '0.8.0'
+down_revision = '14c72daea083'
+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/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/14c72daea083_query_add_column_updated_at.py

@@ -0,0 +1,36 @@
+"""query_add_column_updated_at
+
+Revision ID: 14c72daea083
+Revises: c7c8541fa741
+Create Date: 2025-05-07 14:49:43.993125
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '14c72daea083'
+down_revision = 'c7c8541fa741'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    with op.batch_alter_table('query', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True))
+
+    # Backfill nulls
+    op.execute('UPDATE query SET updated_at = CURRENT_TIMESTAMP WHERE updated_at IS NULL')
+
+    # Make non-nullable
+    with op.batch_alter_table('query', schema=None) as batch_op:
+        batch_op.alter_column('updated_at', nullable=False)
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('query', schema=None) as batch_op:
+        batch_op.drop_column('updated_at')
+
+    # ### end Alembic commands ###

esgpull/migrations/versions/c7c8541fa741_query_add_column_added_at.py

@@ -0,0 +1,37 @@
+"""query_add_column_added_at
+
+Revision ID: c7c8541fa741
+Revises: 0.7.3
+Create Date: 2025-05-05 16:14:57.140262
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'c7c8541fa741'
+down_revision = '0.7.3'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Add as nullable first
+    with op.batch_alter_table('query', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('added_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True))
+
+    # Backfill nulls
+    op.execute('UPDATE query SET added_at = CURRENT_TIMESTAMP WHERE added_at IS NULL')
+
+    # Make non-nullable
+    with op.batch_alter_table('query', schema=None) as batch_op:
+        batch_op.alter_column('added_at', nullable=False)
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('query', schema=None) as batch_op:
+        batch_op.drop_column('added_at')
+
+    # ### 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)

esgpull/models/options.py

@@ -53,7 +53,7 @@ class Options(Base):
     replica: Mapped[Option] = mapped_column(sa.Enum(Option))
     retracted: Mapped[Option] = mapped_column(sa.Enum(Option))
 
-    _distrib_ = Option(False)
+    _distrib_ = Option(True)
     _latest_ = Option(True)
     _replica_ = Option(None)
     _retracted_ = Option(False)

esgpull/models/query.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from collections.abc import Iterator, MutableMapping, Sequence
-from typing import Any, Literal
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Literal
 
 import sqlalchemy as sa
 from rich.console import Console, ConsoleOptions
@@ -11,12 +12,18 @@ from rich.tree import Tree
 from sqlalchemy.orm import Mapped, mapped_column, object_session, relationship
 from typing_extensions import NotRequired, TypedDict
 
+from esgpull import utils
 from esgpull.exceptions import UntrackableQuery
 from esgpull.models.base import Base, Sha
+from esgpull.models.dataset import Dataset
 from esgpull.models.file import FileDict, FileStatus
 from esgpull.models.options import Options
 from esgpull.models.selection import FacetValues, Selection
 from esgpull.models.tag import Tag
+
+if TYPE_CHECKING:
+    from esgpull.models.dataset import Dataset
+
 from esgpull.models.utils import (
     find_int,
     find_str,
@@ -24,7 +31,18 @@ from esgpull.models.utils import (
     rich_measure_impl,
     short_sha,
 )
-from esgpull.utils import format_size
+from esgpull.utils import format_date_iso, format_size
+
+QUERY_DATE_FMT = "%Y-%m-%d %H:%M:%S"
+
+
+def parse_date(d: datetime | str) -> datetime:
+    return utils.parse_date(d, fmt=QUERY_DATE_FMT)
+
+
+def format_date(d: datetime | str) -> str:
+    return utils.format_date(d, fmt=QUERY_DATE_FMT)
+
 
 query_file_proxy = sa.Table(
     "query_file",
@@ -42,9 +60,14 @@ query_tag_proxy = sa.Table(
 
 class File(Base):
     __tablename__ = "file"
+    __table_args__ = (
+        sa.Index("ix_file_dataset_status", "dataset_id", "status"),
+    )
 
     file_id: Mapped[str] = mapped_column(sa.String(255), unique=True)
-    dataset_id: Mapped[str] = mapped_column(sa.String(255))
+    dataset_id: Mapped[str] = mapped_column(
+        sa.String(255), sa.ForeignKey("dataset.dataset_id")
+    )
     master_id: Mapped[str] = mapped_column(sa.String(255))
     url: Mapped[str] = mapped_column(sa.String(255))
     version: Mapped[str] = mapped_column(sa.String(16))
@@ -63,6 +86,11 @@ class File(Base):
         back_populates="files",
         repr=False,
     )
+    dataset: Mapped["Dataset"] = relationship(
+        back_populates="files",
+        init=False,
+        repr=False,
+    )
 
     def _as_bytes(self) -> bytes:
         self_tuple = (self.file_id, self.checksum)
@@ -87,7 +115,7 @@ class File(Base):
             size=source["size"],
         )
         if "status" in source:
-            result.status = FileStatus(source.get("source"))
+            result.status = FileStatus(source.get("status").lower())
         return result
 
     @classmethod
@@ -152,6 +180,8 @@ class QueryDict(TypedDict):
     options: NotRequired[MutableMapping[str, bool | None]]
     selection: NotRequired[MutableMapping[str, FacetValues]]
     files: NotRequired[list[FileDict]]
+    added_at: NotRequired[str]
+    updated_at: NotRequired[str]
 
 
 class Query(Base):
@@ -181,6 +211,14 @@ class Query(Base):
         back_populates="queries",
         repr=False,
     )
+    added_at: Mapped[datetime] = mapped_column(
+        server_default=sa.func.now(),
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        server_default=sa.func.now(),
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
 
     def __init__(
         self,
@@ -191,6 +229,8 @@ class Query(Base):
         options: Options | MutableMapping[str, bool | None] | None = None,
         selection: Selection | MutableMapping[str, FacetValues] | None = None,
         files: list[FileDict] | None = None,
+        added_at: datetime | str | None = None,
+        updated_at: datetime | str | None = None,
     ) -> None:
         self.tracked = tracked
         self.require = require
@@ -219,6 +259,14 @@ class Query(Base):
         if files is not None:
             for file in files:
                 self.files.append(File.fromdict(file))
+        if added_at is not None:
+            self.added_at = parse_date(added_at)
+        else:
+            self.added_at = datetime.now(timezone.utc)
+        if updated_at is not None:
+            self.updated_at = parse_date(updated_at)
+        else:
+            self.updated_at = datetime.now(timezone.utc)
 
     @property
     def has_files(self) -> bool:
@@ -313,6 +361,8 @@ class Query(Base):
             result["options"] = self.options.asdict()
         if self.selection:
             result["selection"] = self.selection.asdict()
+        result["added_at"] = format_date(self.added_at)
+        result["updated_at"] = format_date(self.updated_at)
         return result
 
     def clone(self, compute_sha: bool = True) -> Query:
@@ -360,11 +410,6 @@ class Query(Base):
             self.tags.remove(tag)
         return tag is not None
 
-    def no_require(self) -> Query:
-        cl = self.clone(compute_sha=False)
-        cl._rich_no_require = True  # type: ignore [attr-defined]
-        return cl
-
     def __lshift__(self, child: Query) -> Query:
         result = self.clone(compute_sha=False)
         # if self.name != child.require:
@@ -405,12 +450,16 @@ class Query(Base):
 
     __rich_measure__ = rich_measure_impl
 
-    def _rich_tree(self) -> Tree:
+    def _rich_tree(self, hide_require: bool = False) -> Tree:
         title = Text.from_markup(self.rich_name)
         if not self.tracked:
             title.append(" untracked", style="i red")
+        title.append(
+            f"\n│ added    {format_date_iso(self.added_at)}"
+            f"\n│ updated  {format_date_iso(self.updated_at)}"
+        )
         contents = Table.grid(padding=(0, 1))
-        if not hasattr(self, "_rich_no_require") and self.require is not None:
+        if not hide_require and self.require is not None:
             if len(self.require) == 40:
                 require = Text(short_sha(self.require), style="i green")
             else:
@@ -443,10 +492,44 @@ class Query(Base):
             count_ondisk, size_ondisk = self.files_count_size(FileStatus.Done)
             count_total, size_total = self.files_count_size()
             sizes = f"{format_size(size_ondisk)} / {format_size(size_total)}"
-            lens = f"{count_ondisk}/{count_total}"
-            contents.add_row(
-                "files:", Text(f"{sizes} [{lens}]", style="magenta")
-            )
+            lens = f"{count_ondisk} / {count_total}"
+
+            # Add dataset completion info
+            complete_datasets = 0
+            total_datasets = 0
+            session = object_session(self)
+            orphaned_dataset_count = 0
+
+            if session is not None:
+                from esgpull.models import sql
+
+                dataset_stats = session.execute(
+                    sql.dataset.query_stats(self.sha)
+                ).all()
+
+                # Check for orphaned datasets (dataset_ids from files not in Dataset table)
+                orphaned_dataset_count = (
+                    session.scalar(sql.dataset.orphaned(self.sha)) or 0
+                )
+
+                # Compute counts in Python - simpler and more maintainable
+                total_datasets = len(dataset_stats)
+                complete_datasets = sum(
+                    1 for d in dataset_stats if d.done_count == d.total_files
+                )
+
+            contents.add_row("files:", Text(f"{lens}", style="magenta"))
+            if orphaned_dataset_count > 0:
+                contents.add_row(
+                    "datasets:",
+                    "[magenta]? / ?[/]  [yellow italic]<- update for accurate datasets[/]",
+                )
+            else:
+                contents.add_row(
+                    "datasets:",
+                    f"[magenta]{complete_datasets} / {total_datasets}",
+                )
+            contents.add_row("size:", Text(f"{sizes}", style="magenta"))
         tree = Tree("", hide_root=True, guide_style="dim").add(title)
         if contents.row_count:
             tree.add(contents)

esgpull/models/sql.py

@@ -3,6 +3,7 @@ import functools
 import sqlalchemy as sa
 
 from esgpull.models import Table
+from esgpull.models.dataset import Dataset
 from esgpull.models.facet import Facet
 from esgpull.models.file import FileStatus
 from esgpull.models.query import File, Query, query_file_proxy, query_tag_proxy
@@ -11,15 +12,6 @@ from esgpull.models.synda_file import SyndaFile
 from esgpull.models.tag import Tag
 
 
-def count(item: Table) -> sa.Select[tuple[int]]:
-    table = item.__class__
-    return (
-        sa.select(sa.func.count("*"))
-        .select_from(table)
-        .filter_by(sha=item.sha)
-    )
-
-
 def count_table(table: type[Table]) -> sa.Select[tuple[int]]:
     return sa.select(sa.func.count("*")).select_from(table)
 
@@ -148,6 +140,45 @@ class file:
         return stmt
 
 
+class dataset:
+    @staticmethod
+    @functools.cache
+    def query_stats(query_sha: str) -> sa.Select[tuple[str, int, int]]:
+        return (
+            sa.select(
+                Dataset.dataset_id,
+                Dataset.total_files,
+                sa.func.count(
+                    sa.case((File.status == FileStatus.Done, 1))
+                ).label("done_count"),
+            )
+            .join(File)
+            .join(query_file_proxy)
+            .filter(query_file_proxy.c.query_sha == query_sha)
+            .filter(File.dataset_id.isnot(None))
+            .group_by(Dataset.dataset_id, Dataset.total_files)
+        )
+
+    @staticmethod
+    @functools.cache
+    def orphaned(query_sha: str) -> sa.Select[tuple[int]]:
+        return (
+            sa.select(sa.func.count(sa.distinct(File.dataset_id)))
+            .join(query_file_proxy)
+            .filter(query_file_proxy.c.query_sha == query_sha)
+            .filter(File.dataset_id.isnot(None))
+            .filter(~File.dataset_id.in_(sa.select(Dataset.dataset_id)))
+        )
+
+    @staticmethod
+    @functools.cache
+    def is_complete(dataset: Dataset) -> sa.Select[tuple[bool]]:
+        return sa.select(
+            sa.func.count(sa.case((File.status == FileStatus.Done, 1)))
+            == dataset.total_files
+        ).where(File.dataset_id == dataset.dataset_id)
+
+
 class query:
     @staticmethod
     @functools.cache