1 files changed, 55 insertions, 32 deletions
diff --git a/uploader/population/models.py b/uploader/population/models.py
index 782bc9f..4d95065 100644
--- a/uploader/population/models.py
+++ b/uploader/population/models.py
@@ -26,13 +26,23 @@ def populations_by_species(conn: mdb.Connection, speciesid) -> tuple:
 
     return tuple()
 
+__GENERIC_POPULATION_FAMILIES__ = (
+    "Reference Populations (replicate average, SE, N)",
+    "Crosses and Heterogeneous Stock (individuals)",
+    "Groups Without Genotypes")
 
-def population_families(conn) -> tuple:
+def population_families(conn, species_id: int) -> tuple[str]:
     """Fetch the families under which populations are grouped."""
     with conn.cursor(cursorclass=DictCursor) as cursor:
+        paramstr = ", ".join(["%s"] * len(__GENERIC_POPULATION_FAMILIES__))
         cursor.execute(
-            "SELECT DISTINCT(Family) FROM InbredSet WHERE Family IS NOT NULL")
-        return tuple(row["Family"] for row in cursor.fetchall())
+            "SELECT DISTINCT(Family) FROM InbredSet "
+            "WHERE SpeciesId=%s "
+            "AND Family IS NOT NULL "
+            f"AND Family NOT IN ({paramstr})",
+            (species_id, *__GENERIC_POPULATION_FAMILIES__))
+        return __GENERIC_POPULATION_FAMILIES__ + tuple(
+            row["Family"] for row in cursor.fetchall())
 
 
 def population_genetic_types(conn) -> tuple:
@@ -44,33 +54,46 @@ def population_genetic_types(conn) -> tuple:
         return tuple(row["GeneticType"] for row in cursor.fetchall())
 
 
-def save_population(conn: mdb.Connection, population_details: dict) -> dict:
+def save_population(cursor: mdb.cursors.Cursor, population_details: dict) -> dict:
     """Save the population details to the db."""
-    with conn.cursor(cursorclass=DictCursor) as cursor:
-        #TODO: Handle FamilyOrder here
-        cursor.execute(
-            "INSERT INTO InbredSet("
-            "InbredSetId, InbredSetName, Name, SpeciesId, FullName, "
-            "public, MappingMethodId, GeneticType, Family, MenuOrderId, "
-            "InbredSetCode, Description"
-            ") "
-            "VALUES ("
-            "%(InbredSetId)s, %(InbredSetName)s, %(Name)s, %(SpeciesId)s, "
-            "%(FullName)s, %(public)s, %(MappingMethodId)s, %(GeneticType)s, "
-            "%(Family)s, %(MenuOrderId)s, %(InbredSetCode)s, %(Description)s"
-            ")",
-            {
-                "MenuOrderId": 0,
-                "InbredSetId": 0,
-                "public": 2,
-                **population_details
-            })
-        new_id = cursor.lastrowid
-        cursor.execute("UPDATE InbredSet SET InbredSetId=%s WHERE Id=%s",
-                       (new_id, new_id))
-        return {
-            **population_details,
-            "Id": new_id,
-            "InbredSetId": new_id,
-            "population_id": new_id
-        }
+    cursor.execute("SELECT DISTINCT(Family), FamilyOrder FROM InbredSet "
+                   "WHERE SpeciesId=%s "
+                   "AND Family IS NOT NULL AND Family != '' "
+                   "AND FamilyOrder IS NOT NULL "
+                   "ORDER BY FamilyOrder ASC",
+                   (population_details["SpeciesId"],))
+    _families = {
+        row["Family"]: int(row["FamilyOrder"])
+        for row in cursor.fetchall()
+    }
+    params = {
+        "MenuOrderId": 0,
+        "InbredSetId": 0,
+        "public": 2,
+        **population_details,
+        "FamilyOrder": _families.get(
+            population_details["Family"],
+            max((0,) + tuple(_families.values()))+1)
+    }
+    cursor.execute(
+        "INSERT INTO InbredSet("
+        "InbredSetId, InbredSetName, Name, SpeciesId, FullName, "
+        "public, MappingMethodId, GeneticType, Family, FamilyOrder,"
+        " MenuOrderId, InbredSetCode, Description"
+        ") "
+        "VALUES ("
+        "%(InbredSetId)s, %(InbredSetName)s, %(Name)s, %(SpeciesId)s, "
+        "%(FullName)s, %(public)s, %(MappingMethodId)s, %(GeneticType)s, "
+        "%(Family)s, %(FamilyOrder)s, %(MenuOrderId)s, %(InbredSetCode)s, "
+        "%(Description)s"
+        ")",
+        params)
+    new_id = cursor.lastrowid
+    cursor.execute("UPDATE InbredSet SET InbredSetId=%s WHERE Id=%s",
+                   (new_id, new_id))
+    return {
+        **params,
+        "Id": new_id,
+        "InbredSetId": new_id,
+        "population_id": new_id
+    }