geekorm_core/builder/
columns.rs

1#[cfg(feature = "migrations")]
2use super::alter::{AlterMode, AlterQuery};
3use crate::{ColumnType, ToSqlite};
4use serde::{Deserialize, Serialize};
5
6/// A list of columns in a table
7#[derive(Debug, Clone, Default, Serialize, Deserialize)]
8pub struct Columns {
9    /// List of columns
10    pub columns: Vec<Column>,
11}
12
13impl Columns {
14    /// Create a new instance of Columns
15    pub fn new() -> Self {
16        Columns {
17            columns: Vec::new(),
18        }
19    }
20
21    /// Validate if a column exists
22    pub fn is_valid_column(&self, column: &str) -> bool {
23        for col in &self.columns {
24            if col.name == column {
25                return true;
26            }
27        }
28        false
29    }
30
31    /// Get the Primary Key column of a table
32    pub fn get_primary_key(&self) -> Option<Column> {
33        self.columns
34            .iter()
35            .find(|col| col.column_type.is_primary_key())
36            .cloned()
37    }
38
39    /// Get the Foreign Keys columns of a table
40    pub fn get_foreign_keys(&self) -> Vec<&Column> {
41        self.columns
42            .iter()
43            .filter(|col| matches!(col.column_type, ColumnType::ForeignKey(_)))
44            .collect()
45    }
46
47    /// Get a column by name
48    pub fn get(&self, column: &str) -> Option<&Column> {
49        self.columns
50            .iter()
51            .find(|col| col.name == column || col.alias == column)
52    }
53
54    /// Get the length of the columns
55    pub fn len(&self) -> usize {
56        self.columns.len()
57    }
58
59    /// Check if the columns is empty
60    pub fn is_empty(&self) -> bool {
61        self.columns.is_empty()
62    }
63}
64
65impl Iterator for Columns {
66    type Item = Column;
67
68    fn next(&mut self) -> Option<Self::Item> {
69        self.columns.pop()
70    }
71}
72
73impl From<Vec<Column>> for Columns {
74    fn from(columns: Vec<Column>) -> Self {
75        Columns { columns }
76    }
77}
78
79#[cfg(feature = "migrations")]
80impl quote::ToTokens for Columns {
81    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
82        let columns = &self.columns;
83        tokens.extend(quote::quote! {
84            geekorm::Columns {
85                columns: Vec::from([
86                    #(#columns),*
87                ])
88            }
89        });
90    }
91}
92
93impl ToSqlite for Columns {
94    fn on_create(&self, query: &crate::QueryBuilder) -> Result<String, crate::Error> {
95        let mut sql = Vec::new();
96        for column in &self.columns {
97            match column.on_create(query) {
98                Ok(col) => sql.push(col),
99                Err(crate::Error::ColumnSkipped) => {
100                    // Skip the column
101                    continue;
102                }
103                Err(e) => return Err(e),
104            };
105        }
106
107        for foreign_key in self.get_foreign_keys() {
108            let (ctable, ccolumn) = match &foreign_key.column_type {
109                ColumnType::ForeignKey(opts) => {
110                    let (ctable, ccolumn) = opts
111                        .foreign_key
112                        .split_once('.')
113                        .expect("Invalid foreign key");
114                    (ctable, ccolumn)
115                }
116                _ => unreachable!(),
117            };
118
119            sql.push(format!(
120                "FOREIGN KEY ({parent}) REFERENCES {child}({child_column})",
121                parent = foreign_key.name,
122                child = ctable,
123                child_column = ccolumn
124            ));
125        }
126
127        Ok(format!("({})", sql.join(", ")))
128    }
129
130    fn on_select(&self, query: &crate::QueryBuilder) -> Result<String, crate::Error> {
131        let mut full_query = String::new();
132
133        // Support for WHERE
134        if !query.where_clause.is_empty() {
135            full_query.push_str("WHERE ");
136            for column in &query.where_clause {
137                full_query.push_str(column);
138                full_query.push(' ');
139            }
140        }
141        // Support for ORDER BY
142        let mut order_by = Vec::new();
143        if !query.order_by.is_empty() {
144            for (column, order) in &query.order_by {
145                // TODO(geekmasher): Validate that the column exists in the table
146                order_by.push(format!("{} {}", column, order.to_sqlite()));
147            }
148
149            full_query += format!("ORDER BY {}", order_by.join(", ")).as_str();
150        }
151        Ok(full_query)
152    }
153}
154
155/// A column in a table
156#[derive(Debug, Clone, Serialize, Deserialize)]
157pub struct Column {
158    /// Name of the column
159    pub name: String,
160    /// Type of the column (e.g. TEXT, INTEGER, etc)
161    pub column_type: ColumnType,
162
163    /// Alias for the column
164    pub alias: String,
165    /// Metadata for the column
166    pub skip: bool,
167}
168
169impl Column {
170    /// Create a new instance of Column
171    pub fn new(name: String, column_type: ColumnType) -> Self {
172        Column {
173            name,
174            column_type,
175            alias: String::new(),
176            skip: false,
177        }
178    }
179
180    /// Check if the column is a primary key
181    pub fn is_primary_key(&self) -> bool {
182        self.column_type.is_primary_key()
183    }
184
185    /// Check if the column is nullable
186    pub fn is_not_null(&self) -> bool {
187        self.column_type.is_not_null()
188    }
189    /// If the column unique
190    pub fn is_unique(&self) -> bool {
191        self.column_type.is_unique()
192    }
193}
194
195impl Default for Column {
196    fn default() -> Self {
197        Column {
198            name: String::new(),
199            column_type: ColumnType::Text(Default::default()),
200            alias: String::new(),
201            skip: false,
202        }
203    }
204}
205
206#[cfg(feature = "migrations")]
207impl quote::ToTokens for Column {
208    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
209        let name = &self.name;
210        let coltype = &self.column_type;
211        let alias = &self.alias;
212        let skip = &self.skip;
213
214        tokens.extend(quote::quote! {
215            geekorm::Column {
216                name: String::from(#name),
217                column_type: #coltype,
218                alias: String::from(#alias),
219                skip: #skip,
220            }
221        });
222    }
223}
224
225impl ToSqlite for Column {
226    fn on_create(&self, query: &crate::QueryBuilder) -> Result<String, crate::Error> {
227        if self.skip {
228            return Err(crate::Error::ColumnSkipped);
229        }
230
231        let name = if !&self.alias.is_empty() {
232            self.alias.clone()
233        } else {
234            self.name.clone()
235        };
236        Ok(format!("{} {}", name, self.column_type.on_create(query)?))
237    }
238
239    #[cfg(feature = "migrations")]
240    fn on_alter(&self, query: &AlterQuery) -> Result<String, crate::Error> {
241        Ok(match query.mode {
242            AlterMode::AddTable => {
243                format!("ALTER TABLE {} ADD COLUMN {};", query.table, query.column)
244            }
245            AlterMode::RenameTable => {
246                format!("ALTER TABLE {} RENAME TO {};", query.table, query.column)
247            }
248            AlterMode::DropTable => {
249                format!("ALTER TABLE {} DROP COLUMN {};", query.table, query.column)
250            }
251            AlterMode::AddColumn => {
252                format!(
253                    "ALTER TABLE {} ADD COLUMN {} {};",
254                    query.table,
255                    query.column,
256                    self.column_type.on_alter(query)?
257                )
258            }
259            AlterMode::RenameColumn => {
260                format!(
261                    "ALTER TABLE {} RENAME COLUMN {} TO {};",
262                    query.table,
263                    query.column,
264                    query.rename.as_ref().unwrap_or(&query.column)
265                )
266            }
267            AlterMode::DropColumn => {
268                format!("ALTER TABLE {} DROP COLUMN {};", query.table, query.column)
269            }
270            AlterMode::Skip => {
271                if query.column.is_empty() {
272                    format!("-- Skipping {} this migration", query.table)
273                } else {
274                    format!(
275                        "-- Skipping {}.{} this migration",
276                        query.table, query.column
277                    )
278                }
279            }
280        })
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287    use crate::ColumnTypeOptions;
288
289    fn create_table() -> crate::Table {
290        crate::Table {
291            name: String::from("users"),
292            database: None,
293            columns: Columns::from(vec![
294                Column::new(
295                    String::from("user_id"),
296                    ColumnType::Integer(ColumnTypeOptions::default()),
297                ),
298                Column::new(
299                    String::from("name"),
300                    ColumnType::Text(ColumnTypeOptions::default()),
301                ),
302                Column::new(
303                    String::from("image_id"),
304                    ColumnType::ForeignKey(ColumnTypeOptions {
305                        foreign_key: String::from("images.id"),
306                        ..Default::default()
307                    }),
308                ),
309            ]),
310        }
311    }
312
313    #[test]
314    fn test_column_to_sql() {
315        use super::*;
316        let query = crate::QueryBuilder::default();
317        let column = Column::new(
318            String::from("name"),
319            ColumnType::Text(ColumnTypeOptions::default()),
320        );
321        assert_eq!(column.on_create(&query).unwrap(), "name TEXT");
322
323        let column = Column::new(
324            String::from("age"),
325            ColumnType::Integer(ColumnTypeOptions::default()),
326        );
327        assert_eq!(column.on_create(&query).unwrap(), "age INTEGER");
328
329        // Test renaming the column
330        let column = Column {
331            name: String::from("id"),
332            column_type: ColumnType::Integer(ColumnTypeOptions::default()),
333            alias: String::from("user_id"),
334            ..Default::default()
335        };
336        assert_eq!(column.on_create(&query).unwrap(), "user_id INTEGER");
337    }
338
339    #[test]
340    fn test_foreign_key_to_sql() {
341        let query = crate::QueryBuilder::new().table(create_table());
342
343        let columns = query.table.columns.on_create(&query).unwrap();
344
345        assert_eq!(
346            columns,
347            "(user_id INTEGER, name TEXT, image_id INTEGER, FOREIGN KEY (image_id) REFERENCES images(id))"
348        );
349    }
350
351    #[test]
352    fn test_alter_to_sql() {
353        let query = crate::AlterQuery::new(AlterMode::AddColumn, "Table", "colname");
354
355        let column = Column::new(
356            String::from("name"),
357            ColumnType::Text(ColumnTypeOptions::default()),
358        );
359        assert_eq!(
360            column.on_alter(&query).unwrap(),
361            "ALTER TABLE Table ADD COLUMN colname TEXT;"
362        );
363        let column = Column::new(
364            String::from("name"),
365            ColumnType::Text(ColumnTypeOptions {
366                not_null: true,
367                ..Default::default()
368            }),
369        );
370        assert_eq!(
371            column.on_alter(&query).unwrap(),
372            "ALTER TABLE Table ADD COLUMN colname TEXT NOT NULL DEFAULT '';"
373        );
374    }
375}