esgpull 0.6.3__py3-none-any.whl

esgpull/models/sql.py

@@ -0,0 +1,258 @@
+import functools
+
+import sqlalchemy as sa
+
+from esgpull.models import Table
+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
+from esgpull.models.selection import Selection, selection_facet_proxy
+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)
+
+
+class facet:
+    @staticmethod
+    @functools.cache
+    def all() -> sa.Select[tuple[Facet]]:
+        return sa.select(Facet)
+
+    @staticmethod
+    @functools.cache
+    def shas() -> sa.Select[tuple[str]]:
+        return sa.select(Facet.sha)
+
+    @staticmethod
+    @functools.cache
+    def name_count() -> sa.Select[tuple[str, int]]:
+        return sa.select(Facet.name, sa.func.count("*")).group_by(Facet.name)
+
+    @staticmethod
+    @functools.cache
+    def usage() -> sa.Select[tuple[Facet, int]]:
+        return (
+            sa.select(Facet, sa.func.count("*"))
+            .join(selection_facet_proxy)
+            .group_by(Facet.sha)
+        )
+
+    @staticmethod
+    def known_shas(shas: list[str]) -> sa.Select[tuple[str]]:
+        return sa.select(Facet.sha).where(Facet.sha.in_(shas))
+
+    @staticmethod
+    @functools.cache
+    def names() -> sa.Select[tuple[str]]:
+        return sa.select(Facet.name).distinct()
+
+    @staticmethod
+    def values(name: str) -> sa.Select[tuple[str]]:
+        return sa.select(Facet.value).where(Facet.name == name)
+
+
+class file:
+    @staticmethod
+    @functools.cache
+    def all() -> sa.Select[tuple[File]]:
+        return sa.select(File)
+
+    @staticmethod
+    @functools.cache
+    def shas() -> sa.Select[tuple[str]]:
+        return sa.select(File.sha)
+
+    @staticmethod
+    @functools.cache
+    def orphans() -> sa.Select[tuple[File]]:
+        return (
+            sa.select(File)
+            .outerjoin(query_file_proxy)
+            .filter_by(file_sha=None)
+        )
+
+    @staticmethod
+    @functools.cache
+    def linked() -> sa.Select[tuple[File]]:
+        return sa.select(query_file_proxy.c.file_sha).distinct()
+
+    __dups_cte: sa.CTE = (
+        sa.select(File.master_id)
+        .group_by(File.master_id)
+        .having(sa.func.count("*") > 1)
+        .cte()
+    )
+
+    @staticmethod
+    @functools.cache
+    def duplicates() -> sa.Select[tuple[File]]:
+        return sa.select(File).join(
+            file.__dups_cte,
+            File.master_id == file.__dups_cte.c.master_id,
+        )
+
+    @staticmethod
+    def shas_from_query(query_sha: str) -> sa.Select[tuple[str]]:
+        return sa.select(query_file_proxy.c.file_sha).filter_by(
+            query_sha=query_sha
+        )
+
+    @staticmethod
+    def with_status(*status: FileStatus) -> sa.Select[tuple[File]]:
+        return sa.select(File).where(File.status.in_(status))
+
+    @staticmethod
+    def with_file_id(file_id: str) -> sa.Select[tuple[str]]:
+        return sa.select(File.sha).where(File.file_id == file_id).limit(1)
+
+    @staticmethod
+    def total_size_with_status(
+        *status: FileStatus,
+        query_sha: str | None = None,
+    ) -> sa.Select[tuple[int]]:
+        """
+        This is re-implemented in Query.files_count_size because
+        of cyclic import between query.py and the current file.
+        """
+        stmt = sa.select(sa.func.sum(File.size).where(File.status.in_(status)))
+        if query_sha is not None:
+            stmt = stmt.join_from(query_file_proxy, File).where(
+                query_file_proxy.c.query_sha == query_sha
+            )
+        return stmt
+
+    @staticmethod
+    @functools.cache
+    def status_count_size(
+        all_: bool = False,
+    ) -> sa.Select[tuple[FileStatus, int, int]]:
+        stmt = sa.select(
+            File.status,
+            sa.func.count("*"),
+            sa.func.sum(File.size),
+        ).group_by(File.status)
+        if not all_:
+            stmt = stmt.where(File.status != FileStatus.Done)
+        return stmt
+
+
+class query:
+    @staticmethod
+    @functools.cache
+    def all() -> sa.Select[tuple[Query]]:
+        return sa.select(Query)
+
+    @staticmethod
+    @functools.cache
+    def shas() -> sa.Select[tuple[str]]:
+        return sa.select(Query.sha)
+
+    __tag_query_cte: sa.CTE = (
+        sa.select(Tag.name, query_tag_proxy.c.query_sha)
+        .join_from(query_tag_proxy, Tag)
+        .cte("tag_query_cte")
+    )
+
+    __name_cte: sa.CTE = (
+        sa.select(__tag_query_cte)
+        .group_by(__tag_query_cte.c.name)
+        .having(sa.func.count("*") == 1)
+        .cte("name_cte")
+    )
+
+    __sha_cte: sa.CTE = (
+        sa.select(__tag_query_cte)
+        .group_by(__tag_query_cte.c.query_sha)
+        .having(sa.func.count("*") == 1)
+        .cte("sha_cte")
+    )
+
+    @staticmethod
+    @functools.cache
+    def name_sha() -> sa.Select[tuple[str, str]]:
+        return sa.select(query.__name_cte).join(
+            query.__sha_cte,
+            query.__name_cte.c.query_sha == query.__sha_cte.c.query_sha,
+        )
+
+    @staticmethod
+    def with_shas(*shas: str) -> sa.Select[tuple[Query]]:
+        if not shas:
+            raise ValueError(shas)
+        return sa.select(Query).where(Query.sha.in_(shas))
+
+    @staticmethod
+    def with_tag(tag: str) -> sa.Select[tuple[Query]]:
+        return (
+            sa.select(Query)
+            .join_from(query_tag_proxy, Tag)
+            .join_from(query_tag_proxy, Query)
+            .where(Tag.name == tag)
+        )
+
+    @staticmethod
+    def children(sha: str) -> sa.Select[tuple[Query]]:
+        return sa.select(Query).where(Query.require == sha)
+
+
+class selection:
+    @staticmethod
+    @functools.cache
+    def all() -> sa.Select[tuple[Selection]]:
+        return sa.select(Selection)
+
+    @staticmethod
+    @functools.cache
+    def orphans() -> sa.Select[tuple[Selection]]:
+        return (
+            sa.select(Selection)
+            .outerjoin(Query)
+            .where(Query.sha == None)  # noqa
+        )
+
+
+class tag:
+    @staticmethod
+    @functools.cache
+    def all() -> sa.Select[tuple[Tag]]:
+        return sa.select(Tag)
+
+    @staticmethod
+    @functools.cache
+    def shas() -> sa.Select[tuple[str]]:
+        return sa.select(Tag.sha)
+
+    @staticmethod
+    @functools.cache
+    def orphans() -> sa.Select[tuple[Tag]]:
+        return (
+            sa.select(Tag).outerjoin(query_tag_proxy).filter_by(tag_sha=None)
+        )
+
+
+class synda_file:
+    @staticmethod
+    @functools.cache
+    def all() -> sa.Select[tuple[SyndaFile]]:
+        return sa.select(SyndaFile)
+
+    @staticmethod
+    @functools.cache
+    def ids() -> sa.Select[tuple[int]]:
+        return sa.select(SyndaFile.file_id)
+
+    @staticmethod
+    def with_ids(*ids: int) -> sa.Select[tuple[SyndaFile]]:
+        return sa.select(SyndaFile).where(SyndaFile.file_id.in_(ids))

esgpull/models/synda_file.py

@@ -0,0 +1,85 @@
+from sqlalchemy.orm import (
+    DeclarativeBase,
+    Mapped,
+    MappedAsDataclass,
+    mapped_column,
+)
+
+from esgpull.models.file import FileStatus
+from esgpull.models.query import File
+
+SyndaStatusMap = {
+    "running": FileStatus.Started,
+    "waiting": FileStatus.Queued,
+}
+
+
+class SyndaBase(MappedAsDataclass, DeclarativeBase):
+    pass
+
+
+class SyndaFile(SyndaBase):
+    __tablename__ = "file"
+
+    url: Mapped[str]
+    file_functional_id: Mapped[str]
+    filename: Mapped[str]
+    local_path: Mapped[str]
+    data_node: Mapped[str]
+    checksum: Mapped[str]
+    checksum_type: Mapped[str]
+    duration: Mapped[int]
+    size: Mapped[int]
+    rate: Mapped[int]
+    start_date: Mapped[str]
+    end_date: Mapped[str]
+    crea_date: Mapped[str]
+    status: Mapped[str]
+    error_msg: Mapped[str]
+    sdget_status: Mapped[str]
+    sdget_error_msg: Mapped[str]
+    priority: Mapped[int]
+    tracking_id: Mapped[str]
+    model: Mapped[str]
+    project: Mapped[str]
+    variable: Mapped[str]
+    last_access_date: Mapped[str]
+    dataset_id: Mapped[int]
+    insertion_group_id: Mapped[int]
+    timestamp: Mapped[str]
+    file_id: Mapped[int] = mapped_column(init=False, primary_key=True)
+
+    def get_status(self) -> FileStatus:
+        s = self.status.lower()
+        result: FileStatus
+        if FileStatus.contains(s):
+            result = FileStatus(s)
+        elif s in SyndaStatusMap:
+            result = SyndaStatusMap[s]
+        else:
+            raise ValueError(s)
+        return result
+
+    def to_file(self) -> File:
+        file_id = self.file_functional_id
+        dataset_id = file_id.removesuffix(self.filename).strip(".")
+        dataset_master, version = dataset_id.rsplit(".", 1)
+        master_id = ".".join([dataset_master, self.filename])
+        url = self.url.replace("http://", "https://")
+        local_path = self.local_path.removesuffix(self.filename).strip("/")
+        result = File(
+            file_id=file_id,
+            dataset_id=dataset_id,
+            master_id=master_id,
+            url=url,
+            version=version,
+            filename=self.filename,
+            local_path=local_path,
+            data_node=self.data_node,
+            checksum=self.checksum,
+            checksum_type=self.checksum_type.upper(),
+            size=self.size,
+            status=self.get_status(),
+        )
+        result.compute_sha()
+        return result

esgpull/models/tag.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import sqlalchemy as sa
+from sqlalchemy.orm import Mapped, mapped_column
+
+from esgpull.models.base import Base
+
+
+class Tag(Base):
+    __tablename__ = "tag"
+
+    name: Mapped[str] = mapped_column(sa.String(255))
+    description: Mapped[str | None] = mapped_column(sa.Text, default=None)
+
+    def _as_bytes(self) -> bytes:
+        return self.name.encode()
+
+    def __hash__(self) -> int:
+        return hash(self._as_bytes())

esgpull/models/utils.py

@@ -0,0 +1,54 @@
+from rich.console import Console, ConsoleOptions
+from rich.measure import Measurement, measure_renderables
+
+
+def short_sha(sha: str) -> str:
+    return f"<{sha[:6]}>"
+
+
+def rich_measure_impl(
+    self,
+    console: Console,
+    options: ConsoleOptions,
+) -> Measurement:
+    renderables = list(self.__rich_console__(console, options))
+    return measure_renderables(console, options, renderables)
+
+
+def find_str(container: list | str) -> str:
+    if isinstance(container, list):
+        return find_str(container[0])
+    elif isinstance(container, str):
+        return container
+    else:
+        raise ValueError(container)
+
+
+def find_int(container: list | int) -> int:
+    if isinstance(container, list):
+        return find_int(container[0])
+    elif isinstance(container, int):
+        return container
+    else:
+        raise ValueError(container)
+
+
+def get_local_path(source: dict, version: str) -> str:
+    flat_raw = {}
+    for k, v in source.items():
+        if isinstance(v, list) and len(v) == 1:
+            flat_raw[k] = v[0]
+        else:
+            flat_raw[k] = v
+    template = find_str(flat_raw["directory_format_template_"])
+    # format: "%(a)/%(b)/%(c)/..."
+    template = template.removeprefix("%(root)s/")
+    template = template.replace("%(", "{")
+    template = template.replace(")s", "}")
+    flat_raw.pop("version", None)
+    if "rcm_name" in flat_raw:  # cordex special case
+        institute = flat_raw["institute"]
+        rcm_name = flat_raw["rcm_name"]
+        rcm_model = institute + "-" + rcm_name
+        flat_raw["rcm_model"] = rcm_model
+    return template.format(version=version, **flat_raw)

esgpull/presets.py

@@ -0,0 +1,13 @@
+# from esgpull.models import Query
+
+
+# IPCC_SCENARIOS = Query(
+#     select=dict(
+#         query=(
+#             "experiment:(rcp26 rcp45 rcp60 rcp85) OR "
+#             "experiment_id:(ssp119 ssp126 ssp245 ssp370 ssp460 ssp585)"
+#         )
+#     )
+# )
+
+# TEMPERATURE = Query(select=dict(variable=["tas", "tos", "tasmin", "tasmax"]))

esgpull/processor.py

@@ -0,0 +1,172 @@
+import asyncio
+import ssl
+from collections.abc import AsyncIterator
+from functools import partial
+from typing import TypeAlias
+
+from aiostream.stream import merge
+from httpx import AsyncClient, HTTPError
+
+from esgpull.auth import Auth
+from esgpull.config import Config
+from esgpull.download import DownloadCtx, Simple
+from esgpull.exceptions import DownloadSizeError
+from esgpull.fs import Digest, Filesystem
+from esgpull.models import File
+from esgpull.result import Err, Ok, Result
+from esgpull.tui import logger
+
+# Callback: TypeAlias = Callable[[], None] | partial[None]
+Callback: TypeAlias = partial[None]
+
+default_ssl_context: ssl.SSLContext | bool = False
+default_ssl_context_loaded = False
+
+
+def load_default_ssl_context() -> str:
+    global default_ssl_context
+    global default_ssl_context_loaded
+    if ssl.OPENSSL_VERSION_INFO[0] >= 3:
+        default_ssl_context = ssl.create_default_context()
+        default_ssl_context.options |= 0x4
+        msg = "Using openssl 3 or higher"
+    else:
+        default_ssl_context = True
+        msg = "Using openssl 1"
+    default_ssl_context_loaded = True
+    return msg
+
+
+class Task:
+    def __init__(
+        self,
+        config: Config,
+        auth: Auth,
+        fs: Filesystem,
+        # *,
+        # url: str | None = None,
+        file: File,
+        start_callbacks: list[Callback] | None = None,
+    ) -> None:
+        self.config = config
+        self.auth = auth
+        self.fs = fs
+        self.ctx = DownloadCtx(file)
+        if not self.config.download.disable_checksum:
+            self.ctx.digest = Digest(file)
+        # if file is None and url is not None:
+        #     self.file = self.fetch_file(url)
+        # elif file is not None:
+        #     self.file = file
+        # else:
+        #     raise ValueError("no arguments")
+        self.downloader = Simple()
+        msg: str | None = None
+        if not default_ssl_context_loaded:
+            msg = load_default_ssl_context()
+        self.ssl_context: ssl.SSLContext | bool
+        if self.config.download.disable_ssl:
+            self.ssl_context = False
+        else:
+            if msg is not None:
+                logger.info(msg)
+            self.ssl_context = default_ssl_context
+        if start_callbacks is None:
+            self.start_callbacks = []
+        else:
+            self.start_callbacks = start_callbacks
+
+    @property
+    def file(self) -> File:
+        return self.ctx.file
+
+    # def fetch_file(self, url: str) -> File:
+    #     ctx = Context()
+    #     # [?]TODO: define map data_node->index_node to find url-file
+    #     # ctx.query.index_node = ...
+    #     ctx.query.title = Path(url).name
+    #     results = ctx.search(file=True)
+    #     for res in results:
+    #         file = File.from_dict(res)
+    #         if file.version in url:
+    #             return file
+    #     raise ValueError(f"{url} is not valid")
+
+    async def stream(
+        self, semaphore: asyncio.Semaphore
+    ) -> AsyncIterator[Result]:
+        ctx = self.ctx
+        try:
+            async with (
+                semaphore,
+                self.fs.open(ctx.file) as file_obj,
+                AsyncClient(
+                    follow_redirects=True,
+                    cert=self.auth.cert,
+                    verify=self.ssl_context,
+                    timeout=self.config.download.http_timeout,
+                ) as client,
+            ):
+                for callback in self.start_callbacks:
+                    callback()
+                stream = self.downloader.stream(
+                    client,
+                    ctx,
+                    self.config.download.chunk_size,
+                )
+                async for ctx in stream:
+                    if ctx.chunk is not None:
+                        await file_obj.write(ctx.chunk)
+                    if ctx.error:
+                        err = DownloadSizeError(ctx.completed, ctx.file.size)
+                        yield Err(ctx, err)
+                        await stream.aclose()
+                        break
+                    elif ctx.finished:
+                        await file_obj.to_done()
+                    yield Ok(ctx)
+        except (
+            HTTPError,
+            DownloadSizeError,
+            GeneratorExit,
+            ssl.SSLError
+            # KeyboardInterrupt,
+        ) as err:
+            yield Err(ctx, err)
+
+
+class Processor:
+    def __init__(
+        self,
+        config: Config,
+        auth: Auth,
+        fs: Filesystem,
+        files: list[File],
+        start_callbacks: dict[str, list[Callback]],
+    ) -> None:
+        self.config = config
+        self.fs = fs
+        self.files = list(filter(self.should_download, files))
+        self.tasks = []
+        for file in files:
+            task = Task(
+                config=config,
+                auth=auth,
+                fs=fs,
+                file=file,
+                start_callbacks=start_callbacks[file.sha],
+            )
+            self.tasks.append(task)
+
+    def should_download(self, file: File) -> bool:
+        if self.fs[file].drs.is_file():
+            return False
+        else:
+            return True
+
+    async def process(self) -> AsyncIterator[Result]:
+        semaphore = asyncio.Semaphore(self.config.download.max_concurrent)
+        streams = [task.stream(semaphore) for task in self.tasks]
+        async with merge(*streams).stream() as stream:
+            async for result in stream:
+                yield result

esgpull/result.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass, field
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass
+class Result(Generic[T]):
+    data: T
+    ok: bool
+    err: BaseException | None
+
+
+@dataclass
+class Ok(Result[T]):
+    ok: bool = field(default=True, init=False)
+    err: None = field(default=None, init=False)
+
+
+@dataclass
+class Err(Result[T]):
+    ok: bool = field(default=False, init=False)
+    err: BaseException = field()
+
+
+# class Result:
+#     ok: bool
+#     file: File
+#     completed: int
+#     err: BaseException | None
+
+
+# class Ok(Result):
+#     ok = True
+#     file: File
+#     completed: int
+#     err = None
+
+# def __init__(self, file: File, completed: int) -> None:
+#     self.file = file
+#     self.completed = completed
+
+
+# class Err(Result):
+#     ok = False
+#     file: File
+#     completed: int
+#     err: BaseException
+
+#     def __init__(self, file: File, completed: int, err: BaseException) -> None:
+#         self.file = file
+#         self.completed = completed
+#         self.err = err