db.py

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


class MongoDB(object):

    def __init__(self, host, username, password):
        self.client = MongoClient(host)
        self.db = self.client.chipseq
        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 delete_result(self, result_id):
        # Make sure result exists
        cursor = self.db.results.find({
            "_id": result_id
        })
        if cursor.count() == 1:
            result = cursor.next()
            self.gfs.delete(result["gridfs_id"])
            self.db[result["result_type"]].delete_many({"result_id": result["_id"]})
            self.db.results.delete_one({"_id": result["_id"]})
        else:
            print "result_id %s doesn't exist." % (result_id,)
        return

    def check_result(self, file_name, control_sample_ids, experiment_sample_ids, ref_genome, result_type, params, debug=False):
        try:
            query = {
                "result_type": result_type,
                "ref_genome": ref_genome,
                "timestamp": {"$exists": True},
                "file_name": file_name
            }
            # In the case that there are 0 samples we just want to check for existence.
            query["control_sample_ids"] = {"$all": control_sample_ids} if (len(control_sample_ids) > 0) else {"$exists": True}
            query["experiment_sample_ids"] = {"$all": experiment_sample_ids} if (len(experiment_sample_ids) > 0) else {"$exists": True}
            for job_name in params:
                param_keys = params[job_name].keys()
                if len(param_keys) == 0:
                    query[job_name] = {"$exists": True}
                else:
                    for param_name in param_keys:
                        query[job_name + "." + param_name] = params[job_name][param_name]
            cursor = self.db.results.find(query)
            if debug:
                print "[DB.check_result] Query = %s" % (query,)
                print "[DB.check_result] count = %s" % (cursor.count(),)
            if cursor.count() > 0:
                # Use the latest if available
                return cursor.sort("timestamp", pymongo.DESCENDING).next()
        except pymongo.errors.OperationFailure as e:
            print "Error with experiment_id %s: %s" % (experiment_id, e)
        return None

    def create_result(self, output_file, control_sample_ids, experiment_sample_ids, result_type, additional_data = {}, gfs_attributes = {}):
        # Make sure output_file exists
        if os.path.isfile(output_file):
            # Make sure that all control_sample_ids & experiment_sample_ids are valid
            # REMEMBER, these are ids for control & experiment SAMPLES
            valid_controls = [self.is_valid_sample(cid) for cid in control_sample_ids]
            valid_experiments = [self.is_valid_sample(eid) for eid in experiment_sample_ids]
            if all(valid_controls) and all(valid_experiments):
                # 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_ids": control_sample_ids,
                    "experiment_sample_ids": experiment_sample_ids,
                    "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_ids, experiment_sample_ids], [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 save_bam(self, bam_file, control_sample_ids, experiment_sample_ids, additional_data = {}):
        # Create result entry for bam files.  Since bam is a binary format, the file will only
        # be stored in GridFS
        valid, msg, result_id = self.create_result(bam_file, control_sample_ids, experiment_sample_ids, "bam", additional_data, gfs_attributes = {"file_type": "bam"})
        return (valid, msg, result_id)


    def save_bed(self, bed_file, control_sample_ids, experiment_sample_ids, additional_data = {}):
        # Create result_entry for bed_file
        valid, msg, result_id = self.create_result(bed_file, control_sample_ids, experiment_sample_ids, "bed", additional_data, gfs_attributes = {"file_type": "bed"})
        if valid:
            # Now we load the actual bed data into the bed collection.
            # Data is in a six column format
            # chr, start, end, name, score, strand
            # Load data using a list comprehension over lines,
            # then insert with insert_one()
            # Each document contains "n_lines" number of lines from the
            # result BED file.
            print "loading bed_data..."
            with open(bed_file, "r") as rh:
                msg = "Bed file successfully inserted."
                # Lazy load files in specified line chunk size
                n_lines = chipathlon.conf.result_lines_per_document
                line_set = list(itertools.islice(rh, n_lines))
                while line_set:
                    try:
                        result_lines = []
                        for line in line_set:
                            line_info = line.split()
                            line_record =  {
                                "chr": line_info[0],
                                "start": line_info[1],
                                "end": line_info[2],
                                "name": line_info[3],
                                "score": line_info[4],
                                "strand": line_info[5],
                            }
                            result_lines.append(line_record)

                        self.db.bed.insert_one({"result_id": result_id, "result_lines": result_lines})

                    except pymongo.errors.OperationFailure as e:
                        valid = False
                        msg = "Error inserting bed_file %s: %s" % (bed_file, e)
                    line_set = list(itertools.islice(rh, n_lines))
        return (valid, msg, result_id)

    def save_peak(self, peak_file, control_sample_ids, experiment_sample_ids, additional_data = {}):
        # Create result_entry for peak_file
        valid, msg, result_id = self.create_result(peak_file, control_sample_ids, experiment_sample_ids, "peak", additional_data, gfs_attributes = {"file_type": os.path.splitext(peak_file)[1][1:]})
        if valid:
            # Now we load the actual peak data into the collection
            # Data is in a 10 column format
            # chr, start, end, name, score, strand, signal_value, p_value, q_value, summit
            # Each document contains "n_lines" number of lines from the
            # result peak file.
            with open(peak_file, "r") as rh:
                msg = "Peak file successfully inserted."
                # Lazy load files in specified line chunk size
                n_lines = chipathlon.conf.result_lines_per_document
                line_set = list(itertools.islice(rh, n_lines))
                while line_set:
                    try:
                        result_lines = []
                        for line in line_set:
                            line_info = line.split()
                            line_record =  {
                                "chr": line_info[0],
                                "start": line_info[1],
                                "end": line_info[2],
                                "name": line_info[3],
                                "score": line_info[4],
                                "strand": line_info[5],
                                "signal_value": line_info[6],
                                "p_value": line_info[7],
                                "q_value": line_info[8],
                                "summit": line_info[9]
                            }
                            result_lines.append(line_record)

                        self.db.peak.insert_one({"result_id": result_id, "result_lines": result_lines})

                    except pymongo.errors.OperationFailure as e:
                        valid = False
                        msg = "Error inserting peak_file %s: %s" % (peak_file, e)
                    line_set = list(itertools.islice(rh, n_lines))
        return (valid, msg, result_id)

    def is_valid_sample(self, sample_accession):
        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_id):
        try:
            cursor = self.db.experiments.find({
                "target": {"$exists": True},
                "revoked_files.0": {"$exists": False},
                "@id": "/experiments/%s/" % (experiment_id,)
            })
            if cursor.count() == 1:
                return True
        except pymongo.errors.OperationFailure as e:
            print "Error with experiment_id %s: %s" % (experiment_id, e)
        return False

    def check_valid_samples(self):
        cursor = self.db.experiments.aggregate([
            {
                "$match": {
                    "target": {"$exists": True},
                    "revoked_files.0": {"$exists": False},
                    "assembly.0": {"$exists": True},
                    "assembly.1": {"$exists": False}
                }
            },
            {
                "$lookup": {
                    "from": "samples",
                    "localField": "uuid",
                    "foreignField": "experiment_id",
                    "as": "samples"
                }
            }
        ])
        total = 0
        has_samples = 0
        for document in cursor:
            total += 1
            if len(document["samples"]) > 0:
                has_samples += 1
        return (has_samples, total)

    def get_assembly(self, experiment_id):
        valid = True
        msg = ""
        data = ""
        cursor = self.db.experiments.find({
            "target": {"$exists": True},
            "revoked_files.0": {"$exists": False},
            "assembly.0": {"$exists": True},
            "assembly.1": {"$exists": False},
            "@id": "/experiments/%s/" % (experiment_id,)
        })
        if cursor.count() == 1:
            document = cursor.next()
            data = document["assembly"][0]
            msg = "Succesfully retrieved assembly for experiment with id '%s'.\n" % (experiment_id,)
        else:
            valid = False
            msg = "Experiment with id '%s' does not exist.\n" % (experiment_id,)
        return (valid, msg, data)

    def fetch_from_gridfs(self, gridfs_id, filename, checkmd5=True):
        """
        :param gridfs_id: GridFS _id of file to get.
        :type gridfs_id: bson.objectid.ObjectId
        :param filename: Filename to save file to.
        :type filename: str

        Fetch the file with _id 'gridfs_id' from GridFS and save to the file 'filename'.
        """

        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_samples(self, experiment_id):
        valid = True
        msg = ""
        data = {}
        # First, check to make sure the target experiment is valid
        if self.is_valid_experiment(experiment_id):
            # Next, we check that there is a least 1 possible control
            check3 = self.db.experiments.find({
                "target": {"$exists": True},
                "revoked_files.0": {"$exists": False},
                "assembly.0": {"$exists": True},
                "assembly.1": {"$exists": False},
                "possible_controls.0": {"$exists": True},
                "@id": "/experiments/%s/" % (experiment_id,)
            })
            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},
                            "revoked_files.0": {"$exists": False},
                            "assembly.0": {"$exists": True},
                            "assembly.1": {"$exists": False},
                            "possible_controls.0": {"$exists": True},
                            "@id": "/experiments/%s/" % (experiment_id,)
                        }
                    },
                    {
                        "$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"] == "fastq")]
                experiment_inputs = [sample for sample in document["samples"][0] if ("file_type" in sample and sample["file_type"] == "fastq")]
                if (len(control_inputs) > 0 and len(experiment_inputs) > 0):
                    msg = "Succesfully retrieved input files for experiment with id '%s'.\n" % (experiment_id,)
                    data = {
                        "control": control_inputs,
                        "experiment": experiment_inputs
                    }
                else:
                    valid = False
                    msg = "Experiment with id '%s' has %s possible control inputs, and %s possible experiment inputs.\n" % (experiment_id, len(control_inputs), len(experiment_inputs))
            else:
                valid = False
                msg = "Experiment with id '%s' does not have possible_controls.\n" % (experiment_id,)
        else:
            valid = False
            msg = "Experiment with id '%s' is not valid!  It may not exist, or it may be missing required metadata.\n" % (experiment_id,)
        return (valid, msg, data)