Skip to main content

aegis_client/
query.rs

1//! Aegis Client Query Builder
2//!
3//! Type-safe query building for Aegis database.
4//!
5//! @version 0.1.0
6//! @author AutomataNexus Development Team
7
8use crate::result::Value;
9use serde::{Deserialize, Serialize};
10
11// =============================================================================
12// Query
13// =============================================================================
14
15/// A prepared query.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct Query {
18    pub sql: String,
19    pub params: Vec<Value>,
20}
21
22impl Query {
23    /// Create a new query.
24    pub fn new(sql: impl Into<String>) -> Self {
25        Self {
26            sql: sql.into(),
27            params: Vec::new(),
28        }
29    }
30
31    /// Add a parameter.
32    pub fn param(mut self, value: impl Into<Value>) -> Self {
33        self.params.push(value.into());
34        self
35    }
36
37    /// Add multiple parameters.
38    pub fn params(mut self, values: Vec<Value>) -> Self {
39        self.params.extend(values);
40        self
41    }
42}
43
44// =============================================================================
45// Query Builder
46// =============================================================================
47
48/// Builder for SQL queries.
49#[derive(Debug, Clone, Default)]
50pub struct QueryBuilder {
51    operation: QueryOperation,
52    table: Option<String>,
53    columns: Vec<String>,
54    values: Vec<Vec<Value>>,
55    conditions: Vec<Condition>,
56    order_by: Vec<OrderBy>,
57    limit: Option<u64>,
58    offset: Option<u64>,
59    joins: Vec<Join>,
60    group_by: Vec<String>,
61    having: Vec<Condition>,
62}
63
64#[derive(Debug, Clone, Default)]
65enum QueryOperation {
66    #[default]
67    Select,
68    Insert,
69    Update,
70    Delete,
71}
72
73impl QueryBuilder {
74    /// Create a new query builder.
75    pub fn new() -> Self {
76        Self::default()
77    }
78
79    /// Start a SELECT query.
80    pub fn select(mut self, columns: &[&str]) -> Self {
81        self.operation = QueryOperation::Select;
82        self.columns = columns.iter().map(|s| s.to_string()).collect();
83        self
84    }
85
86    /// Start an INSERT query.
87    pub fn insert_into(mut self, table: &str) -> Self {
88        self.operation = QueryOperation::Insert;
89        self.table = Some(table.to_string());
90        self
91    }
92
93    /// Start an UPDATE query.
94    pub fn update(mut self, table: &str) -> Self {
95        self.operation = QueryOperation::Update;
96        self.table = Some(table.to_string());
97        self
98    }
99
100    /// Start a DELETE query.
101    pub fn delete_from(mut self, table: &str) -> Self {
102        self.operation = QueryOperation::Delete;
103        self.table = Some(table.to_string());
104        self
105    }
106
107    /// Set the table for SELECT.
108    pub fn from(mut self, table: &str) -> Self {
109        self.table = Some(table.to_string());
110        self
111    }
112
113    /// Set columns for INSERT.
114    pub fn columns(mut self, columns: &[&str]) -> Self {
115        self.columns = columns.iter().map(|s| s.to_string()).collect();
116        self
117    }
118
119    /// Add values for INSERT.
120    pub fn values(mut self, values: Vec<Value>) -> Self {
121        self.values.push(values);
122        self
123    }
124
125    /// Add a WHERE condition.
126    pub fn where_eq(mut self, column: &str, value: impl Into<Value>) -> Self {
127        self.conditions
128            .push(Condition::Eq(column.to_string(), value.into()));
129        self
130    }
131
132    /// Add a WHERE != condition.
133    pub fn where_ne(mut self, column: &str, value: impl Into<Value>) -> Self {
134        self.conditions
135            .push(Condition::Ne(column.to_string(), value.into()));
136        self
137    }
138
139    /// Add a WHERE > condition.
140    pub fn where_gt(mut self, column: &str, value: impl Into<Value>) -> Self {
141        self.conditions
142            .push(Condition::Gt(column.to_string(), value.into()));
143        self
144    }
145
146    /// Add a WHERE >= condition.
147    pub fn where_gte(mut self, column: &str, value: impl Into<Value>) -> Self {
148        self.conditions
149            .push(Condition::Gte(column.to_string(), value.into()));
150        self
151    }
152
153    /// Add a WHERE < condition.
154    pub fn where_lt(mut self, column: &str, value: impl Into<Value>) -> Self {
155        self.conditions
156            .push(Condition::Lt(column.to_string(), value.into()));
157        self
158    }
159
160    /// Add a WHERE <= condition.
161    pub fn where_lte(mut self, column: &str, value: impl Into<Value>) -> Self {
162        self.conditions
163            .push(Condition::Lte(column.to_string(), value.into()));
164        self
165    }
166
167    /// Add a WHERE LIKE condition.
168    pub fn where_like(mut self, column: &str, pattern: &str) -> Self {
169        self.conditions
170            .push(Condition::Like(column.to_string(), pattern.to_string()));
171        self
172    }
173
174    /// Add a WHERE IN condition.
175    pub fn where_in(mut self, column: &str, values: Vec<Value>) -> Self {
176        self.conditions
177            .push(Condition::In(column.to_string(), values));
178        self
179    }
180
181    /// Add a WHERE IS NULL condition.
182    pub fn where_null(mut self, column: &str) -> Self {
183        self.conditions.push(Condition::IsNull(column.to_string()));
184        self
185    }
186
187    /// Add a WHERE IS NOT NULL condition.
188    pub fn where_not_null(mut self, column: &str) -> Self {
189        self.conditions
190            .push(Condition::IsNotNull(column.to_string()));
191        self
192    }
193
194    /// Add a JOIN clause.
195    pub fn join(mut self, table: &str, on: &str) -> Self {
196        self.joins.push(Join {
197            join_type: JoinType::Inner,
198            table: table.to_string(),
199            on: on.to_string(),
200        });
201        self
202    }
203
204    /// Add a LEFT JOIN clause.
205    pub fn left_join(mut self, table: &str, on: &str) -> Self {
206        self.joins.push(Join {
207            join_type: JoinType::Left,
208            table: table.to_string(),
209            on: on.to_string(),
210        });
211        self
212    }
213
214    /// Add a RIGHT JOIN clause.
215    pub fn right_join(mut self, table: &str, on: &str) -> Self {
216        self.joins.push(Join {
217            join_type: JoinType::Right,
218            table: table.to_string(),
219            on: on.to_string(),
220        });
221        self
222    }
223
224    /// Add a FULL OUTER JOIN clause.
225    pub fn full_join(mut self, table: &str, on: &str) -> Self {
226        self.joins.push(Join {
227            join_type: JoinType::Full,
228            table: table.to_string(),
229            on: on.to_string(),
230        });
231        self
232    }
233
234    /// Add an ORDER BY clause.
235    pub fn order_by(mut self, column: &str, direction: OrderDirection) -> Self {
236        self.order_by.push(OrderBy {
237            column: column.to_string(),
238            direction,
239        });
240        self
241    }
242
243    /// Add ascending ORDER BY.
244    pub fn order_by_asc(self, column: &str) -> Self {
245        self.order_by(column, OrderDirection::Asc)
246    }
247
248    /// Add descending ORDER BY.
249    pub fn order_by_desc(self, column: &str) -> Self {
250        self.order_by(column, OrderDirection::Desc)
251    }
252
253    /// Add a GROUP BY clause.
254    pub fn group_by(mut self, columns: &[&str]) -> Self {
255        self.group_by = columns.iter().map(|s| s.to_string()).collect();
256        self
257    }
258
259    /// Add a HAVING condition.
260    pub fn having(mut self, condition: Condition) -> Self {
261        self.having.push(condition);
262        self
263    }
264
265    /// Set the LIMIT.
266    pub fn limit(mut self, limit: u64) -> Self {
267        self.limit = Some(limit);
268        self
269    }
270
271    /// Set the OFFSET.
272    pub fn offset(mut self, offset: u64) -> Self {
273        self.offset = Some(offset);
274        self
275    }
276
277    /// Set values for UPDATE.
278    pub fn set(mut self, column: &str, value: impl Into<Value>) -> Self {
279        self.columns.push(column.to_string());
280        if self.values.is_empty() {
281            self.values.push(Vec::new());
282        }
283        self.values[0].push(value.into());
284        self
285    }
286
287    /// Build the query.
288    pub fn build(self) -> Query {
289        let (sql, params) = match self.operation {
290            QueryOperation::Select => self.build_select(),
291            QueryOperation::Insert => self.build_insert(),
292            QueryOperation::Update => self.build_update(),
293            QueryOperation::Delete => self.build_delete(),
294        };
295
296        Query { sql, params }
297    }
298
299    fn build_select(&self) -> (String, Vec<Value>) {
300        let mut sql = String::from("SELECT ");
301        let mut params = Vec::new();
302
303        if self.columns.is_empty() {
304            sql.push('*');
305        } else {
306            sql.push_str(&self.columns.join(", "));
307        }
308
309        if let Some(ref table) = self.table {
310            sql.push_str(" FROM ");
311            sql.push_str(table);
312        }
313
314        for join in &self.joins {
315            sql.push_str(&format!(
316                " {} JOIN {} ON {}",
317                join.join_type.as_str(),
318                join.table,
319                join.on
320            ));
321        }
322
323        self.append_where(&mut sql, &mut params);
324        self.append_group_by(&mut sql);
325        self.append_having(&mut sql, &mut params);
326        self.append_order_by(&mut sql);
327        self.append_limit_offset(&mut sql);
328
329        (sql, params)
330    }
331
332    fn build_insert(&self) -> (String, Vec<Value>) {
333        let mut sql = String::from("INSERT INTO ");
334        let mut params = Vec::new();
335
336        if let Some(ref table) = self.table {
337            sql.push_str(table);
338        }
339
340        if !self.columns.is_empty() {
341            sql.push_str(" (");
342            sql.push_str(&self.columns.join(", "));
343            sql.push(')');
344        }
345
346        sql.push_str(" VALUES ");
347
348        let value_groups: Vec<String> = self
349            .values
350            .iter()
351            .map(|row| {
352                let placeholders: Vec<String> = row
353                    .iter()
354                    .map(|v| {
355                        params.push(v.clone());
356                        format!("${}", params.len())
357                    })
358                    .collect();
359                format!("({})", placeholders.join(", "))
360            })
361            .collect();
362
363        sql.push_str(&value_groups.join(", "));
364
365        (sql, params)
366    }
367
368    fn build_update(&self) -> (String, Vec<Value>) {
369        let mut sql = String::from("UPDATE ");
370        let mut params = Vec::new();
371
372        if let Some(ref table) = self.table {
373            sql.push_str(table);
374        }
375
376        sql.push_str(" SET ");
377
378        let sets: Vec<String> = self
379            .columns
380            .iter()
381            .zip(self.values.first().unwrap_or(&Vec::new()).iter())
382            .map(|(col, val)| {
383                params.push(val.clone());
384                format!("{} = ${}", col, params.len())
385            })
386            .collect();
387
388        sql.push_str(&sets.join(", "));
389
390        self.append_where(&mut sql, &mut params);
391
392        (sql, params)
393    }
394
395    fn build_delete(&self) -> (String, Vec<Value>) {
396        let mut sql = String::from("DELETE FROM ");
397        let mut params = Vec::new();
398
399        if let Some(ref table) = self.table {
400            sql.push_str(table);
401        }
402
403        self.append_where(&mut sql, &mut params);
404
405        (sql, params)
406    }
407
408    fn append_where(&self, sql: &mut String, params: &mut Vec<Value>) {
409        if self.conditions.is_empty() {
410            return;
411        }
412
413        sql.push_str(" WHERE ");
414
415        let conditions: Vec<String> = self.conditions.iter().map(|c| c.to_sql(params)).collect();
416
417        sql.push_str(&conditions.join(" AND "));
418    }
419
420    fn append_group_by(&self, sql: &mut String) {
421        if !self.group_by.is_empty() {
422            sql.push_str(" GROUP BY ");
423            sql.push_str(&self.group_by.join(", "));
424        }
425    }
426
427    fn append_having(&self, sql: &mut String, params: &mut Vec<Value>) {
428        if !self.having.is_empty() {
429            sql.push_str(" HAVING ");
430
431            let conditions: Vec<String> = self.having.iter().map(|c| c.to_sql(params)).collect();
432
433            sql.push_str(&conditions.join(" AND "));
434        }
435    }
436
437    fn append_order_by(&self, sql: &mut String) {
438        if !self.order_by.is_empty() {
439            sql.push_str(" ORDER BY ");
440
441            let orders: Vec<String> = self
442                .order_by
443                .iter()
444                .map(|o| format!("{} {}", o.column, o.direction.as_str()))
445                .collect();
446
447            sql.push_str(&orders.join(", "));
448        }
449    }
450
451    fn append_limit_offset(&self, sql: &mut String) {
452        if let Some(limit) = self.limit {
453            sql.push_str(&format!(" LIMIT {}", limit));
454        }
455
456        if let Some(offset) = self.offset {
457            sql.push_str(&format!(" OFFSET {}", offset));
458        }
459    }
460}
461
462// =============================================================================
463// Condition
464// =============================================================================
465
466/// A query condition.
467#[derive(Debug, Clone)]
468pub enum Condition {
469    Eq(String, Value),
470    Ne(String, Value),
471    Gt(String, Value),
472    Gte(String, Value),
473    Lt(String, Value),
474    Lte(String, Value),
475    Like(String, String),
476    In(String, Vec<Value>),
477    IsNull(String),
478    IsNotNull(String),
479    Raw(String),
480}
481
482impl Condition {
483    fn to_sql(&self, params: &mut Vec<Value>) -> String {
484        match self {
485            Self::Eq(col, val) => {
486                params.push(val.clone());
487                format!("{} = ${}", col, params.len())
488            }
489            Self::Ne(col, val) => {
490                params.push(val.clone());
491                format!("{} != ${}", col, params.len())
492            }
493            Self::Gt(col, val) => {
494                params.push(val.clone());
495                format!("{} > ${}", col, params.len())
496            }
497            Self::Gte(col, val) => {
498                params.push(val.clone());
499                format!("{} >= ${}", col, params.len())
500            }
501            Self::Lt(col, val) => {
502                params.push(val.clone());
503                format!("{} < ${}", col, params.len())
504            }
505            Self::Lte(col, val) => {
506                params.push(val.clone());
507                format!("{} <= ${}", col, params.len())
508            }
509            Self::Like(col, pattern) => {
510                params.push(Value::String(pattern.clone()));
511                format!("{} LIKE ${}", col, params.len())
512            }
513            Self::In(col, vals) => {
514                let placeholders: Vec<String> = vals
515                    .iter()
516                    .map(|v| {
517                        params.push(v.clone());
518                        format!("${}", params.len())
519                    })
520                    .collect();
521                format!("{} IN ({})", col, placeholders.join(", "))
522            }
523            Self::IsNull(col) => format!("{} IS NULL", col),
524            Self::IsNotNull(col) => format!("{} IS NOT NULL", col),
525            Self::Raw(sql) => sql.clone(),
526        }
527    }
528}
529
530// =============================================================================
531// Join
532// =============================================================================
533
534#[derive(Debug, Clone)]
535struct Join {
536    join_type: JoinType,
537    table: String,
538    on: String,
539}
540
541#[derive(Debug, Clone)]
542enum JoinType {
543    Inner,
544    Left,
545    Right,
546    Full,
547}
548
549impl JoinType {
550    fn as_str(&self) -> &'static str {
551        match self {
552            Self::Inner => "INNER",
553            Self::Left => "LEFT",
554            Self::Right => "RIGHT",
555            Self::Full => "FULL",
556        }
557    }
558}
559
560// =============================================================================
561// Order
562// =============================================================================
563
564#[derive(Debug, Clone)]
565struct OrderBy {
566    column: String,
567    direction: OrderDirection,
568}
569
570/// Order direction for ORDER BY.
571#[derive(Debug, Clone, Copy)]
572pub enum OrderDirection {
573    Asc,
574    Desc,
575}
576
577impl OrderDirection {
578    fn as_str(&self) -> &'static str {
579        match self {
580            Self::Asc => "ASC",
581            Self::Desc => "DESC",
582        }
583    }
584}
585
586// =============================================================================
587// Tests
588// =============================================================================
589
590#[cfg(test)]
591mod tests {
592    use super::*;
593
594    #[test]
595    fn test_simple_select() {
596        let query = QueryBuilder::new()
597            .select(&["id", "name"])
598            .from("users")
599            .build();
600
601        assert_eq!(query.sql, "SELECT id, name FROM users");
602    }
603
604    #[test]
605    fn test_select_with_where() {
606        let query = QueryBuilder::new()
607            .select(&["*"])
608            .from("users")
609            .where_eq("id", 1)
610            .build();
611
612        assert_eq!(query.sql, "SELECT * FROM users WHERE id = $1");
613        assert_eq!(query.params.len(), 1);
614    }
615
616    #[test]
617    fn test_select_with_order_limit() {
618        let query = QueryBuilder::new()
619            .select(&["*"])
620            .from("users")
621            .order_by_desc("created_at")
622            .limit(10)
623            .offset(20)
624            .build();
625
626        assert!(query.sql.contains("ORDER BY created_at DESC"));
627        assert!(query.sql.contains("LIMIT 10"));
628        assert!(query.sql.contains("OFFSET 20"));
629    }
630
631    #[test]
632    fn test_insert() {
633        let query = QueryBuilder::new()
634            .insert_into("users")
635            .columns(&["name", "email"])
636            .values(vec![
637                Value::String("Alice".to_string()),
638                Value::String("alice@example.com".to_string()),
639            ])
640            .build();
641
642        assert!(query.sql.starts_with("INSERT INTO users"));
643        assert!(query.sql.contains("(name, email)"));
644        assert!(query.sql.contains("VALUES ($1, $2)"));
645    }
646
647    #[test]
648    fn test_update() {
649        let query = QueryBuilder::new()
650            .update("users")
651            .set("name", "Bob")
652            .set("age", 30)
653            .where_eq("id", 1)
654            .build();
655
656        assert!(query.sql.starts_with("UPDATE users SET"));
657        assert!(query.sql.contains("name = $1"));
658        assert!(query.sql.contains("WHERE id = $3"));
659    }
660
661    #[test]
662    fn test_delete() {
663        let query = QueryBuilder::new()
664            .delete_from("users")
665            .where_eq("id", 1)
666            .build();
667
668        assert_eq!(query.sql, "DELETE FROM users WHERE id = $1");
669    }
670
671    #[test]
672    fn test_join() {
673        let query = QueryBuilder::new()
674            .select(&["users.name", "orders.total"])
675            .from("users")
676            .join("orders", "users.id = orders.user_id")
677            .build();
678
679        assert!(query
680            .sql
681            .contains("INNER JOIN orders ON users.id = orders.user_id"));
682    }
683
684    #[test]
685    fn test_where_in() {
686        let query = QueryBuilder::new()
687            .select(&["*"])
688            .from("users")
689            .where_in("id", vec![Value::Int(1), Value::Int(2), Value::Int(3)])
690            .build();
691
692        assert!(query.sql.contains("id IN ($1, $2, $3)"));
693        assert_eq!(query.params.len(), 3);
694    }
695
696    #[test]
697    fn test_where_null() {
698        let query = QueryBuilder::new()
699            .select(&["*"])
700            .from("users")
701            .where_null("deleted_at")
702            .build();
703
704        assert!(query.sql.contains("deleted_at IS NULL"));
705    }
706}