refine interop query using cursor

Author: SHINTEKI

gene_function/models.py

@@ -50,38 +50,40 @@ def get_by_gene_and_analysis(pairs):
         )
 
     @staticmethod
-    def get_gene_strain_pairs_paginated(skip: int = 0, limit: int = 10000):
+    def get_gene_strain_pairs_paginated(after: str | None = None, limit: int = 10000):
         """
-        Return paginated (gene, genome_id) pairs directly from collection.
-        No $group - assumes data has no duplicates, InteropDB will upsert anyway.
+        Return paginated (gene, genome_id) pairs using cursor-based pagination.
+        Uses _id > after to seek directly, so every page is equally fast.
 
         Args:
-            skip: Number of records to skip
+            after: The _id of the last document from the previous page (None for first page)
             limit: Max records to return
 
         Returns:
             {
                 "pairs": [...],
-                "total": int
+                "next_cursor": str | None  (None means no more data)
             }
         """
-        # Get total count using aggregation
-        count_pipeline = [
-            {"$match": {"gene": {"$ne": None}, "genome_id": {"$ne": None}}},
-            {"$count": "total"}
-        ]
-        count_result = list(GeneInfo.objects.aggregate(count_pipeline))
-        total = count_result[0]["total"] if count_result else 0
+        filter_query = {"gene": {"$ne": None}, "genome_id": {"$ne": None}, "locus_tag": {"$ne": None}}
+        if after is not None:
+            from bson import ObjectId
+            filter_query["_id"] = {"$gt": ObjectId(after)}
 
-        # Direct find with skip/limit (fast)
         cursor = GeneInfo.objects.find(
-            {"gene": {"$ne": None}, "genome_id": {"$ne": None}, "locus_tag": {"$ne": None}},
-            projection={"_id": 0, "gene": 1, "genome_id": 1, "locus_tag": 1}
-        ).skip(skip).limit(limit)
+            filter_query,
+            projection={"_id": 1, "gene": 1, "genome_id": 1, "locus_tag": 1}
+        ).sort("_id", 1).limit(limit)
+
+        pairs = []
+        last_id = None
+        for doc in cursor:
+            pairs.append({"gene": doc["gene"], "strain": doc["genome_id"], "locus_tag": doc["locus_tag"]})
+            last_id = str(doc["_id"])
 
-        pairs = [{"gene": doc["gene"], "strain": doc["genome_id"], "locus_tag": doc["locus_tag"]} for doc in cursor]
+        next_cursor = last_id if len(pairs) == limit else None
 
-        return {"pairs": pairs, "total": total}
+        return {"pairs": pairs, "next_cursor": next_cursor}
 
     def get_gene_info_and_pangenomic_class_pipeline(gene_match): # This is an ugly workaround to make it compatible with Azure Cosmos DB
         return [
@@ -177,16 +179,38 @@ def get_all_strains():
         return [doc["genome_id"] for doc in cursor if doc.get("genome_id")]
 
     @staticmethod
-    def get_all_strains_paginated(skip=0, limit=10000):
+    def get_all_strains_paginated(after: str | None = None, limit: int = 10000):
         """
-        Return paginated distinct genome_id values.
-        Uses distinct() for fast retrieval, then slices in Python.
+        Return paginated genome_id values using cursor-based pagination.
+        Uses _id > after to seek directly, so every page is equally fast.
+
+        Args:
+            after: The _id of the last document from the previous page (None for first page)
+            limit: Max records to return
+
+        Returns:
+            {"strains": [...], "next_cursor": str | None}
         """
-        col = GenomeInfo.objects.collection
-        all_ids = sorted(col.distinct("genome_id"))
-        total = len(all_ids)
-        page = all_ids[skip:skip + limit]
-        return {"strains": page, "total": total}
+        from bson import ObjectId
+
+        filter_query = {"genome_id": {"$ne": None}}
+        if after is not None:
+            filter_query["_id"] = {"$gt": ObjectId(after)}
+
+        cursor = GenomeInfo.objects.find(
+            filter_query,
+            projection={"_id": 1, "genome_id": 1}
+        ).sort("_id", 1).limit(limit)
+
+        strains = []
+        last_id = None
+        for doc in cursor:
+            strains.append(doc["genome_id"])
+            last_id = str(doc["_id"])
+
+        next_cursor = last_id if len(strains) == limit else None
+
+        return {"strains": strains, "next_cursor": next_cursor}
 
 
     def get_genome_and_isolation_info_pipeline(genome_match):

interop_query/views.py

@@ -49,12 +49,12 @@
 )
 @api_view(["GET"])
 def genes(request):
-    """Return all genes with species and PanKB URLs (paginated)."""
+    """Return all genes with species and PanKB URLs (cursor-based pagination)."""
     try:
-        skip = int(request.GET.get("skip", 0))
-        limit = min(int(request.GET.get("limit", 10000)), 50000)
+        after = request.GET.get("after")
+        limit = min(int(request.GET.get("limit", 50000)), 50000)
 
-        result = GeneAnnotations.get_all_genes_paginated(skip=skip, limit=limit)
+        result = GeneAnnotations.get_all_genes_paginated(after=after, limit=limit)
         gene_list = result["genes"]
 
         for gene_data in gene_list:
@@ -67,10 +67,9 @@ def genes(request):
 
         return Response({
             "genes": gene_list,
-            "total": result["total"],
-            "skip": skip,
             "limit": limit,
-            "has_more": skip + len(gene_list) < result["total"],
+            "next_cursor": result["next_cursor"],
+            "has_more": result["next_cursor"] is not None,
         })
     except Exception as e:
         logger.exception("list_all_genes failed")
@@ -112,12 +111,12 @@ def genes(request):
 )
 @api_view(["GET"])
 def strains(request):
-    """Return all strains (genome IDs) with PanKB URLs (paginated)."""
+    """Return all strains (genome IDs) with PanKB URLs (cursor-based pagination)."""
     try:
-        skip = int(request.GET.get("skip", 0))
-        limit = min(int(request.GET.get("limit", 10000)), 50000)
+        after = request.GET.get("after")
+        limit = min(int(request.GET.get("limit", 50000)), 50000)
 
-        result = GenomeInfo.get_all_strains_paginated(skip=skip, limit=limit)
+        result = GenomeInfo.get_all_strains_paginated(after=after, limit=limit)
         strain_ids = result["strains"]
 
         strain_list = []
@@ -129,10 +128,9 @@ def strains(request):
 
         return Response({
             "strains": strain_list,
-            "total": result["total"],
-            "skip": skip,
             "limit": limit,
-            "has_more": skip + len(strain_list) < result["total"],
+            "next_cursor": result["next_cursor"],
+            "has_more": result["next_cursor"] is not None,
         })
     except Exception as e:
         logger.exception("list_all_strains failed")
@@ -176,12 +174,12 @@ def strains(request):
 )
 @api_view(["GET"])
 def gene_strain_pairs(request):
-    """Return distinct (gene, strain, locus_tag) pairs with PanKB URLs."""
+    """Return distinct (gene, strain, locus_tag) pairs with PanKB URLs (cursor-based pagination)."""
     try:
-        skip = int(request.GET.get("skip", 0))
-        limit = min(int(request.GET.get("limit", 10000)), 50000)
+        after = request.GET.get("after")  # _id cursor from previous page
+        limit = min(int(request.GET.get("limit", 50000)), 50000)
 
-        result = GeneInfo.get_gene_strain_pairs_paginated(skip=skip, limit=limit)
+        result = GeneInfo.get_gene_strain_pairs_paginated(after=after, limit=limit)
         pairs = result["pairs"]
 
         for pair in pairs:
@@ -193,10 +191,9 @@ def gene_strain_pairs(request):
 
         return Response({
             "pairs": pairs,
-            "total": result["total"],
-            "skip": skip,
             "limit": limit,
-            "has_more": skip + len(pairs) < result["total"],
+            "next_cursor": result["next_cursor"],
+            "has_more": result["next_cursor"] is not None,
         })
     except Exception as e:
         logger.exception("get_gene_strain_pairs failed")

pangenome_analyses/models.py

@@ -30,21 +30,39 @@ def get_all_genes():
         return result
 
     @staticmethod
-    def get_all_genes_paginated(skip=0, limit=10000):
+    def get_all_genes_paginated(after: str | None = None, limit: int = 10000):
         """
-        Return paginated (gene, pangenome_analysis) pairs.
-        Uses indexed find+sort+skip+limit for performance on large collections.
+        Return paginated (gene, pangenome_analysis) pairs using cursor-based pagination.
+        Uses _id > after to seek directly, so every page is equally fast.
+
+        Args:
+            after: The _id of the last document from the previous page (None for first page)
+            limit: Max records to return
+
+        Returns:
+            {"genes": [...], "next_cursor": str | None}
         """
+        from bson import ObjectId
+
+        filter_query = {"gene": {"$ne": None}, "pangenome_analysis": {"$ne": None}}
+        if after is not None:
+            filter_query["_id"] = {"$gt": ObjectId(after)}
+
         col = GeneAnnotations.objects.collection
         cursor = col.find(
-            {"gene": {"$ne": None}, "pangenome_analysis": {"$ne": None}},
-            {"_id": 0, "gene": 1, "pangenome_analysis": 1},
-        ).sort([("pangenome_analysis", 1), ("gene", 1)]).skip(skip).limit(limit)
-        genes = list(cursor)
+            filter_query,
+            {"_id": 1, "gene": 1, "pangenome_analysis": 1},
+        ).sort("_id", 1).limit(limit)
+
+        genes = []
+        last_id = None
+        for doc in cursor:
+            genes.append({"gene": doc["gene"], "pangenome_analysis": doc["pangenome_analysis"]})
+            last_id = str(doc["_id"])
 
-        total = col.estimated_document_count()
+        next_cursor = last_id if len(genes) == limit else None
 
-        return {"genes": genes, "total": total}
+        return {"genes": genes, "next_cursor": next_cursor}
 
     def get_gene_analysis_pairs(genes):
         """