gen_models/
collection.rs

1use gen_core::traits::Capnp;
2use rusqlite::{Row, params_from_iter};
3use serde::{Deserialize, Serialize};
4
5use crate::{
6    block_group::BlockGroup, db::GraphConnection, gen_models_capnp::collection, traits::*,
7};
8
9#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
10pub struct Collection {
11    pub name: String,
12}
13
14impl<'a> Capnp<'a> for Collection {
15    type Builder = collection::Builder<'a>;
16    type Reader = collection::Reader<'a>;
17
18    fn write_capnp(&self, builder: &mut Self::Builder) {
19        builder.set_name(self.name.clone());
20    }
21
22    fn read_capnp(reader: Self::Reader) -> Self {
23        let name = reader.get_name().unwrap().to_string().unwrap();
24
25        Collection { name }
26    }
27}
28
29impl Query for Collection {
30    type Model = Collection;
31
32    const PRIMARY_KEY: &'static str = "name";
33    const TABLE_NAME: &'static str = "collections";
34
35    fn process_row(row: &Row) -> Self::Model {
36        Collection {
37            name: row.get(0).unwrap(),
38        }
39    }
40}
41
42impl Collection {
43    pub fn exists(conn: &GraphConnection, name: &str) -> bool {
44        let mut stmt = conn
45            .prepare("select name from collections where name = ?1")
46            .unwrap();
47        stmt.exists([name]).unwrap()
48    }
49
50    pub fn create(conn: &GraphConnection, name: &str) -> Collection {
51        let mut stmt = conn
52            .prepare("INSERT INTO collections (name) VALUES (?1) RETURNING *;")
53            .unwrap();
54
55        match stmt.query_row((name,), |_row| {
56            Ok(Collection {
57                name: name.to_string(),
58            })
59        }) {
60            Ok(res) => res,
61            Err(rusqlite::Error::SqliteFailure(err, _details)) => {
62                if err.code == rusqlite::ErrorCode::ConstraintViolation {
63                    Collection {
64                        name: name.to_string(),
65                    }
66                } else {
67                    panic!("something bad happened querying the database")
68                }
69            }
70            Err(_err) => {
71                panic!("something bad happened querying the database")
72            }
73        }
74    }
75
76    pub fn bulk_create(conn: &GraphConnection, names: &Vec<String>) -> Vec<Collection> {
77        let placeholders = names.iter().map(|_| "(?)").collect::<Vec<_>>().join(", ");
78        let q = format!("INSERT INTO collections (name) VALUES {placeholders} RETURNING *",);
79        let mut stmt = conn.prepare(&q).unwrap();
80        let rows = stmt
81            .query_map(params_from_iter(names), |row| {
82                Ok(Collection { name: row.get(0)? })
83            })
84            .unwrap();
85        rows.map(|row| row.unwrap()).collect()
86    }
87
88    pub fn get_block_groups(conn: &GraphConnection, collection_name: &str) -> Vec<BlockGroup> {
89        // Load all block groups that have the given collection_name
90        let mut stmt = conn
91            .prepare("SELECT * FROM block_groups WHERE collection_name = ?1 order by created_on;")
92            .unwrap();
93        let block_group_iter = stmt
94            .query_map([collection_name], |row| Ok(BlockGroup::process_row(row)))
95            .unwrap();
96        block_group_iter.map(|bg| bg.unwrap()).collect()
97    }
98
99    pub fn delete_by_name(conn: &GraphConnection, name: &str) {
100        let mut stmt = conn
101            .prepare("delete from collections where name = ?1")
102            .unwrap();
103        stmt.execute([name]).unwrap();
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use capnp::message::TypedBuilder;
110
111    use super::*;
112    use crate::test_helpers::get_connection;
113
114    #[test]
115    pub fn test_serialization() {
116        let c = Collection {
117            name: "test".to_string(),
118        };
119        let mut message = TypedBuilder::<collection::Owned>::new_default();
120        let mut root = message.init_root();
121        c.write_capnp(&mut root);
122
123        let d = Collection::read_capnp(root.into_reader());
124        assert_eq!(c, d);
125    }
126
127    #[test]
128    pub fn test_delete_by_name() {
129        let conn = &get_connection(None).unwrap();
130        let collection1 = Collection::create(conn, "test1");
131        let collection2 = Collection::create(conn, "test2");
132
133        Collection::delete_by_name(conn, &collection1.name);
134
135        // Verify only the correct one was deleted
136        assert!(!Collection::exists(conn, &collection1.name));
137        assert!(Collection::exists(conn, &collection2.name));
138    }
139}
gen_models/collection.rs

gen_models/
collection.rs
