1use gen_core::traits::Capnp;
2use rusqlite::{Result as SQLResult, Row, params};
3use serde::{Deserialize, Serialize};
4
5use crate::{
6 db::GraphConnection, gen_models_capnp::sample_lineage, lineage::SqlLineage, traits::Query,
7};
8
9#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
10pub struct SampleLineage {
11 pub parent_sample_name: String,
12 pub child_sample_name: String,
13}
14
15impl<'a> Capnp<'a> for SampleLineage {
16 type Builder = sample_lineage::Builder<'a>;
17 type Reader = sample_lineage::Reader<'a>;
18
19 fn write_capnp(&self, builder: &mut Self::Builder) {
20 builder.set_parent_sample_name(&self.parent_sample_name);
21 builder.set_child_sample_name(&self.child_sample_name);
22 }
23
24 fn read_capnp(reader: Self::Reader) -> Self {
25 let parent_sample_name = reader
26 .get_parent_sample_name()
27 .unwrap()
28 .to_string()
29 .unwrap();
30 let child_sample_name = reader.get_child_sample_name().unwrap().to_string().unwrap();
31
32 SampleLineage {
33 parent_sample_name,
34 child_sample_name,
35 }
36 }
37}
38
39impl Query for SampleLineage {
40 type Model = SampleLineage;
41
42 const PRIMARY_KEY: &'static str = "parent_sample_name";
43 const TABLE_NAME: &'static str = "sample_lineage";
44
45 fn process_row(row: &Row) -> Self::Model {
46 SampleLineage {
47 parent_sample_name: row.get(0).unwrap(),
48 child_sample_name: row.get(1).unwrap(),
49 }
50 }
51}
52
53impl SqlLineage for SampleLineage {
54 type Id = String;
55
56 const CHILD_COLUMN: &'static str = "child_sample_name";
57 const CHILD_ID_COLUMN: &'static str = "name";
58 const CHILD_TABLE_NAME: &'static str = "samples";
59 const PARENT_COLUMN: &'static str = "parent_sample_name";
60 const PARENT_ID_COLUMN: &'static str = "name";
61 const PARENT_TABLE_NAME: &'static str = "samples";
62
63 fn parent_id(&self) -> &Self::Id {
64 &self.parent_sample_name
65 }
66
67 fn child_id(&self) -> &Self::Id {
68 &self.child_sample_name
69 }
70}
71
72impl SampleLineage {
73 pub fn get_parents(conn: &GraphConnection, child_sample_name: &str) -> Vec<String> {
74 SampleLineage::query(
75 conn,
76 "SELECT * FROM sample_lineage WHERE child_sample_name = ?1 ORDER BY parent_sample_name;",
77 params![child_sample_name],
78 )
79 .into_iter()
80 .map(|lineage| lineage.parent_sample_name)
81 .collect()
82 }
83
84 pub fn get_children(conn: &GraphConnection, parent_sample_name: &str) -> Vec<String> {
85 SampleLineage::query(
86 conn,
87 "SELECT * FROM sample_lineage WHERE parent_sample_name = ?1 ORDER BY child_sample_name;",
88 params![parent_sample_name],
89 )
90 .into_iter()
91 .map(|lineage| lineage.child_sample_name)
92 .collect()
93 }
94
95 pub fn search_name(conn: &GraphConnection, name: &str) -> Vec<Self> {
96 SampleLineage::query(
97 conn,
98 "SELECT * FROM sample_lineage
99 WHERE instr(lower(parent_sample_name), lower(?1)) > 0
100 OR instr(lower(child_sample_name), lower(?1)) > 0
101 ORDER BY parent_sample_name, child_sample_name;",
102 params![name],
103 )
104 }
105
106 pub fn create(
107 conn: &GraphConnection,
108 parent_sample_name: &str,
109 child_sample_name: &str,
110 ) -> SQLResult<Self> {
111 let query = "INSERT INTO sample_lineage (parent_sample_name, child_sample_name)
112 VALUES (?1, ?2)
113 ON CONFLICT(parent_sample_name, child_sample_name) DO NOTHING;";
114 let mut stmt = conn.prepare(query).unwrap();
115 stmt.execute(params![parent_sample_name, child_sample_name])?;
116
117 Ok(SampleLineage {
118 parent_sample_name: parent_sample_name.to_string(),
119 child_sample_name: child_sample_name.to_string(),
120 })
121 }
122
123 pub fn delete(
124 conn: &GraphConnection,
125 parent_sample_name: &str,
126 child_sample_name: &str,
127 ) -> SQLResult<()> {
128 let query =
129 "DELETE FROM sample_lineage WHERE parent_sample_name = ?1 AND child_sample_name = ?2;";
130 let mut stmt = conn.prepare(query).unwrap();
131 stmt.execute(params![parent_sample_name, child_sample_name])?;
132 Ok(())
133 }
134}
135
136#[cfg(test)]
137mod tests {
138 use capnp::message::TypedBuilder;
139
140 use super::*;
141 use crate::{lineage::SqlLineage, sample::Sample, test_helpers::get_connection};
142
143 #[test]
144 fn test_capnp_serialization() {
145 let lineage = SampleLineage {
146 parent_sample_name: "parent".to_string(),
147 child_sample_name: "child".to_string(),
148 };
149
150 let mut message = TypedBuilder::<sample_lineage::Owned>::new_default();
151 let mut root = message.init_root();
152 lineage.write_capnp(&mut root);
153
154 let deserialized = SampleLineage::read_capnp(root.into_reader());
155 assert_eq!(lineage, deserialized);
156 }
157
158 #[test]
159 fn test_lineage_queries() {
160 let conn = get_connection(None).unwrap();
161
162 for sample in ["root", "left", "right", "leaf", "sibling"] {
163 Sample::get_or_create(&conn, sample);
164 }
165
166 SampleLineage::create(&conn, "root", "left").unwrap();
167 SampleLineage::create(&conn, "root", "right").unwrap();
168 SampleLineage::create(&conn, "left", "leaf").unwrap();
169 SampleLineage::create(&conn, "right", "leaf").unwrap();
170 SampleLineage::create(&conn, "right", "sibling").unwrap();
171
172 let ancestors = SampleLineage::get_ancestors(&conn, &"leaf".to_string(), None);
173 assert_eq!(ancestors, vec!["left", "right", "root"]);
174 assert_eq!(
175 SampleLineage::get_ancestors(&conn, &"leaf".to_string(), Some(1)),
176 vec!["left", "right"]
177 );
178 assert_eq!(
179 SampleLineage::get_ancestors(&conn, &"leaf".to_string(), Some(0)),
180 Vec::<String>::new()
181 );
182
183 assert_eq!(
184 SampleLineage::get_parents(&conn, "leaf"),
185 vec!["left".to_string(), "right".to_string()]
186 );
187 assert_eq!(
188 SampleLineage::get_children(&conn, "right"),
189 vec!["leaf".to_string(), "sibling".to_string()]
190 );
191
192 let descendants = SampleLineage::get_descendants(&conn, &"root".to_string(), None);
193 assert_eq!(descendants, vec!["left", "right", "leaf", "sibling"]);
194 assert_eq!(
195 SampleLineage::get_descendants(&conn, &"root".to_string(), Some(1)),
196 vec!["left", "right"]
197 );
198 assert_eq!(
199 SampleLineage::get_descendants(&conn, &"root".to_string(), Some(0)),
200 Vec::<String>::new()
201 );
202
203 let mut graph = SampleLineage::get_graph(&conn);
204 graph.sort_by(|left, right| {
205 left.parent_sample_name
206 .cmp(&right.parent_sample_name)
207 .then(left.child_sample_name.cmp(&right.child_sample_name))
208 });
209 assert_eq!(
210 graph,
211 vec![
212 SampleLineage {
213 parent_sample_name: "left".to_string(),
214 child_sample_name: "leaf".to_string(),
215 },
216 SampleLineage {
217 parent_sample_name: "right".to_string(),
218 child_sample_name: "leaf".to_string(),
219 },
220 SampleLineage {
221 parent_sample_name: "right".to_string(),
222 child_sample_name: "sibling".to_string(),
223 },
224 SampleLineage {
225 parent_sample_name: "root".to_string(),
226 child_sample_name: "left".to_string(),
227 },
228 SampleLineage {
229 parent_sample_name: "root".to_string(),
230 child_sample_name: "right".to_string(),
231 },
232 ]
233 );
234
235 let path =
236 SampleLineage::get_path_between(&conn, &"leaf".to_string(), &"sibling".to_string());
237 assert_eq!(path, vec!["leaf", "right", "sibling"]);
238
239 let edges = SampleLineage::get_path_edges_between(
240 &conn,
241 &"leaf".to_string(),
242 &"sibling".to_string(),
243 );
244 assert_eq!(
245 edges,
246 vec![
247 SampleLineage {
248 parent_sample_name: "right".to_string(),
249 child_sample_name: "leaf".to_string(),
250 },
251 SampleLineage {
252 parent_sample_name: "right".to_string(),
253 child_sample_name: "sibling".to_string(),
254 },
255 ]
256 );
257 }
258
259 #[test]
260 fn test_lineage_depth_limit() {
261 let conn = get_connection(None).unwrap();
262
263 for sample in ["root", "left", "right", "leaf", "sibling"] {
264 Sample::get_or_create(&conn, sample);
265 }
266
267 SampleLineage::create(&conn, "root", "left").unwrap();
268 SampleLineage::create(&conn, "root", "right").unwrap();
269 SampleLineage::create(&conn, "left", "leaf").unwrap();
270 SampleLineage::create(&conn, "right", "leaf").unwrap();
271 SampleLineage::create(&conn, "right", "sibling").unwrap();
272
273 assert_eq!(
274 SampleLineage::get_ancestors(&conn, &"leaf".to_string(), Some(1)),
275 vec!["left", "right"]
276 );
277 assert_eq!(
278 SampleLineage::get_ancestors(&conn, &"leaf".to_string(), Some(2)),
279 vec!["left", "right", "root"]
280 );
281 assert_eq!(
282 SampleLineage::get_descendants(&conn, &"root".to_string(), Some(1)),
283 vec!["left", "right"]
284 );
285 assert_eq!(
286 SampleLineage::get_descendants(&conn, &"root".to_string(), Some(2)),
287 vec!["left", "right", "leaf", "sibling"]
288 );
289 }
290
291 #[test]
292 fn test_self_references_are_rejected() {
293 let conn = get_connection(None).unwrap();
294 Sample::get_or_create(&conn, "sample");
295
296 let err = SampleLineage::create(&conn, "sample", "sample").unwrap_err();
297 assert!(matches!(
298 err,
299 rusqlite::Error::SqliteFailure(code, _)
300 if code.code == rusqlite::ErrorCode::ConstraintViolation
301 ));
302 }
303
304 #[test]
305 fn test_search_name_returns_partial_matches() {
306 let conn = get_connection(None).unwrap();
307
308 for sample in [
309 "alpha",
310 "BarFooBaz",
311 "child",
312 "foo",
313 "plain-parent",
314 "plain-child",
315 "QuxFood",
316 "zzz",
317 ] {
318 Sample::get_or_create(&conn, sample);
319 }
320
321 SampleLineage::create(&conn, "alpha", "BarFooBaz").unwrap();
322 SampleLineage::create(&conn, "foo", "child").unwrap();
323 SampleLineage::create(&conn, "plain-parent", "plain-child").unwrap();
324 SampleLineage::create(&conn, "zzz", "QuxFood").unwrap();
325
326 let matches = SampleLineage::search_name(&conn, "FoO");
327
328 assert_eq!(
329 matches,
330 vec![
331 SampleLineage {
332 parent_sample_name: "alpha".to_string(),
333 child_sample_name: "BarFooBaz".to_string(),
334 },
335 SampleLineage {
336 parent_sample_name: "foo".to_string(),
337 child_sample_name: "child".to_string(),
338 },
339 SampleLineage {
340 parent_sample_name: "zzz".to_string(),
341 child_sample_name: "QuxFood".to_string(),
342 },
343 ]
344 );
345 }
346}