Get possibly filtered set of markers and count of total records.

Get a possibly filtered set of markers ensuring the markers are
specific to the dataset. Also return the total number of markers in
the dataset, even if we do not return the entire list.

1 files changed, 19 insertions, 6 deletions

diff --git a/uploader/genotypes/models.py b/uploader/genotypes/models.py
index 4c3e634..4f09b02 100644
--- a/uploader/genotypes/models.py
+++ b/uploader/genotypes/models.py
@@ -29,18 +29,31 @@ def genotype_markers_count(conn: mdb.Connection, species_id: int) -> int:
 def genotype_markers(
         conn: mdb.Connection,
         species_id: int,
+        dataset_id: int,
         offset: int = 0,
         limit: Optional[int] = None
-) -> tuple[dict, ...]:
+) -> tuple[tuple[dict, ...], int]:
     """Retrieve markers from the database."""
-    _query = "SELECT * FROM Geno WHERE SpeciesId=%s"
-    if bool(limit) and limit > 0:# type: ignore[operator]
-        _query = _query + f" LIMIT {limit} OFFSET {offset}"
+    _query_template = (
+        "SELECT %%COLS%% FROM GenoXRef AS gxr INNER JOIN Geno AS gno "
+        "ON gxr.GenoId=gno.Id WHERE gxr.GenoFreezeId=%s AND gno.SpeciesId=%s "
+        "%%LIMIT%%")
 
     with conn.cursor(cursorclass=DictCursor) as cursor:
-        cursor.execute(_query, (species_id,))
+        cursor.execute(
+            _query_template.replace("%%LIMIT%%", "").replace(
+                "%%COLS%%", "COUNT(gxr.GenoId) AS total_records"),
+            (species_id, dataset_id))
+        _total_records = cursor.fetchone()["total_records"]
+        cursor.execute(
+            _query_template.replace("%%COLS%%", "gxr.GenoFreezeId, gno.*").replace(
+                "%%LIMIT%%",
+                (f"LIMIT {int(limit)} OFFSET {int(offset)}"
+                 if bool(limit) and limit > 0
+                 else "")),
+            (species_id, dataset_id))
         debug_query(cursor, app.logger)
-        return tuple(dict(row) for row in cursor.fetchall())
+        return tuple(dict(row) for row in cursor.fetchall()), _total_records
 
 
 def genotype_dataset(