litsync 0.0.2__py3-none-any.whl

litsync/sync.py

@@ -0,0 +1,334 @@
+from __future__ import annotations
+
+import contextlib
+import datetime as dt
+import json
+import logging
+import shutil
+import threading
+import zipfile
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Optional
+
+from litsync.config import Config
+from litsync.http import HttpClient
+from litsync.sources import Task
+from litsync.sources.clinicaltrials import ClinicalTrialsSource
+from litsync.sources.fda import FdaSource
+from litsync.sources.pmc import PmcSource
+from litsync.sources.pubmed import PubMedSource
+from litsync.state import FileRecord, StateDB
+from litsync.ui import UI
+from litsync.utils import count_articles, human_bytes, md5_file, new_src_stats, utcnow
+
+LOG = logging.getLogger("litsync")
+
+
+class Syncer:
+    def __init__(self, cfg: Config, ui: UI):
+        self.cfg = cfg
+        self.ui = ui
+        self.http = HttpClient(cfg)
+        self.db = StateDB(cfg.db_path)
+        self._stats_lock = threading.Lock()
+        self.stats = {"skipped": 0, "downloaded": 0, "verified": 0, "failed": 0}
+        self.bytes_downloaded = 0
+        self.articles_downloaded = 0
+        self.per_source: dict[str, dict] = {}
+        self.source_urls: dict[str, str] = {}
+        self.started_at = utcnow()
+
+    # ---- per-file decision logic ---------------------------------------- #
+
+    def _needs_work(self, task: Task) -> tuple[bool, Optional[int]]:
+        row = self.db.get(task.source, task.filename)
+        on_disk = task.dest.exists()
+
+        if (row is not None and row["status"] == "verified" and on_disk
+                and not self.cfg.reverify and task.immutable):
+            recorded = row["remote_size"]
+            if recorded is None or task.dest.stat().st_size == recorded:
+                return False, recorded
+
+        size, mtime, etag = self.http.head(task.url)
+        self.db.mark(task.source, task.filename,
+                     remote_size=size, remote_mtime=mtime, etag=etag)
+
+        if (row is not None and on_disk and not self.cfg.reverify
+                and row["status"] in ("done", "verified")):
+            same_size = size is None or task.dest.stat().st_size == size
+            same_etag = etag is None or row["etag"] == etag
+            if same_size and same_etag:
+                return False, size
+
+        return True, size
+
+    # ---- worker --------------------------------------------------------- #
+
+    def _process(self, task: Task, download_progress) -> None:
+        is_new = self.db.get(task.source, task.filename) is None
+        self.db.upsert_seen(FileRecord(task.source, task.filename, task.url, task.rel_path))
+        self._record_seen(task.source, is_new)
+        try:
+            needs, expected_size = self._needs_work(task)
+        except Exception as exc:
+            self.db.mark(task.source, task.filename, status="failed", error=str(exc))
+            self._bump("failed", task.source)
+            LOG.error("metadata failed for %s: %s", task.filename, exc)
+            return
+
+        if not needs:
+            self._bump("skipped", task.source)
+            return
+
+        if self.cfg.dry_run:
+            self._bump("downloaded", task.source, expected_size or 0)
+            return
+
+        expected_md5 = None
+        if task.md5_url:
+            with contextlib.suppress(Exception):
+                expected_md5 = md5_file_from_url(self.http.get_text(task.md5_url))
+                self.db.mark(task.source, task.filename, md5=expected_md5)
+
+        task_id = download_progress.add_task(task.filename, expected_size)
+
+        def progress_cb(nbytes: int):
+            download_progress.update(task_id, advance=nbytes)
+
+        try:
+            attempts = (self.db.get(task.source, task.filename)["attempts"] or 0) + 1
+            self.db.mark(task.source, task.filename, status="pending", attempts=attempts)
+            written = self.http.download(task.url, task.dest, expected_size, progress_cb)
+            download_progress.update(task_id, completed=written)
+
+            if expected_md5:
+                actual = md5_file(task.dest)
+                if actual.lower() != expected_md5.lower():
+                    task.dest.unlink(missing_ok=True)
+                    raise IOError(f"md5 mismatch: got {actual}, expected {expected_md5}")
+                self.db.mark(task.source, task.filename,
+                             status="verified", local_md5=actual, error=None)
+                self._bump("verified", task.source, written)
+            else:
+                self.db.mark(task.source, task.filename, status="verified", error=None)
+                self._bump("verified", task.source, written)
+
+            if task.extract:
+                self._extract_zip(task)
+            self._count_and_store(task)
+            LOG.info("ok %s", task.rel_path)
+        except Exception as exc:
+            self.db.mark(task.source, task.filename, status="failed", error=str(exc))
+            self._bump("failed", task.source)
+            LOG.error("download failed for %s: %s", task.filename, exc)
+
+    def _bump(self, key: str, source: Optional[str] = None, nbytes: int = 0) -> None:
+        with self._stats_lock:
+            self.stats[key] += 1
+            self.bytes_downloaded += nbytes
+            if source is not None:
+                s = self.per_source.setdefault(source, new_src_stats())
+                s[key] += 1
+                s["bytes"] += nbytes
+
+    def _record_seen(self, source: str, is_new: bool) -> None:
+        with self._stats_lock:
+            s = self.per_source.setdefault(source, new_src_stats())
+            s["new" if is_new else "existing"] += 1
+
+    def _extract_zip(self, task: Task) -> None:
+        if not task.dest.exists():
+            return
+        name = task.dest.name
+        if name.endswith(".json.zip"):
+            extract_dir = task.dest.parent / name[:-9]
+        else:
+            extract_dir = task.dest.with_suffix("")
+        if extract_dir.exists():
+            zip_mtime = task.dest.stat().st_mtime
+            dir_mtime = max(
+                (p.stat().st_mtime for p in extract_dir.rglob("*") if p.is_file()),
+                default=0,
+            )
+            if dir_mtime >= zip_mtime:
+                return
+            shutil.rmtree(extract_dir, ignore_errors=True)
+        self.ui.extract(task.rel_path)
+        extract_dir.parent.mkdir(parents=True, exist_ok=True)
+        with zipfile.ZipFile(task.dest, "r") as zf:
+            zf.extractall(extract_dir)
+
+    def _count_and_store(self, task: Task) -> None:
+        if self.cfg.dry_run:
+            return
+        try:
+            n = count_articles(task.dest, task.source)
+        except Exception as exc:
+            LOG.warning("article count failed for %s: %s", task.filename, exc)
+            return
+        self.db.mark(task.source, task.filename, article_count=n)
+        with self._stats_lock:
+            self.per_source.setdefault(task.source, new_src_stats())["articles"] += n
+            self.articles_downloaded += n
+
+    # ---- prune ---------------------------------------------------------- #
+
+    def _prune(self, tasks_by_source: dict[str, list[Task]]) -> None:
+        remote_by_source = {src: {t.filename for t in tasks}
+                            for src, tasks in tasks_by_source.items()}
+        synced_families = set(self.cfg.sources)
+
+        def family(source: str) -> str:
+            return source.split("_", 1)[0]
+
+        for source in self.db.all_sources():
+            if family(source) not in synced_families:
+                continue
+            remote = remote_by_source.get(source, set())
+            for fname in self.db.known_filenames(source) - remote:
+                row = self.db.get(source, fname)
+                if row and row["rel_path"]:
+                    p = self.cfg.data_root / row["rel_path"]
+                    if p.exists():
+                        LOG.info("prune (no longer remote): %s", row["rel_path"])
+                        if not self.cfg.dry_run:
+                            p.unlink(missing_ok=True)
+
+    # ---- run ------------------------------------------------------------ #
+
+    def run(self) -> int:
+        planners = []
+        if "pubmed" in self.cfg.sources:
+            planners.append(PubMedSource(self.cfg, self.http))
+        if "pmc" in self.cfg.sources:
+            planners.append(PmcSource(self.cfg, self.http))
+        if "fda" in self.cfg.sources:
+            planners.append(FdaSource(self.cfg, self.http))
+        if "clinicaltrials" in self.cfg.sources:
+            planners.append(ClinicalTrialsSource(self.cfg, self.http))
+
+        all_tasks: list[Task] = []
+        by_source: dict[str, list[Task]] = {}
+        for p in planners:
+            self.ui.planning(type(p).__name__)
+            tasks = p.plan()
+            all_tasks.extend(tasks)
+            for t in tasks:
+                by_source.setdefault(t.source, []).append(t)
+                self.source_urls.setdefault(t.source, t.url.rsplit("/", 1)[0] + "/")
+        self.ui.planned(len(all_tasks), len(by_source))
+
+        with self.ui.metadata_progress(len(all_tasks)) as meta, \
+             self.ui.download_progress() as dl:
+            def work(task: Task):
+                self._process(task, dl)
+                meta.advance(1)
+
+            with ThreadPoolExecutor(max_workers=self.cfg.workers) as pool:
+                futures = [pool.submit(work, t) for t in all_tasks]
+                for _ in as_completed(futures):
+                    pass
+
+        if self.cfg.prune:
+            self._prune(by_source)
+
+        self._report()
+        self.db.close()
+        return 1 if self.stats["failed"] else 0
+
+    # ---- article-count backfill (no network) --------------------------- #
+
+    def backfill_counts(self) -> int:
+        rows = self.db.files_missing_counts()
+        self.ui.planned(len(rows), 0)
+        done = self.failed = 0
+
+        def work(item: tuple[str, str, str]) -> None:
+            nonlocal done
+            source, filename, rel_path = item
+            path = self.cfg.data_root / rel_path
+            if not path.exists():
+                return
+            try:
+                n = count_articles(path, source)
+            except Exception as exc:
+                with self._stats_lock:
+                    self.failed += 1
+                LOG.error("count failed for %s: %s", rel_path, exc)
+                return
+            self.db.mark(source, filename, article_count=n)
+            with self._stats_lock:
+                done += 1
+                self.per_source.setdefault(source, new_src_stats())["articles"] += n
+                self.articles_downloaded += n
+
+        with ThreadPoolExecutor(max_workers=self.cfg.workers) as pool:
+            for _ in as_completed([pool.submit(work, it) for it in rows]):
+                pass
+
+        self.stats["failed"] = self.failed
+        self._report()
+        self.db.close()
+        return 1 if self.failed else 0
+
+    # ---- post-sync summary --------------------------------------------- #
+
+    def _report(self) -> None:
+        finished = utcnow()
+        mirror = self.db.summary_by_source()
+        self.write_json_manifest(finished, mirror)
+        self.ui.summary(
+            self.started_at, finished, self.stats, self.per_source, mirror,
+            self.source_urls, self.bytes_downloaded, self.articles_downloaded,
+        )
+
+    def write_json_manifest(self, finished: str, mirror: dict) -> None:
+        sources = sorted(set(mirror) | set(self.per_source))
+        new_dl = self.stats["verified"] + self.stats["downloaded"]
+        manifest = {
+            "started_at": self.started_at,
+            "finished_at": finished,
+            "this_run": {
+                "newly_downloaded": new_dl,
+                "already_current": self.stats["skipped"],
+                "failed": self.stats["failed"],
+                "bytes_downloaded": self.bytes_downloaded,
+                "articles_added": self.articles_downloaded,
+            },
+            "sources": {
+                source: {
+                    "url": self.source_urls.get(source),
+                    "mirror_files": mirror.get(source, {}).get("files", 0),
+                    "mirror_bytes": mirror.get(source, {}).get("bytes", 0),
+                    "mirror_articles": mirror.get(source, {}).get("articles", 0),
+                    "files_counted": mirror.get(source, {}).get("counted", 0),
+                    **self.per_source.get(source, new_src_stats()),
+                }
+                for source in sources
+            },
+            "mirror_total": {
+                "files": sum(m.get("files", 0) for m in mirror.values()),
+                "bytes": sum(m.get("bytes", 0) for m in mirror.values()),
+                "articles": sum(m.get("articles", 0) for m in mirror.values()),
+            },
+        }
+        out = self.cfg.log_dir / f"summary_{dt.date.today().isoformat()}.json"
+        try:
+            out.parent.mkdir(parents=True, exist_ok=True)
+            out.write_text(json.dumps(manifest, indent=2))
+            LOG.info("wrote summary manifest: %s", out)
+        except OSError as exc:
+            LOG.warning("could not write summary manifest: %s", exc)
+
+
+def md5_file_from_url(text: str) -> str:
+    import re
+    m = re.search(r"=\s*([0-9a-fA-F]{32})", text)
+    if m:
+        return m.group(1)
+    m = re.search(r"\b([0-9a-fA-F]{32})\b", text)
+    if not m:
+        raise ValueError(f"could not parse md5 from: {text[:120]!r}")
+    return m.group(1)

litsync/ui.py

@@ -0,0 +1,232 @@
+from __future__ import annotations
+
+import contextlib
+from typing import Optional
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.progress import (
+    BarColumn,
+    DownloadColumn,
+    Progress,
+    SpinnerColumn,
+    TaskID,
+    TextColumn,
+    TimeRemainingColumn,
+    TransferSpeedColumn,
+)
+from rich.table import Table
+from rich.text import Text
+
+from litsync.utils import human_bytes, partial_note
+
+console = Console()
+
+
+class UI:
+    """Plain-text fallback UI (no colors, no progress bars)."""
+
+    def planning(self, name: str):
+        print(f"Planning {name} ...")
+
+    def planned(self, total: int, sources: int):
+        print(f"Planned {total} files across {sources} source groups")
+
+    @contextlib.contextmanager
+    def metadata_progress(self, total: int):
+        yield _NoopProgress(total)
+
+    @contextlib.contextmanager
+    def download_progress(self):
+        yield _NoopDownloadProgress()
+
+    def extract(self, rel_path: str):
+        print(f"Extracting {rel_path}")
+
+    def summary(self, started: str, finished: str, stats: dict, per_source: dict,
+                mirror: dict, source_urls: dict, bytes_downloaded: int, articles_downloaded: int):
+        new_dl = stats.get("verified", 0) + stats.get("downloaded", 0)
+        print("litsync run complete")
+        print(f"  started: {started}")
+        print(f"  finished: {finished}")
+        print(f"  newly downloaded: {new_dl}, skipped: {stats.get('skipped', 0)}, failed: {stats.get('failed', 0)}")
+        print(f"  bytes: {human_bytes(bytes_downloaded)}, articles: {articles_downloaded:,}")
+
+
+class _NoopProgress:
+    def __init__(self, total: int):
+        self.total = total
+
+    def advance(self, n: int = 1):
+        pass
+
+    def set_description(self, desc: str):
+        pass
+
+
+class _NoopDownloadProgress:
+    def add_task(self, description: str, total: Optional[int]) -> int:
+        return 0
+
+    def update(self, task_id: int, advance: int = 0, completed: Optional[int] = None,
+               description: Optional[str] = None):
+        pass
+
+
+class RichUI(UI):
+    """Rich-based progress bars and tables."""
+
+    def planning(self, name: str):
+        console.print(f"[cyan]Planning[/cyan] {name} ...")
+
+    def planned(self, total: int, sources: int):
+        console.print(f"[green]Planned {total} files across {sources} source groups[/green]")
+
+    @contextlib.contextmanager
+    def metadata_progress(self, total: int):
+        progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[bold blue]{task.description}"),
+            BarColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            TextColumn("({task.completed}/{task.total})"),
+            console=console,
+            transient=True,
+        )
+        task = progress.add_task("Checking metadata...", total=total)
+        try:
+            with progress:
+                yield _RichProgress(task, progress)
+        finally:
+            pass
+
+    @contextlib.contextmanager
+    def download_progress(self):
+        progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[bold blue]{task.description}"),
+            BarColumn(bar_width=40),
+            "[progress.percentage]{task.percentage:>3.0f}%",
+            " ",
+            DownloadColumn(),
+            " ",
+            TransferSpeedColumn(),
+            " ",
+            TimeRemainingColumn(),
+            console=console,
+        )
+        try:
+            with progress:
+                yield _RichDownloadProgress(progress)
+        finally:
+            pass
+
+    def extract(self, rel_path: str):
+        console.print(f"[yellow]Extracting[/yellow] {rel_path}")
+
+    def summary(self, started: str, finished: str, stats: dict, per_source: dict,
+                mirror: dict, source_urls: dict, bytes_downloaded: int, articles_downloaded: int):
+        new_dl = stats.get("verified", 0) + stats.get("downloaded", 0)
+        table = Table(title="litsync summary", show_header=True, header_style="bold magenta")
+        table.add_column("Source", style="cyan")
+        table.add_column("Files", justify="right")
+        table.add_column("Size", justify="right")
+        table.add_column("Records", justify="right")
+        table.add_column("Run", justify="left")
+
+        labels = {
+            "pubmed": "PubMed",
+            "pmc": "PubMed Central",
+            "fda": "openFDA",
+            "clinicaltrials": "ClinicalTrials.gov",
+        }
+
+        families: dict[str, list[str]] = {}
+        sources = sorted(set(mirror) | set(per_source))
+        for source in sources:
+            families.setdefault(source.split("_", 1)[0], []).append(source)
+
+        grand_files = grand_bytes = grand_articles = grand_counted = 0
+        for fam in sorted(families):
+            fam_files = fam_bytes = fam_articles = fam_counted = 0
+            for source in families[fam]:
+                m = mirror.get(source, {"files": 0, "bytes": 0, "articles": 0, "counted": 0})
+                r = per_source.get(source, new_src_stats())
+                got = r.get("verified", 0) + r.get("downloaded", 0)
+                if m["counted"]:
+                    art = f"{m['articles']:,}"
+                    if m["counted"] < m["files"]:
+                        art += f" (counted {m['counted']}/{m['files']})"
+                else:
+                    art = "—"
+                run_txt = f"+{r.get('new', 0)} new, {got} fetched, {r.get('skipped', 0)} current, {r.get('failed', 0)} failed"
+                table.add_row(source, str(m["files"]), human_bytes(m["bytes"]), art, run_txt)
+                fam_files += m["files"]
+                fam_bytes += m["bytes"]
+                fam_articles += m["articles"]
+                fam_counted += m["counted"]
+            table.add_row(
+                f"[bold]{labels.get(fam, fam)} subtotal[/bold]",
+                str(fam_files),
+                human_bytes(fam_bytes),
+                f"{fam_articles:,}{partial_note(fam_counted, fam_files)}",
+                "",
+            )
+            grand_files += fam_files
+            grand_bytes += fam_bytes
+            grand_articles += fam_articles
+            grand_counted += fam_counted
+
+        table.add_row(
+            "[bold]TOTAL[/bold]",
+            str(grand_files),
+            human_bytes(grand_bytes),
+            f"{grand_articles:,}{partial_note(grand_counted, grand_files)}",
+            f"+{new_dl} downloaded, {stats.get('skipped', 0)} skipped, {stats.get('failed', 0)} failed",
+            style="bold green",
+        )
+        console.print()
+        console.print(table)
+        console.print(
+            f"[dim]started: {started}  ·  finished: {finished}  · "
+            f" bytes this run: {human_bytes(bytes_downloaded)}  · "
+            f" records this run: {articles_downloaded:,}[/dim]"
+        )
+        if grand_counted < grand_files:
+            console.print(
+                "[dim]Run --count-articles to backfill record counts for already-downloaded files.[/dim]"
+            )
+
+
+class _RichProgress:
+    def __init__(self, task: TaskID, progress: Progress):
+        self.task = task
+        self.progress = progress
+
+    def advance(self, n: int = 1):
+        self.progress.advance(self.task, n)
+
+    def set_description(self, desc: str):
+        self.progress.update(self.task, description=desc)
+
+
+class _RichDownloadProgress:
+    def __init__(self, progress: Progress):
+        self.progress = progress
+
+    def add_task(self, description: str, total: Optional[int]) -> int:
+        return self.progress.add_task(description, total=total)
+
+    def update(self, task_id: int, advance: int = 0, completed: Optional[int] = None,
+               description: Optional[str] = None):
+        kwargs = {"advance": advance}
+        if completed is not None:
+            kwargs["completed"] = completed
+        if description is not None:
+            kwargs["description"] = description
+        self.progress.update(task_id, **kwargs)
+
+
+def new_src_stats() -> dict:
+    return {"new": 0, "existing": 0, "skipped": 0,
+            "downloaded": 0, "verified": 0, "failed": 0, "bytes": 0, "articles": 0}

litsync/utils.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import contextlib
+import datetime as dt
+import fcntl
+import gzip
+import hashlib
+import logging
+import os
+import re
+import tarfile
+from pathlib import Path
+from typing import Optional
+
+LOG = logging.getLogger("litsync")
+CHUNK = 1 << 20  # 1 MiB
+
+
+def utcnow() -> str:
+    return dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds")
+
+
+def new_src_stats() -> dict:
+    return {"new": 0, "existing": 0, "skipped": 0,
+            "downloaded": 0, "verified": 0, "failed": 0, "bytes": 0, "articles": 0}
+
+
+def human_bytes(n) -> str:
+    size = float(n or 0)
+    for unit in ("B", "KB", "MB", "GB", "TB"):
+        if size < 1024:
+            return f"{int(size)} {unit}" if unit == "B" else f"{size:.1f} {unit}"
+        size /= 1024
+    return f"{size:.1f} PB"
+
+
+def partial_note(counted: int, files: int) -> str:
+    return "" if counted >= files else f"  [partial: {counted}/{files} files counted]"
+
+
+def md5_file(path: Path) -> str:
+    h = hashlib.md5()
+    with open(path, "rb") as fh:
+        for chunk in iter(lambda: fh.read(CHUNK), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def parse_md5(text: str) -> str:
+    """NCBI .md5 sidecars look like 'MD5(file.xml.gz)= <hex>' or '<hex>  file'."""
+    m = re.search(r"=\s*([0-9a-fA-F]{32})", text)
+    if m:
+        return m.group(1)
+    m = re.search(r"\b([0-9a-fA-F]{32})\b", text)
+    if not m:
+        raise ValueError(f"could not parse md5 from: {text[:120]!r}")
+    return m.group(1)
+
+
+@contextlib.contextmanager
+def run_lock(path: Path):
+    """Prevent overlapping runs via an exclusive file lock."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fh = open(path, "w")
+    try:
+        try:
+            fcntl.flock(fh, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            raise SystemExit("another litsync run is already in progress; exiting")
+        fh.write(f"{os.getpid()} {utcnow()}\n")
+        fh.flush()
+        yield
+    finally:
+        with contextlib.suppress(Exception):
+            fcntl.flock(fh, fcntl.LOCK_UN)
+        fh.close()
+
+
+def count_articles(path: Path, source: str) -> int:
+    """Cheaply count the articles/records inside one downloaded file."""
+    name = path.name.lower()
+    if source.startswith("pubmed"):
+        if not name.endswith(".xml.gz"):
+            return 0
+        needle = b"<PubmedArticle>"
+        overlap = b""
+        count = 0
+        with gzip.open(path, "rb") as fh:
+            for chunk in iter(lambda: fh.read(CHUNK), b""):
+                buf = overlap + chunk
+                count += buf.count(needle)
+                overlap = buf[-(len(needle) - 1):]
+        return count
+    if source.startswith("pmc"):
+        if not name.endswith(".tar.gz"):
+            return 0
+        count = 0
+        with tarfile.open(path, mode="r|gz") as tar:
+            for member in tar:
+                low = member.name.lower()
+                if member.isfile() and (low.endswith(".xml") or low.endswith(".nxml")):
+                    count += 1
+        return count
+    if source.startswith("fda"):
+        if not name.endswith(".zip"):
+            return 0
+        import json
+        extract_dir = path.parent / name[:-9] if name.endswith(".json.zip") else path.with_suffix("")
+        count = 0
+        for json_file in extract_dir.rglob("*.json"):
+            try:
+                data = json.loads(json_file.read_text(encoding="utf-8"))
+                count += len(data.get("results", []))
+            except Exception:
+                pass
+        return count
+    if source.startswith("clinicaltrials"):
+        if not name.endswith(".zip"):
+            return 0
+        extract_dir = path.with_suffix("")
+        return sum(1 for p in extract_dir.rglob("*.xml") if p.is_file())
+    return 0