esgpull 0.6.3__py3-none-any.whl

esgpull/database.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+from dataclasses import InitVar, dataclass, field
+from pathlib import Path
+from typing import TypeVar
+
+import alembic.command
+import sqlalchemy as sa
+import sqlalchemy.orm
+from alembic.config import Config as AlembicConfig
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from sqlalchemy.orm import Session, joinedload, make_transient
+
+from esgpull import __file__
+from esgpull.config import Config
+from esgpull.models import File, Table, sql
+from esgpull.version import __version__
+
+# from esgpull.exceptions import NoClauseError
+# from esgpull.models import Query
+
+T = TypeVar("T")
+
+
+@dataclass
+class Database:
+    """
+    Main class to interact with esgpull's sqlite db.
+    """
+
+    url: str
+    run_migrations: InitVar[bool] = True
+    _engine: sa.Engine = field(init=False)
+    session: Session = field(init=False)
+    version: str | None = field(init=False, default=None)
+
+    @staticmethod
+    def from_config(config: Config, run_migrations: bool = True) -> Database:
+        url = f"sqlite:///{config.paths.db / config.db.filename}"
+        return Database(url, run_migrations=run_migrations)
+
+    def __post_init__(self, run_migrations: bool) -> None:
+        self._engine = sa.create_engine(self.url)
+        self.session = Session(self._engine)
+        if run_migrations:
+            self._update()
+
+    def _update(self) -> None:
+        alembic_config = AlembicConfig()
+        migrations_path = Path(__file__).parent / "migrations"
+        alembic_config.set_main_option("script_location", str(migrations_path))
+        alembic_config.attributes["connection"] = self._engine
+        script = ScriptDirectory.from_config(alembic_config)
+        head = script.get_current_head()
+        with self._engine.begin() as conn:
+            opts = {"version_table": "version"}
+            ctx = MigrationContext.configure(conn, opts=opts)
+            self.version = ctx.get_current_revision()
+        if self.version != head:
+            alembic.command.upgrade(alembic_config, __version__)
+            self.version = head
+        if self.version != __version__:
+            alembic.command.revision(
+                alembic_config,
+                message="update tables",
+                autogenerate=True,
+                rev_id=__version__,
+            )
+            self.version = __version__
+
+    @property
+    @contextmanager
+    def safe(self) -> Iterator[None]:
+        try:
+            yield
+        except (sa.exc.SQLAlchemyError, KeyboardInterrupt):
+            self.session.rollback()
+            raise
+
+    def get(
+        self,
+        table: type[Table],
+        sha: str,
+        lazy: bool = True,
+        detached: bool = False,
+    ) -> Table | None:
+        if lazy:
+            result = self.session.get(table, sha)
+        else:
+            stmt = sa.select(table).filter_by(sha=sha)
+            match self.scalars(stmt.options(joinedload("*")), unique=True):
+                case [result]:
+                    ...
+                case []:
+                    result = None
+                case [*many]:
+                    raise ValueError(f"{len(many)} found, expected 1.")
+        if detached and result is not None:
+            result = table(**result.asdict())
+        return result
+
+    def scalars(
+        self, statement: sa.Select[tuple[T]], unique: bool = False
+    ) -> Sequence[T]:
+        with self.safe:
+            result = self.session.scalars(statement)
+            if unique:
+                result = result.unique()
+            return result.all()
+
+    SomeTuple = TypeVar("SomeTuple", bound=tuple)
+
+    def rows(self, statement: sa.Select[SomeTuple]) -> list[sa.Row[SomeTuple]]:
+        with self.safe:
+            return list(self.session.execute(statement).all())
+
+    def add(self, *items: Table) -> None:
+        with self.safe:
+            self.session.add_all(items)
+            self.session.commit()
+            for item in items:
+                self.session.refresh(item)
+
+    def delete(self, *items: Table) -> None:
+        with self.safe:
+            for item in items:
+                self.session.delete(item)
+            self.session.commit()
+        for item in items:
+            make_transient(item)
+
+    def __contains__(self, item: Table) -> bool:
+        return self.scalars(sql.count(item))[0] > 0
+
+    def has_file_id(self, file: File) -> bool:
+        return len(self.scalars(sql.file.with_file_id(file.file_id))) == 1
+
+    def merge(self, item: Table, commit: bool = False) -> Table:
+        with self.safe:
+            result = self.session.merge(item)
+            if commit:
+                self.session.commit()
+        return result
+
+    def get_deprecated_files(self) -> list[File]:
+        duplicates = self.scalars(sql.file.duplicates())
+        duplicates_dict: dict[str, list[File]] = {}
+        for file in duplicates:
+            duplicates_dict.setdefault(file.master_id, [])
+            duplicates_dict[file.master_id].append(file)
+        deprecated: list[File] = []
+        for files in duplicates_dict.values():
+            versions = [int(f.version[1:]) for f in files]
+            latest_version = "v" + str(max(versions))
+            for file in files:
+                if file.version != latest_version:
+                    deprecated.append(file)
+        return deprecated

esgpull/download.py

@@ -0,0 +1,162 @@
+# from math import ceil
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+
+from httpx import AsyncClient
+
+from esgpull.fs import Digest
+from esgpull.models import File
+
+# import asyncio
+# from urllib.parse import urlsplit
+# from esgpull.auth import Auth
+# from esgpull.context import Context
+
+
+@dataclass
+class DownloadCtx:
+    file: File
+    completed: int = 0
+    chunk: bytes | None = None
+    digest: Digest | None = None
+
+    @property
+    def finished(self) -> bool:
+        return self.completed == self.file.size
+
+    @property
+    def error(self) -> bool:
+        return self.completed > self.file.size
+
+    def update_digest(self) -> None:
+        if self.digest is not None and self.chunk is not None:
+            self.digest.update(self.chunk)
+
+
+class BaseDownloader:
+    def stream(
+        self,
+        client: AsyncClient,
+        ctx: DownloadCtx,
+        chunk_size: int,
+    ) -> AsyncGenerator[DownloadCtx, None]:
+        raise NotImplementedError
+
+
+class Simple(BaseDownloader):
+    """
+    Simple chunked async downloader.
+    """
+
+    async def stream(
+        self,
+        client: AsyncClient,
+        ctx: DownloadCtx,
+        chunk_size: int,
+    ) -> AsyncGenerator[DownloadCtx, None]:
+        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):
+                ctx.completed += len(chunk)
+                ctx.chunk = chunk
+                ctx.update_digest()
+                yield ctx
+
+
+# class Distributed(BaseDownloader):
+#     """
+#     Distributed chunked async downloader.
+#     Fetches chunks from multiple URLs pointing to the same file.
+#     """
+
+#     def __init__(
+#         self,
+#         auth: Auth,
+#         *,
+#         file: File | None = None,
+#         url: str | None = None,
+#         config: Config | None = None,
+#         max_ping: float = 5.0,
+#     ) -> None:
+#         super().__init__(auth, file=file, url=url, config=config)
+#         self.max_ping = max_ping
+
+#     async def try_url(self, url: str, client: AsyncClient) -> str | None:
+#         result = None
+#         node = urlsplit(url).netloc
+#         print(f"trying url on '{node}'")
+#         try:
+#             resp = await client.head(url)
+#             print(f"got response on '{node}'")
+#             resp.raise_for_status()
+#             accept_ranges = resp.headers.get("Accept-Ranges")
+#             content_length = resp.headers.get("Content-Length")
+#             if (
+#                 accept_ranges == "bytes"
+#                 and int(content_length) == self.file.size
+#             ):
+#                 result = str(resp.url)
+#             else:
+#                 print(dict(resp.headers))
+#         except HTTPError as err:
+#             print(type(err))
+#             print(err.request.headers)
+#         return result
+
+#     async def process_queue(
+#         self, url: str, queue: asyncio.Queue
+#     ) -> tuple[list[tuple[int, bytes]], str]:
+#         node = urlsplit(url).netloc
+#         print(f"starting process on '{node}'")
+#         chunks: list[tuple[int, bytes]] = []
+#         async with self.make_client() as client:
+#             final_url = await self.try_url(url, client)
+#             if final_url is None:
+#                 print(f"no url found for '{node}'")
+#                 return chunks, url
+#             else:
+#                 url = final_url
+#             while not queue.empty():
+#                 chunk_idx = await queue.get()
+#                 print(f"processing chunk {chunk_idx} on '{node}'")
+#                 start = chunk_idx * self.config.download.chunk_size
+#                 end = min(
+#                     self.file.size,
+#                     (chunk_idx + 1) * self.config.download.chunk_size - 1,
+#                 )
+#                 headers = {"Range": f"bytes={start}-{end}"}
+#                 resp = await client.get(url, headers=headers)
+#                 queue.task_done()
+#                 if resp.status_code == 206:
+#                     chunks.append((chunk_idx, resp.content))
+#                 else:
+#                     await queue.put(chunk_idx)
+#                     print(f"error status {resp.status_code} on '{node}'")
+#                     break
+#         return chunks, url
+
+#     async def fetch_urls(self) -> list[str]:
+#         ctx = Context(distrib=True)
+#         ctx.query.instance_id = self.file.file_id
+#         results = await ctx._search(file=True)
+#         files = [File.from_dict(item) for item in results]
+#         return [file.url for file in files]
+
+# async def aget(self) -> bytes:
+#     nb_chunks = ceil(self.file.size / self.config.download.chunk_size)
+#     queue: asyncio.Queue[int] = asyncio.Queue(nb_chunks)
+#     for chunk_idx in range(nb_chunks):
+#         queue.put_nowait(chunk_idx)
+#     completed: list[bool] = [False for _ in range(nb_chunks)]
+#     chunks: list[bytes] = [bytes() for _ in range(nb_chunks)]
+#     urls = await self.fetch_urls()
+#     workers = [self.process_queue(url, queue) for url in urls]
+#     for future in asyncio.as_completed(workers):
+#         some_chunks, url = await future
+#         print(f"got {len(some_chunks)} chunks from {url}")
+#         for chunk_idx, chunk in some_chunks:
+#             completed[chunk_idx] = True
+#             chunks[chunk_idx] = chunk
+#     if not all(completed):
+#         raise ValueError("TODO: progressive write (with .part file)")
+#     return b"".join(chunks)