vibesql_executor/
session.rs

1//! Session - Prepared Statement Execution Context
2//!
3//! Provides a session-based API for prepared statement execution, offering
4//! significant performance benefits for repeated query patterns by caching
5//! parsed SQL and performing AST-level parameter binding.
6//!
7//! ## Performance Benefits
8//!
9//! For repeated queries, prepared statements can provide 10-100x speedup by:
10//! - Parsing SQL once and caching the AST
11//! - Binding parameters at the AST level (no re-parsing)
12//! - Reusing the same cached statement for different parameter values
13//!
14//! ## Usage
15//!
16//! ```text
17//! use vibesql_executor::Session;
18//! use vibesql_storage::Database;
19//! use vibesql_types::SqlValue;
20//!
21//! let mut db = Database::new();
22//! // ... create tables and insert data ...
23//!
24//! // Create a session for prepared statement execution
25//! let session = Session::new(&db);
26//!
27//! // Prepare a statement (parses SQL once)
28//! let stmt = session.prepare("SELECT * FROM users WHERE id = ?")?;
29//!
30//! // Execute with different parameters (reuses parsed AST - no re-parsing!)
31//! let result1 = session.execute_prepared(&stmt, &[SqlValue::Integer(1)])?;
32//! let result2 = session.execute_prepared(&stmt, &[SqlValue::Integer(2)])?;
33//!
34//! // For DML statements that modify data, use execute_prepared_mut
35//! let mut session = Session::new(&mut db);
36//! let insert_stmt = session.prepare("INSERT INTO users (id, name) VALUES (?, ?)")?;
37//! session.execute_prepared_mut(&insert_stmt, &[SqlValue::Integer(3), SqlValue::Varchar("Alice".into())])?;
38//! ```
39//!
40//! ## Thread Safety
41//!
42//! `PreparedStatement` is `Clone` and can be shared across threads via `Arc`.
43//! The `PreparedStatementCache` uses internal locking for thread-safe access.
44
45use std::sync::Arc;
46
47use vibesql_ast::Statement;
48use vibesql_storage::{Database, Row};
49use vibesql_types::SqlValue;
50
51use crate::cache::{
52    ArenaParseError, ArenaPreparedStatement, CachedPlan, PkPointLookupPlan, PreparedStatement,
53    PreparedStatementCache, PreparedStatementError, ProjectionPlan,
54};
55use crate::errors::ExecutorError;
56use crate::{DeleteExecutor, InsertExecutor, SelectExecutor, SelectResult, UpdateExecutor};
57
58/// Execution result for prepared statements
59#[derive(Debug)]
60pub enum PreparedExecutionResult {
61    /// Result from a SELECT query
62    Select(SelectResult),
63    /// Number of rows affected by INSERT/UPDATE/DELETE
64    RowsAffected(usize),
65    /// DDL or other statement that doesn't return rows
66    Ok,
67}
68
69impl PreparedExecutionResult {
70    /// Get rows if this is a SELECT result
71    pub fn rows(&self) -> Option<&[Row]> {
72        match self {
73            PreparedExecutionResult::Select(result) => Some(&result.rows),
74            _ => None,
75        }
76    }
77
78    /// Get rows affected if this is a DML result
79    pub fn rows_affected(&self) -> Option<usize> {
80        match self {
81            PreparedExecutionResult::RowsAffected(n) => Some(*n),
82            _ => None,
83        }
84    }
85
86    /// Convert to SelectResult if this is a SELECT result
87    pub fn into_select_result(self) -> Option<SelectResult> {
88        match self {
89            PreparedExecutionResult::Select(result) => Some(result),
90            _ => None,
91        }
92    }
93}
94
95/// Error type for session operations
96#[derive(Debug)]
97pub enum SessionError {
98    /// Error during prepared statement operations
99    PreparedStatement(PreparedStatementError),
100    /// Error during query execution
101    Execution(ExecutorError),
102    /// Statement type not supported for this operation
103    UnsupportedStatement(String),
104}
105
106impl std::fmt::Display for SessionError {
107    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        match self {
109            SessionError::PreparedStatement(e) => write!(f, "Prepared statement error: {}", e),
110            SessionError::Execution(e) => write!(f, "Execution error: {:?}", e),
111            SessionError::UnsupportedStatement(msg) => write!(f, "Unsupported statement: {}", msg),
112        }
113    }
114}
115
116impl std::error::Error for SessionError {}
117
118impl From<PreparedStatementError> for SessionError {
119    fn from(e: PreparedStatementError) -> Self {
120        SessionError::PreparedStatement(e)
121    }
122}
123
124impl From<ExecutorError> for SessionError {
125    fn from(e: ExecutorError) -> Self {
126        SessionError::Execution(e)
127    }
128}
129
130/// Session for executing prepared statements
131///
132/// A session holds a reference to the database and a cache of prepared statements.
133/// Use this for executing repeated queries with different parameters.
134pub struct Session<'a> {
135    db: &'a Database,
136    cache: Arc<PreparedStatementCache>,
137}
138
139impl<'a> Session<'a> {
140    /// Create a new session with a reference to the database
141    ///
142    /// Uses a default cache size of 1000 prepared statements.
143    pub fn new(db: &'a Database) -> Self {
144        Self {
145            db,
146            cache: Arc::new(PreparedStatementCache::default_cache()),
147        }
148    }
149
150    /// Create a new session with a custom cache size
151    pub fn with_cache_size(db: &'a Database, cache_size: usize) -> Self {
152        Self {
153            db,
154            cache: Arc::new(PreparedStatementCache::new(cache_size)),
155        }
156    }
157
158    /// Create a new session with a shared cache
159    ///
160    /// This allows multiple sessions to share the same prepared statement cache,
161    /// which is useful for connection pooling scenarios.
162    pub fn with_shared_cache(db: &'a Database, cache: Arc<PreparedStatementCache>) -> Self {
163        Self { db, cache }
164    }
165
166    /// Get the underlying database reference
167    pub fn database(&self) -> &Database {
168        self.db
169    }
170
171    /// Get the prepared statement cache
172    pub fn cache(&self) -> &PreparedStatementCache {
173        &self.cache
174    }
175
176    /// Get the shared cache Arc (for sharing with other sessions)
177    pub fn shared_cache(&self) -> Arc<PreparedStatementCache> {
178        Arc::clone(&self.cache)
179    }
180
181    /// Prepare a SQL statement for execution
182    ///
183    /// Parses the SQL and caches the result. Subsequent calls with the same
184    /// SQL string will return the cached statement without re-parsing.
185    ///
186    /// Supports `?` placeholders for parameter binding.
187    ///
188    /// # Example
189    ///
190    /// ```text
191    /// let stmt = session.prepare("SELECT * FROM users WHERE id = ?")?;
192    /// assert_eq!(stmt.param_count(), 1);
193    /// ```
194    pub fn prepare(&self, sql: &str) -> Result<Arc<PreparedStatement>, SessionError> {
195        self.cache.get_or_prepare(sql).map_err(SessionError::from)
196    }
197
198    /// Prepare a SQL SELECT statement using arena allocation
199    ///
200    /// This is optimized for SELECT statements and provides better cache locality.
201    /// Arena-based statements store the parsed AST in contiguous memory, which
202    /// can improve performance for frequently executed queries.
203    ///
204    /// For non-SELECT statements, this will return an error - use `prepare()` instead.
205    ///
206    /// # Performance Benefits
207    ///
208    /// Arena allocation provides:
209    /// - Better cache locality (contiguous memory layout)
210    /// - Lower allocation overhead (single arena vs multiple heap allocations)
211    /// - Potential for zero-copy parameter binding in future phases
212    ///
213    /// # Example
214    ///
215    /// ```text
216    /// let stmt = session.prepare_arena("SELECT * FROM users WHERE id = ?")?;
217    /// assert_eq!(stmt.param_count(), 1);
218    /// ```
219    pub fn prepare_arena(&self, sql: &str) -> Result<Arc<ArenaPreparedStatement>, ArenaParseError> {
220        self.cache.get_or_prepare_arena(sql)
221    }
222
223    /// Execute a prepared SELECT statement with parameters
224    ///
225    /// Binds the parameters to the prepared statement and executes it.
226    /// This is the fast path for repeated queries - no SQL parsing occurs.
227    ///
228    /// For simple PK point lookups, uses cached execution plan to bypass
229    /// the full query execution pipeline.
230    ///
231    /// # Example
232    ///
233    /// ```text
234    /// let stmt = session.prepare("SELECT * FROM users WHERE id = ?")?;
235    /// let result = session.execute_prepared(&stmt, &[SqlValue::Integer(42)])?;
236    /// ```
237    pub fn execute_prepared(
238        &self,
239        stmt: &PreparedStatement,
240        params: &[SqlValue],
241    ) -> Result<PreparedExecutionResult, SessionError> {
242        // Check parameter count first
243        if params.len() != stmt.param_count() {
244            return Err(SessionError::PreparedStatement(
245                PreparedStatementError::ParameterCountMismatch {
246                    expected: stmt.param_count(),
247                    actual: params.len(),
248                },
249            ));
250        }
251
252        // Try fast-path execution using cached plan
253        if let CachedPlan::PkPointLookup(plan) = stmt.cached_plan() {
254            if let Some(result) = self.try_execute_pk_lookup(plan, params)? {
255                return Ok(result);
256            }
257            // Fall through to standard execution if fast path fails
258        }
259
260        // Bind parameters to get executable statement
261        let bound_stmt = stmt.bind(params)?;
262
263        // Execute based on statement type
264        self.execute_statement(&bound_stmt)
265    }
266
267    /// Try to execute a PK point lookup using the cached plan
268    ///
269    /// Returns `Ok(Some(result))` if execution succeeded via fast path,
270    /// `Ok(None)` if we need to fall back to standard execution,
271    /// or `Err` if a real error occurred.
272    fn try_execute_pk_lookup(
273        &self,
274        plan: &PkPointLookupPlan,
275        params: &[SqlValue],
276    ) -> Result<Option<PreparedExecutionResult>, SessionError> {
277        // Get the table
278        let table = match self.db.get_table(&plan.table_name) {
279            Some(t) => t,
280            None => return Ok(None), // Table doesn't exist or it's a view - fall back
281        };
282
283        // Verify PK columns match what we expected
284        let actual_pk_columns = match &table.schema.primary_key {
285            Some(cols) if cols.len() == plan.pk_columns.len() => cols,
286            _ => return Ok(None), // PK structure changed - fall back
287        };
288
289        // Build PK values in the correct order
290        let mut pk_values = Vec::with_capacity(plan.pk_columns.len());
291        for (param_idx, pk_col_idx) in &plan.param_to_pk_col {
292            if *param_idx >= params.len() || *pk_col_idx >= plan.pk_columns.len() {
293                return Ok(None); // Invalid mapping - fall back
294            }
295
296            // Verify the column name still matches
297            let expected_col = &plan.pk_columns[*pk_col_idx];
298            let actual_col = &actual_pk_columns[*pk_col_idx];
299            if !expected_col.eq_ignore_ascii_case(actual_col) {
300                return Ok(None); // Column mismatch - fall back
301            }
302
303            pk_values.push(params[*param_idx].clone());
304        }
305
306        // Perform the PK lookup
307        let row = if pk_values.len() == 1 {
308            self.db
309                .get_row_by_pk(&plan.table_name, &pk_values[0])
310                .map_err(|e| SessionError::Execution(ExecutorError::StorageError(e.to_string())))?
311        } else {
312            self.db
313                .get_row_by_composite_pk(&plan.table_name, &pk_values)
314                .map_err(|e| SessionError::Execution(ExecutorError::StorageError(e.to_string())))?
315        };
316
317        let rows = match row {
318            Some(r) => vec![r.clone()],
319            None => vec![],
320        };
321
322        // Apply projection if needed
323        let (columns, result_rows) = match &plan.projection {
324            ProjectionPlan::Wildcard => {
325                // Return all columns
326                let columns: Vec<String> = table.schema.columns.iter().map(|c| c.name.clone()).collect();
327                (columns, rows)
328            }
329            ProjectionPlan::Columns(projections) => {
330                // Build column indices for projection
331                let mut col_indices = Vec::with_capacity(projections.len());
332                let mut column_names = Vec::with_capacity(projections.len());
333
334                for proj in projections {
335                    let idx = table
336                        .schema
337                        .columns
338                        .iter()
339                        .position(|c| c.name.eq_ignore_ascii_case(&proj.column_name));
340
341                    match idx {
342                        Some(i) => {
343                            col_indices.push(i);
344                            column_names.push(proj.alias.clone().unwrap_or_else(|| proj.column_name.clone()));
345                        }
346                        None => return Ok(None), // Column not found - fall back
347                    }
348                }
349
350                // Project rows
351                let projected_rows: Vec<Row> = rows
352                    .into_iter()
353                    .map(|row| {
354                        let projected_values: Vec<SqlValue> =
355                            col_indices.iter().map(|&i| row.values[i].clone()).collect();
356                        Row::new(projected_values)
357                    })
358                    .collect();
359
360                (column_names, projected_rows)
361            }
362        };
363
364        Ok(Some(PreparedExecutionResult::Select(SelectResult {
365            columns,
366            rows: result_rows,
367        })))
368    }
369
370    /// Execute a bound statement (internal helper)
371    fn execute_statement(&self, stmt: &Statement) -> Result<PreparedExecutionResult, SessionError> {
372        match stmt {
373            Statement::Select(select_stmt) => {
374                let executor = SelectExecutor::new(self.db);
375                let result = executor.execute_with_columns(select_stmt)?;
376                Ok(PreparedExecutionResult::Select(result))
377            }
378            _ => Err(SessionError::UnsupportedStatement(
379                "Only SELECT is supported for read-only sessions. Use SessionMut for DML.".into(),
380            )),
381        }
382    }
383}
384
385/// Mutable session for executing prepared statements that modify data
386///
387/// Use this session type when you need to execute INSERT, UPDATE, or DELETE
388/// statements in addition to SELECT.
389pub struct SessionMut<'a> {
390    db: &'a mut Database,
391    cache: Arc<PreparedStatementCache>,
392}
393
394impl<'a> SessionMut<'a> {
395    /// Create a new mutable session with a reference to the database
396    pub fn new(db: &'a mut Database) -> Self {
397        Self {
398            db,
399            cache: Arc::new(PreparedStatementCache::default_cache()),
400        }
401    }
402
403    /// Create a new mutable session with a custom cache size
404    pub fn with_cache_size(db: &'a mut Database, cache_size: usize) -> Self {
405        Self {
406            db,
407            cache: Arc::new(PreparedStatementCache::new(cache_size)),
408        }
409    }
410
411    /// Create a new mutable session with a shared cache
412    pub fn with_shared_cache(db: &'a mut Database, cache: Arc<PreparedStatementCache>) -> Self {
413        Self { db, cache }
414    }
415
416    /// Get the underlying database reference (immutable)
417    pub fn database(&self) -> &Database {
418        self.db
419    }
420
421    /// Get the underlying database reference (mutable)
422    pub fn database_mut(&mut self) -> &mut Database {
423        self.db
424    }
425
426    /// Get the prepared statement cache
427    pub fn cache(&self) -> &PreparedStatementCache {
428        &self.cache
429    }
430
431    /// Get the shared cache Arc
432    pub fn shared_cache(&self) -> Arc<PreparedStatementCache> {
433        Arc::clone(&self.cache)
434    }
435
436    /// Prepare a SQL statement for execution
437    pub fn prepare(&self, sql: &str) -> Result<Arc<PreparedStatement>, SessionError> {
438        self.cache.get_or_prepare(sql).map_err(SessionError::from)
439    }
440
441    /// Prepare a SQL SELECT statement using arena allocation
442    ///
443    /// See [`Session::prepare_arena`] for details.
444    pub fn prepare_arena(&self, sql: &str) -> Result<Arc<ArenaPreparedStatement>, ArenaParseError> {
445        self.cache.get_or_prepare_arena(sql)
446    }
447
448    /// Execute a prepared statement with parameters (read-only)
449    ///
450    /// Use this for SELECT queries.
451    pub fn execute_prepared(
452        &self,
453        stmt: &PreparedStatement,
454        params: &[SqlValue],
455    ) -> Result<PreparedExecutionResult, SessionError> {
456        let bound_stmt = stmt.bind(params)?;
457        self.execute_statement_readonly(&bound_stmt)
458    }
459
460    /// Execute a prepared statement with parameters (read-write)
461    ///
462    /// Use this for INSERT, UPDATE, DELETE statements.
463    pub fn execute_prepared_mut(
464        &mut self,
465        stmt: &PreparedStatement,
466        params: &[SqlValue],
467    ) -> Result<PreparedExecutionResult, SessionError> {
468        let bound_stmt = stmt.bind(params)?;
469        self.execute_statement_mut(&bound_stmt)
470    }
471
472    /// Execute a read-only statement
473    fn execute_statement_readonly(
474        &self,
475        stmt: &Statement,
476    ) -> Result<PreparedExecutionResult, SessionError> {
477        match stmt {
478            Statement::Select(select_stmt) => {
479                let executor = SelectExecutor::new(self.db);
480                let result = executor.execute_with_columns(select_stmt)?;
481                Ok(PreparedExecutionResult::Select(result))
482            }
483            _ => Err(SessionError::UnsupportedStatement(
484                "Use execute_prepared_mut for DML statements".into(),
485            )),
486        }
487    }
488
489    /// Execute a statement that may modify data
490    fn execute_statement_mut(
491        &mut self,
492        stmt: &Statement,
493    ) -> Result<PreparedExecutionResult, SessionError> {
494        match stmt {
495            Statement::Select(select_stmt) => {
496                let executor = SelectExecutor::new(self.db);
497                let result = executor.execute_with_columns(select_stmt)?;
498                Ok(PreparedExecutionResult::Select(result))
499            }
500            Statement::Insert(insert_stmt) => {
501                let rows_affected = InsertExecutor::execute(self.db, insert_stmt)?;
502                // Invalidate cache for affected table
503                self.cache.invalidate_table(&insert_stmt.table_name);
504                Ok(PreparedExecutionResult::RowsAffected(rows_affected))
505            }
506            Statement::Update(update_stmt) => {
507                let rows_affected = UpdateExecutor::execute(update_stmt, self.db)?;
508                // Invalidate cache for affected table
509                self.cache.invalidate_table(&update_stmt.table_name);
510                Ok(PreparedExecutionResult::RowsAffected(rows_affected))
511            }
512            Statement::Delete(delete_stmt) => {
513                let rows_affected = DeleteExecutor::execute(delete_stmt, self.db)?;
514                // Invalidate cache for affected table
515                self.cache.invalidate_table(&delete_stmt.table_name);
516                Ok(PreparedExecutionResult::RowsAffected(rows_affected))
517            }
518            _ => Err(SessionError::UnsupportedStatement(format!(
519                "Statement type {:?} not supported for prepared execution",
520                std::mem::discriminant(stmt)
521            ))),
522        }
523    }
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529    use vibesql_catalog::{ColumnSchema, TableSchema};
530    use vibesql_types::DataType;
531
532    fn create_test_db() -> Database {
533        let mut db = Database::new();
534        // Enable case-insensitive identifiers (default MySQL behavior)
535        db.catalog.set_case_sensitive_identifiers(false);
536
537        // Create users table
538        let columns = vec![
539            ColumnSchema::new("id".to_string(), DataType::Integer, false),
540            ColumnSchema::new("name".to_string(), DataType::Varchar { max_length: Some(100) }, true),
541        ];
542        let schema = TableSchema::with_primary_key(
543            "users".to_string(),
544            columns,
545            vec!["id".to_string()],
546        );
547        db.create_table(schema).unwrap();
548
549        // Insert test data
550        let row1 = Row::new(vec![SqlValue::Integer(1), SqlValue::Varchar("Alice".into())]);
551        let row2 = Row::new(vec![SqlValue::Integer(2), SqlValue::Varchar("Bob".into())]);
552        let row3 = Row::new(vec![SqlValue::Integer(3), SqlValue::Varchar("Charlie".into())]);
553
554        db.insert_row("users", row1).unwrap();
555        db.insert_row("users", row2).unwrap();
556        db.insert_row("users", row3).unwrap();
557
558        db
559    }
560
561    #[test]
562    fn test_session_prepare() {
563        let db = create_test_db();
564        let session = Session::new(&db);
565
566        let stmt = session.prepare("SELECT * FROM users WHERE id = ?").unwrap();
567        assert_eq!(stmt.param_count(), 1);
568    }
569
570    #[test]
571    fn test_session_execute_prepared() {
572        let db = create_test_db();
573        let session = Session::new(&db);
574
575        let stmt = session.prepare("SELECT * FROM users WHERE id = ?").unwrap();
576
577        // Execute with id = 1
578        let result = session
579            .execute_prepared(&stmt, &[SqlValue::Integer(1)])
580            .unwrap();
581
582        if let PreparedExecutionResult::Select(select_result) = result {
583            assert_eq!(select_result.rows.len(), 1);
584            assert_eq!(select_result.rows[0].values[0], SqlValue::Integer(1));
585            assert_eq!(
586                select_result.rows[0].values[1],
587                SqlValue::Varchar("Alice".into())
588            );
589        } else {
590            panic!("Expected Select result");
591        }
592    }
593
594    #[test]
595    fn test_session_reuse_prepared() {
596        let db = create_test_db();
597        let session = Session::new(&db);
598
599        let stmt = session.prepare("SELECT * FROM users WHERE id = ?").unwrap();
600
601        // Execute multiple times with different parameters
602        let result1 = session
603            .execute_prepared(&stmt, &[SqlValue::Integer(1)])
604            .unwrap();
605        let result2 = session
606            .execute_prepared(&stmt, &[SqlValue::Integer(2)])
607            .unwrap();
608        let result3 = session
609            .execute_prepared(&stmt, &[SqlValue::Integer(3)])
610            .unwrap();
611
612        // Verify each returned the correct row
613        assert_eq!(
614            result1.rows().unwrap()[0].values[1],
615            SqlValue::Varchar("Alice".into())
616        );
617        assert_eq!(
618            result2.rows().unwrap()[0].values[1],
619            SqlValue::Varchar("Bob".into())
620        );
621        assert_eq!(
622            result3.rows().unwrap()[0].values[1],
623            SqlValue::Varchar("Charlie".into())
624        );
625
626        // Verify cache was used (should have 1 miss, then hits)
627        let stats = session.cache().stats();
628        assert_eq!(stats.misses, 1);
629        // hits is always >= 0 since it's a u64, just ensure it exists
630        let _hits = stats.hits;
631    }
632
633    #[test]
634    fn test_session_param_count_mismatch() {
635        let db = create_test_db();
636        let session = Session::new(&db);
637
638        let stmt = session.prepare("SELECT * FROM users WHERE id = ?").unwrap();
639
640        // Try to execute with wrong number of parameters
641        let result = session.execute_prepared(&stmt, &[]);
642        assert!(result.is_err());
643
644        let result = session.execute_prepared(&stmt, &[SqlValue::Integer(1), SqlValue::Integer(2)]);
645        assert!(result.is_err());
646    }
647
648    #[test]
649    fn test_session_mut_insert() {
650        let mut db = create_test_db();
651        let mut session = SessionMut::new(&mut db);
652
653        let stmt = session
654            .prepare("INSERT INTO users (id, name) VALUES (?, ?)")
655            .unwrap();
656
657        let result = session
658            .execute_prepared_mut(
659                &stmt,
660                &[SqlValue::Integer(4), SqlValue::Varchar("David".into())],
661            )
662            .unwrap();
663
664        assert_eq!(result.rows_affected(), Some(1));
665
666        // Verify the row was inserted
667        let select_stmt = session.prepare("SELECT * FROM users WHERE id = ?").unwrap();
668        let select_result = session
669            .execute_prepared(&select_stmt, &[SqlValue::Integer(4)])
670            .unwrap();
671
672        assert_eq!(select_result.rows().unwrap().len(), 1);
673        assert_eq!(
674            select_result.rows().unwrap()[0].values[1],
675            SqlValue::Varchar("David".into())
676        );
677    }
678
679    #[test]
680    fn test_session_mut_update() {
681        let mut db = create_test_db();
682        let mut session = SessionMut::new(&mut db);
683
684        let stmt = session
685            .prepare("UPDATE users SET name = ? WHERE id = ?")
686            .unwrap();
687
688        let result = session
689            .execute_prepared_mut(
690                &stmt,
691                &[SqlValue::Varchar("Alicia".into()), SqlValue::Integer(1)],
692            )
693            .unwrap();
694
695        assert_eq!(result.rows_affected(), Some(1));
696
697        // Verify the row was updated
698        let select_stmt = session.prepare("SELECT * FROM users WHERE id = ?").unwrap();
699        let select_result = session
700            .execute_prepared(&select_stmt, &[SqlValue::Integer(1)])
701            .unwrap();
702
703        assert_eq!(
704            select_result.rows().unwrap()[0].values[1],
705            SqlValue::Varchar("Alicia".into())
706        );
707    }
708
709    #[test]
710    fn test_session_mut_delete() {
711        let mut db = create_test_db();
712        let mut session = SessionMut::new(&mut db);
713
714        let stmt = session.prepare("DELETE FROM users WHERE id = ?").unwrap();
715
716        let result = session
717            .execute_prepared_mut(&stmt, &[SqlValue::Integer(1)])
718            .unwrap();
719
720        assert_eq!(result.rows_affected(), Some(1));
721
722        // Verify the row was deleted
723        let select_stmt = session.prepare("SELECT * FROM users WHERE id = ?").unwrap();
724        let select_result = session
725            .execute_prepared(&select_stmt, &[SqlValue::Integer(1)])
726            .unwrap();
727
728        assert_eq!(select_result.rows().unwrap().len(), 0);
729    }
730
731    #[test]
732    fn test_shared_cache() {
733        let db = create_test_db();
734
735        // Create first session and prepare a statement
736        let session1 = Session::new(&db);
737        let stmt = session1
738            .prepare("SELECT * FROM users WHERE id = ?")
739            .unwrap();
740
741        // Get the shared cache
742        let shared_cache = session1.shared_cache();
743        let initial_misses = session1.cache().stats().misses;
744
745        // Create second session with shared cache
746        let session2 = Session::with_shared_cache(&db, shared_cache);
747
748        // Prepare the same statement - should hit cache
749        let _stmt2 = session2
750            .prepare("SELECT * FROM users WHERE id = ?")
751            .unwrap();
752
753        // Verify cache was shared (no additional miss)
754        assert_eq!(session2.cache().stats().misses, initial_misses);
755
756        // Execute on both sessions
757        let result1 = session1
758            .execute_prepared(&stmt, &[SqlValue::Integer(1)])
759            .unwrap();
760        let result2 = session2
761            .execute_prepared(&stmt, &[SqlValue::Integer(2)])
762            .unwrap();
763
764        assert_eq!(
765            result1.rows().unwrap()[0].values[1],
766            SqlValue::Varchar("Alice".into())
767        );
768        assert_eq!(
769            result2.rows().unwrap()[0].values[1],
770            SqlValue::Varchar("Bob".into())
771        );
772    }
773
774    #[test]
775    fn test_no_params_statement() {
776        let db = create_test_db();
777        let session = Session::new(&db);
778
779        let stmt = session.prepare("SELECT * FROM users").unwrap();
780        assert_eq!(stmt.param_count(), 0);
781
782        let result = session.execute_prepared(&stmt, &[]).unwrap();
783        assert_eq!(result.rows().unwrap().len(), 3);
784    }
785
786    #[test]
787    fn test_multiple_placeholders() {
788        let db = create_test_db();
789        let session = Session::new(&db);
790
791        let stmt = session
792            .prepare("SELECT * FROM users WHERE id >= ? AND id <= ?")
793            .unwrap();
794        assert_eq!(stmt.param_count(), 2);
795
796        let result = session
797            .execute_prepared(&stmt, &[SqlValue::Integer(1), SqlValue::Integer(2)])
798            .unwrap();
799
800        assert_eq!(result.rows().unwrap().len(), 2);
801    }
802
803    #[test]
804    fn test_session_prepare_arena() {
805        let db = create_test_db();
806        let session = Session::new(&db);
807
808        // Test prepare_arena for SELECT statement
809        let stmt = session
810            .prepare_arena("SELECT * FROM users WHERE id = ?")
811            .unwrap();
812        assert_eq!(stmt.param_count(), 1);
813
814        // Verify tables were extracted (arena parser uses uppercase)
815        assert!(stmt.tables().contains("USERS"));
816
817        // Test caching - second call should hit cache
818        let stmt2 = session
819            .prepare_arena("SELECT * FROM users WHERE id = ?")
820            .unwrap();
821        assert_eq!(stmt2.param_count(), 1);
822
823        // Verify it's the same cached statement
824        assert!(std::sync::Arc::ptr_eq(&stmt, &stmt2));
825    }
826
827    #[test]
828    fn test_session_prepare_arena_no_params() {
829        let db = create_test_db();
830        let session = Session::new(&db);
831
832        let stmt = session
833            .prepare_arena("SELECT * FROM users")
834            .unwrap();
835        assert_eq!(stmt.param_count(), 0);
836    }
837
838    #[test]
839    fn test_session_prepare_arena_join() {
840        let db = create_test_db();
841        let session = Session::new(&db);
842
843        // Create orders table first
844        use vibesql_catalog::{ColumnSchema, TableSchema};
845        use vibesql_types::DataType;
846
847        let orders_columns = vec![
848            ColumnSchema::new("id".to_string(), DataType::Integer, false),
849            ColumnSchema::new("user_id".to_string(), DataType::Integer, false),
850        ];
851        let _orders_schema = TableSchema::with_primary_key(
852            "orders".to_string(),
853            orders_columns,
854            vec!["id".to_string()],
855        );
856
857        // We need a mutable db for this test, so we'll just test the parse works
858        // without actually querying
859        let stmt = session
860            .prepare_arena("SELECT u.id FROM users u JOIN orders o ON u.id = o.user_id")
861            .unwrap();
862
863        // Both tables should be tracked
864        let tables = stmt.tables();
865        assert!(tables.contains("USERS"), "Expected USERS in {:?}", tables);
866        assert!(tables.contains("ORDERS"), "Expected ORDERS in {:?}", tables);
867    }
868
869    #[test]
870    fn test_session_mut_prepare_arena() {
871        let mut db = create_test_db();
872        let session = SessionMut::new(&mut db);
873
874        // Test prepare_arena for SELECT statement
875        let stmt = session
876            .prepare_arena("SELECT * FROM users WHERE id = ?")
877            .unwrap();
878        assert_eq!(stmt.param_count(), 1);
879    }
880}
vibesql_executor/session.rs

vibesql_executor/
session.rs
