#! /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: <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
"""
RIF_CACHE_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: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()
