Reorganise the database code

Issue:
https://github.com/genenetwork/gn-gemtext-threads/blob/main/topics/gn1-migration-to-gn2/clustering.gmi

* Reorganise the code to separate the datasets from the traits, and to more
  closely conform to the same flow as that in GN1

4 files changed, 485 insertions, 266 deletions

diff --git a/gn3/db/datasets.py b/gn3/db/datasets.py
new file mode 100644
index 0000000..3ad50f6
--- /dev/null
+++ b/gn3/db/datasets.py
@@ -0,0 +1,251 @@
+from typing import Any, Dict, Union
+
+def retrieve_probeset_trait_dataset_name(
+        threshold: int, name: str, connection: Any):
+    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):
+    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):
+    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):
+    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):
+    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():
+    return {
+        "searchfield": ["name", "description"],
+        "disfield": ["name", "description"],
+        "type": "Temp",
+        "dataset_id": 1,
+        "fullname": "Temporary Storage",
+        "shortname": "Temp"
+    }
+
+def retrieve_geno_trait_dataset():
+    return {
+        "searchfield": ["name","chr"],
+	"disfield": ["name","chr","mb", "source2", "sequence"],
+	"type": "Geno"
+    }
+
+def retrieve_publish_trait_dataset():
+    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():
+    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):
+    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 9f89510..85cccfa 100644
--- a/gn3/db/traits.py
+++ b/gn3/db/traits.py
@@ -1,6 +1,7 @@
 """This class contains functions relating to trait data manipulation"""
 from typing import Any, Dict, Union
 from gn3.function_helpers import compose
+from gn3.db.datasets import retrieve_trait_dataset
 
 
 def get_trait_csv_sample_data(conn: Any,
@@ -77,41 +78,6 @@ 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
-    """
-    table_map = {
-        "ProbeSet": "ProbeSetFreeze",
-        "Publish": "PublishFreeze",
-        "Geno": "GenoFreeze",
-        "Temp": "TempFreeze"}
-    columns = "Id, Name, FullName, ShortName{}".format(
-        ", DataScale" if trait_type == "ProbeSet" else "")
-    query = (
-        "SELECT %(columns)s "
-        "FROM %(table)s "
-        "WHERE "
-        "public > %(threshold)s "
-        "AND "
-        "(Name = %(name)s OR FullName = %(name)s OR ShortName = %(name)s)")
-    with connection.cursor() as cursor:
-        cursor.execute(
-            query,
-            {
-                "table": table_map[trait_type],
-                "columns": columns,
-                "threshold": threshold,
-                "name": name
-            })
-        return cursor.fetchone()
-
-
 def retrieve_publish_trait_info(trait_data_source: Dict[str, Any], conn: Any):
     """Retrieve trait information for type `Publish` traits.
 
@@ -156,11 +122,11 @@ def retrieve_publish_trait_info(trait_data_source: Dict[str, Any], conn: Any):
             })
         return dict(zip([k.lower() for k in keys], cursor.fetchone()))
 
-def set_confidential_field(trait_info):
+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_info["type"] == "Publish":
+    if trait_type == "Publish":
         return {
             **trait_info,
             "confidential": 1 if (
@@ -271,7 +237,7 @@ def set_homologene_id_field_probeset(trait_info, conn):
             return {**trait_info, "homologeneid": res[0]}
     return {**trait_info, "homologeneid": None}
 
-def set_homologene_id_field(trait_info, conn):
+def set_homologene_id_field(trait_type, trait_info, conn):
     """
     Common postprocessing function for all trait types.
 
@@ -283,84 +249,83 @@ def set_homologene_id_field(trait_info, conn):
         "Publish": set_to_null,
         "ProbeSet": lambda ti: set_homologene_id_field_probeset(ti, conn)
     }
-    return functions_table[trait_info["type"]](trait_info)
+    return functions_table[trait_type](trait_info)
 
-def set_geno_riset_fields(name, conn):
-    """
-    Retrieve the RISet, and RISetID values for various Geno trait types.
-    """
+def load_publish_qtl_info(trait_info, conn):
     query = (
-        "SELECT InbredSet.Name, InbredSet.Id "
-        "FROM InbredSet, GenoFreeze "
-        "WHERE GenoFreeze.InbredSetId = InbredSet.Id "
-        "AND GenoFreeze.Name = %(name)s")
+        "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:
-        return cursor.execute(query, {"name": name})
+        cursor.execute()
+        return dict(zip(["locus", "lrs", "additive"], cursor.fetchone()))
+    return {"locus": "", "lrs": "", "additive": ""}
 
-def set_publish_riset_fields(name, conn):
-    """
-    Retrieve the RISet, and RISetID values for various Publish trait types.
-    """
+def load_probeset_qtl_info(trait_info, conn):
     query = (
-        "SELECT InbredSet.Name, InbredSet.Id "
-        "FROM InbredSet, PublishFreeze "
-        "WHERE PublishFreeze.InbredSetId = InbredSet.Id "
-        "AND PublishFreeze.Name = %(name)s")
+        "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:
-        return cursor.execute(query, {"name": name})
+        cursor.execute()
+        return dict(zip(
+            ["locus", "lrs", "pvalue", "mean", "additive"], cursor.fetchone()))
+    return {"locus": "", "lrs": "", "pvalue": "", "mean": "", "additive": ""}
 
-def set_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:
-        return cursor.execute(query, {"name": name})
-
-def set_riset_fields(trait_info, conn):
-    """
-    Retrieve the RISet, and RISetID values for various trait types.
-    """
-    riset_functions_map = {
-        "Temp": lambda ti, con: (None, None),
-        "Geno": set_geno_riset_fields,
-        "Publish": set_publish_riset_fields,
-        "ProbeSet": set_probeset_riset_fields
+def load_qtl_info(qtl, trait_type, trait_info, conn):
+    if not qtl:
+        return trait_info
+    qtl_info_functions = {
+        "Publish": load_publish_qtl_info,
+        "ProbeSet": load_probeset_qtl_info
     }
-    if not trait_info.get("haveinfo", None):
+    if trait_inf["name"] not in qtl_info_functions.keys():
         return trait_info
 
-    riset, riid = riset_functions_map[trait_info["type"]](
-        trait_info["name"], conn)
-    return {
-        **trait_info, "risetid": riid,
-        "riset": "BXD" if riset == "BXD300" else riset}
+    return qtl_info_functions[trait_type](trait_info, conn)
 
 def build_trait_name(trait_fullname):
     name_parts = trait_fullname.split("::")
     assert len(name_parts) >= 2, "Name format error"
     return {
-        "trait_db": name_parts[0],
+        "db": {"dataset_name": name_parts[0]},
         "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):
+    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_full_name: str, trait_dataset_id: int,
-        trait_dataset_name: str, conn: Any, qtl=None):
+        trait_type: str, 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."""
-    # pylint: disable=[R0913]
     trait = build_trait_name(trait_full_name)
     trait_info_function_table = {
         "Publish": retrieve_publish_trait_info,
@@ -370,15 +335,19 @@ def retrieve_trait_info(
     }
 
     common_post_processing_fn = compose(
-        lambda ti: set_riset_fields(ti, conn),
-        lambda ti: set_homologene_id_field(ti, conn),
-        lambda ti: {"type": trait_type, **ti},
-        lambda ti: {**ti, **trait},
+        lambda ti: load_qtl_info(qtl, trait_type, ti, conn),
+        lambda ti: set_homologene_id_field(trait_type, ti, conn),
+        lambda ti: {"trait_type": trait_type, **ti},
+        lambda ti: {**trait, **ti},
         set_haveinfo_field)
 
     trait_post_processing_functions_table = {
-        "Publish": compose(set_confidential_field, common_post_processing_fn),
-        "ProbeSet": compose(common_post_processing_fn),
+        "Publish": compose(
+            lambda ti: set_confidential_field(trait_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
     }
@@ -387,10 +356,16 @@ def retrieve_trait_info(
         trait_post_processing_functions_table[trait_type],
         trait_info_function_table[trait_type])
 
-    return retrieve_info(
+    trait_dataset = retrieve_trait_dataset(trait_type, trait, threshold, conn)
+    trait_info = retrieve_info(
         {
             "trait_name": trait["trait_name"],
-            "trait_dataset_id": trait_dataset_id,
-            "trait_dataset_name":trait_dataset_name
+            "trait_dataset_id": trait_dataset["dataset_id"],
+            "trait_dataset_name": trait_dataset["dataset_name"]
         },
         conn)
+    return {
+        **trait_info,
+        "db": {**trait["db"], **trait_dataset},
+        "riset": trait_dataset["riset"]
+    }
diff --git a/tests/unit/db/test_datasets.py b/tests/unit/db/test_datasets.py
new file mode 100644
index 0000000..34fe7f0
--- /dev/null
+++ b/tests/unit/db/test_datasets.py
@@ -0,0 +1,133 @@
+from unittest import mock, TestCase
+
+class TestDatasetsDBFunctions(TestCase):
+
+    def test_retrieve_trait_dataset_name(self):
+        """Test that the function is called correctly."""
+        for trait_type, thresh, trait_dataset_name, columns, table in [
+                ["ProbeSet", 9, "testName",
+                 "Id, Name, FullName, ShortName, DataScale", "ProbeSetFreeze"],
+                ["Geno", 3, "genoTraitName", "Id, Name, FullName, ShortName",
+                 "GenoFreeze"],
+                ["Publish", 6, "publishTraitName",
+                 "Id, Name, FullName, ShortName", "PublishFreeze"],
+                ["Temp", 4, "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 = (
+                        "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 %(columns)s "
+                        "FROM %(table)s "
+                        "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,
+                         "table": table, "columns": columns})
+
+    def test_set_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(
+                        set_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_set_riset_fields(self):
+        """
+        Test that the riset fields are set up correctly for the different trait
+        types.
+        """
+        for trait_info, expected in [
+                [{}, {}],
+                [{"haveinfo": 0, "type": "Publish"},
+                 {"haveinfo": 0, "type": "Publish"}],
+                [{"haveinfo": 0, "type": "ProbeSet"},
+                 {"haveinfo": 0, "type": "ProbeSet"}],
+                [{"haveinfo": 0, "type": "Geno"},
+                 {"haveinfo": 0, "type": "Geno"}],
+                [{"haveinfo": 0, "type": "Temp"},
+                 {"haveinfo": 0, "type": "Temp"}],
+                [{"haveinfo": 1, "type": "Publish", "name": "test"},
+                 {"haveinfo": 1, "type": "Publish", "name": "test",
+                  "riset": "riset_name", "risetid": 0}],
+                [{"haveinfo": 1, "type": "ProbeSet", "name": "test"},
+                 {"haveinfo": 1, "type": "ProbeSet", "name": "test",
+                  "riset": "riset_name", "risetid": 0}],
+                [{"haveinfo": 1, "type": "Geno", "name": "test"},
+                 {"haveinfo": 1, "type": "Geno", "name": "test",
+                  "riset": "riset_name", "risetid": 0}],
+                [{"haveinfo": 1, "type": "Temp", "name": "test"},
+                 {"haveinfo": 1, "type": "Temp", "name": "test", "riset": None,
+                  "risetid": None}]
+        ]:
+            db_mock = mock.MagicMock()
+            with self.subTest(trait_info=trait_info, expected=expected):
+                with db_mock.cursor() as cursor:
+                    cursor.execute.return_value = ("riset_name", 0)
+                    self.assertEqual(
+                        set_riset_fields(trait_info, db_mock), expected)
+
+    def test_set_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(
+                        set_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_set_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(
+                        set_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 39d7a31..7d161bf 100644
--- a/tests/unit/db/test_traits.py
+++ b/tests/unit/db/test_traits.py
@@ -2,55 +2,19 @@
 from unittest import mock, TestCase
 from gn3.db.traits import (
     build_trait_name,
-    set_riset_fields,
     set_haveinfo_field,
     update_sample_data,
     retrieve_trait_info,
-    set_geno_riset_fields,
     set_confidential_field,
     set_homologene_id_field,
     retrieve_geno_trait_info,
     retrieve_temp_trait_info,
-    set_publish_riset_fields,
-    set_probeset_riset_fields,
-    retrieve_trait_dataset_name,
     retrieve_publish_trait_info,
     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, table in [
-                ["ProbeSet", 9, "testName",
-                 "Id, Name, FullName, ShortName, DataScale", "ProbeSetFreeze"],
-                ["Geno", 3, "genoTraitName", "Id, Name, FullName, ShortName",
-                 "GenoFreeze"],
-                ["Publish", 6, "publishTraitName",
-                 "Id, Name, FullName, ShortName", "PublishFreeze"],
-                ["Temp", 4, "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 = (
-                        "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 %(columns)s "
-                        "FROM %(table)s "
-                        "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,
-                         "table": table, "columns": columns})
-
     def test_retrieve_publish_trait_info(self):
         """Test retrieval of type `Publish` traits."""
         db_mock = mock.MagicMock()
@@ -159,10 +123,10 @@ class TestTraitsDBFunctions(TestCase):
     def test_build_trait_name_with_good_fullnames(self):
         for fullname, expected in [
                 ["testdb::testname",
-                 {"trait_db": "testdb", "trait_name": "testname", "cellid": "",
-                  "trait_fullname": "testdb::testname"}],
+                 {"db": {"dataset_name": "testdb"}, "trait_name": "testname",
+                  "cellid": "", "trait_fullname": "testdb::testname"}],
                 ["testdb::testname::testcell",
-                 {"trait_db": "testdb", "trait_name": "testname",
+                 {"db": {"dataset_name": "testdb"}, "trait_name": "testname",
                   "cellid": "testcell",
                   "trait_fullname": "testdb::testname::testcell"}]]:
             with self.subTest(fullname=fullname):
@@ -176,26 +140,26 @@ class TestTraitsDBFunctions(TestCase):
 
     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, expected in [
-                ["Publish", "pubDb::PublishTraitName::pubCell", 1,
-                 "PublishDatasetTraitName",
+        for trait_type, threshold, trait_fullname, expected in [
+                ["Publish", 9, "pubDb::PublishTraitName::pubCell",
                  {"haveinfo": 0, "homologeneid": None, "type": "Publish",
-                  "confidential": 0, "trait_db": "pubDb",
+                  "confidential": 0, "db": {"dataset_name": "pubDb"},
                   "trait_name": "PublishTraitName", "cellid": "pubCell",
                   "trait_fullname": "pubDb::PublishTraitName::pubCell"}],
-                ["ProbeSet", "prbDb::ProbeSetTraitName::prbCell", 2,
-                 "ProbeSetDatasetTraitName",
+                ["ProbeSet", 5, "prbDb::ProbeSetTraitName::prbCell",
                  {"haveinfo": 0, "homologeneid": None, "type": "ProbeSet",
                   "trait_fullname": "prbDb::ProbeSetTraitName::prbCell",
-                  "trait_db": "prbDb", "trait_name": "ProbeSetTraitName",
-                  "cellid": "prbCell"}],
-                ["Geno", "genDb::GenoTraitName", 3, "GenoDatasetTraitName",
+                  "db": {"dataset_name": "prbDb"},
+                  "trait_name": "ProbeSetTraitName", "cellid": "prbCell"}],
+                ["Geno", 12, "genDb::GenoTraitName",
                  {"haveinfo": 0, "homologeneid": None, "type": "Geno",
-                  "trait_fullname": "genDb::GenoTraitName", "trait_db": "genDb",
+                  "trait_fullname": "genDb::GenoTraitName",
+                  "db": {"dataset_name": "genDb"},
                   "trait_name": "GenoTraitName", "cellid": ""}],
-                ["Temp", "tmpDb::TempTraitName", 4, "TempDatasetTraitName",
+                ["Temp", 6, "tmpDb::TempTraitName",
                  {"haveinfo": 0, "homologeneid": None, "type": "Temp",
-                  "trait_fullname": "tmpDb::TempTraitName", "trait_db": "tmpDb",
+                  "trait_fullname": "tmpDb::TempTraitName",
+                  "db": {"dataset_name": "tmpDb"},
                   "trait_name": "TempTraitName", "cellid": ""}]]:
             db_mock = mock.MagicMock()
             with self.subTest(trait_type=trait_type):
@@ -203,8 +167,7 @@ class TestTraitsDBFunctions(TestCase):
                     cursor.fetchone.return_value = tuple()
                     self.assertEqual(
                         retrieve_trait_info(
-                            trait_type, trait_name, trait_dataset_id,
-                            trait_dataset_name, db_mock),
+                            trait_type, threshold, trait_fullname, db_mock),
                         expected)
 
     def test_update_sample_data(self):
@@ -246,128 +209,25 @@ class TestTraitsDBFunctions(TestCase):
 
     def test_set_homologene_id_field(self):
         """Test that the `homologene_id` field is set up correctly"""
-        for trait_info, expected in [
-                [{"type": "Publish"},
-                 {"type": "Publish", "homologeneid": None}],
-                [{"type": "ProbeSet"},
-                 {"type": "ProbeSet", "homologeneid": None}],
-                [{"type": "Geno"}, {"type": "Geno", "homologeneid": None}],
-                [{"type": "Temp"}, {"type": "Temp", "homologeneid": None}]]:
+        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_info, db_mock), expected)
+                        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_info, expected in [
-                [{"type": "Publish"}, {"type": "Publish", "confidential": 0}],
-                [{"type": "ProbeSet"}, {"type": "ProbeSet"}],
-                [{"type": "Geno"}, {"type": "Geno"}],
-                [{"type": "Temp"}, {"type": "Temp"}]]:
+        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_info), expected)
-
-    def test_set_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(
-                        set_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})
-
-
-    def test_set_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(
-                        set_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_set_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(
-                        set_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_set_riset_fields(self):
-        """
-        Test that the riset fields are set up correctly for the different trait
-        types.
-        """
-        for trait_info, expected in [
-                [{}, {}],
-                [{"haveinfo": 0, "type": "Publish"},
-                 {"haveinfo": 0, "type": "Publish"}],
-                [{"haveinfo": 0, "type": "ProbeSet"},
-                 {"haveinfo": 0, "type": "ProbeSet"}],
-                [{"haveinfo": 0, "type": "Geno"},
-                 {"haveinfo": 0, "type": "Geno"}],
-                [{"haveinfo": 0, "type": "Temp"},
-                 {"haveinfo": 0, "type": "Temp"}],
-                [{"haveinfo": 1, "type": "Publish", "name": "test"},
-                 {"haveinfo": 1, "type": "Publish", "name": "test",
-                  "riset": "riset_name", "risetid": 0}],
-                [{"haveinfo": 1, "type": "ProbeSet", "name": "test"},
-                 {"haveinfo": 1, "type": "ProbeSet", "name": "test",
-                  "riset": "riset_name", "risetid": 0}],
-                [{"haveinfo": 1, "type": "Geno", "name": "test"},
-                 {"haveinfo": 1, "type": "Geno", "name": "test",
-                  "riset": "riset_name", "risetid": 0}],
-                [{"haveinfo": 1, "type": "Temp", "name": "test"},
-                 {"haveinfo": 1, "type": "Temp", "name": "test", "riset": None,
-                  "risetid": None}]
-        ]:
-            db_mock = mock.MagicMock()
-            with self.subTest(trait_info=trait_info, expected=expected):
-                with db_mock.cursor() as cursor:
-                    cursor.execute.return_value = ("riset_name", 0)
-                    self.assertEqual(
-                        set_riset_fields(trait_info, db_mock), expected)
+                    set_confidential_field(trait_type, trait_info), expected)