geekorm_core/builder/
table.rs

1use serde::{Deserialize, Serialize};
2use std::fmt::Display;
3
4use crate::{Columns, QueryBuilder, ToSqlite, Values};
5
6/// The Table struct for defining a table
7#[derive(Debug, Clone, Default, Serialize, Deserialize)]
8pub struct Table {
9    /// Name of the table
10    pub name: String,
11    /// Columns in the table
12    pub columns: Columns,
13    /// Database name (optional)
14    ///
15    /// This is used to support multiple databases in the same project
16    #[serde(default, skip_serializing_if = "Option::is_none")]
17    pub database: Option<String>,
18}
19
20impl Table {
21    /// Function to check if a column name is valid
22    pub fn is_valid_column(&self, column: &str) -> bool {
23        if let Some((table, column)) = column.split_once('.') {
24            if table != self.name {
25                return false;
26            }
27            self.columns.is_valid_column(column)
28        } else {
29            self.columns.is_valid_column(column)
30        }
31    }
32
33    /// Get the name of the primary key column
34    pub fn get_primary_key(&self) -> String {
35        self.columns
36            .columns
37            .iter()
38            .find(|col| col.column_type.is_primary_key())
39            .map(|col| col.name.clone())
40            .unwrap_or_else(|| String::from("id"))
41    }
42
43    /// Get the foreign key by table name
44    pub fn get_foreign_key(&self, table_name: String) -> &crate::Column {
45        for column in self.columns.get_foreign_keys() {
46            if column.column_type.is_foreign_key_table(&table_name) {
47                return column;
48            }
49        }
50        panic!("No foreign key found for column: {}", table_name);
51    }
52
53    /// Get the full name of a column (table.column)
54    pub fn get_fullname(&self, column: &str) -> Result<String, crate::Error> {
55        let column = self.columns.get(column).ok_or_else(|| {
56            crate::Error::ColumnNotFound(self.name.to_string(), column.to_string())
57        })?;
58        let name = if column.alias.is_empty() {
59            column.name.clone()
60        } else {
61            column.alias.clone()
62        };
63        Ok(format!("{}.{}", self.name, name))
64    }
65
66    /// Get dependencies for the table
67    ///
68    /// This is a list of tables that the table depends on
69    pub fn get_dependencies(&self) -> Vec<String> {
70        let mut dependencies = Vec::new();
71        for column in &self.columns.columns {
72            if let Some(ftable) = column.column_type.foreign_key_table_name() {
73                dependencies.push(ftable);
74            }
75        }
76        dependencies
77    }
78}
79
80/// Implement the `ToTokens` trait for the `Table` struct
81#[cfg(feature = "migrations")]
82impl quote::ToTokens for Table {
83    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
84        let name = &self.name;
85        let columns = &self.columns;
86        let database = self
87            .database
88            .clone()
89            .unwrap_or_else(|| "Database".to_string());
90
91        tokens.extend(quote::quote! {
92            geekorm::Table {
93                name: String::from(#name),
94                columns: #columns,
95                database: Some(String::from(#database)),
96            }
97        });
98    }
99}
100
101impl ToSqlite for Table {
102    fn on_create(&self, query: &QueryBuilder) -> Result<String, crate::Error> {
103        Ok(format!(
104            "CREATE TABLE IF NOT EXISTS {} {};",
105            self.name,
106            self.columns.on_create(query)?
107        ))
108    }
109
110    fn on_select(&self, qb: &QueryBuilder) -> Result<String, crate::Error> {
111        let mut full_query = String::new();
112
113        // Resolve the rest of the query, and append if necessary
114        let columns = self.columns.on_select(qb);
115
116        if let Ok(ref columns) = columns {
117            if qb.count {
118                // If the query is a count query, return the count query
119                full_query = String::from("SELECT COUNT(1)");
120            } else {
121                // Select selective columns
122                let mut select_columns: Vec<String> = Vec::new();
123
124                let scolumns: Vec<String> = if !qb.columns.is_empty() {
125                    qb.columns.clone()
126                } else {
127                    self.columns
128                        .columns
129                        .iter()
130                        .filter(|col| !col.skip)
131                        .map(|col| col.name.clone())
132                        .collect()
133                };
134
135                for column in scolumns {
136                    // TODO(geekmasher): Validate that the column exists in the table
137                    if qb.joins.is_empty() {
138                        // If the query does not join multiple tables, we can use the column name directly
139                        select_columns.push(column);
140                    } else {
141                        // We have to use the full column name
142                        if column.contains('.') {
143                            // Table.column
144                            select_columns.push(column);
145                        } else {
146                            // Lookup the column in the table
147                            let fullname = qb.table.get_fullname(&column)?;
148                            select_columns.push(fullname);
149                        }
150                    }
151                }
152                full_query = format!("SELECT {}", select_columns.join(", "));
153            }
154
155            // FROM {table}
156            full_query.push_str(" FROM ");
157            full_query.push_str(&self.name);
158
159            // JOIN
160            if !qb.joins.is_empty() {
161                full_query.push(' ');
162                full_query.push_str(qb.joins.on_select(qb)?.as_str());
163            }
164
165            // WHERE {where_clause} ORDER BY {order_by}
166            if !columns.is_empty() {
167                full_query.push(' ');
168                full_query.push_str(columns);
169            }
170
171            // LIMIT {limit} OFFSET {offset}
172            if let Some(limit) = qb.limit {
173                // TODO(geekmasher): Check offset
174                full_query.push_str(" LIMIT ");
175                full_query.push_str(&limit.to_string());
176                if let Some(offset) = qb.offset {
177                    full_query.push_str(" OFFSET ");
178                    full_query.push_str(&offset.to_string());
179                }
180            }
181
182            // End
183            full_query = full_query.trim().to_string();
184            full_query.push(';');
185        }
186        Ok(full_query)
187    }
188
189    fn on_insert(&self, query: &QueryBuilder) -> Result<(String, Values), crate::Error> {
190        let mut full_query = format!("INSERT INTO {} ", self.name);
191
192        let mut columns: Vec<String> = Vec::new();
193        let mut values: Vec<String> = Vec::new();
194        let mut parameters = Values::new();
195
196        for (cname, value) in query.values.values.iter() {
197            let column = query.table.columns.get(cname.as_str()).unwrap();
198
199            // Get the column (might be an alias)
200            let mut column_name = column.name.clone();
201            if !column.alias.is_empty() {
202                column_name = column.alias.to_string();
203            }
204
205            // Skip auto increment columns
206            if column.column_type.is_auto_increment() {
207                continue;
208            }
209
210            columns.push(column_name.clone());
211
212            // Add to Values
213            match value {
214                crate::Value::Identifier(_) | crate::Value::Text(_) | crate::Value::Json(_) => {
215                    // Security: String values should never be directly inserted into the query
216                    // This is to prevent SQL injection attacks
217                    values.push(String::from("?"));
218                    parameters.push(column_name, value.clone());
219                }
220                crate::Value::Blob(value) => {
221                    // Security: Blods should never be directly inserted into the query
222                    values.push(String::from("?"));
223                    parameters.push(column_name, value.clone());
224                }
225                crate::Value::Integer(value) => values.push(value.to_string()),
226                crate::Value::Boolean(value) => values.push(value.to_string()),
227                crate::Value::Null => values.push("NULL".to_string()),
228            }
229        }
230
231        // Generate the column names
232        full_query.push('(');
233        full_query.push_str(&columns.join(", "));
234        full_query.push(')');
235
236        // Generate values
237        full_query.push_str(" VALUES (");
238        full_query.push_str(&values.join(", "));
239        full_query.push(')');
240        full_query.push(';');
241
242        Ok((full_query, parameters))
243    }
244
245    fn on_update(&self, query: &QueryBuilder) -> Result<(String, Values), crate::Error> {
246        let mut full_query = format!("UPDATE {} SET ", self.name);
247
248        let mut columns: Vec<String> = Vec::new();
249        let mut parameters = Values::new();
250
251        for (cname, value) in query.values.values.iter() {
252            let column = query.table.columns.get(cname.as_str()).unwrap();
253
254            // Skip if primary key
255            if column.column_type.is_primary_key() || cname == "id" {
256                continue;
257            }
258            // Get the column (might be an alias)
259            let mut column_name = column.name.clone();
260            if !column.alias.is_empty() {
261                column_name = column.alias.to_string();
262            }
263
264            // Add to Values
265            match value {
266                crate::Value::Identifier(_)
267                | crate::Value::Text(_)
268                | crate::Value::Blob(_)
269                | crate::Value::Json(_) => {
270                    // Security: String values should never be directly inserted into the query
271                    // This is to prevent SQL injection attacks
272                    columns.push(format!("{} = ?", column_name));
273                    parameters.push(column_name, value.clone());
274                }
275                crate::Value::Integer(value) => {
276                    columns.push(format!("{} = {}", column_name, value))
277                }
278                crate::Value::Boolean(value) => {
279                    columns.push(format!("{} = {}", column_name, value))
280                }
281                crate::Value::Null => columns.push(format!("{} = NULL", column_name)),
282            }
283        }
284
285        // Generate the column names
286        full_query.push_str(&columns.join(", "));
287
288        // WHERE
289        // TODO(geekmasher): We only support updating by primary key
290        let primary_key_name = query.table.get_primary_key();
291        let primary_key = query.values.get(&primary_key_name).unwrap();
292        let where_clause = format!(" WHERE {} = {}", primary_key_name, primary_key);
293        full_query.push_str(&where_clause);
294        full_query.push(';');
295
296        Ok((full_query, parameters))
297    }
298
299    /// Function to delete a row from the table
300    ///
301    /// Only supports deleting by primary key
302    fn on_delete(&self, query: &QueryBuilder) -> Result<(String, Values), crate::Error> {
303        let mut full_query = format!("DELETE FROM {}", self.name);
304        let mut parameters = Values::new();
305
306        // Delete by primary key
307        let primary_key_name = self.get_primary_key();
308        let primary_key = query.values.get(&primary_key_name).unwrap();
309
310        parameters.push(primary_key_name.to_string(), primary_key.clone());
311
312        full_query.push_str(&format!(" WHERE {} = ?;", primary_key_name));
313
314        Ok((full_query, parameters))
315    }
316
317    #[cfg(feature = "migrations")]
318    fn on_alter(&self, query: &crate::AlterQuery) -> Result<String, crate::Error> {
319        for column in &self.columns.columns {
320            if column.name == query.column {
321                return column.clone().on_alter(query);
322            }
323        }
324
325        Err(crate::Error::ColumnNotFound(
326            self.name.clone(),
327            query.column.clone(),
328        ))
329    }
330}
331
332impl Display for Table {
333    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
334        write!(f, "Table('{}')", self.name)
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341
342    fn table() -> Table {
343        use crate::{Column, ColumnType, ColumnTypeOptions};
344
345        Table {
346            name: "Test".to_string(),
347            database: None,
348            columns: vec![
349                Column::new(
350                    "id".to_string(),
351                    ColumnType::Integer(ColumnTypeOptions::primary_key()),
352                ),
353                Column::new(
354                    "name".to_string(),
355                    ColumnType::Text(ColumnTypeOptions::default()),
356                ),
357            ]
358            .into(),
359        }
360    }
361
362    #[test]
363    fn test_table_to_sql() {
364        let table = table();
365
366        let query = crate::QueryBuilder::select().table(table.clone());
367        // Basic CREATE and SELECT
368        assert_eq!(
369            table.on_create(&query).unwrap(),
370            "CREATE TABLE IF NOT EXISTS Test (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT);"
371        );
372        assert_eq!(
373            table.on_select(&query).unwrap(),
374            "SELECT id, name FROM Test;"
375        );
376
377        let query = crate::QueryBuilder::select()
378            .table(table.clone())
379            .where_eq("name", "this");
380        assert_eq!(
381            table.on_select(&query).unwrap(),
382            "SELECT id, name FROM Test WHERE name = ?;"
383        );
384    }
385
386    #[test]
387    fn test_count() {
388        let table = table();
389
390        let query = crate::QueryBuilder::select().table(table.clone()).count();
391        assert_eq!(
392            table.on_select(&query).unwrap(),
393            "SELECT COUNT(1) FROM Test;"
394        );
395
396        let query = crate::QueryBuilder::select()
397            .table(table.clone())
398            .count()
399            .where_eq("name", "this");
400        assert_eq!(
401            table.on_select(&query).unwrap(),
402            "SELECT COUNT(1) FROM Test WHERE name = ?;"
403        );
404
405        let query = crate::QueryBuilder::select()
406            .table(table.clone())
407            .count()
408            .where_ne("name", "this");
409        assert_eq!(
410            table.on_select(&query).unwrap(),
411            "SELECT COUNT(1) FROM Test WHERE name != ?;"
412        );
413    }
414
415    #[test]
416    fn test_row_delete() {
417        let table = table();
418
419        let query = crate::QueryBuilder::delete()
420            .table(table.clone())
421            .where_eq("id", 1);
422        let (delete_query, _) = table.on_delete(&query).unwrap();
423
424        assert_eq!(delete_query, "DELETE FROM Test WHERE id = ?;");
425    }
426
427    #[test]
428    fn test_is_valid_column() {
429        let table = table();
430
431        assert!(table.is_valid_column("id"));
432        assert!(table.is_valid_column("name"));
433        assert!(!table.is_valid_column("name2"));
434        // Test with table name
435        assert!(table.is_valid_column("Test.name"));
436        assert!(!table.is_valid_column("Tests.name"));
437    }
438}
geekorm_core/builder/table.rs

geekorm_core/builder/
table.rs
