#! /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, Counter 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 re import shutil import sys import hashlib import tempfile from typing import Callable, Dict, 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 = 100_000 # Running the script in prod consumers ~1GB per process when handling 100_000 Documents per chunk. # To prevent running out of RAM, we set this as the upper bound for total concurrent processes PROCESS_COUNT_LIMIT = 67 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")]) 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:GNWikiEntry ; gnt:belongsToSpecies ?species ; rdfs:comment ?comment . ?species gnt:shortName ?speciesName . } GROUP BY ?speciesName ?symbolName """ 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: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, remove_common_words: bool = False): 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"] count: Counter[str] = Counter() words_regex = re.compile(r"\w+") for entry in bindings : x = (entry["speciesName"]["value"], entry["symbolName"]["value"],) value = entry["comment"]["value"] value = " ".join(words_regex.findall(value)) # remove punctuation cache[x] = value count.update(Counter(value.lower().strip().split())) if not remove_common_words: return cache words_to_drop = set() for word, cnt in count.most_common(1000): if len(word) < 4 or cnt > 3000: words_to_drop.add(word) smaller_cache = {} for entry, value in cache.items(): new_value = set(word for word in value.lower().split() if word not in words_to_drop) smaller_cache[entry] = " ".join(new_value) return smaller_cache def md5hash_ttl_dir(ttl_dir: pathlib.Path) -> str: if not ttl_dir.exists(): return "-1" ttl_hash = hashlib.new("md5") for ttl_file in ttl_dir.glob("*.ttl"): with open(ttl_file, encoding="utf-8") as f_: ttl_hash.update(f_.read().encode()) return ttl_hash.hexdigest() # 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")) termgenerator.set_stopper_strategy(xapian.TermGenerator.STOP_ALL) termgenerator.set_stopper(xapian.SimpleStopper()) 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_without_positions(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], 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 default_no_of_workers = os.cpu_count() or 1 no_of_workers = min(default_no_of_workers, PROCESS_COUNT_LIMIT) try: with worker_queue(no_of_workers) 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 parallel_xapian_compact(combined_index: pathlib.Path, indices: List[pathlib.Path]) -> None: # We found that compacting 50 files of ~600MB has decent performance no_of_workers = 20 file_groupings = 50 with temporary_directory("parallel_combine", str(combined_index)) as parallel_combine: parallel_combine.mkdir(parents=True, exist_ok=True) with worker_queue(no_of_workers) as spawn_worker: i = 0 while i < len(indices): end_index = (i + file_groupings) files = indices[i:end_index] last_item_idx = i + len(files) spawn_worker(xapian_compact, (parallel_combine / f"{i}_{last_item_idx}", files)) logging.debug("Spawned worker to compact files from %s to %s", i, last_item_idx) i = end_index logging.debug("Completed parallel xapian compacts") xapian_compact(combined_index, list(parallel_combine.iterdir())) 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)) combined_index.mkdir(parents=True, exist_ok=True) start = time.monotonic() 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() logging.debug("Removing databases that were compacted into %s", combined_index.name) for folder in indices: shutil.rmtree(folder) logging.debug("Completed xapian-compact %s; handled %s files in %s minutes", combined_index.name, len(indices), (time.monotonic() - start) / 60) @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. generif = pathlib.Path("/var/lib/data/") if (db.get_metadata("generif-checksum").decode() == md5hash_ttl_dir(generif) 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') 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 logging.info("Building wiki cache") wiki_cache = build_rdf_cache(sparql_uri, WIKI_CACHE_QUERY, remove_common_words=True) logging.info("Building rif cache") rif_cache = build_rdf_cache(sparql_uri, RIF_CACHE_QUERY, remove_common_words=True) 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") parallel_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", md5hash_ttl_dir(pathlib.Path("/var/lib/data/")).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()