Skip to main content

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 multiple database backends.
6//!
7//! ## Supported Database Drivers
8//!
9//! - **PostgreSQL**: Full support with advanced features like UUID, JSONB, arrays
10//! - **MySQL**: Complete support for standard MySQL/MariaDB features
11//! - **SQLite**: In-memory and file-based SQLite databases
12//!
13//! ## Features
14//!
15//! - **Connection Pooling**: Automatic connection pool management via sqlx
16//! - **Driver Detection**: Automatic database driver detection from connection URL
17//! - **Schema Management**: Table creation with indexes, constraints, and foreign keys
18//! - **Type Safety**: Type-safe operations across different database backends
19//!
20//! ## Example Usage
21//!
22//! ```rust,ignore
23//! use bottle_orm::Database;
24//!
25//! // Connect to PostgreSQL
26//! let db = Database::connect("postgres://user:pass@localhost/mydb").await?;
27//!
28//! // Connect to SQLite
29//! let db = Database::connect("sqlite::memory:").await?;
30//!
31//! // Connect to MySQL
32//! let db = Database::connect("mysql://user:pass@localhost/mydb").await?;
33//!
34//! // Create table for a model
35//! db.create_table::<User>().await?;
36//!
37//! // Assign foreign keys
38//! db.assign_foreign_keys::<Post>().await?;
39//!
40//! // Start building queries
41//! let users = db.model::<User>().scan().await?;
42//! ```
43
44// ============================================================================
45// External Crate Imports
46// ============================================================================
47
48use heck::ToSnakeCase;
49use sqlx::{
50    any::{AnyArguments, AnyPoolOptions},
51    AnyPool, Arguments, Error, Row,
52};
53use std::time::Duration;
54
55// ============================================================================
56// Internal Crate Imports
57// ============================================================================
58
59use crate::{migration::Migrator, model::Model, query_builder::QueryBuilder, Transaction};
60
61// ============================================================================
62// Database Driver Enumeration
63// ============================================================================
64
65/// Supported database driver types.
66///
67/// This enum represents the different database backends that Bottle ORM can work with.
68/// The driver type is automatically detected from the connection URL and used to
69/// generate appropriate SQL syntax for each database system.
70///
71/// # Variants
72///
73/// * `Postgres` - PostgreSQL database (9.5+)
74/// * `SQLite` - SQLite database (3.x)
75/// * `MySQL` - MySQL or MariaDB database (5.7+/10.2+)
76///
77/// # SQL Dialect Differences
78///
79/// Different drivers use different SQL syntax:
80///
81/// - **Placeholders**:
82///   - PostgreSQL: `$1, $2, $3` (numbered)
83///   - SQLite/MySQL: `?, ?, ?` (positional)
84///
85/// - **Type Casting**:
86///   - PostgreSQL: `$1::UUID`, `$2::TIMESTAMPTZ`
87///   - SQLite/MySQL: Automatic type inference
88///
89/// - **Schema Queries**:
90///   - PostgreSQL: `information_schema` tables
91///   - SQLite: `sqlite_master` system table
92///   - MySQL: `information_schema` tables
93///
94/// # Example
95///
96/// ```rust,ignore
97/// match db.driver {
98///     Drivers::Postgres => println!("Using PostgreSQL"),
99///     Drivers::SQLite => println!("Using SQLite"),
100///     Drivers::MySQL => println!("Using MySQL"),
101/// }
102/// ```
103#[derive(Clone, Debug, Copy)]
104pub enum Drivers {
105    /// PostgreSQL driver.
106    ///
107    /// Used for PostgreSQL databases. Supports advanced features like:
108    /// - UUID native type
109    /// - JSONB for JSON data
110    /// - Array types
111    /// - Full-text search
112    /// - Advanced indexing (GiST, GIN, etc.)
113    Postgres,
114
115    /// SQLite driver.
116    ///
117    /// Used for SQLite databases (both in-memory and file-based). Characteristics:
118    /// - Lightweight and embedded
119    /// - Single-file database
120    /// - Limited concurrent write support
121    /// - Good for development and small applications
122    SQLite,
123
124    /// MySQL driver.
125    ///
126    /// Used for MySQL and MariaDB databases. Features:
127    /// - Wide compatibility
128    /// - Good performance for read-heavy workloads
129    /// - Mature ecosystem
130    /// - ACID compliance (with InnoDB)
131    MySQL,
132}
133
134// ============================================================================
135// Database Builder
136// ============================================================================
137
138/// A builder for creating a `Database` connection with custom options.
139///
140/// Allows configuration of connection pool settings such as maximum connections,
141/// timeouts, and lifetimes.
142///
143/// # Example
144///
145/// ```rust,ignore
146/// let db = Database::builder()
147///     .max_connections(10)
148///     .min_connections(2)
149///     .acquire_timeout(std::time::Duration::from_secs(5))
150///     .connect("postgres://...")
151///     .await?;
152/// ```
153pub struct DatabaseBuilder {
154    options: AnyPoolOptions,
155}
156
157impl DatabaseBuilder {
158    /// Sets the maximum number of connections that this pool should maintain.
159    pub fn max_connections(mut self, max: u32) -> Self {
160        self.options = self.options.max_connections(max);
161        self
162    }
163
164    /// Sets the minimum number of connections that this pool should maintain.
165    pub fn min_connections(mut self, min: u32) -> Self {
166        self.options = self.options.min_connections(min);
167        self
168    }
169
170    /// Sets the maximum amount of time to spend waiting for a connection.
171    pub fn acquire_timeout(mut self, timeout: Duration) -> Self {
172        self.options = self.options.acquire_timeout(timeout);
173        self
174    }
175
176    /// Sets the maximum amount of time a connection may be idle.
177    pub fn idle_timeout(mut self, timeout: Duration) -> Self {
178        self.options = self.options.idle_timeout(Some(timeout));
179        self
180    }
181
182    /// Sets the maximum lifetime of a connection.
183    pub fn max_lifetime(mut self, lifetime: Duration) -> Self {
184        self.options = self.options.max_lifetime(Some(lifetime));
185        self
186    }
187
188    /// Connects to the database using the configured options.
189    pub async fn connect(self, url: &str) -> Result<Database, Error> {
190        // Install default drivers for sqlx::Any
191        sqlx::any::install_default_drivers();
192
193        let pool = self.options.connect(url).await?;
194
195        // Detect driver type from URL scheme
196        let (driver_str, _) = url.split_once(':').unwrap_or(("sqlite", ""));
197        let driver = match driver_str {
198            "postgresql" | "postgres" => Drivers::Postgres,
199            "mysql" => Drivers::MySQL,
200            _ => Drivers::SQLite,
201        };
202
203        Ok(Database { pool, driver })
204    }
205}
206
207// ============================================================================
208// Database Connection and Management
209// ============================================================================
210
211/// The main entry point for database connection and management.
212///
213/// `Database` handles connection pooling, driver detection, and provides methods
214/// for schema operations and query building. It uses sqlx's `AnyPool` to support
215/// multiple database backends with a unified interface.
216///
217/// # Fields
218///
219/// * `pool` - The sqlx connection pool for executing queries
220/// * `driver` - The detected database driver type
221///
222/// # Thread Safety
223///
224/// `Database` implements `Clone` and can be safely shared across threads.
225/// The underlying connection pool is thread-safe and handles connection
226/// distribution automatically.
227///
228/// # Example
229///
230/// ```rust,ignore
231/// use bottle_orm::Database;
232///
233/// #[tokio::main]
234/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
235///     // Connect to database
236///     let db = Database::connect("postgres://localhost/mydb").await?;
237///
238///     // Create migrator
239///     let migrator = db.migrator();
240///
241///     // Build queries
242///     let query = db.model::<User>();
243///
244///     // Database can be cloned and shared
245///     let db_clone = db.clone();
246///     tokio::spawn(async move {
247///         let users = db_clone.model::<User>().scan().await;
248///     });
249///
250///     Ok(())
251/// }
252/// ```
253#[derive(Clone)]
254pub struct Database {
255    /// The sqlx connection pool for executing database queries.
256    ///
257    /// This pool manages a set of database connections that can be reused
258    /// across multiple queries, improving performance by avoiding the overhead
259    /// of creating new connections for each operation.
260    pub(crate) pool: AnyPool,
261
262    /// The detected database driver type.
263    ///
264    /// Used to generate driver-specific SQL syntax (e.g., placeholders,
265    /// type casting, schema queries).
266    pub(crate) driver: Drivers,
267}
268
269// ============================================================================
270// Database Implementation
271// ============================================================================
272
273impl Database {
274    // ========================================================================
275    // Connection Management
276    // ========================================================================
277
278    /// Creates a builder to configure the database connection options.
279    ///
280    /// Returns a `DatabaseBuilder` which allows setting pool options like
281    /// `max_connections`, timeouts, etc.
282    ///
283    /// # Example
284    ///
285    /// ```rust,ignore
286    /// let db = Database::builder()
287    ///     .max_connections(20)
288    ///     .connect("postgres://...")
289    ///     .await?;
290    /// ```
291    pub fn builder() -> DatabaseBuilder {
292        DatabaseBuilder { options: AnyPoolOptions::new() }
293    }
294
295    /// Connects to the database using a connection string (Database URL).
296    ///
297    /// This method establishes a connection pool to the specified database and
298    /// automatically detects the driver type based on the URL scheme. The connection
299    /// pool is configured with a default maximum of 5 connections.
300    ///
301    /// # Arguments
302    ///
303    /// * `url` - The database connection string with the format:
304    ///   `<scheme>://<user>:<password>@<host>:<port>/<database>`
305    ///
306    /// # Supported URL Schemes
307    ///
308    /// - **PostgreSQL**: `postgres://` or `postgresql://`
309    /// - **MySQL**: `mysql://`
310    /// - **SQLite**: `sqlite://` or `sqlite::memory:` (for in-memory databases)
311    ///
312    /// # Connection Pool Configuration
313    ///
314    /// - Maximum connections: 5
315    /// - Connection timeout: Default (30 seconds)
316    /// - Idle timeout: Default (10 minutes)
317    ///
318    /// # Returns
319    ///
320    /// * `Ok(Database)` - Successfully connected database instance
321    /// * `Err(sqlx::Error)` - Connection error (invalid URL, authentication failure, etc.)
322    ///
323    /// # Examples
324    ///
325    /// ```rust,ignore
326    /// // PostgreSQL connection
327    /// let db = Database::connect("postgres://user:password@localhost:5432/mydb").await?;
328    ///
329    /// // PostgreSQL with SSL
330    /// let db = Database::connect("postgres://user:password@localhost/mydb?sslmode=require").await?;
331    ///
332    /// // SQLite in-memory database (great for testing)
333    /// let db = Database::connect("sqlite::memory:").await?;
334    ///
335    /// // SQLite file-based database
336    /// let db = Database::connect("sqlite://./database.db").await?;
337    ///
338    /// // MySQL connection
339    /// let db = Database::connect("mysql://user:password@localhost:3306/mydb").await?;
340    /// ```
341    ///
342    /// # Error Handling
343    ///
344    /// ```rust,ignore
345    /// match Database::connect("postgres://localhost/mydb").await {
346    ///     Ok(db) => println!("Connected successfully"),
347    ///     Err(e) => eprintln!("Connection failed: {}", e),
348    /// }
349    /// ```
350    pub async fn connect(url: &str) -> Result<Self, Error> {
351        Self::builder().max_connections(5).connect(url).await
352    }
353
354    // ========================================================================
355    // Schema Management
356    // ========================================================================
357
358    /// Creates a `Migrator` instance to manage schema migrations.
359    ///
360    /// The migrator allows you to register multiple models and execute
361    /// all necessary schema changes (table creation, foreign keys) in the
362    /// correct order.
363    ///
364    /// # Returns
365    ///
366    /// A new `Migrator` instance associated with this database connection
367    ///
368    /// # Example
369    ///
370    /// ```rust,ignore
371    /// let db = Database::connect("sqlite::memory:").await?;
372    ///
373    /// db.migrator()
374    ///     .register::<User>()
375    ///     .register::<Post>()
376    ///     .register::<Comment>()
377    ///     .run()
378    ///     .await?;
379    /// ```
380    ///
381    /// # See Also
382    ///
383    /// * [`Migrator`] - For detailed migration documentation
384    /// * [`Migrator::register()`] - For registering models
385    /// * [`Migrator::run()`] - For executing migrations
386    pub fn migrator(&self) -> Migrator<'_> {
387        Migrator::new(self)
388    }
389
390    // ========================================================================
391    // Query Building
392    // ========================================================================
393
394    /// Starts building a query for a specific Model.
395    ///
396    /// This method creates a new `QueryBuilder` instance configured for the
397    /// specified model type. The query builder provides a fluent interface
398    /// for constructing SELECT and INSERT queries.
399    ///
400    /// # Type Parameters
401    ///
402    /// * `T` - The Model type to query. Must implement `Model + Send + Sync + Unpin`
403    ///
404    /// # Returns
405    ///
406    /// A new `QueryBuilder` instance ready for query construction
407    ///
408    /// # Example
409    ///
410    /// ```rust,ignore
411    /// // Simple query
412    /// let users: Vec<User> = db.model::<User>().scan().await?;
413    ///
414    /// // Filtered query
415    /// let adults: Vec<User> = db.model::<User>()
416    ///     .filter("age", ">=", 18)
417    ///     .scan()
418    ///     .await?;
419    ///
420    /// // Insert operation
421    /// let new_user = User { /* ... */ };
422    /// db.model::<User>().insert(&new_user).await?;
423    /// ```
424    ///
425    /// # See Also
426    ///
427    /// * [`QueryBuilder`] - For detailed query building documentation
428    /// * [`QueryBuilder::filter()`] - For adding WHERE clauses
429    /// * [`QueryBuilder::scan()`] - For executing SELECT queries
430    /// * [`QueryBuilder::insert()`] - For INSERT operations
431    pub fn model<T: Model + Send + Sync + Unpin>(&self) -> QueryBuilder<'_, T, Self> {
432        // Get active column names from the model
433        let active_columns = T::active_columns();
434        let mut columns: Vec<String> = Vec::with_capacity(active_columns.capacity());
435
436        // Convert column names to snake_case and strip 'r#' prefix if present
437        for col in active_columns {
438            columns.push(col.strip_prefix("r#").unwrap_or(col).to_snake_case());
439        }
440
441        // Create and return the query builder
442        QueryBuilder::new(self.clone(), self.driver, T::table_name(), T::columns(), columns)
443    }
444
445    /// Creates a raw SQL query builder.
446    ///
447    /// This provides a "safety hatch" to execute raw SQL queries when the fluent
448    /// QueryBuilder is not sufficient (e.g., complex joins, CTEs, specific DB features).
449    ///
450    /// # Arguments
451    ///
452    /// * `sql` - The raw SQL query string (use `$1`, `?`, etc. for placeholders based on driver)
453    ///
454    /// # Returns
455    ///
456    /// A `RawQuery` builder that allows binding parameters and executing the query.
457    ///
458    /// # Example
459    ///
460    /// ```rust,ignore
461    /// // Fetching complex data
462    /// let users: Vec<User> = db.raw("SELECT * FROM users WHERE age > $1")
463    ///     .bind(18)
464    ///     .fetch_all()
465    ///     .await?;
466    ///
467    /// // Executing an update
468    /// let affected = db.raw("UPDATE users SET active = $1 WHERE last_login < $2")
469    ///     .bind(false)
470    ///     .bind(one_year_ago)
471    ///     .execute()
472    ///     .await?;
473    /// ```
474    pub fn raw<'a>(&'a self, sql: &'a str) -> RawQuery<'a, Self> {
475        RawQuery::new(self.clone(), sql)
476    }
477
478    // ========================================================================
479    // Table Creation
480    // ========================================================================
481
482    /// Creates the table for model `T` if it does not exist.
483    ///
484    /// This method generates and executes SQL to create a table based on the
485    /// model's structure. It handles column definitions, primary keys, unique
486    /// constraints, default values, and indexes.
487    ///
488    /// # Type Parameters
489    ///
490    /// * `T` - The Model type representing the table
491    ///
492    /// # Returns
493    ///
494    /// * `Ok(&Self)` - Reference to self for method chaining
495    /// * `Err(sqlx::Error)` - Database error during table creation
496    ///
497    /// # Generated SQL Features
498    ///
499    /// - **Primary Keys**: Automatically marked with `PRIMARY KEY`
500    /// - **NOT NULL**: Non-nullable fields get `NOT NULL` constraint
501    /// - **UNIQUE**: Fields marked with `#[orm(unique)]` get `UNIQUE` constraint
502    /// - **DEFAULT**: Fields marked with `#[orm(create_time)]` get `DEFAULT CURRENT_TIMESTAMP`
503    /// - **Indexes**: Fields marked with `#[orm(index)]` get database indexes
504    ///
505    /// # Example
506    ///
507    /// ```rust,ignore
508    /// use bottle_orm::Model;
509    /// use uuid::Uuid;
510    /// use chrono::{DateTime, Utc};
511    ///
512    /// #[derive(Model)]
513    /// struct User {
514    ///     #[orm(primary_key)]
515    ///     id: Uuid,
516    ///     #[orm(size = 50, unique)]
517    ///     username: String,
518    ///     #[orm(index)]
519    ///     email: String,
520    ///     age: i32,
521    ///     #[orm(create_time)]
522    ///     created_at: DateTime<Utc>,
523    /// }
524    ///
525    /// // Creates table with:
526    /// // - UUID primary key
527    /// // - Unique username constraint
528    /// // - Index on email
529    /// // - created_at with DEFAULT CURRENT_TIMESTAMP
530    /// db.create_table::<User>().await?;
531    /// ```
532    ///
533    /// # Generated SQL Example (PostgreSQL)
534    ///
535    /// ```sql
536    /// CREATE TABLE IF NOT EXISTS "user" (
537    ///     "id" UUID PRIMARY KEY,
538    ///     "username" VARCHAR(50) NOT NULL UNIQUE,
539    ///     "email" TEXT NOT NULL,
540    ///     "age" INTEGER NOT NULL,
541    ///     "created_at" TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
542    /// );
543    /// CREATE INDEX IF NOT EXISTS "idx_user_email" ON "user" ("email");
544    /// ```
545    pub async fn create_table<T: Model>(&self) -> Result<&Self, Error> {
546        // Get table name in snake_case format
547        let table_name = T::table_name().to_snake_case();
548        let columns = T::columns();
549
550        let mut column_defs = Vec::new();
551        let mut index_statements = Vec::new();
552
553        // Build column definitions
554        for col in &columns {
555            // Strip 'r#' prefix if present (for Rust keywords used as column names)
556            let col_name = col.name.strip_prefix("r#").unwrap_or(col.name).to_snake_case();
557            let mut def = format!("\"{}\" {}", col_name, col.sql_type);
558
559            // Add PRIMARY KEY constraint
560            if col.is_primary_key {
561                def.push_str(" PRIMARY KEY");
562            }
563
564            // Add NOT NULL constraint (except for primary keys, which are implicitly NOT NULL)
565            if !col.is_nullable && !col.is_primary_key {
566                def.push_str(" NOT NULL");
567            }
568
569            // Add DEFAULT CURRENT_TIMESTAMP for create_time fields
570            if col.create_time {
571                def.push_str(" DEFAULT CURRENT_TIMESTAMP");
572            }
573
574            // Add UNIQUE constraint
575            if col.unique {
576                def.push_str(" UNIQUE");
577            }
578
579            column_defs.push(def);
580
581            // Generate index creation statement if needed
582            if col.index {
583                let index_type = if col.unique { "UNIQUE INDEX" } else { "INDEX" };
584                let index_name = format!("idx_{}_{}", table_name, col_name);
585
586                let index_query = format!(
587                    "CREATE {} IF NOT EXISTS \"{}\" ON \"{}\" (\"{}\")",
588                    index_type, index_name, table_name, col_name,
589                );
590
591                index_statements.push(index_query);
592            }
593        }
594
595        // Add SQLite Foreign Keys inline (SQLite doesn't support ADD CONSTRAINT)
596        if let Drivers::SQLite = self.driver {
597            for col in &columns {
598                if let (Some(f_table), Some(f_key)) = (col.foreign_table, col.foreign_key) {
599                    let col_name = col.name.strip_prefix("r#").unwrap_or(col.name).to_snake_case();
600                    let f_table_clean = f_table.to_snake_case();
601                    let f_key_clean = f_key.to_snake_case();
602
603                    let fk_def = format!(
604                        "FOREIGN KEY (\"{}\") REFERENCES \"{}\" (\"{}\")",
605                        col_name, f_table_clean, f_key_clean
606                    );
607                    column_defs.push(fk_def);
608                }
609            }
610        }
611
612        // Build and execute CREATE TABLE statement
613        let create_table_query =
614            format!("CREATE TABLE IF NOT EXISTS \"{}\" ({})", table_name.to_snake_case(), column_defs.join(", "));
615        log::info!("{}", create_table_query);
616
617        sqlx::query(&create_table_query).execute(&self.pool).await?;
618
619        // Create indexes
620        for idx_stmt in index_statements {
621            sqlx::query(&idx_stmt).execute(&self.pool).await?;
622        }
623
624        Ok(self)
625    }
626
627    /// Starts a new database transaction.
628    ///
629    /// Returns a `Transaction` wrapper that can be used to execute multiple
630    /// queries atomically. The transaction must be explicitly committed
631    /// using `commit()`, otherwise it will be rolled back when dropped.
632    ///
633    /// # Returns
634    ///
635    /// * `Ok(Transaction)` - A new transaction instance
636    /// * `Err(sqlx::Error)` - Database error starting transaction
637    ///
638    /// # Example
639    ///
640    /// ```rust,ignore
641    /// let mut tx = db.begin().await?;
642    /// // ... perform operations ...
643    /// tx.commit().await?;
644    /// ```
645    pub async fn begin<'a>(&self) -> Result<Transaction<'a>, sqlx::Error> {
646        let tx = self.pool.begin().await?;
647        Ok(Transaction { tx, driver: self.driver })
648    }
649
650    // ========================================================================
651    // Foreign Key Management
652    // ========================================================================
653
654    /// Checks for and assigns Foreign Keys for model `T`.
655    ///
656    /// This method examines all columns marked with `#[orm(foreign_key = "Table::Column")]`
657    /// and creates the appropriate foreign key constraints. It verifies that constraints
658    /// don't already exist before attempting to create them, preventing duplication errors.
659    ///
660    /// # Type Parameters
661    ///
662    /// * `T` - The Model type to process for foreign keys
663    ///
664    /// # Returns
665    ///
666    /// * `Ok(&Self)` - Reference to self for method chaining
667    /// * `Err(sqlx::Error)` - Database error during foreign key creation
668    ///
669    /// # Constraint Naming
670    ///
671    /// Foreign key constraints are named using the pattern:
672    /// `fk_{table_name}_{column_name}`
673    ///
674    /// # Example
675    ///
676    /// ```rust,ignore
677    /// use bottle_orm::Model;
678    /// use uuid::Uuid;
679    ///
680    /// #[derive(Model)]
681    /// struct User {
682    ///     #[orm(primary_key)]
683    ///     id: Uuid,
684    ///     username: String,
685    /// }
686    ///
687    /// #[derive(Model)]
688    /// struct Post {
689    ///     #[orm(primary_key)]
690    ///     id: Uuid,
691    ///     #[orm(foreign_key = "User::id")]
692    ///     user_id: Uuid,
693    ///     title: String,
694    /// }
695    ///
696    /// // Create tables first
697    /// db.create_table::<User>().await?;
698    /// db.create_table::<Post>().await?;
699    ///
700    /// // Then assign foreign keys
701    /// db.assign_foreign_keys::<Post>().await?;
702    /// ```
703    ///
704    /// # Generated SQL Example
705    ///
706    /// ```sql
707    /// ALTER TABLE "post"
708    /// ADD CONSTRAINT "fk_post_user_id"
709    /// FOREIGN KEY ("user_id")
710    /// REFERENCES "user" ("id");
711    /// ```
712    ///
713    /// # Important Notes
714    ///
715    /// - Foreign key assignment should be done **after** all tables are created
716    /// - The referenced table and column must exist before creating the foreign key
717    /// - Use the `Migrator` to handle the correct order automatically
718    /// - Currently optimized for PostgreSQL (uses `information_schema`)
719    ///
720    /// # See Also
721    ///
722    /// * [`Migrator`] - For automatic migration order management
723    pub async fn assign_foreign_keys<T: Model>(&self) -> Result<&Self, Error> {
724        // SQLite handles FKs in create_table, so we skip here
725        if let Drivers::SQLite = self.driver {
726            return Ok(self);
727        }
728
729        // Get table name in snake_case format
730        let table_name = T::table_name().to_snake_case();
731        let columns = T::columns();
732
733        // Process each column that has a foreign key definition
734        for col in columns {
735            if let (Some(f_table), Some(f_key)) = (col.foreign_table, col.foreign_key) {
736                // Clean up column and reference names
737                let col_name = col.name.strip_prefix("r#").unwrap_or(col.name).to_snake_case();
738                let f_table_clean = f_table.to_snake_case();
739                let f_key_clean = f_key.to_snake_case();
740
741                // Generate constraint name
742                let constraint_name = format!("fk_{}_{}", table_name, col_name);
743
744                // Check if constraint already exists
745                let count: i64 = match self.driver {
746                    Drivers::Postgres => {
747                        let check_query =
748                            "SELECT count(*) FROM information_schema.table_constraints WHERE constraint_name = $1";
749                        let row = sqlx::query(check_query).bind(&constraint_name).fetch_one(&self.pool).await?;
750                        row.try_get(0).unwrap_or(0)
751                    }
752                    Drivers::MySQL => {
753                        let check_query = "SELECT count(*) FROM information_schema.TABLE_CONSTRAINTS WHERE CONSTRAINT_NAME = ? AND TABLE_SCHEMA = DATABASE()";
754                        let row = sqlx::query(check_query).bind(&constraint_name).fetch_one(&self.pool).await?;
755                        row.try_get(0).unwrap_or(0)
756                    }
757                    Drivers::SQLite => 0, // Unreachable
758                };
759
760                // Skip if constraint already exists
761                if count > 0 {
762                    continue;
763                }
764
765                // Create foreign key constraint
766                let alter_query = format!(
767                    "ALTER TABLE \"{}\" ADD CONSTRAINT \"{}\" FOREIGN KEY (\"{}\") REFERENCES \"{}\" (\"{}\")",
768                    table_name, constraint_name, col_name, f_table_clean, f_key_clean
769                );
770
771                sqlx::query(&alter_query).execute(&self.pool).await?;
772            }
773        }
774
775        Ok(self)
776    }
777}
778
779/// A trait representing a database connection or transaction.
780///
781/// This trait abstracts over `Database` (pool) and `Transaction` types, allowing
782/// the `QueryBuilder` to work seamlessly with both. It uses Generic Associated Types (GATs)
783/// to handle the lifetimes of the executor references correctly.
784pub trait Connection {
785    /// The type of the executor returned by this connection.
786    ///
787    /// This uses GATs to bind the lifetime of the executor (`'c`) to the lifetime
788    /// of the borrow of the connection (`&'c mut self`).
789    type Exec<'c>: sqlx::Executor<'c, Database = sqlx::Any>
790    where
791        Self: 'c;
792
793    /// Returns a mutable reference to the SQLx executor.
794    ///
795    /// # Returns
796    ///
797    /// An executor capable of running SQL queries (either a Pool or a Transaction).
798    fn executor<'c>(&'c mut self) -> Self::Exec<'c>;
799}
800
801/// Implementation of Connection for the main Database struct.
802///
803/// Uses the internal connection pool to execute queries.
804impl Connection for Database {
805    type Exec<'c> = &'c sqlx::Pool<sqlx::Any>;
806
807    fn executor<'c>(&'c mut self) -> Self::Exec<'c> {
808        &self.pool
809    }
810}
811
812/// Implementation of Connection for a mutable reference to Database.
813impl<'a> Connection for &'a mut Database {
814    type Exec<'c>
815        = &'c sqlx::Pool<sqlx::Any>
816    where
817        Self: 'c;
818
819    fn executor<'c>(&'c mut self) -> Self::Exec<'c> {
820        &self.pool
821    }
822}
823
824/// Implementation of Connection for a mutable reference to sqlx::Transaction.
825impl<'a> Connection for &mut sqlx::Transaction<'a, sqlx::Any> {
826    type Exec<'c>
827        = &'c mut sqlx::AnyConnection
828    where
829        Self: 'c;
830
831    fn executor<'c>(&'c mut self) -> Self::Exec<'c> {
832        &mut **self
833    }
834}
835
836// ============================================================================
837
838// Raw SQL Query Builder
839
840// ============================================================================
841
842/// A builder for executing raw SQL queries with parameter binding.
843
844///
845
846/// Returned by `Database::raw()` or `Transaction::raw()`. Allows constructing safe, parameterized
847
848/// SQL queries that can bypass the standard model-based QueryBuilder when
849
850/// complex SQL features (CTEs, Window Functions, etc.) are needed.
851
852pub struct RawQuery<'a, C> {
853    conn: C,
854
855    sql: &'a str,
856
857    args: AnyArguments<'a>,
858}
859
860impl<'a, C> RawQuery<'a, C>
861where
862    C: Connection + Send,
863{
864    /// Creates a new RawQuery instance.
865
866    pub(crate) fn new(conn: C, sql: &'a str) -> Self {
867        Self { conn, sql, args: AnyArguments::default() }
868    }
869
870    /// Binds a parameter to the query.
871
872    ///
873
874    /// # Arguments
875
876    ///
877
878    /// * `value` - The value to bind. Must implement `sqlx::Encode` and `sqlx::Type`.
879
880    pub fn bind<T>(mut self, value: T) -> Self
881    where
882        T: 'a + sqlx::Encode<'a, sqlx::Any> + sqlx::Type<sqlx::Any> + Send + Sync,
883    {
884        let _ = self.args.add(value);
885
886        self
887    }
888
889    /// Executes the query and returns all matching rows mapped to type `T`.
890
891    pub async fn fetch_all<T>(mut self) -> Result<Vec<T>, Error>
892    where
893        T: for<'r> sqlx::FromRow<'r, sqlx::any::AnyRow> + Send + Unpin,
894    {
895        sqlx::query_as_with::<_, T, _>(self.sql, self.args).fetch_all(self.conn.executor()).await
896    }
897
898    /// Executes the query and returns the first matching row mapped to type `T`.
899
900    pub async fn fetch_one<T>(mut self) -> Result<T, Error>
901    where
902        T: for<'r> sqlx::FromRow<'r, sqlx::any::AnyRow> + Send + Unpin,
903    {
904        sqlx::query_as_with::<_, T, _>(self.sql, self.args).fetch_one(self.conn.executor()).await
905    }
906
907    /// Executes the query and returns the first matching row, or None if not found.
908
909    pub async fn fetch_optional<T>(mut self) -> Result<Option<T>, Error>
910    where
911        T: for<'r> sqlx::FromRow<'r, sqlx::any::AnyRow> + Send + Unpin,
912    {
913        sqlx::query_as_with::<_, T, _>(self.sql, self.args).fetch_optional(self.conn.executor()).await
914    }
915
916    /// Executes the query and returns a single scalar value.
917
918    ///
919
920    /// Useful for queries like `SELECT count(*) ...` or `SELECT id ...`.
921
922    pub async fn fetch_scalar<O>(mut self) -> Result<O, Error>
923    where
924        O: for<'r> sqlx::Decode<'r, sqlx::Any> + sqlx::Type<sqlx::Any> + Send + Unpin,
925    {
926        sqlx::query_scalar_with::<_, O, _>(self.sql, self.args).fetch_one(self.conn.executor()).await
927    }
928
929    /// Executes the query and returns a single scalar value, or None if not found.
930
931    pub async fn fetch_scalar_optional<O>(mut self) -> Result<Option<O>, Error>
932    where
933        O: for<'r> sqlx::Decode<'r, sqlx::Any> + sqlx::Type<sqlx::Any> + Send + Unpin,
934    {
935        sqlx::query_scalar_with::<_, O, _>(self.sql, self.args).fetch_optional(self.conn.executor()).await
936    }
937
938    /// Executes the query (INSERT, UPDATE, DELETE) and returns the number of affected rows.
939
940    pub async fn execute(mut self) -> Result<u64, Error> {
941        let result = sqlx::query_with(self.sql, self.args).execute(self.conn.executor()).await?;
942
943        Ok(result.rows_affected())
944    }
945}