Merge pull request #32 from genenetwork/heatmap_decompose_db_retrieval

Heatmap decompose db retrieval

7 files changed, 1305 insertions, 155 deletions

diff --git a/gn3/computations/heatmap.py b/gn3/computations/heatmap.py
new file mode 100644
index 0000000..a0e778a
--- /dev/null
+++ b/gn3/computations/heatmap.py
@@ -0,0 +1,177 @@
+"""
+This module will contain functions to be used in computation of the data used to
+generate various kinds of heatmaps.
+"""
+
+from functools import reduce
+from typing import Any, Dict, Sequence
+from gn3.computations.slink import slink
+from gn3.db.traits import retrieve_trait_data, retrieve_trait_info
+from gn3.computations.correlations2 import compute_correlation
+
+def export_trait_data(
+        trait_data: dict, strainlist: Sequence[str], dtype: str = "val",
+        var_exists: bool = False, n_exists: bool = False):
+    """
+    Export data according to `strainlist`. Mostly used in calculating
+    correlations.
+
+    DESCRIPTION:
+    Migrated from
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L166-L211
+
+    PARAMETERS
+    trait: (dict)
+      The dictionary of key-value pairs representing a trait
+    strainlist: (list)
+      A list of strain names
+    type: (str)
+      ... verify what this is ...
+    var_exists: (bool)
+      A flag indicating existence of variance
+    n_exists: (bool)
+      A flag indicating existence of ndata
+    """
+    def __export_all_types(tdata, strain):
+        sample_data = []
+        if tdata[strain]["val"]:
+            sample_data.append(tdata[strain]["val"])
+            if var_exists:
+                if tdata[strain].var:
+                    sample_data.append(tdata[strain]["var"])
+                else:
+                    sample_data.append(None)
+            if n_exists:
+                if tdata[strain]["ndata"]:
+                    sample_data.append(tdata[strain]["ndata"])
+                else:
+                    sample_data.append(None)
+        else:
+            if var_exists and n_exists:
+                sample_data += [None, None, None]
+            elif var_exists or n_exists:
+                sample_data += [None, None]
+            else:
+                sample_data.append(None)
+
+        return tuple(sample_data)
+
+    def __exporter(accumulator, strain):
+        # pylint: disable=[R0911]
+        if trait_data.has_key(strain):
+            if dtype == "val":
+                return accumulator + (trait_data[strain]["val"], )
+            if dtype == "var":
+                return accumulator + (trait_data[strain]["var"], )
+            if dtype == "N":
+                return trait_data[strain]["ndata"]
+            if dtype == "all":
+                return accumulator + __export_all_types(trait_data, strain)
+            raise KeyError("Type `%s` is incorrect" % dtype)
+        if var_exists and n_exists:
+            return accumulator + (None, None, None)
+        if var_exists or n_exists:
+            return accumulator + (None, None)
+        return accumulator + (None,)
+
+    return reduce(__exporter, strainlist, tuple())
+
+def trait_display_name(trait: Dict):
+    """
+    Given a trait, return a name to use to display the trait on a heatmap.
+
+    DESCRIPTION
+    Migrated from
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L141-L157
+    """
+    if trait.get("db", None) and trait.get("trait_name", None):
+        if trait["db"]["dataset_type"] == "Temp":
+            desc = trait["description"]
+            if desc.find("PCA") >= 0:
+                return "%s::%s" % (
+                    trait["db"]["displayname"],
+                    desc[desc.rindex(':')+1:].strip())
+            return "%s::%s" % (
+                trait["db"]["displayname"],
+                desc[:desc.index('entered')].strip())
+        prefix = "%s::%s" % (
+            trait["db"]["dataset_name"], trait["trait_name"])
+        if trait["cellid"]:
+            return "%s::%s" % (prefix, trait["cellid"])
+        return prefix
+    return trait["description"]
+
+def cluster_traits(traits_data_list: Sequence[Dict]):
+    """
+    Clusters the trait values.
+
+    DESCRIPTION
+    Attempts to replicate the clustering of the traits, as done at
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/heatmap/Heatmap.py#L138-L162
+    """
+    def __compute_corr(tdata_i, tdata_j):
+        if tdata_j[0] < tdata_i[0]:
+            corr_vals = compute_correlation(tdata_i, tdata_j)
+            corr = corr_vals[0]
+            if (1 - corr) < 0:
+                return 0.0
+            return 1 - corr
+        return 0.0
+
+    def __cluster(tdata_i):
+        return tuple(
+            __compute_corr(tdata_i, tdata_j)
+            for tdata_j in enumerate(traits_data_list))
+
+    return tuple(__cluster(tdata_i) for tdata_i in enumerate(traits_data_list))
+
+def heatmap_data(formd, search_result, conn: Any):
+    """
+    heatmap function
+
+    DESCRIPTION
+    This function is an attempt to reproduce the initialisation at
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/heatmap/Heatmap.py#L46-L64
+    and also the clustering and slink computations at
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/heatmap/Heatmap.py#L138-L165
+    with the help of the `gn3.computations.heatmap.cluster_traits` function.
+
+    It does not try to actually draw the heatmap image.
+
+    PARAMETERS:
+    TODO: Elaborate on the parameters here...
+    """
+    threshold = 0 # webqtlConfig.PUBLICTHRESH
+    cluster_checked = formd.formdata.getvalue("clusterCheck", "")
+    strainlist = [
+        strain for strain in formd.strainlist if strain not in formd.parlist]
+    genotype = formd.genotype
+
+    def __retrieve_traitlist_and_datalist(threshold, fullname):
+        trait = retrieve_trait_info(threshold, fullname, conn)
+        return (
+            trait,
+            export_trait_data(retrieve_trait_data(trait, conn), strainlist))
+
+    traits_details = [
+        __retrieve_traitlist_and_datalist(threshold, fullname)
+        for fullname in search_result]
+    traits_list = map(lambda x: x[0], traits_details)
+    traits_data_list = map(lambda x: x[1], traits_details)
+
+    return {
+        "target_description_checked": formd.formdata.getvalue(
+            "targetDescriptionCheck", ""),
+        "cluster_checked": cluster_checked,
+        "slink_data": (
+            slink(cluster_traits(traits_data_list))
+            if cluster_checked else False),
+        "sessionfile": formd.formdata.getvalue("session"),
+        "genotype": genotype,
+        "nLoci": sum(map(len, genotype)),
+        "strainlist": strainlist,
+        "ppolar": formd.ppolar,
+        "mpolar":formd.mpolar,
+        "traits_list": traits_list,
+        "traits_data_list": traits_data_list
+    }
diff --git a/gn3/computations/slink.py b/gn3/computations/slink.py
index 23d3d88..5953e6b 100644
--- a/gn3/computations/slink.py
+++ b/gn3/computations/slink.py
@@ -7,7 +7,7 @@ slink:
     TODO: Describe what the function does...
 """
 import logging
-from typing import List, Tuple, Union, Sequence
+from typing import Union, Sequence
 
 NumType = Union[int, float]
 SeqOfNums = Sequence[NumType]
diff --git a/gn3/db/datasets.py b/gn3/db/datasets.py
new file mode 100644
index 0000000..53d6811
--- /dev/null
+++ b/gn3/db/datasets.py
@@ -0,0 +1,291 @@
+"""
+This module contains functions relating to specific trait dataset manipulation
+"""
+from typing import Any
+
+def retrieve_probeset_trait_dataset_name(
+        threshold: int, name: str, connection: Any):
+    """
+    Get the ID, DataScale and various name formats for a `ProbeSet` trait.
+    """
+    query = (
+        "SELECT Id, Name, FullName, ShortName, DataScale "
+        "FROM ProbeSetFreeze "
+        "WHERE "
+        "public > %(threshold)s "
+        "AND "
+        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)")
+    with connection.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                "threshold": threshold,
+                "name": name
+            })
+        return dict(zip(
+            ["dataset_id", "dataset_name", "dataset_fullname",
+             "dataset_shortname", "dataset_datascale"],
+            cursor.fetchone))
+
+def retrieve_publish_trait_dataset_name(
+        threshold: int, name: str, connection: Any):
+    """
+    Get the ID, DataScale and various name formats for a `Publish` trait.
+    """
+    query = (
+        "SELECT Id, Name, FullName, ShortName "
+        "FROM PublishFreeze "
+        "WHERE "
+        "public > %(threshold)s "
+        "AND "
+        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)")
+    with connection.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                "threshold": threshold,
+                "name": name
+            })
+        return dict(zip(
+            ["dataset_id", "dataset_name", "dataset_fullname",
+             "dataset_shortname"],
+            cursor.fetchone))
+
+def retrieve_geno_trait_dataset_name(
+        threshold: int, name: str, connection: Any):
+    """
+    Get the ID, DataScale and various name formats for a `Geno` trait.
+    """
+    query = (
+        "SELECT Id, Name, FullName, ShortName "
+        "FROM GenoFreeze "
+        "WHERE "
+        "public > %(threshold)s "
+        "AND "
+        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)")
+    with connection.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                "threshold": threshold,
+                "name": name
+            })
+        return dict(zip(
+            ["dataset_id", "dataset_name", "dataset_fullname",
+             "dataset_shortname"],
+            cursor.fetchone))
+
+def retrieve_temp_trait_dataset_name(
+        threshold: int, name: str, connection: Any):
+    """
+    Get the ID, DataScale and various name formats for a `Temp` trait.
+    """
+    query = (
+        "SELECT Id, Name, FullName, ShortName "
+        "FROM TempFreeze "
+        "WHERE "
+        "public > %(threshold)s "
+        "AND "
+        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)")
+    with connection.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                "threshold": threshold,
+                "name": name
+            })
+        return dict(zip(
+            ["dataset_id", "dataset_name", "dataset_fullname",
+             "dataset_shortname"],
+            cursor.fetchone))
+
+def retrieve_dataset_name(
+        trait_type: str, threshold: int, trait_name: str, dataset_name: str,
+        conn: Any):
+    """
+    Retrieve the name of a trait given the trait's name
+
+    This is extracted from the `webqtlDataset.retrieveName` function as is
+    implemented at
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlDataset.py#L140-L169
+    """
+    fn_map = {
+        "ProbeSet": retrieve_probeset_trait_dataset_name,
+        "Publish": retrieve_publish_trait_dataset_name,
+        "Geno": retrieve_geno_trait_dataset_name,
+        "Temp": retrieve_temp_trait_dataset_name}
+    if trait_type == "Temp":
+        return retrieve_temp_trait_dataset_name(threshold, trait_name, conn)
+    return fn_map[trait_type](threshold, dataset_name, conn)
+
+
+def retrieve_geno_riset_fields(name, conn):
+    """
+    Retrieve the RISet, and RISetID values for various Geno trait types.
+    """
+    query = (
+        "SELECT InbredSet.Name, InbredSet.Id "
+        "FROM InbredSet, GenoFreeze "
+        "WHERE GenoFreeze.InbredSetId = InbredSet.Id "
+        "AND GenoFreeze.Name = %(name)s")
+    with conn.cursor() as cursor:
+        cursor.execute(query, {"name": name})
+        return dict(zip(["riset", "risetid"], cursor.fetchone()))
+    return {}
+
+def retrieve_publish_riset_fields(name, conn):
+    """
+    Retrieve the RISet, and RISetID values for various Publish trait types.
+    """
+    query = (
+        "SELECT InbredSet.Name, InbredSet.Id "
+        "FROM InbredSet, PublishFreeze "
+        "WHERE PublishFreeze.InbredSetId = InbredSet.Id "
+        "AND PublishFreeze.Name = %(name)s")
+    with conn.cursor() as cursor:
+        cursor.execute(query, {"name": name})
+        return dict(zip(["riset", "risetid"], cursor.fetchone()))
+    return {}
+
+def retrieve_probeset_riset_fields(name, conn):
+    """
+    Retrieve the RISet, and RISetID values for various ProbeSet trait types.
+    """
+    query = (
+        "SELECT InbredSet.Name, InbredSet.Id "
+        "FROM InbredSet, ProbeSetFreeze, ProbeFreeze "
+        "WHERE ProbeFreeze.InbredSetId = InbredSet.Id "
+        "AND ProbeFreeze.Id = ProbeSetFreeze.ProbeFreezeId "
+        "AND ProbeSetFreeze.Name = %(name)s")
+    with conn.cursor() as cursor:
+        cursor.execute(query, {"name": name})
+        return dict(zip(["riset", "risetid"], cursor.fetchone()))
+    return {}
+
+def retrieve_temp_riset_fields(name, conn):
+    """
+    Retrieve the RISet, and RISetID values for `Temp` trait types.
+    """
+    query = (
+        "SELECT InbredSet.Name, InbredSet.Id "
+        "FROM InbredSet, Temp "
+        "WHERE Temp.InbredSetId = InbredSet.Id "
+        "AND Temp.Name = %(name)s")
+    with conn.cursor() as cursor:
+        cursor.execute(query, {"name": name})
+        return dict(zip(["riset", "risetid"], cursor.fetchone()))
+    return {}
+
+def retrieve_riset_fields(trait_type, trait_name, dataset_info, conn):
+    """
+    Retrieve the RISet, and RISetID values for various trait types.
+    """
+    riset_fns_map = {
+        "Geno": retrieve_geno_riset_fields,
+        "Publish": retrieve_publish_riset_fields,
+        "ProbeSet": retrieve_probeset_riset_fields
+    }
+
+    if trait_type == "Temp":
+        riset_info = retrieve_temp_riset_fields(trait_name, conn)
+    else:
+        riset_info = riset_fns_map[trait_type](dataset_info["dataset_name"], conn)
+
+    return {
+        **dataset_info,
+        **riset_info,
+        "riset": (
+            "BXD" if riset_info.get("riset") == "BXD300"
+            else riset_info.get("riset", ""))
+    }
+
+def retrieve_temp_trait_dataset():
+    """
+    Retrieve the dataset that relates to `Temp` traits
+    """
+    # pylint: disable=[C0330]
+    return {
+        "searchfield": ["name", "description"],
+        "disfield": ["name", "description"],
+        "type": "Temp",
+        "dataset_id": 1,
+        "fullname": "Temporary Storage",
+        "shortname": "Temp"
+    }
+
+def retrieve_geno_trait_dataset():
+    """
+    Retrieve the dataset that relates to `Geno` traits
+    """
+    # pylint: disable=[C0330]
+    return {
+        "searchfield": ["name", "chr"],
+	"disfield": ["name", "chr", "mb", "source2", "sequence"],
+	"type": "Geno"
+    }
+
+def retrieve_publish_trait_dataset():
+    """
+    Retrieve the dataset that relates to `Publish` traits
+    """
+    # pylint: disable=[C0330]
+    return {
+        "searchfield": [
+            "name", "post_publication_description", "abstract", "title",
+            "authors"],
+        "disfield": [
+            "name", "pubmed_id", "pre_publication_description",
+            "post_publication_description", "original_description",
+	    "pre_publication_abbreviation", "post_publication_abbreviation",
+	    "lab_code", "submitter", "owner", "authorized_users",
+	    "authors", "title", "abstract", "journal", "volume", "pages",
+            "month", "year", "sequence", "units", "comments"],
+        "type": "Publish"
+    }
+
+def retrieve_probeset_trait_dataset():
+    """
+    Retrieve the dataset that relates to `ProbeSet` traits
+    """
+    # pylint: disable=[C0330]
+    return {
+        "searchfield": [
+            "name", "description", "probe_target_description", "symbol",
+            "alias", "genbankid", "unigeneid", "omim", "refseq_transcriptid",
+            "probe_set_specificity", "probe_set_blat_score"],
+	"disfield": [
+            "name", "symbol", "description", "probe_target_description", "chr",
+            "mb", "alias", "geneid", "genbankid", "unigeneid", "omim",
+            "refseq_transcriptid", "blatseq", "targetseq", "chipid", "comments",
+            "strand_probe", "strand_gene", "probe_set_target_region",
+            "proteinid", "probe_set_specificity", "probe_set_blat_score",
+            "probe_set_blat_mb_start", "probe_set_blat_mb_end",
+            "probe_set_strand", "probe_set_note_by_rw", "flag"],
+	"type": "ProbeSet"
+    }
+
+def retrieve_trait_dataset(trait_type, trait, threshold, conn):
+    """
+    Retrieve the dataset that relates to a specific trait.
+    """
+    dataset_fns = {
+        "Temp": retrieve_temp_trait_dataset,
+        "Geno": retrieve_geno_trait_dataset,
+        "Publish": retrieve_publish_trait_dataset,
+        "ProbeSet": retrieve_probeset_trait_dataset
+    }
+    dataset_name_info = {
+        "dataset_id": None,
+        "dataset_name": trait["db"]["dataset_name"],
+        **retrieve_dataset_name(
+            trait_type, threshold, trait["trait_name"],
+            trait["db"]["dataset_name"], conn)
+    }
+    riset = retrieve_riset_fields(
+        trait_type, trait["trait_name"], dataset_name_info, conn)
+    return {
+        "display_name": dataset_name_info["dataset_name"],
+        **dataset_name_info,
+        **dataset_fns[trait_type](),
+        **riset
+    }
diff --git a/gn3/db/traits.py b/gn3/db/traits.py
index 61bc94d..6ea24be 100644
--- a/gn3/db/traits.py
+++ b/gn3/db/traits.py
@@ -1,5 +1,7 @@
 """This class contains functions relating to trait data manipulation"""
-from typing import Any, Dict, Union
+from typing import Any, Dict, Union, Sequence
+from gn3.function_helpers import compose
+from gn3.db.datasets import retrieve_trait_dataset
 
 
 def get_trait_csv_sample_data(conn: Any,
@@ -41,6 +43,7 @@ def update_sample_data(conn: Any,
                        count: Union[int, str]):
     """Given the right parameters, update sample-data from the relevant
     table."""
+    # pylint: disable=[R0913, R0914]
     STRAIN_ID_SQL: str = "UPDATE Strain SET Name = %s WHERE Id = %s"
     PUBLISH_DATA_SQL: str = ("UPDATE PublishData SET value = %s "
                              "WHERE StrainId = %s AND Id = %s")
@@ -76,159 +79,590 @@ def update_sample_data(conn: Any,
     return (updated_strains, updated_published_data,
             updated_se_data, updated_n_strains)
 
-
-def retrieve_trait_dataset_name(
-        trait_type: str, threshold: int, name: str, connection: Any):
-    """
-    Retrieve the name of a trait given the trait's name
-
-    This is extracted from the `webqtlDataset.retrieveName` function as is
-    implemented at
-    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlDataset.py#L140-L169
-    """
-    columns = "Id, Name, FullName, ShortName{}".format(
-        ", DataScale" if trait_type == "ProbeSet" else "")
-    query = (
-        "SELECT {columns} "
-        "FROM {trait_type}Freeze "
-        "WHERE "
-        "public > %(threshold)s "
-        "AND "
-        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)").format(
-            columns=columns, trait_type=trait_type)
-    with connection.cursor() as cursor:
-        cursor.execute(query, {"threshold": threshold, "name": name})
-        return cursor.fetchone()
-
-PUBLISH_TRAIT_INFO_QUERY = (
-    "SELECT "
-    "PublishXRef.Id, Publication.PubMed_ID, "
-    "Phenotype.Pre_publication_description, "
-    "Phenotype.Post_publication_description, "
-    "Phenotype.Original_description, "
-    "Phenotype.Pre_publication_abbreviation, "
-    "Phenotype.Post_publication_abbreviation, "
-    "Phenotype.Lab_code, Phenotype.Submitter, Phenotype.Owner, "
-    "Phenotype.Authorized_Users, CAST(Publication.Authors AS BINARY), "
-    "Publication.Title, Publication.Abstract, Publication.Journal, "
-    "Publication.Volume, Publication.Pages, Publication.Month, "
-    "Publication.Year, PublishXRef.Sequence, Phenotype.Units, "
-    "PublishXRef.comments "
-    "FROM "
-    "PublishXRef, Publication, Phenotype, PublishFreeze "
-    "WHERE "
-    "PublishXRef.Id = %(trait_name)s AND "
-    "Phenotype.Id = PublishXRef.PhenotypeId AND "
-    "Publication.Id = PublishXRef.PublicationId AND "
-    "PublishXRef.InbredSetId = PublishFreeze.InbredSetId AND "
-    "PublishFreeze.Id =%(trait_dataset_id)s")
-
-
 def retrieve_publish_trait_info(trait_data_source: Dict[str, Any], conn: Any):
     """Retrieve trait information for type `Publish` traits.
 
     https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L399-L421"""
+    keys = (
+        "Id", "PubMed_ID", "Pre_publication_description",
+        "Post_publication_description", "Original_description",
+        "Pre_publication_abbreviation", "Post_publication_abbreviation",
+        "Lab_code", "Submitter", "Owner", "Authorized_Users", "Authors",
+        "Title", "Abstract", "Journal", "Volume", "Pages", "Month", "Year",
+        "Sequence", "Units", "comments")
+    columns = (
+        "PublishXRef.Id, Publication.PubMed_ID, "
+        "Phenotype.Pre_publication_description, "
+        "Phenotype.Post_publication_description, "
+        "Phenotype.Original_description, "
+        "Phenotype.Pre_publication_abbreviation, "
+        "Phenotype.Post_publication_abbreviation, "
+        "Phenotype.Lab_code, Phenotype.Submitter, Phenotype.Owner, "
+        "Phenotype.Authorized_Users, CAST(Publication.Authors AS BINARY), "
+        "Publication.Title, Publication.Abstract, Publication.Journal, "
+        "Publication.Volume, Publication.Pages, Publication.Month, "
+        "Publication.Year, PublishXRef.Sequence, Phenotype.Units, "
+        "PublishXRef.comments")
+    query = (
+        "SELECT "
+        "{columns} "
+        "FROM "
+        "PublishXRef, Publication, Phenotype, PublishFreeze "
+        "WHERE "
+        "PublishXRef.Id = %(trait_name)s AND "
+        "Phenotype.Id = PublishXRef.PhenotypeId AND "
+        "Publication.Id = PublishXRef.PublicationId AND "
+        "PublishXRef.InbredSetId = PublishFreeze.InbredSetId AND "
+        "PublishFreeze.Id =%(trait_dataset_id)s").format(columns=columns)
     with conn.cursor() as cursor:
         cursor.execute(
-            PUBLISH_TRAIT_INFO_QUERY,
+            query,
             {
                 k:v for k, v in trait_data_source.items()
                 if k in ["trait_name", "trait_dataset_id"]
             })
-        return cursor.fetchone()
-
-PROBESET_TRAIT_INFO_QUERY = (
-    "SELECT "
-    "ProbeSet.name, ProbeSet.symbol, ProbeSet.description, "
-    "ProbeSet.probe_target_description, ProbeSet.chr, ProbeSet.mb, "
-    "ProbeSet.alias, ProbeSet.geneid, ProbeSet.genbankid, ProbeSet.unigeneid, "
-    "ProbeSet.omim, ProbeSet.refseq_transcriptid, ProbeSet.blatseq, "
-    "ProbeSet.targetseq, ProbeSet.chipid, ProbeSet.comments, "
-    "ProbeSet.strand_probe, ProbeSet.strand_gene, "
-    "ProbeSet.probe_set_target_region, ProbeSet.proteinid, "
-    "ProbeSet.probe_set_specificity, ProbeSet.probe_set_blat_score, "
-    "ProbeSet.probe_set_blat_mb_start, ProbeSet.probe_set_blat_mb_end, "
-    "ProbeSet.probe_set_strand, ProbeSet.probe_set_note_by_rw, "
-    "ProbeSet.flag "
-    "FROM "
-    "ProbeSet, ProbeSetFreeze, ProbeSetXRef "
-    "WHERE "
-    "ProbeSetXRef.ProbeSetFreezeId = ProbeSetFreeze.Id AND "
-    "ProbeSetXRef.ProbeSetId = ProbeSet.Id AND "
-    "ProbeSetFreeze.Name = %(trait_dataset_name)s AND "
-    "ProbeSet.Name = %(trait_name)s")
+        return dict(zip([k.lower() for k in keys], cursor.fetchone()))
+
+def set_confidential_field(trait_type, trait_info):
+    """Post processing function for 'Publish' trait types.
+
+    It sets the value for the 'confidential' key."""
+    if trait_type == "Publish":
+        return {
+            **trait_info,
+            "confidential": 1 if (
+                trait_info.get("pre_publication_description", None)
+                and not trait_info.get("pubmed_id", None)) else 0}
+    return trait_info
 
 def retrieve_probeset_trait_info(trait_data_source: Dict[str, Any], conn: Any):
     """Retrieve trait information for type `ProbeSet` traits.
 
     https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L424-L435"""
+    keys = (
+        "name", "symbol", "description", "probe_target_description", "chr",
+        "mb", "alias", "geneid", "genbankid", "unigeneid", "omim",
+        "refseq_transcriptid", "blatseq", "targetseq", "chipid", "comments",
+        "strand_probe", "strand_gene", "probe_set_target_region", "proteinid",
+        "probe_set_specificity", "probe_set_blat_score",
+        "probe_set_blat_mb_start", "probe_set_blat_mb_end", "probe_set_strand",
+        "probe_set_note_by_rw", "flag")
+    query = (
+        "SELECT "
+        "{columns} "
+        "FROM "
+        "ProbeSet, ProbeSetFreeze, ProbeSetXRef "
+        "WHERE "
+        "ProbeSetXRef.ProbeSetFreezeId = ProbeSetFreeze.Id AND "
+        "ProbeSetXRef.ProbeSetId = ProbeSet.Id AND "
+        "ProbeSetFreeze.Name = %(trait_dataset_name)s AND "
+        "ProbeSet.Name = %(trait_name)s").format(
+            columns=", ".join(["ProbeSet.{}".format(x) for x in keys]))
     with conn.cursor() as cursor:
         cursor.execute(
-            PROBESET_TRAIT_INFO_QUERY,
+            query,
             {
                 k:v for k, v in trait_data_source.items()
                 if k in ["trait_name", "trait_dataset_name"]
             })
-        return cursor.fetchone()
-
-GENO_TRAIT_INFO_QUERY = (
-    "SELECT "
-    "Geno.name, Geno.chr, Geno.mb, Geno.source2, Geno.sequence "
-    "FROM "
-    "Geno, GenoFreeze, GenoXRef "
-    "WHERE "
-    "GenoXRef.GenoFreezeId = GenoFreeze.Id AND GenoXRef.GenoId = Geno.Id AND "
-    "GenoFreeze.Name = %(trait_dataset_name)s AND Geno.Name = %(trait_name)s")
+        return dict(zip(keys, cursor.fetchone()))
 
 def retrieve_geno_trait_info(trait_data_source: Dict[str, Any], conn: Any):
     """Retrieve trait information for type `Geno` traits.
 
     https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L438-L449"""
+    keys = ("name", "chr", "mb", "source2", "sequence")
+    query = (
+        "SELECT "
+        "{columns} "
+        "FROM "
+        "Geno, GenoFreeze, GenoXRef "
+        "WHERE "
+        "GenoXRef.GenoFreezeId = GenoFreeze.Id AND GenoXRef.GenoId = Geno.Id AND "
+        "GenoFreeze.Name = %(trait_dataset_name)s AND "
+        "Geno.Name = %(trait_name)s").format(
+            columns=", ".join(["Geno.{}".format(x) for x in keys]))
     with conn.cursor() as cursor:
         cursor.execute(
-            GENO_TRAIT_INFO_QUERY,
+            query,
             {
                 k:v for k, v in trait_data_source.items()
                 if k in ["trait_name", "trait_dataset_name"]
             })
-        return cursor.fetchone()
-
-TEMP_TRAIT_INFO_QUERY = (
-    "SELECT name, description FROM Temp "
-    "WHERE Name = %(trait_name)s")
+        return dict(zip(keys, cursor.fetchone()))
 
 def retrieve_temp_trait_info(trait_data_source: Dict[str, Any], conn: Any):
     """Retrieve trait information for type `Temp` traits.
 
     https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L450-452"""
+    keys = ("name", "description")
+    query = (
+        "SELECT {columns} FROM Temp "
+        "WHERE Name = %(trait_name)s").format(columns=", ".join(keys))
     with conn.cursor() as cursor:
         cursor.execute(
-            TEMP_TRAIT_INFO_QUERY,
+            query,
             {
                 k:v for k, v in trait_data_source.items()
                 if k in ["trait_name"]
             })
-        return cursor.fetchone()
+        return dict(zip(keys, cursor.fetchone()))
+
+def set_haveinfo_field(trait_info):
+    """
+    Common postprocessing function for all trait types.
+
+    Sets the value for the 'haveinfo' field."""
+    return {**trait_info, "haveinfo": 1 if trait_info else 0}
+
+def set_homologene_id_field_probeset(trait_info, conn):
+    """
+    Postprocessing function for 'ProbeSet' traits.
+
+    Sets the value for the 'homologene' key.
+    """
+    query = (
+        "SELECT HomologeneId FROM Homologene, Species, InbredSet"
+        " WHERE Homologene.GeneId = %(geneid)s AND InbredSet.Name = %(riset)s"
+        " AND InbredSet.SpeciesId = Species.Id AND"
+        " Species.TaxonomyId = Homologene.TaxonomyId")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                k:v for k, v in trait_info.items()
+                if k in ["geneid", "riset"]
+            })
+        res = cursor.fetchone()
+        if res:
+            return {**trait_info, "homologeneid": res[0]}
+    return {**trait_info, "homologeneid": None}
+
+def set_homologene_id_field(trait_type, trait_info, conn):
+    """
+    Common postprocessing function for all trait types.
+
+    Sets the value for the 'homologene' key."""
+    set_to_null = lambda ti: {**ti, "homologeneid": None}
+    functions_table = {
+        "Temp": set_to_null,
+        "Geno": set_to_null,
+        "Publish": set_to_null,
+        "ProbeSet": lambda ti: set_homologene_id_field_probeset(ti, conn)
+    }
+    return functions_table[trait_type](trait_info)
+
+def load_publish_qtl_info(trait_info, conn):
+    """
+    Load extra QTL information for `Publish` traits
+    """
+    query = (
+        "SELECT PublishXRef.Locus, PublishXRef.LRS, PublishXRef.additive "
+        "FROM PublishXRef, PublishFreeze "
+        "WHERE PublishXRef.Id = %(trait_name)s "
+        "AND PublishXRef.InbredSetId = PublishFreeze.InbredSetId "
+        "AND PublishFreeze.Id = %(dataset_id)s")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                "trait_name": trait_info["trait_name"],
+                "dataset_id": trait_info["db"]["dataset_id"]
+            })
+        return dict(zip(["locus", "lrs", "additive"], cursor.fetchone()))
+    return {"locus": "", "lrs": "", "additive": ""}
+
+def load_probeset_qtl_info(trait_info, conn):
+    """
+    Load extra QTL information for `ProbeSet` traits
+    """
+    query = (
+        "SELECT ProbeSetXRef.Locus, ProbeSetXRef.LRS, ProbeSetXRef.pValue, "
+        "ProbeSetXRef.mean, ProbeSetXRef.additive "
+        "FROM ProbeSetXRef, ProbeSet "
+        "WHERE ProbeSetXRef.ProbeSetId = ProbeSet.Id "
+        " AND ProbeSet.Name = %(trait_name)s "
+        "AND ProbeSetXRef.ProbeSetFreezeId = %(dataset_id)s")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                "trait_name": trait_info["trait_name"],
+                "dataset_id": trait_info["db"]["dataset_id"]
+            })
+        return dict(zip(
+            ["locus", "lrs", "pvalue", "mean", "additive"], cursor.fetchone()))
+    return {"locus": "", "lrs": "", "pvalue": "", "mean": "", "additive": ""}
+
+def load_qtl_info(qtl, trait_type, trait_info, conn):
+    """
+    Load extra QTL information for traits
+
+    DESCRIPTION:
+    Migrated from
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L500-L534
+
+    PARAMETERS:
+    qtl: boolean
+    trait_type: string
+      The type of the trait in consideration
+    trait_info: map/dictionary
+      A dictionary of the trait's key-value pairs
+    conn:
+      A database connection object
+    """
+    if not qtl:
+        return trait_info
+    qtl_info_functions = {
+        "Publish": load_publish_qtl_info,
+        "ProbeSet": load_probeset_qtl_info
+    }
+    if trait_info["name"] not in qtl_info_functions.keys():
+        return trait_info
+
+    return qtl_info_functions[trait_type](trait_info, conn)
+
+def build_trait_name(trait_fullname):
+    """
+    Initialises the trait's name, and other values from the search data provided
+    """
+    def dataset_type(dset_name):
+        if dset_name.find('Temp') >= 0:
+            return "Temp"
+        if dset_name.find('Geno') >= 0:
+            return "Geno"
+        if dset_name.find('Publish') >= 0:
+            return "Publish"
+        return "ProbeSet"
+
+    name_parts = trait_fullname.split("::")
+    assert len(name_parts) >= 2, "Name format error"
+    dataset_name = name_parts[0]
+    dataset_type = dataset_type(dataset_name)
+    return {
+        "db": {
+            "dataset_name": dataset_name,
+            "dataset_type": dataset_type},
+        "trait_fullname": trait_fullname,
+        "trait_name": name_parts[1],
+        "cellid": name_parts[2] if len(name_parts) == 3 else ""
+    }
+
+def retrieve_probeset_sequence(trait, conn):
+    """
+    Retrieve a 'ProbeSet' trait's sequence information
+    """
+    query = (
+        "SELECT ProbeSet.BlatSeq "
+        "FROM ProbeSet, ProbeSetFreeze, ProbeSetXRef "
+        "WHERE ProbeSet.Id=ProbeSetXRef.ProbeSetId "
+        "AND ProbeSetFreeze.Id = ProbeSetXRef.ProbeSetFreezeId "
+        "AND ProbeSet.Name = %(trait_name)s "
+        "AND ProbeSetFreeze.Name = %(dataset_name)s")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {
+                "trait_name": trait["trait_name"],
+                "dataset_name": trait["db"]["dataset_name"]
+            })
+        seq = cursor.fetchone()
+        return {**trait, "sequence": seq[0] if seq else ""}
 
 def retrieve_trait_info(
-        trait_type: str, trait_name: str, trait_dataset_id: int,
-        trait_dataset_name: str, conn: Any):
+        threshold: int, trait_full_name: str, conn: Any,
+        qtl=None):
     """Retrieves the trait information.
 
     https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L397-L456
 
     This function, or the dependent functions, might be incomplete as they are
     currently."""
+    trait = build_trait_name(trait_full_name)
+    trait_dataset_type = trait["db"]["dataset_type"]
     trait_info_function_table = {
         "Publish": retrieve_publish_trait_info,
         "ProbeSet": retrieve_probeset_trait_info,
         "Geno": retrieve_geno_trait_info,
         "Temp": retrieve_temp_trait_info
     }
-    return trait_info_function_table[trait_type](
+
+    common_post_processing_fn = compose(
+        lambda ti: load_qtl_info(qtl, trait_dataset_type, ti, conn),
+        lambda ti: set_homologene_id_field(trait_dataset_type, ti, conn),
+        lambda ti: {"trait_type": trait_dataset_type, **ti},
+        lambda ti: {**trait, **ti})
+
+    trait_post_processing_functions_table = {
+        "Publish": compose(
+            lambda ti: set_confidential_field(trait_dataset_type, ti),
+            common_post_processing_fn),
+        "ProbeSet": compose(
+            lambda ti: retrieve_probeset_sequence(ti, conn),
+            common_post_processing_fn),
+        "Geno": common_post_processing_fn,
+        "Temp": common_post_processing_fn
+    }
+
+    retrieve_info = compose(
+        set_haveinfo_field, trait_info_function_table[trait_dataset_type])
+
+    trait_dataset = retrieve_trait_dataset(
+        trait_dataset_type, trait, threshold, conn)
+    trait_info = retrieve_info(
         {
-            "trait_name": trait_name,
-            "trait_dataset_id": trait_dataset_id,
-            "trait_dataset_name":trait_dataset_name
+            "trait_name": trait["trait_name"],
+            "trait_dataset_id": trait_dataset["dataset_id"],
+            "trait_dataset_name": trait_dataset["dataset_name"]
         },
         conn)
+    if trait_info["haveinfo"]:
+        return {
+            **trait_post_processing_functions_table[trait_dataset_type](trait_info),
+            "db": {**trait["db"], **trait_dataset},
+            "riset": trait_dataset["riset"]
+        }
+    return trait_info
+
+def retrieve_temp_trait_data(trait_info: dict, conn: Any):
+    """
+    Retrieve trait data for `Temp` traits.
+    """
+    query = (
+        "SELECT "
+        "Strain.Name, TempData.value, TempData.SE, TempData.NStrain, "
+        "TempData.Id "
+        "FROM TempData, Temp, Strain "
+        "WHERE TempData.StrainId = Strain.Id "
+        "AND TempData.Id = Temp.DataId "
+        "AND Temp.name = %(trait_name)s "
+        "ORDER BY Strain.Name")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {"trait_name": trait_info["trait_name"]})
+        return [dict(zip(
+            ["strain_name", "value", "se_error", "nstrain", "id"], row))
+                for row in cursor.fetchall()]
+    return []
+
+def retrieve_species_id(riset, conn: Any):
+    """
+    Retrieve a species id given the RISet value
+    """
+    with conn.cursor as cursor:
+        cursor.execute(
+            "SELECT SpeciesId from InbredSet WHERE Name = %(riset)s",
+            {"riset": riset})
+        return cursor.fetchone()[0]
+    return None
+
+def retrieve_geno_trait_data(trait_info: Dict, conn: Any):
+    """
+    Retrieve trait data for `Geno` traits.
+    """
+    query = (
+        "SELECT Strain.Name, GenoData.value, GenoSE.error, GenoData.Id "
+        "FROM (GenoData, GenoFreeze, Strain, Geno, GenoXRef) "
+        "LEFT JOIN GenoSE ON "
+        "(GenoSE.DataId = GenoData.Id AND GenoSE.StrainId = GenoData.StrainId) "
+        "WHERE Geno.SpeciesId = %(species_id)s "
+        "AND Geno.Name = %(trait_name)s AND GenoXRef.GenoId = Geno.Id "
+        "AND GenoXRef.GenoFreezeId = GenoFreeze.Id "
+        "AND GenoFreeze.Name = %(dataset_name)s "
+        "AND GenoXRef.DataId = GenoData.Id "
+        "AND GenoData.StrainId = Strain.Id "
+        "ORDER BY Strain.Name")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {"trait_name": trait_info["trait_name"],
+             "dataset_name": trait_info["db"]["dataset_name"],
+             "species_id": retrieve_species_id(
+                 trait_info["db"]["riset"], conn)})
+        return [dict(zip(
+            ["strain_name", "value", "se_error", "id"], row))
+                for row in cursor.fetchall()]
+    return []
+
+def retrieve_publish_trait_data(trait_info: Dict, conn: Any):
+    """
+    Retrieve trait data for `Publish` traits.
+    """
+    query = (
+        "SELECT "
+        "Strain.Name, PublishData.value, PublishSE.error, NStrain.count, "
+        "PublishData.Id "
+        "FROM (PublishData, Strain, PublishXRef, PublishFreeze) "
+        "LEFT JOIN PublishSE ON "
+        "(PublishSE.DataId = PublishData.Id "
+        "AND PublishSE.StrainId = PublishData.StrainId) "
+        "LEFT JOIN NStrain ON "
+        "(NStrain.DataId = PublishData.Id "
+        "AND NStrain.StrainId = PublishData.StrainId) "
+        "WHERE PublishXRef.InbredSetId = PublishFreeze.InbredSetId "
+        "AND PublishData.Id = PublishXRef.DataId "
+        "AND PublishXRef.Id = %(trait_name)s "
+        "AND PublishFreeze.Id = %(dataset_id)s "
+        "AND PublishData.StrainId = Strain.Id "
+        "ORDER BY Strain.Name")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {"trait_name": trait_info["trait_name"],
+             "dataset_id": trait_info["db"]["dataset_id"]})
+        return [dict(zip(
+            ["strain_name", "value", "se_error", "nstrain", "id"], row))
+                for row in cursor.fetchall()]
+    return []
+
+def retrieve_cellid_trait_data(trait_info: Dict, conn: Any):
+    """
+    Retrieve trait data for `Probe Data` types.
+    """
+    query = (
+        "SELECT "
+        "Strain.Name, ProbeData.value, ProbeSE.error, ProbeData.Id "
+        "FROM (ProbeData, ProbeFreeze, ProbeSetFreeze, ProbeXRef, Strain,"
+        " Probe, ProbeSet) "
+        "LEFT JOIN ProbeSE ON "
+        "(ProbeSE.DataId = ProbeData.Id "
+        " AND ProbeSE.StrainId = ProbeData.StrainId) "
+        "WHERE Probe.Name = %(cellid)s "
+        "AND ProbeSet.Name = %(trait_name)s "
+        "AND Probe.ProbeSetId = ProbeSet.Id "
+        "AND ProbeXRef.ProbeId = Probe.Id "
+        "AND ProbeXRef.ProbeFreezeId = ProbeFreeze.Id "
+        "AND ProbeSetFreeze.ProbeFreezeId = ProbeFreeze.Id "
+        "AND ProbeSetFreeze.Name = %(dataset_name)s "
+        "AND ProbeXRef.DataId = ProbeData.Id "
+        "AND ProbeData.StrainId = Strain.Id "
+        "ORDER BY Strain.Name")
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {"cellid": trait_info["cellid"],
+             "trait_name": trait_info["trait_name"],
+             "dataset_id": trait_info["db"]["dataset_id"]})
+        return [dict(zip(
+            ["strain_name", "value", "se_error", "id"], row))
+                for row in cursor.fetchall()]
+    return []
+
+def retrieve_probeset_trait_data(trait_info: Dict, conn: Any):
+    """
+    Retrieve trait data for `ProbeSet` traits.
+    """
+    query = (
+        "SELECT Strain.Name, ProbeSetData.value, ProbeSetSE.error, "
+        "ProbeSetData.Id "
+        "FROM (ProbeSetData, ProbeSetFreeze, Strain, ProbeSet, ProbeSetXRef) "
+        "LEFT JOIN ProbeSetSE ON "
+        "(ProbeSetSE.DataId = ProbeSetData.Id "
+        "AND ProbeSetSE.StrainId = ProbeSetData.StrainId) "
+        "WHERE ProbeSet.Name = %(trait_name)s "
+        "AND ProbeSetXRef.ProbeSetId = ProbeSet.Id "
+        "AND ProbeSetXRef.ProbeSetFreezeId = ProbeSetFreeze.Id "
+        "AND ProbeSetFreeze.Name = %(dataset_name)s "
+        "AND ProbeSetXRef.DataId = ProbeSetData.Id "
+        "AND ProbeSetData.StrainId = Strain.Id "
+        "ORDER BY Strain.Name")
+
+    with conn.cursor() as cursor:
+        cursor.execute(
+            query,
+            {"trait_name": trait_info["trait_name"],
+             "dataset_name": trait_info["db"]["dataset_name"]})
+        return [dict(zip(
+            ["strain_name", "value", "se_error", "id"], row))
+                for row in cursor.fetchall()]
+    return []
+
+def with_strainlist_data_setup(strainlist: Sequence[str]):
+    """
+    Build function that computes the trait data from provided list of strains.
+
+    PARAMETERS
+    strainlist: (list)
+      A list of strain names
+
+    RETURNS:
+      Returns a function that given some data from the database, computes the
+      strain's value, variance and ndata values, only if the strain is present
+      in the provided `strainlist` variable.
+    """
+    def setup_fn(tdata):
+        if tdata["strain_name"] in strainlist:
+            val = tdata["value"]
+            if val is not None:
+                return {
+                    "strain_name": tdata["strain_name"],
+                    "value": val,
+                    "variance": tdata["se_error"],
+                    "ndata": tdata.get("nstrain", None)
+                }
+        return None
+    return setup_fn
+
+def without_strainlist_data_setup():
+    """
+    Build function that computes the trait data.
+
+    RETURNS:
+      Returns a function that given some data from the database, computes the
+      strain's value, variance and ndata values.
+    """
+    def setup_fn(tdata):
+        val = tdata["value"]
+        if val is not None:
+            return {
+                "strain_name": tdata["strain_name"],
+                "value": val,
+                "variance": tdata["se_error"],
+                "ndata": tdata.get("nstrain", None)
+            }
+        return None
+    return setup_fn
+
+def retrieve_trait_data(trait: dict, conn: Any, strainlist: Sequence[str] = tuple()):
+    """
+    Retrieve trait data
+
+    DESCRIPTION
+    Retrieve trait data as is done in
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/base/webqtlTrait.py#L258-L386
+    """
+    # I do not like this section, but it retains the flow in the old codebase
+    if trait["db"]["dataset_type"] == "Temp":
+        results = retrieve_temp_trait_data(trait, conn)
+    elif trait["db"]["dataset_type"] == "Publish":
+        results = retrieve_publish_trait_data(trait, conn)
+    elif trait["cellid"]:
+        results = retrieve_cellid_trait_data(trait, conn)
+    elif trait["db"]["dataset_type"] == "ProbeSet":
+        results = retrieve_probeset_trait_data(trait, conn)
+    else:
+        results = retrieve_geno_trait_data(trait, conn)
+
+    if results:
+        # do something with mysqlid
+        mysqlid = results[0]["id"]
+        if strainlist:
+            data = [
+                item for item in
+                map(with_strainlist_data_setup(strainlist), results)
+                if item is not None]
+        else:
+            data = [
+                item for item in
+                map(without_strainlist_data_setup(), results)
+                if item is not None]
+
+        return {
+            "mysqlid": mysqlid,
+            "data": dict(map(
+                lambda x: (
+                    x["strain_name"],
+                    {k:v for k, v in x.items() if x != "strain_name"}),
+                data))}
+    return {}
diff --git a/gn3/function_helpers.py b/gn3/function_helpers.py
new file mode 100644
index 0000000..397b2da
--- /dev/null
+++ b/gn3/function_helpers.py
@@ -0,0 +1,36 @@
+"""
+This module will contain helper functions that should assist in maintaining a
+mostly functional way of programming.
+
+It will also contain miscellaneous functions that can be used globally, and thus
+do not fit well in any other module.
+
+FUNCTIONS:
+compose: This function is used to compose multiple functions into a single
+    function. It passes the results of calling one function to the other until
+    all the functions to be composed are called.
+"""
+from functools import reduce
+
+def compose(*functions):
+    """Compose multiple functions into a single function.
+
+    The utility in this function is not specific to this module, and as such,
+    this function can, and probably should, be moved to a more global module.
+
+    DESCRIPTION:
+    Given `cfn = compose(f_1, f_2, ... f_(n-1), f_n )`, calling
+    `cfn(arg_1, arg_2, ..., arg_m)` should call `f_n` with the arguments passed
+    to `cfn` and the results of that should be passed as arguments to `f_(n-1)`
+    and so on until `f_1` is called with the results of the cumulative calls and
+    that is the result of the entire chain of calls.
+
+    PARAMETERS:
+    functions: a variable argument list of function.
+    """
+    def composed_function(*args, **kwargs):
+        return reduce(
+            lambda res, fn: fn(res),
+            reversed(functions[:-1]),
+            functions[-1](*args, **kwargs))
+    return composed_function
diff --git a/tests/unit/db/test_datasets.py b/tests/unit/db/test_datasets.py
new file mode 100644
index 0000000..38de0e2
--- /dev/null
+++ b/tests/unit/db/test_datasets.py
@@ -0,0 +1,133 @@
+"""Tests for gn3/db/datasets.py"""
+
+from unittest import mock, TestCase
+from gn3.db.datasets import (
+    retrieve_dataset_name,
+    retrieve_riset_fields,
+    retrieve_geno_riset_fields,
+    retrieve_publish_riset_fields,
+    retrieve_probeset_riset_fields)
+
+class TestDatasetsDBFunctions(TestCase):
+    """Test cases for datasets functions."""
+
+    def test_retrieve_dataset_name(self):
+        """Test that the function is called correctly."""
+        for trait_type, thresh, trait_name, dataset_name, columns, table in [
+                ["ProbeSet", 9, "probesetTraitName", "probesetDatasetName",
+                 "Id, Name, FullName, ShortName, DataScale", "ProbeSetFreeze"],
+                ["Geno", 3, "genoTraitName", "genoDatasetName",
+                 "Id, Name, FullName, ShortName", "GenoFreeze"],
+                ["Publish", 6, "publishTraitName", "publishDatasetName",
+                 "Id, Name, FullName, ShortName", "PublishFreeze"],
+                ["Temp", 4, "tempTraitName", "tempTraitName",
+                 "Id, Name, FullName, ShortName", "TempFreeze"]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_type=trait_type):
+                with db_mock.cursor() as cursor:
+                    cursor.fetchone.return_value = {}
+                    self.assertEqual(
+                        retrieve_dataset_name(
+                            trait_type, thresh, trait_name, dataset_name, db_mock),
+                        {})
+                    cursor.execute.assert_called_once_with(
+                        "SELECT {cols} "
+                        "FROM {table} "
+                        "WHERE public > %(threshold)s AND "
+                        "(Name = %(name)s "
+                        "OR FullName = %(name)s "
+                        "OR ShortName = %(name)s)".format(
+                            table=table, cols=columns),
+                        {"threshold": thresh, "name": dataset_name})
+
+    def test_retrieve_probeset_riset_fields(self):
+        """
+        Test that the `riset` and `riset_id` fields are retrieved appropriately
+        for the 'ProbeSet' trait type.
+        """
+        for trait_name, expected in [
+                ["testProbeSetName", {}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_name=trait_name, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ()
+                    self.assertEqual(
+                        retrieve_probeset_riset_fields(trait_name, db_mock),
+                        expected)
+                    cursor.execute.assert_called_once_with(
+                        (
+                            "SELECT InbredSet.Name, InbredSet.Id"
+                            " FROM InbredSet, ProbeSetFreeze, ProbeFreeze"
+                            " WHERE ProbeFreeze.InbredSetId = InbredSet.Id"
+                            " AND ProbeFreeze.Id = ProbeSetFreeze.ProbeFreezeId"
+                            " AND ProbeSetFreeze.Name = %(name)s"),
+                        {"name": trait_name})
+
+    def test_retrieve_riset_fields(self):
+        """
+        Test that the riset fields are set up correctly for the different trait
+        types.
+        """
+        for trait_type, trait_name, dataset_info, expected in [
+                ["Publish", "pubTraitName01", {"dataset_name": "pubDBName01"},
+                 {"dataset_name": "pubDBName01", "riset": ""}],
+                ["ProbeSet", "prbTraitName01", {"dataset_name": "prbDBName01"},
+                 {"dataset_name": "prbDBName01", "riset": ""}],
+                ["Geno", "genoTraitName01", {"dataset_name": "genoDBName01"},
+                 {"dataset_name": "genoDBName01", "riset": ""}],
+                ["Temp", "tempTraitName01", {}, {"riset": ""}],
+                ]:
+            db_mock = mock.MagicMock()
+            with self.subTest(
+                    trait_type=trait_type, trait_name=trait_name,
+                    dataset_info=dataset_info):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ("riset_name", 0)
+                    self.assertEqual(
+                        retrieve_riset_fields(
+                            trait_type, trait_name, dataset_info, db_mock),
+                        expected)
+
+    def test_retrieve_publish_riset_fields(self):
+        """
+        Test that the `riset` and `riset_id` fields are retrieved appropriately
+        for the 'Publish' trait type.
+        """
+        for trait_name, expected in [
+                ["testPublishName", {}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_name=trait_name, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ()
+                    self.assertEqual(
+                        retrieve_publish_riset_fields(trait_name, db_mock),
+                        expected)
+                    cursor.execute.assert_called_once_with(
+                        (
+                            "SELECT InbredSet.Name, InbredSet.Id"
+                            " FROM InbredSet, PublishFreeze"
+                            " WHERE PublishFreeze.InbredSetId = InbredSet.Id"
+                            " AND PublishFreeze.Name = %(name)s"),
+                        {"name": trait_name})
+
+    def test_retrieve_geno_riset_fields(self):
+        """
+        Test that the `riset` and `riset_id` fields are retrieved appropriately
+        for the 'Geno' trait type.
+        """
+        for trait_name, expected in [
+                ["testGenoName", {}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_name=trait_name, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ()
+                    self.assertEqual(
+                        retrieve_geno_riset_fields(trait_name, db_mock),
+                        expected)
+                    cursor.execute.assert_called_once_with(
+                        (
+                            "SELECT InbredSet.Name, InbredSet.Id"
+                            " FROM InbredSet, GenoFreeze"
+                            " WHERE GenoFreeze.InbredSetId = InbredSet.Id"
+                            " AND GenoFreeze.Name = %(name)s"),
+                        {"name": trait_name})
diff --git a/tests/unit/db/test_traits.py b/tests/unit/db/test_traits.py
index c8f28b5..ee98893 100644
--- a/tests/unit/db/test_traits.py
+++ b/tests/unit/db/test_traits.py
@@ -1,51 +1,20 @@
 """Tests for gn3/db/traits.py"""
 from unittest import mock, TestCase
 from gn3.db.traits import (
-    GENO_TRAIT_INFO_QUERY,
-    TEMP_TRAIT_INFO_QUERY,
-    PUBLISH_TRAIT_INFO_QUERY,
-    PROBESET_TRAIT_INFO_QUERY)
-from gn3.db.traits import (
+    build_trait_name,
+    set_haveinfo_field,
+    update_sample_data,
     retrieve_trait_info,
+    set_confidential_field,
+    set_homologene_id_field,
     retrieve_geno_trait_info,
     retrieve_temp_trait_info,
-    retrieve_trait_dataset_name,
     retrieve_publish_trait_info,
-    retrieve_probeset_trait_info,
-    update_sample_data)
-
+    retrieve_probeset_trait_info)
 
 class TestTraitsDBFunctions(TestCase):
     "Test cases for traits functions"
 
-    def test_retrieve_trait_dataset_name(self):
-        """Test that the function is called correctly."""
-        for trait_type, thresh, trait_dataset_name, columns in [
-                ["ProbeSet", 9, "testName",
-                 "Id, Name, FullName, ShortName, DataScale"],
-                ["Geno", 3, "genoTraitName", "Id, Name, FullName, ShortName"],
-                ["Publish", 6, "publishTraitName",
-                 "Id, Name, FullName, ShortName"],
-                ["Temp", 4, "tempTraitName", "Id, Name, FullName, ShortName"]]:
-            db_mock = mock.MagicMock()
-            with self.subTest(trait_type=trait_type):
-                with db_mock.cursor() as cursor:
-                    cursor.fetchone.return_value = (
-                        "testName", "testNameFull", "testNameShort",
-                        "dataScale")
-                    self.assertEqual(
-                        retrieve_trait_dataset_name(
-                            trait_type, thresh, trait_dataset_name, db_mock),
-                        ("testName", "testNameFull", "testNameShort",
-                         "dataScale"))
-                    cursor.execute.assert_called_once_with(
-                        "SELECT {cols} "
-                        "FROM {ttype}Freeze "
-                        "WHERE public > %(threshold)s AND "
-                        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)".format(
-                            cols=columns, ttype=trait_type),
-                        {"threshold": thresh, "name": trait_dataset_name})
-
     def test_retrieve_publish_trait_info(self):
         """Test retrieval of type `Publish` traits."""
         db_mock = mock.MagicMock()
@@ -54,12 +23,32 @@ class TestTraitsDBFunctions(TestCase):
             trait_source = {
                 "trait_name": "PublishTraitName", "trait_dataset_id": 1}
             self.assertEqual(
-                retrieve_publish_trait_info(
-                    trait_source,
-                    db_mock),
-                tuple())
+                retrieve_publish_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                PUBLISH_TRAIT_INFO_QUERY, trait_source)
+                ("SELECT "
+                 "PublishXRef.Id, Publication.PubMed_ID,"
+                 " Phenotype.Pre_publication_description,"
+                 " Phenotype.Post_publication_description,"
+                 " Phenotype.Original_description,"
+                 " Phenotype.Pre_publication_abbreviation,"
+                 " Phenotype.Post_publication_abbreviation,"
+                 " Phenotype.Lab_code, Phenotype.Submitter, Phenotype.Owner,"
+                 " Phenotype.Authorized_Users,"
+                 " CAST(Publication.Authors AS BINARY),"
+                 " Publication.Title, Publication.Abstract,"
+                 " Publication.Journal,"
+                 " Publication.Volume, Publication.Pages, Publication.Month,"
+                 " Publication.Year, PublishXRef.Sequence, Phenotype.Units,"
+                 " PublishXRef.comments"
+                 " FROM"
+                 " PublishXRef, Publication, Phenotype, PublishFreeze"
+                 " WHERE"
+                 " PublishXRef.Id = %(trait_name)s"
+                 " AND Phenotype.Id = PublishXRef.PhenotypeId"
+                 " AND Publication.Id = PublishXRef.PublicationId"
+                 " AND PublishXRef.InbredSetId = PublishFreeze.InbredSetId"
+                 " AND PublishFreeze.Id =%(trait_dataset_id)s"),
+                trait_source)
 
     def test_retrieve_probeset_trait_info(self):
         """Test retrieval of type `Probeset` traits."""
@@ -70,9 +59,31 @@ class TestTraitsDBFunctions(TestCase):
                 "trait_name": "ProbeSetTraitName",
                 "trait_dataset_name": "ProbeSetDatasetTraitName"}
             self.assertEqual(
-                retrieve_probeset_trait_info(trait_source, db_mock), tuple())
+                retrieve_probeset_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                PROBESET_TRAIT_INFO_QUERY, trait_source)
+                (
+                    "SELECT "
+                    "ProbeSet.name, ProbeSet.symbol, ProbeSet.description, "
+                    "ProbeSet.probe_target_description, ProbeSet.chr, "
+                    "ProbeSet.mb, ProbeSet.alias, ProbeSet.geneid, "
+                    "ProbeSet.genbankid, ProbeSet.unigeneid, ProbeSet.omim, "
+                    "ProbeSet.refseq_transcriptid, ProbeSet.blatseq, "
+                    "ProbeSet.targetseq, ProbeSet.chipid, ProbeSet.comments, "
+                    "ProbeSet.strand_probe, ProbeSet.strand_gene, "
+                    "ProbeSet.probe_set_target_region, ProbeSet.proteinid, "
+                    "ProbeSet.probe_set_specificity, "
+                    "ProbeSet.probe_set_blat_score, "
+                    "ProbeSet.probe_set_blat_mb_start, "
+                    "ProbeSet.probe_set_blat_mb_end, "
+                    "ProbeSet.probe_set_strand, ProbeSet.probe_set_note_by_rw, "
+                    "ProbeSet.flag "
+                    "FROM "
+                    "ProbeSet, ProbeSetFreeze, ProbeSetXRef "
+                    "WHERE "
+                    "ProbeSetXRef.ProbeSetFreezeId = ProbeSetFreeze.Id "
+                    "AND ProbeSetXRef.ProbeSetId = ProbeSet.Id "
+                    "AND ProbeSetFreeze.Name = %(trait_dataset_name)s "
+                    "AND ProbeSet.Name = %(trait_name)s"), trait_source)
 
     def test_retrieve_geno_trait_info(self):
         """Test retrieval of type `Geno` traits."""
@@ -83,9 +94,19 @@ class TestTraitsDBFunctions(TestCase):
                 "trait_name": "GenoTraitName",
                 "trait_dataset_name": "GenoDatasetTraitName"}
             self.assertEqual(
-                retrieve_geno_trait_info(trait_source, db_mock), tuple())
+                retrieve_geno_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                GENO_TRAIT_INFO_QUERY, trait_source)
+                (
+                    "SELECT "
+                    "Geno.name, Geno.chr, Geno.mb, Geno.source2, Geno.sequence "
+                    "FROM "
+                    "Geno, GenoFreeze, GenoXRef "
+                    "WHERE "
+                    "GenoXRef.GenoFreezeId = GenoFreeze.Id "
+                    "AND GenoXRef.GenoId = Geno.Id "
+                    "AND GenoFreeze.Name = %(trait_dataset_name)s "
+                    "AND Geno.Name = %(trait_name)s"),
+                trait_source)
 
     def test_retrieve_temp_trait_info(self):
         """Test retrieval of type `Temp` traits."""
@@ -94,26 +115,51 @@ class TestTraitsDBFunctions(TestCase):
             cursor.fetchone.return_value = tuple()
             trait_source = {"trait_name": "TempTraitName"}
             self.assertEqual(
-                retrieve_temp_trait_info(trait_source, db_mock), tuple())
+                retrieve_temp_trait_info(trait_source, db_mock), {})
             cursor.execute.assert_called_once_with(
-                TEMP_TRAIT_INFO_QUERY, trait_source)
+                "SELECT name, description FROM Temp WHERE Name = %(trait_name)s",
+                trait_source)
+
+    def test_build_trait_name_with_good_fullnames(self):
+        """
+        Check that the name is built correctly.
+        """
+        for fullname, expected in [
+                ["testdb::testname",
+                 {"db": {"dataset_name": "testdb", "dataset_type": "ProbeSet"},
+                  "trait_name": "testname", "cellid": "",
+                  "trait_fullname": "testdb::testname"}],
+                ["testdb::testname::testcell",
+                 {"db": {"dataset_name": "testdb", "dataset_type": "ProbeSet"},
+                  "trait_name": "testname", "cellid": "testcell",
+                  "trait_fullname": "testdb::testname::testcell"}]]:
+            with self.subTest(fullname=fullname):
+                self.assertEqual(build_trait_name(fullname), expected)
+
+    def test_build_trait_name_with_bad_fullnames(self):
+        """
+        Check that an exception is raised if the full name format is wrong.
+        """
+        for fullname in ["", "test", "test:test"]:
+            with self.subTest(fullname=fullname):
+                with self.assertRaises(AssertionError, msg="Name format error"):
+                    build_trait_name(fullname)
 
     def test_retrieve_trait_info(self):
         """Test that information on traits is retrieved as appropriate."""
-        for trait_type, trait_name, trait_dataset_id, trait_dataset_name, in [
-                ["Publish", "PublishTraitName", 1, "PublishDatasetTraitName"],
-                ["ProbeSet", "ProbeSetTraitName", 2, "ProbeSetDatasetTraitName"],
-                ["Geno", "GenoTraitName", 3, "GenoDatasetTraitName"],
-                ["Temp", "TempTraitName", 4, "TempDatasetTraitName"]]:
+        for threshold, trait_fullname, expected in [
+                [9, "pubDb::PublishTraitName::pubCell", {"haveinfo": 0}],
+                [5, "prbDb::ProbeSetTraitName::prbCell", {"haveinfo": 0}],
+                [12, "genDb::GenoTraitName", {"haveinfo": 0}],
+                [6, "tmpDb::TempTraitName", {"haveinfo": 0}]]:
             db_mock = mock.MagicMock()
-            with self.subTest(trait_type=trait_type):
+            with self.subTest(trait_fullname=trait_fullname):
                 with db_mock.cursor() as cursor:
                     cursor.fetchone.return_value = tuple()
                     self.assertEqual(
                         retrieve_trait_info(
-                            trait_type, trait_name, trait_dataset_id,
-                            trait_dataset_name, db_mock),
-                        tuple())
+                            threshold, trait_fullname, db_mock),
+                        expected)
 
     def test_update_sample_data(self):
         """Test that the SQL queries when calling update_sample_data are called with
@@ -143,3 +189,36 @@ class TestTraitsDBFunctions(TestCase):
                  mock.call(PUBLISH_SE_SQL, (2.3, 10, 8967049)),
                  mock.call(N_STRAIN_SQL, (2, 10, 8967049))]
             )
+
+    def test_set_haveinfo_field(self):
+        """Test that the `haveinfo` field is set up correctly"""
+        for trait_info, expected in [
+                [{}, {"haveinfo": 0}],
+                [{"k1": "v1"}, {"k1": "v1", "haveinfo": 1}]]:
+            with self.subTest(trait_info=trait_info, expected=expected):
+                self.assertEqual(set_haveinfo_field(trait_info), expected)
+
+    def test_set_homologene_id_field(self):
+        """Test that the `homologene_id` field is set up correctly"""
+        for trait_type, trait_info, expected in [
+                ["Publish", {}, {"homologeneid": None}],
+                ["ProbeSet", {}, {"homologeneid": None}],
+                ["Geno", {}, {"homologeneid": None}],
+                ["Temp", {}, {"homologeneid": None}]]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_info=trait_info, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.fetchone.return_value = ()
+                    self.assertEqual(
+                        set_homologene_id_field(trait_type, trait_info, db_mock), expected)
+
+    def test_set_confidential_field(self):
+        """Test that the `confidential` field is set up correctly"""
+        for trait_type, trait_info, expected in [
+                ["Publish", {}, {"confidential": 0}],
+                ["ProbeSet", {}, {}],
+                ["Geno", {}, {}],
+                ["Temp", {}, {}]]:
+            with self.subTest(trait_info=trait_info, expected=expected):
+                self.assertEqual(
+                    set_confidential_field(trait_type, trait_info), expected)