about summary refs log tree commit diff
path: root/uploader/population/models.py
diff options
context:
space:
mode:
Diffstat (limited to 'uploader/population/models.py')
-rw-r--r--uploader/population/models.py22
1 files changed, 17 insertions, 5 deletions
diff --git a/uploader/population/models.py b/uploader/population/models.py
index d78a821..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:
@@ -47,9 +57,11 @@ def population_genetic_types(conn) -> tuple:
 def save_population(cursor: mdb.cursors.Cursor, population_details: dict) -> dict:
     """Save the population details to the db."""
     cursor.execute("SELECT DISTINCT(Family), FamilyOrder FROM InbredSet "
-                   "WHERE Family IS NOT NULL AND Family != '' "
+                   "WHERE SpeciesId=%s "
+                   "AND Family IS NOT NULL AND Family != '' "
                    "AND FamilyOrder IS NOT NULL "
-                   "ORDER BY FamilyOrder ASC")
+                   "ORDER BY FamilyOrder ASC",
+                   (population_details["SpeciesId"],))
     _families = {
         row["Family"]: int(row["FamilyOrder"])
         for row in cursor.fetchall()