db.py

from pymongo import MongoClient
import pymongo.errors
import gridfs
import sys
import traceback
import os
import itertools
import time
import collections
import chipathlon.conf
from pprint import pprint
import hashlib
from chipathlon.utils import progress


class MongoDB(object):
    """
    :param host: The host address of the MongoDB database.
    :type host: str
    :param username: The username of the account for the MongoDB database.
    :type username: str
    :param password: The password for the user.
    :type password: str
    :param debug: A flag for printing additional messages.
    :type debug: bool

    This class is used to manage all interactions with the encode metadata.
    The metadata can be very unruly and difficult to deal with.  There
    are several helper functions within this class to make some database
    operations much easier.
    """

    def __init__(self, host="localhost", username=None, password=None, debug=False):
        self.debug = debug
        self.host = host
        self.username = username
        self.password = password
        self.client = MongoClient(host)
        self.db = self.client.chipseq
        self.cache = collections.defaultdict(dict)
        self.cache = {}
        if username and password:
            try:
                self.db.authenticate(username, password, mechanism="SCRAM-SHA-1")
            except:
                print("Could not authenticate to db %s!" % (host,))
                print traceback.format_exc()
                sys.exit(1)
        self.gfs = gridfs.GridFS(self.db)
        return

    def add_cache(self, accession, file_type, data):
        """
        :param accession: The accession of the file to store.
        :type accession: str
        :param file_type: The type of file to store.
        :type file_type: str
        :param data: The data to add to the cache.
        :type data: Object

        Adds a data result to the internal cache.  This is used to speed up
        requests that are identical.  We may have multiple runs that use
        identical control / signal files but change around the alignment or
        peak calling tools.  In these cases we don't want to request info
        from the database multiple times for the same data.
        """
        self.cache[accession][file_type] = data
        return

    def get_cache(self, accession, file_type):
        """
        :param accession: The accession of the file to retrieve.
        :type accession: str
        :param file_type: The type of file to retrieve.
        :type file_type: str

        Gets a data item from the internal cache.
        """
        if accession in self.cache:
            return self.cache[accession].get(file_type)
        return None

    def delete_result(self, result, genome):
        """
        :param result: The result to delete
        :type result: :py:class:`~chipathlon.result.Result`
        :param genome: The genome to find information from.
        :type genome: :py:class:`~chipathlon.genome.Genome`

        Deletes a result and it's corresponding gridfs entry.
        """
        result_id = self.get_reuslt_id(result, genome)
        cursor = self.db.results.find({
            "_id": result_id
        })
        if cursor.count() == 1:
            result = cursor.next()
            self.gfs.delete(result["gridfs_id"])
            self.db.results.delete_one({"_id": result["_id"]})
        else:
            print "result_id %s doesn't exist." % (result_id,)
        return

    def _get_result_query(self, result, genome):
        query = {
            "result_type": result.file_type,
            "assembly": genome.assembly,
            "timestamp": {"$exists": True},
            "file_name": result.full_name
        }
        # In the case that there are 0 samples we just want to check for existence.
        control_sample_accessions = result.get_accessions("control")
        signal_sample_accessions = result.get_accessions("signal")
        query["control_sample_accessions"] = {"$all": control_sample_accessions} if (len(control_sample_accessions) > 0) else {"$exists": True}
        query["signal_sample_accessions"] = {"$all": signal_sample_accessions} if (len(signal_sample_accessions) > 0) else {"$exists": True}
        for job in result.all_jobs:
            job_args = job.get_db_arguments()
            arg_keys = job_args.keys()
            if len(arg_keys) == 0:
                query[job.job_name] = {"$exists": True}
            else:
                for arg_name in arg_keys:
                    query[job.job_name + "." + arg_name] = job_args[arg_name]
        if self.debug:
            print "Result query: %s" % (query,)
        return query

    def result_exists(self, result, genome):
        """
        :param result: The result to check.
        :type result: :py:meth:`~chipathlon.result.Result`
        :param genome: The genome to find information from.
        :type genome: :py:meth:`~chipathlon.genome.Genome`

        Check if a result exists in the database.  The genome parameter
        is required since some files have been aligned or use individual
        chromsome fasta or size files for peak calling.
        """
        try:
            cursor = self.db.results.find(self._get_result_query(result, genome))
            return cursor.count() > 0
        except pymongo.errors.OperationFailure as e:
            print "Error checking result [%s]: %s" % (file_name, e)
        return False

    def get_result_id(self, result, genome):
        """
        :param result: The result to check.
        :type result: :py:meth:`~chipathlon.result.Result`
        :param genome: The genome to find information from.
        :type genome: :py:meth:`~chipathlon.genome.Genome`
        :returns: The id found or None

        Get the id of a result in the database.
        """
        try:
            cursor = self.db.results.find(self._get_result_query(result, genome))
            if cursor.count() == 1:
                return cursor._id
        except pymongo.errors.OperationFailure as e:
            print "Error getting result id [%s]: %s" % (file_name, e)
        return None

    def get_result(self, result, genome):
        """
        :param result: The result to check.
        :type result: :py:meth:~chipathlon.result.Result
        :param genome: The genome to find information from.
        :type genome: :py:meth:~chipathlon.genome.Genome

        Get the metadata for the result from the database.  If multiple results
        exist, the most recently saved result is returned.
        """
        try:
            cursor = self.db.results.find(self._get_result_query(result, genome))
            if cursor.count() > 0:
                return cursor.sort("timestamp", pymongo.DESCENDING).next()
        except pymongo.errors.OperationFailure as e:
            print "Error checking result [%s]: %s" % (file_name, e)
        return None

    def save_result(self, output_file, control_sample_accessions, signal_sample_accessions, result_type, additional_data = {}, gfs_attributes = {}):
        """
        :param output_file: The path to the result to save.
        :type output_file: str
        :param control_sample_accessions: A list of control accessions.
        :type control_sample_accessions: list
        :param signal_sample_accessions: A list of signal accessions.
        :type signal_sample_accessions: list
        :param result_type: Useful file type info
        :type result_type: str
        :param additional_data: Additional metadata to store in mongo.
        :type additional_data: dict
        :param gfs_attributes: Additional metadata to store in gridfs.
        :type gfs_attributes: dict

        Saves a result entry into MongodDB and uploads the file into gridfs.
        The only difference between additional_data and gfs_attributes is the
        location the metadata is stored.  Both just store key value pairs of
        information, the additional_data information is stored in the result
        entry, the gfs_attributes information is stored in gridfs.
        """
        # Make sure output_file exists
        if os.path.isfile(output_file):
            # Make sure that all control_sample_accessions & signal_sample_accessions are valid
            # REMEMBER, these are ids for control & experiment SAMPLES
            valid_controls = [self.is_valid_sample(cid) for cid in control_sample_accessions]
            valid_experiments = [self.is_valid_sample(eid) for eid in signal_sample_accessions]
            if all(valid_controls) and all(valid_experiments):
                gfs_attributes["file_type"] = result_type
                # First, we load the output file into gfs
                with open(output_file, "r") as rh:
                    # Calling put returns the gfs id
                    gridfs_id = self.gfs.put(rh, filename=os.path.basename(output_file), **gfs_attributes)
                # Now, we create the actual result entry by combining all necessary info
                result_entry = {
                    "gridfs_id": gridfs_id,
                    "control_sample_accessions": control_sample_accessions,
                    "signal_sample_accessions": signal_sample_accessions,
                    "result_type": result_type,
                    "file_name": output_file,
                    "timestamp": time.time()
                }
                # Add additional attributes into the result_entry
                result_entry.update(additional_data)
                # Insert the entry into the database, and return the id
                result = self.db.results.insert_one(result_entry)
                return (True, "Result created successfully.", result.inserted_id)
            else:
                msg = "Not all input ids are valid.  The following are invalid: "
                for id_list, valid_list in zip([control_sample_accessions, signal_sample_accessions], [valid_controls, valid_experiments]):
                    msg += ", ".join([id_list[i] for i, valid in enumerate(valid_list) if not valid])
        else:
            msg = "Specified output_file %s does not exist." % (output_file,)
        return (False, msg, None)

    def is_valid_sample(self, sample_accession):
        """
        :param sample_accession: The accession number to check.
        :type sample_accession: str
        :returns: Whether or not the sample is valid.

        Ensures that a sample with the accession specified actually exists.
        """
        try:
            cursor = self.db.samples.find({
                "accession": sample_accession
            })
            if cursor.count() == 1:
                return True
        except pymongo.errors.OperationFailure as e:
            print "Error with sample_accession %s: %s" % (sample_accession, e)
        return False

    def is_valid_experiment(self, experiment_accession):
        """
        :param experiment_accession: The accession number to check.
        :type experiment_accession: str
        :returns: Whether or not the experiment is valid

        Ensures that an experiment with the accession specified actually exists.
        """
        try:
            cursor = self.db.experiments.find({
                "target": {"$exists": True},
                "@id": "/experiments/%s/" % (experiment_accession,)
            })
            if cursor.count() == 1:
                return True
        except pymongo.errors.OperationFailure as e:
            print "Error with experiment_accession %s: %s" % (experiment_accession, e)
        return False

    def fetch_from_gridfs(self, gridfs_id, filename, checkmd5=True):
        """
        :param gridfs_id: GridFS _id of file to get.
        :type gridfs_id: :py:class:`bson.objectid.ObjectId`
        :param filename: Filename to save file to.
        :type filename: str
        :param checkmd5: Whether or not to validate the md5 of the result
        :type checkmd5: bool

        Fetch the file with the corresponding id and save it under the
        specified 'filename'.  If checkmd5 is specified, validate that the
        saved file has a correct md5 value.
        """
        try:
            gridfs_file = self.gfs.get(gridfs_id)
            gridfs_md5 = gridfs_file.md5
        except gridfs.errors.NoFile as e:
            print "Error fetching file from GridFS!\nNo file with ID '%s'" % (gridfs_id)
            print e
            sys.exit(1)

        try:
            output_fh = open(filename,'wb')
        except IOError as e:
            print "Error creating GridFS output file '%s':" % (filename)
            print (e.errno,e.strerror)
            sys.exit(1)

        hash_md5 = hashlib.md5()
        for chunk in gridfs_file:
            output_fh.write(chunk)
            hash_md5.update(chunk)

        output_fh.close()
        gridfs_file.close()

        if checkmd5:
            if gridfs_md5 == hash_md5.hexdigest():
                return True
            else:
                print "MD5 mismatch saving file from GridFS to '%s'" % (filename)
                return False
        else:
            return True

    def get_sample(self, accession, file_type):
        """
        :param accession: The accession number of the target sample
        :type accession: string
        :param file_type: The file type of the target sample.
        :type file_type: string
        :returns: A tuple (valid, msg, data)

        Gets the associated sample based on accession number and file_type.
        For loading input files for workflows the file_type should be fastq
        or bam.  Other file types can be specified for loading additional files
        saved in the experiment metadata.
        """
        valid = True
        msg = ""
        data = {}
        check_cache = self.get_cache(accession, file_type)
        if check_cache is not None:
            msg = "Retrieved data from cache."
            data = check_cache
        else:
            cursor = self.db.samples.find({
                "accession": accession,
                "file_type": file_type
            })
            if cursor.count() == 1:
                data = cursor.next()
                self.add_cache(accession, file_type, data)
            else:
                valid = False
                msg = "Found %s files with accession: %s, file_type: %s. Should only be 1." % (
                    cursor.count(),
                    accession,
                    file_type
                )
        return (valid, msg, data)

    def clean_gfs(self):
        """
        This function finds all files stored in gridfs that are not currently
        referenced by any result file and removes them.
        A clean database is a happy database.
        """
        cursor = self.db.results.aggregate([
            {
                "$group": {
                    "_id": 1,
                    "valid_ids": {"$push": "$gridfs_id"}
                }
            }
        ])
        # Doc contains all our valid ids
        id_doc = cursor.next()
        # Find all fs.files documents
        gfs_cursor = self.db.fs.files.find({
            "_id": {
                "$nin": id_doc["valid_ids"]
            }
        })
        # Iterate through file, delete fs.chunks then fs.files
        total_files = gfs_cursor.count()
        print "Found %s unused gridfs files.  Preparing to delete...." % (total_files,)
        for i, fs_file in enumerate(gfs_cursor):
            progress(i, total_files)
            self.db.fs.chunks.remove({
                "files_id": fs_file["_id"]
            })
            self.db.fs.files.remove({
                "_id": fs_file["_id"]
            })
        progress(total_files, total_files)
        return

    def get_samples(self, experiment_accession, file_type):
        """
        :param experiment_accession: Accession number of the experiment to grab samples from.
        :type experiment_accession: str
        :param file_type: File type of samples to grab usually fastq or bam
        :type file_type: str
        :returns: A tuple (valid, msg, data)

        Validates and gets samples for the given experiment.  Experiments must
        have control and signal samples of the provided file_type to be
        considered valid.  Returns a tuple with three values (valid, msg, data)
        valid -- Whether or not the accession / file_type combo is a valid exp
        msg -- Why it is or is not valid
        data -- A dictionary containing a list of all control / sample documents.

        The data dictionary has two keys, "control" and "signal", each one containing
        a list of all metadata related to the experiment samples.  The sample metadata
        is taken directly from Mongo.
        """
        valid = True
        msg = ""
        data = {}
        # First, check to make sure the target experiment is valid
        if self.is_valid_experiment(experiment_accession):
            # Next, we check that there is a least 1 possible control
            check3 = self.db.experiments.find({
                "target": {"$exists": True},
                "possible_controls.0": {"$exists": True},
                "@id": "/experiments/%s/" % (experiment_accession,)
            })
            if check3.count() == 1:
                # Complicated aggregtaion pipeline does the following steps:
                # 1. Find the experiment that matches the given id
                # 2. Join samples into the collection by exp_id
                # 3. Iterate through possible_controls
                # 4. Join possible_control data into control_exps
                # 5. Iterate through control_exps
                # 6. Join samples into the control_exps by exp_id
                # 7. Re-aggregate all data into arrays
                pipeline = [
                    {
                        "$match": {
                            "target": {"$exists": True},
                            "possible_controls.0": {"$exists": True},
                            "@id": "/experiments/%s/" % (experiment_accession,)
                        }
                    },
                    {
                        "$lookup": {
                            "from": "samples",
                            "localField": "uuid",
                            "foreignField": "experiment_id",
                            "as": "samples"
                        }
                    },
                    {
                        "$unwind": "$possible_controls"
                    },
                    {
                        "$lookup": {
                            "from": "samples",
                            "localField": "possible_controls.uuid",
                            "foreignField": "experiment_id",
                            "as": "possible_controls.samples"
                        }
                    },
                    {
                        "$group": {
                            "_id": "$_id",
                            "possible_controls": {"$push": "$possible_controls"},
                            "samples": {"$push": "$samples"}
                        }
                    }
                ]
                cursor = self.db.experiments.aggregate(pipeline)
                # We should have only 1 document
                document = cursor.next()
                control_inputs = [sample for control in document["possible_controls"] for sample in control["samples"] if ("file_type" in sample and sample["file_type"] == file_type)]
                signal_inputs = [sample for sample in document["samples"][0] if ("file_type" in sample and sample["file_type"] == file_type)]
                if (len(control_inputs) > 0 and len(signal_inputs) > 0):
                    msg = "Succesfully retrieved input files for experiment with id '%s'.\n" % (experiment_accession,)
                    data = {
                        "control": control_inputs,
                        "signal": signal_inputs
                    }
                else:
                    valid = False
                    msg = "Experiment with id '%s' has %s possible control inputs, and %s possible signal inputs.\n" % (experiment_accession, len(control_inputs), len(signal_inputs))
            else:
                valid = False
                msg = "Experiment with id '%s' does not have possible_controls.\n" % (experiment_accession,)
        else:
            valid = False
            msg = "Experiment with id '%s' is not valid!  It may not exist, or it may be missing required metadata.\n" % (experiment_accession,)
        return (valid, msg, data)