55class GeneInfo :
66 objects = database .MongoDBObjects ("pankb_gene_info" )
77
8+ def get_by_gene_analysis_genome (query_args ):
9+ """
10+ Return every GeneInfo document whose (gene, pangenome_analysis, genome_id)
11+ matches any tuple in *query_args*.
12+ """
13+ if not query_args :
14+ return []
15+
16+ or_conditions = [
17+ {
18+ "gene" : t ["gene" ],
19+ "pangenome_analysis" : t ["pangenome_analysis" ],
20+ "genome_id" : t ["genome_id" ],
21+ }
22+ for t in query_args
23+ ]
24+
25+ return list (
26+ GeneInfo .objects .find (
27+ {"$or" : or_conditions },
28+ projection = {"_id" : 0 }
29+ )
30+ )
31+
32+ def get_by_gene_and_analysis (pairs ):
33+ """
34+ Return every GeneInfo document whose (gene, pangenome_analysis)
35+ matches any tuple in *pairs*.
36+ """
37+ if not pairs :
38+ return []
39+
40+ or_conditions = [
41+ {"gene" : g , "pangenome_analysis" : a }
42+ for g , a in pairs
43+ ]
44+
45+ return list (
46+ GeneInfo .objects .find (
47+ {"$or" : or_conditions },
48+ projection = {"_id" : 0 }
49+ )
50+ )
51+
52+ @staticmethod
53+ def get_gene_strain_pairs_paginated (skip : int = 0 , limit : int = 10000 ):
54+ """
55+ Return paginated (gene, genome_id) pairs directly from collection.
56+ No $group - assumes data has no duplicates, InteropDB will upsert anyway.
57+
58+ Args:
59+ skip: Number of records to skip
60+ limit: Max records to return
61+
62+ Returns:
63+ {
64+ "pairs": [...],
65+ "total": int
66+ }
67+ """
68+ # Get total count using aggregation
69+ count_pipeline = [
70+ {"$match" : {"gene" : {"$ne" : None }, "genome_id" : {"$ne" : None }}},
71+ {"$count" : "total" }
72+ ]
73+ count_result = list (GeneInfo .objects .aggregate (count_pipeline ))
74+ total = count_result [0 ]["total" ] if count_result else 0
75+
76+ # Direct find with skip/limit (fast)
77+ cursor = GeneInfo .objects .find (
78+ {"gene" : {"$ne" : None }, "genome_id" : {"$ne" : None }, "locus_tag" : {"$ne" : None }},
79+ projection = {"_id" : 0 , "gene" : 1 , "genome_id" : 1 , "locus_tag" : 1 }
80+ ).skip (skip ).limit (limit )
81+
82+ pairs = [{"gene" : doc ["gene" ], "strain" : doc ["genome_id" ], "locus_tag" : doc ["locus_tag" ]} for doc in cursor ]
83+
84+ return {"pairs" : pairs , "total" : total }
85+
86+ def get_gene_info_and_pangenomic_class_pipeline (gene_match ): # This is an ugly workaround to make it compatible with Azure Cosmos DB
87+ return [
88+ {"$match" : gene_match },
89+ {
90+ "$unionWith" : {
91+ "coll" : "pankb_gene_annotations" ,
92+ "pipeline" : [
93+ {
94+ "$match" : {
95+ "pangenome_analysis" : gene_match ["pangenome_analysis" ]
96+ }
97+ },
98+ {"$project" : {"_id" : 1 , "gene" : 1 , "pangenomic_class" : 1 }},
99+ ],
100+ }
101+ },
102+ {"$group" : {"_id" : "$gene" , "doc" : {"$mergeObjects" : "$$ROOT" }}},
103+ {"$replaceRoot" : {"newRoot" : "$doc" }},
104+ {"$match" : {"locus_tag" : {"$exists" : True }}},
105+ {"$fill" : {"output" : {"pangenomic_class" : {"value" : "-" }}}}
106+ ]
107+
108+ def get_gene_info_and_pangenomic_class_pipeline_mongodb_only (gene_match ):
109+ return [
110+ {"$match" : gene_match },
111+ {
112+ "$lookup" : {
113+ "from" : "pankb_gene_annotations" ,
114+ "let" : {
115+ "q_gene" : "$gene" ,
116+ "q_pangenome_analysis" : "$pangenome_analysis" ,
117+ },
118+ "pipeline" : [
119+ {
120+ "$match" : {
121+ "$expr" : {
122+ "$and" : [
123+ {"$eq" : ["$gene" , "$$q_gene" ]},
124+ {
125+ "$eq" : [
126+ "$pangenome_analysis" ,
127+ "$$q_pangenome_analysis" ,
128+ ]
129+ },
130+ ]
131+ }
132+ }
133+ }
134+ ],
135+ "as" : "pangenomic_class" ,
136+ }
137+ },
138+ {
139+ "$set" : {
140+ "pangenomic_class" : {
141+ "$ifNull" : [
142+ {"$arrayElemAt" : ["$pangenomic_class.pangenomic_class" , 0 ]},
143+ "-" ,
144+ ]
145+ }
146+ }
147+ },
148+ ]
149+
150+ def get_gene_info_and_pangenomic_class (genome_match , projection = None ):
151+ pipeline = GeneInfo .get_gene_info_and_pangenomic_class_pipeline (genome_match )
152+ if isinstance (projection , list ):
153+ projection = {p : 1 for p in projection }
154+ if not "_id" in projection :
155+ projection ["_id" ] = 0
156+ if projection :
157+ pipeline .append ({"$project" : projection })
158+ return GeneInfo .objects .aggregate (pipeline )
159+
8160
9161# Model for the Genome Info table content
10162class GenomeInfo :
11163 objects = database .MongoDBObjects ("pankb_genome_info" )
12164
165+ @staticmethod
166+ def get_all_strains ():
167+ """
168+ Return a sorted list of distinct genome_id values.
169+ """
170+ pipeline = [
171+ {"$group" : {"_id" : "$genome_id" }},
172+ {"$sort" : {"_id" : 1 }},
173+ {"$project" : {"_id" : 0 , "genome_id" : "$_id" }},
174+ ]
175+
176+ cursor = GenomeInfo .objects .aggregate (pipeline )
177+ return [doc ["genome_id" ] for doc in cursor if doc .get ("genome_id" )]
178+
179+
13180 def get_genome_and_isolation_info_pipeline (genome_match ):
14181 return [
15182 {"$match" : genome_match },
@@ -30,6 +197,23 @@ def get_genome_and_isolation_info_pipeline(genome_match):
30197 {"$project" : {"_id" : 0 , "isolation_info" : 0 }},
31198 ]
32199
200+ def get_by_genome_ids (genome_ids , projection = None ):
201+ """
202+ Fetch one or more genomes by ID.
203+
204+ :param genome_ids: list[str] – genome_id values to look up
205+ :param include_isolation: bool – whether to perform the $lookup join
206+ :param projection: list[str] | dict | None – optional projection
207+ :return: list[dict]
208+ """
209+ genome_match = {"genome_id" : {"$in" : genome_ids }}
210+
211+ cursor = GenomeInfo .get_genome_and_isolation_info (
212+ genome_match , projection = projection
213+ )
214+
215+ return list (cursor )
216+
33217 def get_genome_and_isolation_info (genome_match , projection = None ):
34218 pipeline = GenomeInfo .get_genome_and_isolation_info_pipeline (genome_match )
35219 if isinstance (projection , list ):
0 commit comments