1use std::{collections::HashMap, rc::Rc};
2
3use gen_core::{HashId, calculate_hash, traits::Capnp};
4use rusqlite::{self, Row, params, types::Value};
5use serde::{Deserialize, Serialize};
6
7use crate::{
8 db::GraphConnection,
9 edge::{Edge, EdgeData},
10 gen_models_capnp::block_group_edge,
11 traits::*,
12};
13
14#[derive(Clone, Debug, Deserialize, Serialize, Eq, Hash, PartialEq, Ord, PartialOrd)]
15pub struct BlockGroupEdge {
16 pub id: HashId,
17 pub block_group_id: HashId,
18 pub edge_id: HashId,
19 pub chromosome_index: i64,
20 pub phased: i64,
21 pub created_on: i64,
22}
23
24impl<'a> Capnp<'a> for BlockGroupEdge {
25 type Builder = block_group_edge::Builder<'a>;
26 type Reader = block_group_edge::Reader<'a>;
27
28 fn write_capnp(&self, builder: &mut Self::Builder) {
29 builder.set_id(&self.id.0).unwrap();
30 builder.set_block_group_id(&self.block_group_id.0).unwrap();
31 builder.set_edge_id(&self.edge_id.0).unwrap();
32 builder.set_chromosome_index(self.chromosome_index);
33 builder.set_phased(self.phased);
34 builder.set_created_on(self.created_on);
35 }
36
37 fn read_capnp(reader: Self::Reader) -> Self {
38 let id = reader
39 .get_id()
40 .unwrap()
41 .as_slice()
42 .unwrap()
43 .try_into()
44 .unwrap();
45 let block_group_id = reader
46 .get_block_group_id()
47 .unwrap()
48 .as_slice()
49 .unwrap()
50 .try_into()
51 .unwrap();
52 let edge_id = reader
53 .get_edge_id()
54 .unwrap()
55 .as_slice()
56 .unwrap()
57 .try_into()
58 .unwrap();
59 let chromosome_index = reader.get_chromosome_index();
60 let phased = reader.get_phased();
61 let created_on = reader.get_created_on();
62
63 BlockGroupEdge {
64 id,
65 block_group_id,
66 edge_id,
67 chromosome_index,
68 phased,
69 created_on,
70 }
71 }
72}
73
74#[derive(Clone, Debug, Eq, Hash, PartialEq, Ord, PartialOrd)]
75pub struct BlockGroupEdgeData {
76 pub block_group_id: HashId,
77 pub edge_id: HashId,
78 pub chromosome_index: i64,
79 pub phased: i64,
80}
81
82impl BlockGroupEdgeData {
83 pub fn id_hash(&self) -> HashId {
84 HashId(calculate_hash(&format!(
85 "{}:{}:{}:{}",
86 self.block_group_id, self.edge_id, self.chromosome_index, self.phased
87 )))
88 }
89}
90
91impl From<&BlockGroupEdge> for BlockGroupEdgeData {
92 fn from(item: &BlockGroupEdge) -> Self {
93 BlockGroupEdgeData {
94 block_group_id: item.block_group_id,
95 edge_id: item.edge_id,
96 chromosome_index: item.chromosome_index,
97 phased: item.phased,
98 }
99 }
100}
101
102#[derive(Clone, Debug, Eq, Hash, PartialEq, Ord, PartialOrd)]
103pub struct AugmentedEdge {
104 pub edge: Edge,
105 pub chromosome_index: i64,
106 pub phased: i64,
107 pub created_on: i64,
108}
109
110#[derive(Clone, Debug, Eq, Hash, PartialEq, Ord, PartialOrd)]
111pub struct AugmentedEdgeData {
112 pub edge_data: EdgeData,
113 pub chromosome_index: i64,
114 pub phased: i64,
115}
116
117impl Query for BlockGroupEdge {
118 type Model = BlockGroupEdge;
119
120 const TABLE_NAME: &'static str = "block_group_edges";
121
122 fn process_row(row: &Row) -> Self::Model {
123 BlockGroupEdge {
124 id: row.get(0).unwrap(),
125 block_group_id: row.get(1).unwrap(),
126 edge_id: row.get(2).unwrap(),
127 chromosome_index: row.get(3).unwrap(),
128 phased: row.get(4).unwrap(),
129 created_on: row.get(5).unwrap(),
130 }
131 }
132}
133
134impl BlockGroupEdge {
135 pub fn bulk_create(conn: &GraphConnection, block_group_edges: &[BlockGroupEdgeData]) {
136 let batch_size = max_rows_per_batch(conn, 6);
137
138 for chunk in block_group_edges.chunks(batch_size) {
139 let mut sql = String::from(
140 "INSERT OR IGNORE INTO block_group_edges
141 (id, block_group_id, edge_id, chromosome_index, phased, created_on) VALUES ",
142 );
143 let mut rows_to_insert = vec![];
144 let mut params: Vec<Box<dyn rusqlite::ToSql>> = Vec::new();
145 let timestamp = chrono::Utc::now().timestamp_nanos_opt().unwrap();
146 for block_group_edge in chunk {
147 rows_to_insert.push("(?, ?, ?, ?, ?, ?)".to_string());
148 let hash = block_group_edge.id_hash();
149 params.push(Box::new(hash));
150 params.push(Box::new(block_group_edge.block_group_id));
151 params.push(Box::new(block_group_edge.edge_id));
152 params.push(Box::new(block_group_edge.chromosome_index));
153 params.push(Box::new(block_group_edge.phased));
154 params.push(Box::new(timestamp));
155 }
156
157 sql.push_str(&rows_to_insert.join(", "));
158
159 let mut stmt = conn.prepare(&sql).unwrap();
160 stmt.execute(rusqlite::params_from_iter(params)).unwrap();
161 }
162 }
163
164 pub fn bulk_delete(conn: &GraphConnection, block_group_edges: &[BlockGroupEdgeData]) {
165 let hashes = block_group_edges
166 .iter()
167 .map(|bge| bge.id_hash())
168 .collect::<Vec<_>>();
169 BlockGroupEdge::delete_by_ids(conn, &hashes);
170 }
171
172 pub fn edges_for_block_group(
173 conn: &GraphConnection,
174 block_group_id: &HashId,
175 ) -> Vec<AugmentedEdge> {
176 let block_group_edges = BlockGroupEdge::query(
177 conn,
178 "select * from block_group_edges where block_group_id = ?1 ORDER BY created_on DESC;",
179 params![block_group_id],
180 );
181 let edge_ids = block_group_edges
182 .iter()
183 .map(|block_group_edge| block_group_edge.edge_id)
184 .collect::<Vec<_>>();
185 let edges = Edge::query_by_ids(conn, &edge_ids);
186 let edge_map = edges
187 .iter()
188 .map(|edge| (&edge.id, edge))
189 .collect::<HashMap<_, &Edge>>();
190 block_group_edges
191 .into_iter()
192 .map(|bge| {
193 let edge_info = *edge_map.get(&bge.edge_id).unwrap();
194 AugmentedEdge {
195 edge: edge_info.clone(),
196 chromosome_index: bge.chromosome_index,
197 phased: bge.phased,
198 created_on: bge.created_on,
199 }
200 })
201 .collect()
202 }
203
204 pub fn specific_edges_for_block_group(
205 conn: &GraphConnection,
206 block_group_id: &HashId,
207 edge_ids: &[HashId],
208 ) -> Vec<AugmentedEdge> {
209 let block_group_edges = BlockGroupEdge::query(
210 conn,
211 "SELECT * FROM block_group_edges WHERE block_group_id = ?1 AND edge_id in rarray(?2);",
212 params![
213 block_group_id,
214 Rc::new(
215 edge_ids
216 .iter()
217 .map(|x| Value::from(*x))
218 .collect::<Vec<Value>>()
219 )
220 ],
221 );
222 let edge_ids = block_group_edges
223 .iter()
224 .map(|block_group_edge| block_group_edge.edge_id)
225 .collect::<Vec<_>>();
226 let edges = Edge::query_by_ids(conn, &edge_ids);
227
228 let edge_map = edges
229 .iter()
230 .map(|edge| (&edge.id, edge))
231 .collect::<HashMap<_, &Edge>>();
232 block_group_edges
233 .into_iter()
234 .map(|bge| {
235 let edge_info = *edge_map.get(&bge.edge_id).unwrap();
236 AugmentedEdge {
237 edge: edge_info.clone(),
238 chromosome_index: bge.chromosome_index,
239 phased: bge.phased,
240 created_on: bge.created_on,
241 }
242 })
243 .collect()
244 }
245}
246
247#[cfg(test)]
248mod tests {
249 use capnp::message::TypedBuilder;
250 use chrono::Utc;
251
252 use super::*;
253 use crate::gen_models_capnp::block_group_edge;
254
255 #[test]
256 fn test_block_group_edge_capnp_serialization() {
257 let block_group_edge = BlockGroupEdge {
258 id: "0000000000000000000000000000030000000000000000000000000000000000"
259 .try_into()
260 .unwrap(),
261 block_group_id: "0000000000000000000000000000030000000020000000000000000000000000"
262 .try_into()
263 .unwrap(),
264 edge_id: "0000000000000000000000000000030000000000000000000000000000000000"
265 .try_into()
266 .unwrap(),
267 chromosome_index: 1,
268 phased: 0,
269 created_on: Utc::now().timestamp_nanos_opt().unwrap(),
270 };
271
272 let mut message = TypedBuilder::<block_group_edge::Owned>::new_default();
273 let mut root = message.init_root();
274 block_group_edge.write_capnp(&mut root);
275
276 let deserialized = BlockGroupEdge::read_capnp(root.into_reader());
277 assert_eq!(block_group_edge, deserialized);
278 }
279}