bottle_orm/
database.rs

1//! # Database Module
2//!
3//! This module provides the core database connection and management functionality for Bottle ORM.
4//! It handles connection pooling, driver detection, table creation, and foreign key management
5//! across PostgreSQL, MySQL, and SQLite.
6
7// ============================================================================
8// External Crate Imports
9// ============================================================================
10
11use futures::future::BoxFuture;
12use heck::ToSnakeCase;
13use sqlx::{any::AnyArguments, AnyPool, Row, Arguments};
14use std::sync::Arc;
15
16// ============================================================================
17// Internal Crate Imports
18// ============================================================================
19
20use crate::{migration::Migrator, Error, Model, QueryBuilder};
21
22// ============================================================================
23// Database Driver Enum
24// ============================================================================
25
26/// Supported database drivers for Bottle ORM.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum Drivers {
29    /// PostgreSQL driver
30    Postgres,
31    /// MySQL driver
32    MySQL,
33    /// SQLite driver
34    SQLite,
35}
36
37// ============================================================================
38// Database Struct
39// ============================================================================
40
41/// The main entry point for Bottle ORM database operations.
42///
43/// `Database` manages a connection pool and provides methods for starting
44/// transactions, creating tables, and building queries for models.
45///
46/// It is designed to be thread-safe and easily shared across an application
47/// (internally uses an `Arc` for the connection pool).
48#[derive(Debug, Clone)]
49pub struct Database {
50    /// The underlying SQLx connection pool
51    pub(crate) pool: AnyPool,
52    /// The detected database driver
53    pub(crate) driver: Drivers,
54}
55
56// ============================================================================
57// Database Implementation
58// ============================================================================
59
60impl Database {
61    /// Creates a new DatabaseBuilder for configuring the connection.
62    pub fn builder() -> DatabaseBuilder {
63        DatabaseBuilder::new()
64    }
65
66    /// Connects to a database using the provided connection string.
67    ///
68    /// This is a convenience method that uses default builder settings.
69    ///
70    /// # Arguments
71    ///
72    /// * `url` - A database connection URL (e.g., "postgres://user:pass@localhost/db")
73    pub async fn connect(url: &str) -> Result<Self, Error> {
74        DatabaseBuilder::new().connect(url).await
75    }
76
77    /// Returns a new Migrator instance for managing schema changes.
78    pub fn migrator(&self) -> Migrator<'_> {
79        Migrator::new(self)
80    }
81
82    /// Starts building a query for the specified model.
83    ///
84    /// # Type Parameters
85    ///
86    /// * `T` - The Model type to query.
87    pub fn model<T: Model + Send + Sync + Unpin + crate::AnyImpl>(&self) -> QueryBuilder<T, Self> {
88        let active_columns = T::active_columns();
89        let mut columns: Vec<String> = Vec::with_capacity(active_columns.capacity());
90
91        for col in active_columns {
92            columns.push(col.strip_prefix("r#").unwrap_or(col).to_snake_case());
93        }
94
95        QueryBuilder::new(self.clone(), self.driver, T::table_name(), <T as Model>::columns(), columns)
96    }
97
98    /// Creates a raw SQL query builder.
99    pub fn raw<'a>(&self, sql: &'a str) -> RawQuery<'a, Self> {
100        RawQuery::new(self.clone(), sql)
101    }
102
103    /// Starts a new database transaction.
104    pub async fn begin(&self) -> Result<crate::transaction::Transaction<'_>, Error> {
105        let tx = self.pool.begin().await?;
106        Ok(crate::transaction::Transaction {
107            tx: Arc::new(tokio::sync::Mutex::new(Some(tx))),
108            driver: self.driver,
109        })
110    }
111
112    /// Checks if a table exists in the database.
113    pub async fn table_exists(&self, table_name: &str) -> Result<bool, Error> {
114        let table_name_snake = table_name.to_snake_case();
115        let query = match self.driver {
116            Drivers::Postgres => {
117                "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = $1 AND table_schema = 'public')".to_string()
118            }
119            Drivers::MySQL => {
120                "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = ? AND table_schema = DATABASE())".to_string()
121            }
122            Drivers::SQLite => {
123                "SELECT count(*) FROM sqlite_master WHERE type='table' AND name=?".to_string()
124            }
125        };
126
127        let row = sqlx::query(&query).bind(&table_name_snake).fetch_one(&self.pool).await?;
128
129        match self.driver {
130            Drivers::SQLite => {
131                let count: i64 = row.try_get(0)?;
132                Ok(count > 0)
133            }
134            _ => {
135                let exists: bool = row.try_get(0)?;
136                Ok(exists)
137            }
138        }
139    }
140
141    /// Creates a table based on the provided Model metadata.
142    pub async fn create_table<T: Model>(&self) -> Result<(), Error> {
143        let table_name = T::table_name().to_snake_case();
144        let columns = T::columns();
145
146        let mut query = format!("CREATE TABLE IF NOT EXISTS \"{}\" (", table_name);
147        let mut column_defs = Vec::new();
148        let mut indexes = Vec::new();
149
150        // Identify primary key columns
151        let pk_columns: Vec<String> = columns.iter()
152            .filter(|c| c.is_primary_key)
153            .map(|c| format!("\"{}\"", c.name.strip_prefix("r#").unwrap_or(c.name).to_snake_case()))
154            .collect();
155
156        for col in columns {
157            let col_name_clean = col.name.strip_prefix("r#").unwrap_or(col.name).to_snake_case();
158            let mut def = format!("\"{}\" {}", col_name_clean, col.sql_type);
159
160            // If it's a single primary key, we can keep it inline for simplicity
161            // If it's composite, we MUST define it as a table constraint
162            if col.is_primary_key && pk_columns.len() == 1 {
163                def.push_str(" PRIMARY KEY");
164            } else if !col.is_nullable || col.is_primary_key {
165                def.push_str(" NOT NULL");
166            }
167
168            if col.unique && !col.is_primary_key {
169                def.push_str(" UNIQUE");
170            }
171
172            if col.index && !col.is_primary_key && !col.unique {
173                indexes.push(format!(
174                    "CREATE INDEX IF NOT EXISTS \"idx_{}_{}\" ON \"{}\" (\"{}\")",
175                    table_name, col_name_clean, table_name, col_name_clean
176                ));
177            }
178
179            column_defs.push(def);
180        }
181
182        // Add composite primary key if multiple columns are specified
183        if pk_columns.len() > 1 {
184            column_defs.push(format!("PRIMARY KEY ({})", pk_columns.join(", ")));
185        }
186
187        query.push_str(&column_defs.join(", "));
188        query.push(')');
189
190        sqlx::query(&query).execute(&self.pool).await?;
191
192        for idx_query in indexes {
193            sqlx::query(&idx_query).execute(&self.pool).await?;
194        }
195
196        Ok(())
197    }
198
199    /// Synchronizes a table schema by adding missing columns or indexes.
200    pub async fn sync_table<T: Model>(&self) -> Result<(), Error> {
201        if !self.table_exists(T::table_name()).await? {
202            return self.create_table::<T>().await;
203        }
204
205        let table_name = T::table_name().to_snake_case();
206        let model_columns = T::columns();
207        let existing_columns = self.get_table_columns(&table_name).await?;
208
209        for col in model_columns {
210            let col_name_clean = col.name.strip_prefix("r#").unwrap_or(col.name).to_snake_case();
211            if !existing_columns.contains(&col_name_clean) {
212                let mut alter_query = format!("ALTER TABLE \"{}\" ADD COLUMN \"{}\" {}", table_name, col_name_clean, col.sql_type);
213                if !col.is_nullable {
214                    alter_query.push_str(" DEFAULT ");
215                    match col.sql_type {
216                        "INTEGER" | "INT" | "BIGINT" => alter_query.push('0'),
217                        "BOOLEAN" | "BOOL" => alter_query.push_str("FALSE"),
218                        _ => alter_query.push_str("''"),
219                    }
220                }
221                sqlx::query(&alter_query).execute(&self.pool).await?;
222            }
223
224            if col.index || col.unique {
225                let existing_indexes = self.get_table_indexes(&table_name).await?;
226                let idx_name = format!("idx_{}_{}", table_name, col_name_clean);
227                let uniq_name = format!("unique_{}_{}", table_name, col_name_clean);
228
229                if col.unique && !existing_indexes.contains(&uniq_name) {
230                    let mut query = format!("CREATE UNIQUE INDEX \"{}\" ON \"{}\" (\"{}\")", uniq_name, table_name, col_name_clean);
231                    if matches!(self.driver, Drivers::SQLite) {
232                        query = format!("CREATE UNIQUE INDEX IF NOT EXISTS \"{}\" ON \"{}\" (\"{}\")", uniq_name, table_name, col_name_clean);
233                    }
234                    sqlx::query(&query).execute(&self.pool).await?;
235                } else if col.index && !existing_indexes.contains(&idx_name) && !col.unique {
236                    let mut query = format!("CREATE INDEX \"{}\" ON \"{}\" (\"{}\")", idx_name, table_name, col_name_clean);
237                    if matches!(self.driver, Drivers::SQLite) {
238                        query = format!("CREATE INDEX IF NOT EXISTS \"{}\" ON \"{}\" (\"{}\")", idx_name, table_name, col_name_clean);
239                    }
240                    sqlx::query(&query).execute(&self.pool).await?;
241                }
242            }
243        }
244
245        Ok(())
246    }
247
248    /// Returns the current columns of a table.
249    pub async fn get_table_columns(&self, table_name: &str) -> Result<Vec<String>, Error> {
250        let table_name_snake = table_name.to_snake_case();
251        let query = match self.driver {
252            Drivers::Postgres => "SELECT column_name::TEXT FROM information_schema.columns WHERE table_name = $1 AND table_schema = 'public'".to_string(),
253            Drivers::MySQL => "SELECT column_name FROM information_schema.columns WHERE table_name = ? AND table_schema = DATABASE()".to_string(),
254            Drivers::SQLite => format!("PRAGMA table_info(\"{}\")", table_name_snake),
255        };
256
257        let rows = if let Drivers::SQLite = self.driver {
258            sqlx::query(&query).fetch_all(&self.pool).await?
259        } else {
260            sqlx::query(&query).bind(&table_name_snake).fetch_all(&self.pool).await?
261        };
262
263        let mut columns = Vec::new();
264        for row in rows {
265            let col_name: String = if let Drivers::SQLite = self.driver {
266                row.try_get("name")?
267            } else {
268                row.try_get(0)?
269            };
270            columns.push(col_name);
271        }
272        Ok(columns)
273    }
274
275    /// Returns the current indexes of a table.
276    pub async fn get_table_indexes(&self, table_name: &str) -> Result<Vec<String>, Error> {
277        let table_name_snake = table_name.to_snake_case();
278        let query = match self.driver {
279            Drivers::Postgres => "SELECT indexname::TEXT FROM pg_indexes WHERE tablename = $1 AND schemaname = 'public'".to_string(),
280            Drivers::MySQL => "SELECT INDEX_NAME FROM information_schema.STATISTICS WHERE TABLE_NAME = ? AND TABLE_SCHEMA = DATABASE()".to_string(),
281            Drivers::SQLite => format!("PRAGMA index_list(\"{}\")", table_name_snake),
282        };
283
284        let rows = if let Drivers::SQLite = self.driver {
285            sqlx::query(&query).fetch_all(&self.pool).await?
286        } else {
287            sqlx::query(&query).bind(&table_name_snake).fetch_all(&self.pool).await?
288        };
289
290        let mut indexes = Vec::new();
291        for row in rows {
292            let idx_name: String = if let Drivers::SQLite = self.driver {
293                row.try_get("name")?
294            } else {
295                row.try_get(0)?
296            };
297            indexes.push(idx_name);
298        }
299        Ok(indexes)
300    }
301
302    /// Assigns foreign keys to a table.
303    pub async fn assign_foreign_keys<T: Model>(&self) -> Result<(), Error> {
304        let table_name = T::table_name().to_snake_case();
305        let columns = T::columns();
306
307        for col in columns {
308            if let (Some(f_table), Some(f_key)) = (col.foreign_table, col.foreign_key) {
309                if matches!(self.driver, Drivers::SQLite) { continue; }
310                let constraint_name = format!("fk_{}_{}_{}", table_name, f_table.to_snake_case(), col.name.to_snake_case());
311                let query = format!(
312                    "ALTER TABLE \"{}\" ADD CONSTRAINT \"{}\" FOREIGN KEY (\"{}\") REFERENCES \"{}\"(\"{}\")",
313                    table_name, constraint_name, col.name.to_snake_case(), f_table.to_snake_case(), f_key.to_snake_case()
314                );
315                let _ = sqlx::query(&query).execute(&self.pool).await;
316            }
317        }
318        Ok(())
319    }
320}
321
322// ============================================================================
323// DatabaseBuilder Struct
324// ============================================================================
325
326pub struct DatabaseBuilder {
327    max_connections: u32,
328}
329
330impl DatabaseBuilder {
331    /// Creates a new DatabaseBuilder with default settings.
332    ///
333    /// # Example
334    ///
335    /// ```rust,ignore
336    /// let builder = DatabaseBuilder::new();
337    /// ```
338    pub fn new() -> Self { Self { max_connections: 5 } }
339
340    /// Sets the maximum number of connections for the database pool.
341    ///
342    /// # Arguments
343    ///
344    /// * `max` - The maximum number of connections.
345    ///
346    /// # Example
347    ///
348    /// ```rust,ignore
349    /// let db = Database::builder()
350    ///     .max_connections(10)
351    ///     .connect("sqlite::memory:")
352    ///     .await?;
353    /// ```
354    pub fn max_connections(mut self, max: u32) -> Self { self.max_connections = max; self }
355
356    /// Connects to the database using the configured settings.
357    ///
358    /// # Arguments
359    ///
360    /// * `url` - The database connection string.
361    ///
362    /// # Example
363    ///
364    /// ```rust,ignore
365    /// let db = Database::builder()
366    ///     .connect("sqlite::memory:")
367    ///     .await?;
368    /// ```
369    pub async fn connect(self, url: &str) -> Result<Database, Error> {
370        // Ensure sqlx drivers are registered for Any driver support
371        let _ = sqlx::any::install_default_drivers();
372
373        let pool = sqlx::any::AnyPoolOptions::new().max_connections(self.max_connections).connect(url).await?;
374        let driver = if url.starts_with("postgres") { Drivers::Postgres }
375                    else if url.starts_with("mysql") { Drivers::MySQL }
376                    else { Drivers::SQLite };
377        Ok(Database { pool, driver })
378    }
379}
380
381// ============================================================================
382// Connection Trait
383// ============================================================================
384
385pub trait Connection: Send + Sync {
386    fn execute<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<sqlx::any::AnyQueryResult, sqlx::Error>>;
387    fn fetch_all<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<Vec<sqlx::any::AnyRow>, sqlx::Error>>;
388    fn fetch_one<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<sqlx::any::AnyRow, sqlx::Error>>;
389    fn fetch_optional<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<Option<sqlx::any::AnyRow>, sqlx::Error>>;
390}
391
392impl Connection for Database {
393    fn execute<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<sqlx::any::AnyQueryResult, sqlx::Error>> {
394        Box::pin(async move { sqlx::query_with(sql, args).execute(&self.pool).await })
395    }
396    fn fetch_all<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<Vec<sqlx::any::AnyRow>, sqlx::Error>> {
397        Box::pin(async move { sqlx::query_with(sql, args).fetch_all(&self.pool).await })
398    }
399    fn fetch_one<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<sqlx::any::AnyRow, sqlx::Error>> {
400        Box::pin(async move { sqlx::query_with(sql, args).fetch_one(&self.pool).await })
401    }
402    fn fetch_optional<'a, 'q: 'a>(&'a self, sql: &'q str, args: AnyArguments<'q>) -> BoxFuture<'a, Result<Option<sqlx::any::AnyRow>, sqlx::Error>> {
403        Box::pin(async move { sqlx::query_with(sql, args).fetch_optional(&self.pool).await })
404    }
405}
406
407// ============================================================================
408// Raw SQL Query Builder
409// ============================================================================
410
411pub struct RawQuery<'a, C> {
412    conn: C,
413    sql: &'a str,
414    args: AnyArguments<'a>,
415}
416
417impl<'a, C> RawQuery<'a, C> where C: Connection {
418    pub(crate) fn new(conn: C, sql: &'a str) -> Self {
419        Self { conn, sql, args: AnyArguments::default() }
420    }
421
422    /// Binds a value to the SQL query.
423    ///
424    /// # Type Parameters
425    ///
426    /// * `T` - The type of the value to bind.
427    ///
428    /// # Example
429    ///
430    /// ```rust,ignore
431    /// let user: User = db.raw("SELECT * FROM users WHERE id = ?")
432    ///     .bind(1)
433    ///     .fetch_one()
434    ///     .await?;
435    /// ```
436    pub fn bind<T>(mut self, value: T) -> Self where T: 'a + sqlx::Encode<'a, sqlx::Any> + sqlx::Type<sqlx::Any> + Send + Sync {
437        let _ = self.args.add(value);
438        self
439    }
440
441    /// Executes the query and returns all matching rows.
442    ///
443    /// # Type Parameters
444    ///
445    /// * `T` - The type to map the rows to.
446    ///
447    /// # Example
448    ///
449    /// ```rust,ignore
450    /// let users: Vec<User> = db.raw("SELECT * FROM users")
451    ///     .fetch_all()
452    ///     .await?;
453    /// ```
454    pub async fn fetch_all<T>(self) -> Result<Vec<T>, Error> where T: for<'r> sqlx::FromRow<'r, sqlx::any::AnyRow> + Send + Unpin {
455        let rows = self.conn.fetch_all(self.sql, self.args).await?;
456        Ok(rows.iter().map(|r| T::from_row(r).unwrap()).collect())
457    }
458
459    /// Executes the query and returns exactly one row.
460    ///
461    /// # Type Parameters
462    ///
463    /// * `T` - The type to map the row to.
464    ///
465    /// # Example
466    ///
467    /// ```rust,ignore
468    /// let user: User = db.raw("SELECT * FROM users WHERE id = 1")
469    ///     .fetch_one()
470    ///     .await?;
471    /// ```
472    pub async fn fetch_one<T>(self) -> Result<T, Error> where T: for<'r> sqlx::FromRow<'r, sqlx::any::AnyRow> + Send + Unpin {
473        let row = self.conn.fetch_one(self.sql, self.args).await?;
474        Ok(T::from_row(&row)?)
475    }
476
477    /// Executes the query and returns an optional row.
478    ///
479    /// # Type Parameters
480    ///
481    /// * `T` - The type to map the row to.
482    ///
483    /// # Example
484    ///
485    /// ```rust,ignore
486    /// let user: Option<User> = db.raw("SELECT * FROM users WHERE id = 1")
487    ///     .fetch_optional()
488    ///     .await?;
489    /// ```
490    pub async fn fetch_optional<T>(self) -> Result<Option<T>, Error> where T: for<'r> sqlx::FromRow<'r, sqlx::any::AnyRow> + Send + Unpin {
491        let row = self.conn.fetch_optional(self.sql, self.args).await?;
492        Ok(row.map(|r| T::from_row(&r).unwrap()))
493    }
494
495    /// Executes the query and returns the number of affected rows.
496    ///
497    /// Useful for UPDATE, DELETE or INSERT queries.
498    ///
499    /// # Example
500    ///
501    /// ```rust,ignore
502    /// let affected = db.raw("DELETE FROM users WHERE id = 1")
503    ///     .execute()
504    ///     .await?;
505    /// ```
506    pub async fn execute(self) -> Result<u64, Error> {
507        let result = self.conn.execute(self.sql, self.args).await?;
508        Ok(result.rows_affected())
509    }
510}
bottle_orm/database.rs

bottle_orm/
database.rs
