path: root/scripts/index-genenetwork

#! /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 collections import deque, namedtuple
import contextlib
import time
import datetime
from functools import partial
import itertools
import json
import logging
from multiprocessing import Lock, Process
import os
import pathlib
import resource
import shutil
import sys
import tempfile
from typing import Callable, Generator, 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")])


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_rif_cache(sparql_uri: str):
    cache = {}
    sparql = SPARQLWrapper(sparql_uri)
    sparql.setReturnFormat(JSON)
    query = """
PREFIX rdf: <http://www.w3.org/1999/02/22-rdf-syntax-ns#>
PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
PREFIX gnt: <http://genenetwork.org/term/>
PREFIX gnc: <http://genenetwork.org/category/>

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
"""
    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: <http://www.w3.org/1999/02/22-rdf-syntax-ns#>
PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
PREFIX gnt: <http://genenetwork.org/term/>
PREFIX gnc: <http://genenetwork.org/category/>

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(2)
def index_rif_comments(species, symbol):
    key = (species, symbol,)
    entry = rdfcache.get(key)
    if entry:
        termgenerator.index_text(entry, 0, "XRF")


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.
data: Iterable
rdfcache: 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 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_fullname"]:
                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_fullname"].bind(index_dataset)
            trait["symbol"].bind(index_symbol)
            trait["chr"].bind(index_chr)
            trait["geno_chr"].bind(index_peakchr)

            Maybe.apply(
                index_rif_comments
            ).to_arguments(
                trait["species"], trait["symbol"]
            )

            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 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_fullname"]:
                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_fullname"].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, 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:
            global rdfcache
            rdfcache = build_rif_cache(sparql_uri)
            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):
                    # pylint: disable=global-statement
                    global data
                    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 = " ".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"),
                        filename=os.environ.get("INDEX_LOG_FILE"),
                        filemode="a",
                        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:
            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()