#! /usr/bin/env python3 # pylint: disable=invalid-name """This script must be run each time the database is updated. It runs queries against the SQL database, indexes the results and builds a xapian index. This xapian index is later used in providing search through the web interface. """ from dataclasses import dataclass from collections import deque, namedtuple import contextlib import time import datetime from functools import partial import itertools import json import logging from multiprocessing import Lock, Manager, Process, managers import os import pathlib import resource import shutil import sys import tempfile from typing import Callable, Generator, Hashable, Iterable, List from SPARQLWrapper import SPARQLWrapper, JSON import MySQLdb import click from pymonad.maybe import Just, Maybe, Nothing from pymonad.tools import curry import xapian from gn3.db_utils import database_connection from gn3.monads import query_sql DOCUMENTS_PER_CHUNK = 100000 SQLQuery = namedtuple("SQLQuery", ["fields", "tables", "where", "offset", "limit"], defaults=[Nothing, 0, Nothing]) SQLTableClause = namedtuple("SQLTableClause", ["join_type", "table", "condition"]) # FIXME: Some Max LRS values in the DB are wrongly listed as 0.000, # but shouldn't be displayed. Make them NULLs in the database. genes_query = SQLQuery( ["ProbeSet.Name AS name", "ProbeSet.Symbol AS symbol", "ProbeSet.description AS description", "ProbeSet.Chr AS chr", "ProbeSet.Mb as mb", "ProbeSet.alias AS alias", "ProbeSet.GenbankId AS genbankid", "ProbeSet.UniGeneId AS unigeneid", "ProbeSet.Probe_Target_Description AS probe_target_description", "ProbeSetFreeze.Name AS dataset", "ProbeSetFreeze.FullName AS dataset_fullname", "Species.Name AS species", "InbredSet.Name AS `group`", "Tissue.Name AS tissue", "ProbeSetXRef.Mean AS mean", "ProbeSetXRef.LRS AS lrs", "ProbeSetXRef.additive AS additive", "Geno.Chr AS geno_chr", "Geno.Mb as geno_mb"], ["Species", SQLTableClause("INNER JOIN", "InbredSet", "InbredSet.SpeciesId = Species.Id"), SQLTableClause("INNER JOIN", "ProbeFreeze", "ProbeFreeze.InbredSetId = InbredSet.Id"), SQLTableClause("INNER JOIN", "Tissue", "ProbeFreeze.TissueId = Tissue.Id"), SQLTableClause("INNER JOIN", "ProbeSetFreeze", "ProbeSetFreeze.ProbeFreezeId = ProbeFreeze.Id"), SQLTableClause("INNER JOIN", "ProbeSetXRef", "ProbeSetXRef.ProbeSetFreezeId = ProbeSetFreeze.Id"), SQLTableClause("INNER JOIN", "ProbeSet", "ProbeSet.Id = ProbeSetXRef.ProbeSetId"), SQLTableClause("LEFT JOIN", "Geno", "ProbeSetXRef.Locus = Geno.Name AND Geno.SpeciesId = Species.Id")], Just("ProbeSetFreeze.confidentiality < 1 AND ProbeSetFreeze.public > 0")) # FIXME: Some years are blank strings or strings that contain text # other than the year. These should be fixed in the database and the # year field must be made an integer. phenotypes_query = SQLQuery( ["Species.Name AS species", "InbredSet.Name AS `group`", "PublishFreeze.Name AS dataset", "PublishFreeze.FullName AS dataset_fullname", "PublishXRef.Id AS name", """COALESCE(Phenotype.Post_publication_abbreviation, Phenotype.Pre_publication_abbreviation) AS abbreviation""", """COALESCE(Phenotype.Post_publication_description, Phenotype.Pre_publication_description) AS description""", "Phenotype.Lab_code", "Publication.Abstract", "Publication.Title", "Publication.Authors AS authors", """IF(CONVERT(Publication.Year, UNSIGNED)=0, NULL, CONVERT(Publication.Year, UNSIGNED)) AS year""", "Publication.PubMed_ID AS pubmed_id", "PublishXRef.LRS as lrs", "PublishXRef.additive", "InbredSet.InbredSetCode AS inbredsetcode", "PublishXRef.mean", "Geno.Chr as geno_chr", "Geno.Mb as geno_mb"], ["Species", SQLTableClause("INNER JOIN", "InbredSet", "InbredSet.SpeciesId = Species.Id"), SQLTableClause("INNER JOIN", "PublishFreeze", "PublishFreeze.InbredSetId = InbredSet.Id"), SQLTableClause("INNER JOIN", "PublishXRef", "PublishXRef.InbredSetId = InbredSet.Id"), SQLTableClause("INNER JOIN", "Phenotype", "PublishXRef.PhenotypeId = Phenotype.Id"), SQLTableClause("INNER JOIN", "Publication", "PublishXRef.PublicationId = Publication.Id"), SQLTableClause("LEFT JOIN", "Geno", "PublishXRef.Locus = Geno.Name AND Geno.SpeciesId = Species.Id")]) RIF_CACHE_QUERY = """ PREFIX rdf: PREFIX rdfs: PREFIX gnt: PREFIX gnc: SELECT ?symbolName ?speciesName GROUP_CONCAT(DISTINCT ?comment ; separator=\"\\n\") AS ?comment WHERE { ?symbol rdfs:comment _:node ; rdfs:label ?symbolName . _:node rdf:type gnc:GNWikiEntry ; gnt:belongsToSpecies ?species ; rdfs:comment ?comment . ?species gnt:shortName ?speciesName . } GROUP BY ?speciesName ?symbolName """ WIKI_CACHE_QUERY = """ PREFIX rdf: PREFIX rdfs: PREFIX gnt: PREFIX gnc: SELECT ?symbolName ?speciesName GROUP_CONCAT(DISTINCT ?comment ; separator=\"\\n\") AS ?comment WHERE { ?symbol rdfs:comment _:node ; rdfs:label ?symbolName . _:node rdf:type gnc:NCBIWikiEntry ; gnt:belongsToSpecies ?species ; rdfs:comment ?comment . ?species gnt:shortName ?speciesName . } GROUP BY ?speciesName ?symbolName """ def serialize_sql(query: SQLQuery) -> str: """Serialize SQLQuery object to a string.""" table_clauses = [clause if isinstance(clause, str) else f"{clause.join_type} {clause.table} ON {clause.condition}" for clause in query.tables] sql = f"SELECT {', '.join(query.fields)} FROM {' '.join(table_clauses)}" def append_to_sql(appendee): nonlocal sql sql += appendee query.where.bind(lambda where: append_to_sql(f" WHERE {where}")) query.limit.bind(lambda limit: append_to_sql(f" LIMIT {limit}")) if query.offset != 0: sql += f" OFFSET {query.offset}" return sql @contextlib.contextmanager def locked_xapian_writable_database(path: pathlib.Path) -> xapian.WritableDatabase: """Open xapian database for writing. When a process is writing to a xapian database opened by this function, no other process may do so. This avoids I/O contention between processes. """ # pylint: disable-next=invalid-name if not path.exists(): os.makedirs(path) db = xapian.WritableDatabase(str(path)) db.begin_transaction() try: yield db except Exception as exception: db.cancel_transaction() raise exception else: xapian_lock.acquire() try: db.commit_transaction() finally: xapian_lock.release() finally: db.close() def build_rdf_cache(sparql_uri: str, query: str): cache = {} sparql = SPARQLWrapper(sparql_uri) sparql.setReturnFormat(JSON) sparql.setQuery(query) results = sparql.queryAndConvert() if not isinstance(results, dict): raise TypeError(f"Expected results to be a dict but found {type(results)}") bindings = results["results"]["bindings"] for entry in bindings : x = (entry["speciesName"]["value"], entry["symbolName"]["value"],) cache[x] = entry["comment"]["value"] return cache def hash_generif_graph(sparql_uri: str): sparql = SPARQLWrapper(sparql_uri) sparql.setReturnFormat(JSON) query = """ PREFIX rdf: PREFIX rdfs: PREFIX gnt: PREFIX gnc: SELECT SHA256(GROUP_CONCAT(?entries ; separator=\"\\n\")) AS ?hash WHERE { {{ SELECT ?type CONCAT(?symbolName, ",", ?speciesName, \"\\n\",GROUP_CONCAT(?comment ; separator=\"\\n\")) AS ?entries WHERE { ?symbol rdfs:comment _:node ; rdfs:label ?symbolName . _:node rdf:type gnc:GNWikiEntry ; rdf:type ?type ; gnt:belongsToSpecies ?species ; rdfs:comment ?comment . ?species gnt:shortName ?speciesName . } GROUP BY ?speciesName ?symbolName ?type }} } GROUP BY ?type """ sparql.setQuery(query) results = sparql.queryAndConvert() if not isinstance(results, dict): raise TypeError(f"Expected results to be a dict but found {type(results)}") bindings = results["results"]["bindings"] return bindings[0]["hash"]["value"] # pylint: disable=invalid-name def write_document(db: xapian.WritableDatabase, identifier: str, doctype: str, doc: xapian.Document) -> None: """Write document into xapian database.""" # We use the XT and Q prefixes to indicate the type and idterm # respectively. idterm = f"Q{doctype}:{identifier.lower()}" doc.add_boolean_term(f"XT{doctype}") doc.add_boolean_term(idterm) db.replace_document(idterm, doc) termgenerator = xapian.TermGenerator() termgenerator.set_stemmer(xapian.Stem("en")) def index_text(text: str) -> None: """Index text and increase term position.""" termgenerator.index_text(text) termgenerator.increase_termpos() @curry(3) def index_from_dictionary(keys: Hashable, prefix: str, dictionary: dict): entry = dictionary.get(keys) if not entry: return termgenerator.index_text(entry, 0, prefix) index_text_without_positions = lambda text: termgenerator.index_text_without_positions(text) index_authors = lambda authors: termgenerator.index_text(authors, 0, "A") index_species = lambda species: termgenerator.index_text_without_positions(species, 0, "XS") index_group = lambda group: termgenerator.index_text_without_positions(group, 0, "XG") index_tissue = lambda tissue: termgenerator.index_text(tissue, 0, "XI") index_dataset = lambda dataset: termgenerator.index_text(dataset, 0, "XDS") index_symbol = lambda symbol: termgenerator.index_text_without_positions(symbol, 0, "XY") index_chr = lambda chr: termgenerator.index_text_without_positions(chr, 0, "XC") index_peakchr = lambda peakchr: termgenerator.index_text_without_positions(peakchr, 0, "XPC") add_mean = lambda doc, mean: doc.add_value(0, xapian.sortable_serialise(mean)) add_peak = lambda doc, peak: doc.add_value(1, xapian.sortable_serialise(peak)) add_mb = lambda doc, mb: doc.add_value(2, xapian.sortable_serialise(mb)) add_peakmb = lambda doc, peakmb: doc.add_value(3, xapian.sortable_serialise(peakmb)) add_additive = lambda doc, additive: doc.add_value(4, xapian.sortable_serialise(additive)) add_year = lambda doc, year: doc.add_value(5, xapian.sortable_serialise(float(year))) # When a child process is forked, it inherits a copy of the memory of # its parent. We use this to pass data retrieved from SQL from parent # to child. Specifically, we use this global variable. # This is copy-on-write so make sure child processes don't modify this data mysql_data: Iterable rif_cache: Iterable wiki_cache: Iterable # We use this lock to ensure that only one process writes its Xapian # index to disk at a time. xapian_lock = Lock() def index_genes(xapian_build_directory: pathlib.Path, chunk_index: int) -> None: """Index genes data into a Xapian index.""" with locked_xapian_writable_database(xapian_build_directory / f"genes-{chunk_index:04d}") as db: for trait in mysql_data: # pylint: disable=cell-var-from-loop doc = xapian.Document() termgenerator.set_document(doc) # Add values. trait["mean"].bind(partial(add_mean, doc)) trait["lrs"].bind(partial(add_peak, doc)) trait["mb"].bind(partial(add_mb, doc)) trait["geno_mb"].bind(partial(add_peakmb, doc)) trait["additive"].bind(partial(add_additive, doc)) # Index free text. for key in ["description", "tissue", "dataset"]: trait[key].bind(index_text) trait.pop("probe_target_description").bind(index_text) for key in ["name", "symbol", "species", "group"]: trait[key].bind(index_text_without_positions) for key in ["alias", "genbankid", "unigeneid"]: trait.pop(key).bind(index_text_without_positions) # Index text with prefixes. trait["species"].bind(index_species) trait["group"].bind(index_group) trait["tissue"].bind(index_tissue) trait["dataset"].bind(index_dataset) trait["symbol"].bind(index_symbol) trait["chr"].bind(index_chr) trait["geno_chr"].bind(index_peakchr) Maybe.apply(index_from_dictionary).to_arguments( Just((trait["species"].value, trait["symbol"].value)), Just("XRF"), Just(rif_cache) ) Maybe.apply(index_from_dictionary).to_arguments( Just((trait["species"].value, trait["symbol"].value)), Just("XWK"), Just(wiki_cache) ) doc.set_data(json.dumps(trait.data)) (Maybe.apply(curry(2, lambda name, dataset: f"{name}:{dataset}")) .to_arguments(trait["name"], trait["dataset"]) .bind(lambda idterm: write_document(db, idterm, "gene", doc))) def index_phenotypes(xapian_build_directory: pathlib.Path, chunk_index: int) -> None: """Index phenotypes data into a Xapian index.""" with locked_xapian_writable_database( xapian_build_directory / f"phenotypes-{chunk_index:04d}") as db: for trait in mysql_data: # pylint: disable=cell-var-from-loop doc = xapian.Document() termgenerator.set_document(doc) # Add values. trait["mean"].bind(partial(add_mean, doc)) trait["lrs"].bind(partial(add_peak, doc)) trait["geno_mb"].bind(partial(add_peakmb, doc)) trait["additive"].bind(partial(add_additive, doc)) trait["year"].bind(partial(add_year, doc)) # Index free text. for key in ["description", "authors", "dataset"]: trait[key].bind(index_text) for key in ["Abstract", "Title"]: trait.pop(key).bind(index_text) for key in ["species", "group", "inbredsetcode"]: trait[key].bind(index_text_without_positions) for key in ["abbreviation", "Lab_code"]: trait.pop(key).bind(index_text_without_positions) # Index text with prefixes. trait["species"].bind(index_species) trait["group"].bind(index_group) trait["authors"].bind(index_authors) trait["geno_chr"].bind(index_peakchr) trait["dataset"].bind(index_dataset) # Convert name from integer to string. trait["name"] = trait["name"].map(str) # Split comma-separated authors into a list. trait["authors"] = trait["authors"].map( lambda s: [author.strip() for author in s.split(",")]) doc.set_data(json.dumps(trait.data)) (Maybe.apply(curry(2, lambda name, dataset: f"{name}:{dataset}")) .to_arguments(trait["name"], trait["dataset"]) .bind(lambda idterm: write_document(db, idterm, "phenotype", doc))) def group(generator: Iterable, chunk_size: int) -> Iterable: """Group elements of generator into chunks.""" return iter(lambda: tuple(itertools.islice(generator, chunk_size)), ()) @contextlib.contextmanager def worker_queue(number_of_workers: int = os.cpu_count() or 1) -> Generator: """Manage a pool of worker processes returning a function to spawn them.""" processes: deque = deque() def spawn(target, args): if len(processes) == number_of_workers: processes.popleft().join() process = Process(target=target, args=args) process.start() processes.append(process) yield spawn for process in processes: process.join() def index_query(index_function: Callable[[pathlib.Path, int, managers.Namespace], None], query: SQLQuery, xapian_build_directory: pathlib.Path, sql_uri: str, sparql_uri: str, start: int = 0) -> None: """Run SQL query, and index its results for Xapian.""" i = start try: with worker_queue() as spawn_worker: with database_connection(sql_uri) as conn: for chunk in group(query_sql(conn, serialize_sql( # KLUDGE: MariaDB does not allow an offset # without a limit. So, set limit to a "high" # value. query._replace(limit=Just(2**64 - 1), offset=start*DOCUMENTS_PER_CHUNK)), server_side=True), DOCUMENTS_PER_CHUNK): global mysql_data mysql_data = chunk spawn_worker(index_function, (xapian_build_directory, i)) logging.debug("Spawned worker process on chunk %s", i) i += 1 # In the event of an operational error, open a new connection and # resume indexing. # pylint: disable=protected-access except MySQLdb._exceptions.OperationalError: logging.warning("Reopening connection to recovering from SQL operational error", exc_info=True) index_query(index_function, query, xapian_build_directory, sql_uri, sparql_uri, i) @contextlib.contextmanager def temporary_directory(prefix: str, parent_directory: str) -> Generator: """Create temporary directory returning it as a PosixPath.""" with tempfile.TemporaryDirectory(prefix=prefix, dir=parent_directory) as tmpdirname: yield pathlib.Path(tmpdirname) def xapian_compact(combined_index: pathlib.Path, indices: List[pathlib.Path]) -> None: """Compact and combine several Xapian indices.""" # xapian-compact opens all indices simultaneously. So, raise the limit on # the number of open files. soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, min(10*len(indices), hard)), hard)) db = xapian.Database() try: for index in indices: db.add_database(xapian.Database(str(index))) db.compact(str(combined_index), xapian.DBCOMPACT_MULTIPASS | xapian.Compactor.FULLER) finally: db.close() @click.command(help="Verify checksums and return True when the data has been changed.") @click.argument("xapian_directory") @click.argument("sql_uri") @click.argument("sparql_uri") def is_data_modified(xapian_directory: str, sql_uri: str, sparql_uri: str) -> None: dir_ = pathlib.Path(xapian_directory) with locked_xapian_writable_database(dir_) as db, database_connection(sql_uri) as conn: checksums = "-1" if db.get_metadata('tables'): checksums = " ".join([ str(result["Checksum"].value) for result in query_sql( conn, f"CHECKSUM TABLE {', '.join(db.get_metadata('tables').decode().split())}") ]) # Return a zero exit status code when the data has changed; # otherwise exit with a 1 exit status code. if (db.get_metadata("generif-checksum").decode() == hash_generif_graph(sparql_uri) and db.get_metadata("checksums").decode() == checksums): sys.exit(1) sys.exit(0) @click.command(help="Index GeneNetwork data and build Xapian search index in XAPIAN_DIRECTORY.") @click.argument("xapian_directory") @click.argument("sql_uri") @click.argument("sparql_uri") # pylint: disable=missing-function-docstring def create_xapian_index(xapian_directory: str, sql_uri: str, sparql_uri: str) -> None: logging.basicConfig(level=os.environ.get("LOGLEVEL", "DEBUG"), format='%(asctime)s %(levelname)s: %(message)s', datefmt='%Y-%m-%d %H:%M:%S %Z') logging.info("Verifying the checksums") if not pathlib.Path(xapian_directory).exists(): pathlib.Path(xapian_directory).mkdir() # Ensure no other build process is running. if any(pathlib.Path(xapian_directory).iterdir()): logging.error("Build directory %s has build files; " "perhaps another build process is running.", xapian_directory) sys.exit(1) start_time = time.perf_counter() with temporary_directory("combined", xapian_directory) as combined_index: with temporary_directory("build", xapian_directory) as xapian_build_directory: global rif_cache global wiki_cache rif_cache = build_rdf_cache(sparql_uri, RIF_CACHE_QUERY) wiki_cache = build_rdf_cache(sparql_uri, WIKI_CACHE_QUERY) logging.info("Indexing genes") index_query(index_genes, genes_query, xapian_build_directory, sql_uri, sparql_uri) logging.info("Indexing phenotypes") index_query(index_phenotypes, phenotypes_query, xapian_build_directory, sql_uri, sparql_uri) logging.info("Combining and compacting indices") xapian_compact(combined_index, list(xapian_build_directory.iterdir())) logging.info("Writing table checksums into index") with locked_xapian_writable_database(combined_index) as db: # Build a (deduplicated) set of all tables referenced in # queries. tables = set(clause if isinstance(clause, str) else clause.table for clause in genes_query.tables + phenotypes_query.tables) with database_connection(sql_uri) as conn: checksums = [ result["Checksum"].bind(str) # type: ignore for result in query_sql(conn, f"CHECKSUM TABLE {', '.join(tables)}") ] db.set_metadata("tables", " ".join(tables)) db.set_metadata("checksums", " ".join(checksums)) logging.info("Writing generif checksums into index") db.set_metadata("generif-checksum", hash_generif_graph(sparql_uri).encode()) for child in combined_index.iterdir(): shutil.move(child, xapian_directory) logging.info("Index built") end_time = time.perf_counter() index_time = datetime.timedelta(seconds=end_time - start_time) logging.info(f"Time to Index: {index_time}") @click.group() def cli(): pass cli.add_command(is_data_modified) cli.add_command(create_xapian_index) if __name__ == "__main__": # pylint: disable=no-value-for-parameter cli()