Skip to main content

ri/database/orm/
mod.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18//! # ORM Module
19//!
20//! This module provides ORM-like database operations for Ri.
21
22use crate::core::{RiResult, RiError};
23use serde::{Deserialize, Serialize};
24use std::collections::HashMap as FxHashMap;
25use crate::database::RiDatabase;
26
27pub mod repository;
28
29/// Validates a SQL identifier (table name, column name) to prevent SQL injection.
30///
31/// Only allows alphanumeric characters and underscores. Must start with a letter or underscore.
32/// Maximum length is 128 characters to prevent buffer overflow attacks.
33fn validate_identifier(identifier: &str) -> RiResult<()> {
34    if identifier.is_empty() {
35        return Err(RiError::Other("Identifier cannot be empty".to_string()));
36    }
37    
38    if identifier.len() > 128 {
39        return Err(RiError::Other("Identifier too long (max 128 characters)".to_string()));
40    }
41    
42    let chars: Vec<char> = identifier.chars().collect();
43    
44    if !chars[0].is_ascii_alphabetic() && chars[0] != '_' {
45        return Err(RiError::Other(
46            "Identifier must start with a letter or underscore".to_string()
47        ));
48    }
49    
50    for c in &chars {
51        if !c.is_ascii_alphanumeric() && *c != '_' {
52            return Err(RiError::Other(
53                "Identifier can only contain alphanumeric characters and underscores".to_string()
54            ));
55        }
56    }
57    
58    let lower = identifier.to_lowercase();
59    let dangerous_keywords = [
60        "drop", "delete", "truncate", "update", "insert", "alter", "create",
61        "exec", "execute", "xp_", "sp_", "union", "select", "from", "where",
62        "having", "group", "order", "limit", "offset", "join", "inner", "outer",
63        "left", "right", "full", "cross", "natural", "using", "on", "as",
64    ];
65    
66    for keyword in &dangerous_keywords {
67        if lower == *keyword {
68            return Err(RiError::Other(
69                format!("Identifier cannot be a reserved SQL keyword: {}", identifier)
70            ));
71        }
72    }
73    
74    Ok(())
75}
76
77/// Validates multiple identifiers at once.
78fn validate_identifiers(identifiers: &[&str]) -> RiResult<()> {
79    for id in identifiers {
80        validate_identifier(id)?;
81    }
82    Ok(())
83}
84
85pub use repository::{RiORMSimpleRepository, RiORMCrudRepository, RiORMRepository};
86
87#[cfg(feature = "pyo3")]
88pub mod py_repository;
89
90#[cfg(feature = "pyo3")]
91pub use py_repository::RiPyORMRepository;
92
93#[derive(Debug, Clone, PartialEq, Eq, Hash)]
94#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
95pub struct ColumnDefinition {
96    pub name: String,
97    pub column_type: String,
98    pub is_primary_key: bool,
99    pub is_nullable: bool,
100    pub default_value: Option<String>,
101    pub is_unique: bool,
102    pub max_length: Option<usize>,
103}
104
105impl Default for ColumnDefinition {
106    fn default() -> Self {
107        Self {
108            name: String::new(),
109            column_type: "TEXT".to_string(),
110            is_primary_key: false,
111            is_nullable: true,
112            default_value: None,
113            is_unique: false,
114            max_length: None,
115        }
116    }
117}
118
119#[derive(Debug, Clone, PartialEq, Eq)]
120#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
121pub struct IndexDefinition {
122    pub name: String,
123    pub columns: Vec<String>,
124    pub is_unique: bool,
125    pub is_full_text: bool,
126}
127
128impl Default for IndexDefinition {
129    fn default() -> Self {
130        Self {
131            name: String::new(),
132            columns: Vec::new(),
133            is_unique: false,
134            is_full_text: false,
135        }
136    }
137}
138
139#[derive(Debug, Clone, PartialEq, Eq)]
140#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
141pub struct ForeignKeyDefinition {
142    pub name: String,
143    pub column: String,
144    pub referenced_table: String,
145    pub referenced_column: String,
146    pub on_delete: String,
147    pub on_update: String,
148}
149
150impl Default for ForeignKeyDefinition {
151    fn default() -> Self {
152        Self {
153            name: String::new(),
154            column: String::new(),
155            referenced_table: String::new(),
156            referenced_column: String::new(),
157            on_delete: "CASCADE".to_string(),
158            on_update: "CASCADE".to_string(),
159        }
160    }
161}
162
163#[derive(Debug, Clone)]
164#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
165pub struct TableDefinition {
166    pub table_name: String,
167    pub columns: FxHashMap<String, ColumnDefinition>,
168    pub primary_key: Vec<String>,
169    pub indexes: Vec<IndexDefinition>,
170    pub foreign_keys: Vec<ForeignKeyDefinition>,
171    pub engine: Option<String>,
172    pub charset: Option<String>,
173}
174
175impl TableDefinition {
176    pub fn new(table_name: &str) -> Self {
177        Self {
178            table_name: table_name.to_string(),
179            columns: FxHashMap::default(),
180            primary_key: Vec::new(),
181            indexes: Vec::new(),
182            foreign_keys: Vec::new(),
183            engine: None,
184            charset: None,
185        }
186    }
187
188    pub fn add_column(&mut self, column: ColumnDefinition) -> RiResult<()> {
189        validate_identifier(&column.name)?;
190        self.columns.insert(column.name.clone(), column);
191        Ok(())
192    }
193
194    pub fn set_primary_key(&mut self, columns: Vec<String>) -> RiResult<()> {
195        validate_identifiers(&columns.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
196        self.primary_key = columns;
197        Ok(())
198    }
199
200    pub fn add_index(&mut self, index: IndexDefinition) -> RiResult<()> {
201        validate_identifier(&index.name)?;
202        validate_identifiers(&index.columns.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
203        self.indexes.push(index);
204        Ok(())
205    }
206
207    pub fn add_foreign_key(&mut self, fk: ForeignKeyDefinition) -> RiResult<()> {
208        validate_identifier(&fk.name)?;
209        validate_identifier(&fk.column)?;
210        validate_identifier(&fk.referenced_table)?;
211        validate_identifier(&fk.referenced_column)?;
212        self.foreign_keys.push(fk);
213        Ok(())
214    }
215
216    pub fn get_create_sql(&self) -> RiResult<String> {
217        validate_identifier(&self.table_name)?;
218        
219        let mut sql = format!("CREATE TABLE IF NOT EXISTS {} (", self.table_name);
220
221        let mut column_defs = Vec::with_capacity(self.columns.len());
222
223        for (name, col) in &self.columns {
224            validate_identifier(name)?;
225            validate_identifier(&col.column_type)?;
226            
227            let mut def = format!("{} {}", name, col.column_type);
228
229            if !col.is_nullable {
230                def.push_str(" NOT NULL");
231            }
232
233            if col.is_primary_key {
234                def.push_str(" PRIMARY KEY");
235            }
236
237            if let Some(default) = &col.default_value {
238                if !default.starts_with('\'') && !default.chars().all(|c| c.is_numeric() || c == '.' || c == '-') {
239                    return Err(RiError::Other("Invalid default value format".to_string()));
240                }
241                def.push_str(&format!(" DEFAULT {}", default));
242            }
243
244            if col.is_unique {
245                def.push_str(" UNIQUE");
246            }
247
248            column_defs.push(def);
249        }
250
251        if !self.primary_key.is_empty() {
252            validate_identifiers(&self.primary_key.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
253            column_defs.push(format!("PRIMARY KEY ({})", self.primary_key.join(", ")));
254        }
255
256        sql.push_str(&column_defs.join(", "));
257
258        for fk in &self.foreign_keys {
259            validate_identifier(&fk.column)?;
260            validate_identifier(&fk.referenced_table)?;
261            validate_identifier(&fk.referenced_column)?;
262            
263            sql.push_str(&format!(
264                ", FOREIGN KEY ({}) REFERENCES {}({}) ON DELETE {} ON UPDATE {}",
265                fk.column, fk.referenced_table, fk.referenced_column, fk.on_delete, fk.on_update
266            ));
267        }
268
269        sql.push_str(")");
270
271        if let Some(engine) = &self.engine {
272            validate_identifier(engine)?;
273            sql.push_str(&format!(" ENGINE={}", engine));
274        }
275
276        if let Some(charset) = &self.charset {
277            validate_identifier(charset)?;
278            sql.push_str(&format!(" DEFAULT CHARSET={}", charset));
279        }
280
281        sql.push_str(";");
282
283        Ok(sql)
284    }
285}
286
287#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
288#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
289pub enum ComparisonOperator {
290    Equal,
291    NotEqual,
292    GreaterThan,
293    GreaterThanOrEqual,
294    LessThan,
295    LessThanOrEqual,
296    Like,
297    ILike,
298    In,
299    NotIn,
300    IsNull,
301    IsNotNull,
302    Between,
303}
304
305#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
306#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
307pub enum LogicalOperator {
308    And,
309    Or,
310    Not,
311}
312
313#[derive(Debug, Clone)]
314#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
315pub struct Criteria {
316    pub column: String,
317    pub operator: ComparisonOperator,
318    pub value: serde_json::Value,
319}
320
321impl Criteria {
322    pub fn new(column: &str, operator: ComparisonOperator, value: serde_json::Value) -> RiResult<Self> {
323        validate_identifier(column)?;
324        Ok(Self {
325            column: column.to_string(),
326            operator,
327            value,
328        })
329    }
330
331    pub fn to_sql(&self) -> RiResult<(String, Vec<serde_json::Value>)> {
332        validate_identifier(&self.column)?;
333        
334        let (op_str, value) = match self.operator {
335            ComparisonOperator::Equal => ("=".to_string(), vec![self.value.clone()]),
336            ComparisonOperator::NotEqual => ("!=".to_string(), vec![self.value.clone()]),
337            ComparisonOperator::GreaterThan => (">".to_string(), vec![self.value.clone()]),
338            ComparisonOperator::GreaterThanOrEqual => (">=".to_string(), vec![self.value.clone()]),
339            ComparisonOperator::LessThan => ("<".to_string(), vec![self.value.clone()]),
340            ComparisonOperator::LessThanOrEqual => ("<=".to_string(), vec![self.value.clone()]),
341            ComparisonOperator::Like => ("LIKE".to_string(), vec![self.value.clone()]),
342            ComparisonOperator::ILike => ("ILIKE".to_string(), vec![self.value.clone()]),
343            ComparisonOperator::In => {
344                if let serde_json::Value::Array(arr) = &self.value {
345                    let placeholders = (0..arr.len()).map(|_| "?").collect::<Vec<_>>().join(", ");
346                    (format!("IN ({})", placeholders), arr.clone())
347                } else {
348                    ("= ?".to_string(), vec![self.value.clone()])
349                }
350            }
351            ComparisonOperator::NotIn => {
352                if let serde_json::Value::Array(arr) = &self.value {
353                    let placeholders = (0..arr.len()).map(|_| "?").collect::<Vec<_>>().join(", ");
354                    (format!("NOT IN ({})", placeholders), arr.clone())
355                } else {
356                    ("!= ?".to_string(), vec![self.value.clone()])
357                }
358            }
359            ComparisonOperator::IsNull => ("IS NULL".to_string(), Vec::new()),
360            ComparisonOperator::IsNotNull => ("IS NOT NULL".to_string(), Vec::new()),
361            ComparisonOperator::Between => {
362                if let serde_json::Value::Array(arr) = &self.value {
363                    if arr.len() == 2 {
364                        ("BETWEEN ? AND ?".to_string(), arr.clone())
365                    } else {
366                        ("BETWEEN ? AND ?".to_string(), vec![self.value.clone()])
367                    }
368                } else {
369                    ("BETWEEN ? AND ?".to_string(), vec![self.value.clone()])
370                }
371            }
372        };
373
374        let placeholder = if matches!(self.operator, ComparisonOperator::IsNull | ComparisonOperator::IsNotNull) {
375            "".to_string()
376        } else {
377            "?".to_string()
378        };
379
380        Ok((format!("{} {} {}", self.column, op_str, placeholder), value))
381    }
382}
383
384#[derive(Debug, Clone, PartialEq, Eq)]
385#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
386pub struct SortOrder {
387    pub column: String,
388    pub ascending: bool,
389}
390
391impl SortOrder {
392    pub fn new(column: &str, ascending: bool) -> RiResult<Self> {
393        validate_identifier(column)?;
394        Ok(Self {
395            column: column.to_string(),
396            ascending,
397        })
398    }
399
400    pub fn asc(column: &str) -> RiResult<Self> {
401        Self::new(column, true)
402    }
403
404    pub fn desc(column: &str) -> RiResult<Self> {
405        Self::new(column, false)
406    }
407}
408
409#[derive(Debug, Clone)]
410#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
411pub struct Pagination {
412    pub page: u64,
413    pub page_size: u64,
414}
415
416impl Default for Pagination {
417    fn default() -> Self {
418        Self {
419            page: 1,
420            page_size: 20,
421        }
422    }
423}
424
425impl Pagination {
426    pub fn new(page: u64, page_size: u64) -> Self {
427        Self { page, page_size }
428    }
429
430    pub fn offset(&self) -> u64 {
431        (self.page - 1) * self.page_size
432    }
433
434    pub fn limit(&self) -> u64 {
435        self.page_size
436    }
437}
438
439#[derive(Debug, Clone)]
440#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
441pub struct QueryBuilder {
442    pub table_name: String,
443    pub criteria: Vec<Criteria>,
444    pub sort_orders: Vec<SortOrder>,
445    pub pagination: Option<Pagination>,
446    pub select_columns: Option<Vec<String>>,
447    pub group_by_columns: Option<Vec<String>>,
448    pub having_criteria: Vec<Criteria>,
449    pub distinct: bool,
450    pub joins: Vec<JoinClause>,
451}
452
453#[derive(Debug, Clone)]
454#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
455pub struct JoinClause {
456    pub join_type: JoinType,
457    pub table_name: String,
458    pub on_column: String,
459    pub referenced_column: String,
460}
461
462#[derive(Debug, Clone, PartialEq, Eq)]
463#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
464pub enum JoinType {
465    Inner,
466    Left,
467    Right,
468    Full,
469}
470
471impl QueryBuilder {
472    pub fn new(table_name: &str) -> RiResult<Self> {
473        validate_identifier(table_name)?;
474        Ok(Self {
475            table_name: table_name.to_string(),
476            criteria: Vec::new(),
477            sort_orders: Vec::new(),
478            pagination: None,
479            select_columns: None,
480            group_by_columns: None,
481            having_criteria: Vec::new(),
482            distinct: false,
483            joins: Vec::new(),
484        })
485    }
486
487    pub fn select(&mut self, columns: Vec<&str>) -> RiResult<&mut Self> {
488        validate_identifiers(&columns)?;
489        self.select_columns = Some(columns.iter().map(|s| s.to_string()).collect());
490        Ok(self)
491    }
492
493    pub fn where_criteria(&mut self, criteria: Criteria) -> &mut Self {
494        self.criteria.push(criteria);
495        self
496    }
497
498    pub fn and_where(&mut self, criteria: Criteria) -> &mut Self {
499        self.where_criteria(criteria)
500    }
501
502    pub fn or_where(&mut self, criteria: Criteria) -> &mut Self {
503        self.criteria.push(criteria);
504        self
505    }
506
507    pub fn order_by(&mut self, sort_order: SortOrder) -> &mut Self {
508        self.sort_orders.push(sort_order);
509        self
510    }
511
512    pub fn paginate(&mut self, page: u64, page_size: u64) -> &mut Self {
513        self.pagination = Some(Pagination::new(page, page_size));
514        self
515    }
516
517    pub fn distinct(&mut self) -> &mut Self {
518        self.distinct = true;
519        self
520    }
521
522    pub fn group_by(&mut self, columns: Vec<&str>) -> RiResult<&mut Self> {
523        validate_identifiers(&columns)?;
524        self.group_by_columns = Some(columns.iter().map(|s| s.to_string()).collect());
525        Ok(self)
526    }
527
528    pub fn inner_join(&mut self, table_name: &str, on_column: &str, referenced_column: &str) -> RiResult<&mut Self> {
529        validate_identifier(table_name)?;
530        validate_identifier(on_column)?;
531        validate_identifier(referenced_column)?;
532        self.joins.push(JoinClause {
533            join_type: JoinType::Inner,
534            table_name: table_name.to_string(),
535            on_column: on_column.to_string(),
536            referenced_column: referenced_column.to_string(),
537        });
538        Ok(self)
539    }
540
541    pub fn left_join(&mut self, table_name: &str, on_column: &str, referenced_column: &str) -> RiResult<&mut Self> {
542        validate_identifier(table_name)?;
543        validate_identifier(on_column)?;
544        validate_identifier(referenced_column)?;
545        self.joins.push(JoinClause {
546            join_type: JoinType::Left,
547            table_name: table_name.to_string(),
548            on_column: on_column.to_string(),
549            referenced_column: referenced_column.to_string(),
550        });
551        Ok(self)
552    }
553
554    pub fn build(&self) -> RiResult<(String, Vec<serde_json::Value>)> {
555        validate_identifier(&self.table_name)?;
556        
557        let mut sql = String::new();
558        let mut params = Vec::with_capacity(8);
559
560        sql.push_str("SELECT ");
561
562        if self.distinct {
563            sql.push_str("DISTINCT ");
564        }
565
566        if let Some(columns) = &self.select_columns {
567            validate_identifiers(&columns.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
568            sql.push_str(&columns.join(", "));
569        } else {
570            sql.push_str("*");
571        }
572
573        sql.push_str(&format!(" FROM {}", self.table_name));
574
575        for join in &self.joins {
576            validate_identifier(&join.table_name)?;
577            validate_identifier(&join.on_column)?;
578            validate_identifier(&join.referenced_column)?;
579            
580            let join_type = match join.join_type {
581                JoinType::Inner => "INNER JOIN",
582                JoinType::Left => "LEFT JOIN",
583                JoinType::Right => "RIGHT JOIN",
584                JoinType::Full => "FULL JOIN",
585            };
586            sql.push_str(&format!(
587                " {} {} ON {}.{} = {}.{}",
588                join_type, join.table_name, self.table_name, join.on_column, join.table_name, join.referenced_column
589            ));
590        }
591
592        if !self.criteria.is_empty() {
593            sql.push_str(" WHERE 1=1");
594            for criteria in &self.criteria {
595                sql.push_str(" AND ");
596                let (clause, values) = criteria.to_sql()?;
597                sql.push_str(&clause);
598                params.extend(values);
599            }
600        }
601
602        if let Some(group_by) = &self.group_by_columns {
603            validate_identifiers(&group_by.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
604            sql.push_str(&format!(" GROUP BY {}", group_by.join(", ")));
605        }
606
607        if !self.having_criteria.is_empty() {
608            sql.push_str(" HAVING 1=1");
609            for criteria in &self.having_criteria {
610                sql.push_str(" AND ");
611                let (clause, values) = criteria.to_sql()?;
612                sql.push_str(&clause);
613                params.extend(values);
614            }
615        }
616
617        if !self.sort_orders.is_empty() {
618            for o in &self.sort_orders {
619                validate_identifier(&o.column)?;
620            }
621            let orders: Vec<String> = self.sort_orders.iter()
622                .map(|o| format!("{} {}", o.column, if o.ascending { "ASC" } else { "DESC" }))
623                .collect();
624            sql.push_str(&format!(" ORDER BY {}", orders.join(", ")));
625        }
626
627        if let Some(pagination) = &self.pagination {
628            if pagination.page_size > 10000 {
629                return Err(RiError::Other("Page size too large (max 10000)".to_string()));
630            }
631            sql.push_str(&format!(" LIMIT {} OFFSET {}", pagination.limit(), pagination.offset()));
632        }
633
634        Ok((sql, params))
635    }
636}