refresh_assemblies.py

"""Add the taxon-to-accession table the atlas needs to open Database results.

The taxonomy snapshot stores counts, not accessions, so a group in the atlas had
no way to name a single assembly the Database tab could look up. This writes an
``assemblies`` table into the existing ``data/taxonomy.sqlite`` holding one row
per *annotated* eukaryotic assembly: the accession the Database tab searches and
the taxon it belongs to.

Only annotated assemblies are kept. The atlas offers this as a jump into
published annotations, so an accession with nothing behind it would be a dead
end, and the restriction keeps the table at roughly 33,000 rows.

Counts are deliberately untouched. Rebuilding them here would re-read a newer
assembly summary and silently shift every total in the interface away from the
published snapshot; see refresh_taxonomy.py for that, which is a separate
decision.
"""
import argparse
from datetime import datetime, timezone
import json
from pathlib import Path
import sqlite3
import tarfile

ROOT = Path(__file__).resolve().parent


def merged_taxids(taxdump):
    with tarfile.open(taxdump, "r:gz") as archive, archive.extractfile("merged.dmp") as stream:
        return {int(parts[0]): int(parts[1]) for parts in
                ([p.strip() for p in line.decode().split("|")] for line in stream)}


def build(summary, taxdump, database, coverage):
    annotated = set(json.loads(Path(coverage).read_text())["accessions"])
    merged = merged_taxids(taxdump)
    with sqlite3.connect(database) as conn:
        known = {row[0] for row in conn.execute("SELECT taxid FROM taxa")}

    rows, skipped, unresolved = [], 0, 0
    with open(summary) as stream:
        columns = None
        for line in stream:
            if line.startswith("##"):
                continue
            if line.startswith("#"):
                header = line.lstrip("# ").rstrip("\n").split("\t")
                columns = {name: header.index(name) for name in ("assembly_accession", "taxid", "version_status")}
                continue
            row = line.rstrip("\n").split("\t")
            accession = row[columns["assembly_accession"]]
            if accession not in annotated or row[columns["version_status"]] != "latest":
                skipped += 1
                continue
            try:
                taxid = int(row[columns["taxid"]])
            except ValueError:
                unresolved += 1
                continue
            seen = set()
            while taxid in merged and taxid not in seen:
                seen.add(taxid)
                taxid = merged[taxid]
            if taxid not in known:
                unresolved += 1
                continue
            rows.append((accession, taxid))

    if not rows:
        raise ValueError("No annotated assemblies matched the summary; check the inputs.")
    with sqlite3.connect(database) as conn:
        conn.executescript("DROP TABLE IF EXISTS assemblies;"
                           "CREATE TABLE assemblies (accession TEXT PRIMARY KEY, taxid INTEGER);")
        conn.executemany("INSERT OR REPLACE INTO assemblies VALUES (?,?)", rows)
        conn.execute("CREATE INDEX assembly_taxa ON assemblies(taxid)")
        metadata = json.loads(conn.execute("SELECT value FROM metadata").fetchone()[0])
        metadata["assemblies"] = {"created_at": datetime.now(timezone.utc).isoformat(),
                                 "rows": len(rows), "scope": "annotated eukaryotic assemblies only",
                                 "source": "https://ftp.ncbi.nlm.nih.gov/genomes/ASSEMBLY_REPORTS/assembly_summary_genbank.txt"}
        conn.execute("UPDATE metadata SET value=?", (json.dumps(metadata),))
    print(json.dumps({"rows": len(rows), "annotated_accessions": len(annotated),
                      "unresolved_taxids": unresolved, "rows_skipped": skipped}, indent=2))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--summary", type=Path, default=ROOT / ".cache/coverage/assembly_summary_genbank.txt")
    parser.add_argument("--taxdump", type=Path, default=ROOT / ".cache/coverage/taxdump.tar.gz")
    parser.add_argument("--database", type=Path, default=ROOT / "data/taxonomy.sqlite")
    parser.add_argument("--coverage", type=Path, default=ROOT / "data/coverage.json")
    args = parser.parse_args()
    build(args.summary, args.taxdump, args.database, args.coverage)


if __name__ == "__main__":
    main()