sql_cli/data/
temp_table_registry.rs

1/// Temporary table registry for script execution
2/// Manages temporary tables (prefixed with #) that persist across query batches
3/// within a single script execution. Tables are dropped when the script completes.
4use anyhow::{anyhow, Result};
5use std::collections::HashMap;
6use std::sync::Arc;
7use tracing::{debug, info};
8
9use crate::data::datatable::DataTable;
10
11/// Registry for temporary tables that persist across GO-separated batches
12/// within a single script execution
13#[derive(Debug, Clone)]
14pub struct TempTableRegistry {
15    tables: HashMap<String, Arc<DataTable>>,
16}
17
18impl TempTableRegistry {
19    /// Create a new empty temp table registry
20    pub fn new() -> Self {
21        debug!("Creating new TempTableRegistry");
22        Self {
23            tables: HashMap::new(),
24        }
25    }
26
27    /// Store a temporary table
28    /// Returns error if table already exists (to prevent accidental overwrites)
29    pub fn insert(&mut self, name: String, table: Arc<DataTable>) -> Result<()> {
30        if !name.starts_with('#') {
31            return Err(anyhow!(
32                "Temporary table name must start with #, got: {}",
33                name
34            ));
35        }
36
37        if self.tables.contains_key(&name) {
38            return Err(anyhow!(
39                "Temporary table {} already exists. Drop it first or use a different name.",
40                name
41            ));
42        }
43
44        let row_count = table.row_count();
45        let col_count = table.column_count();
46
47        info!(
48            "Storing temporary table {} ({} rows, {} columns)",
49            name, row_count, col_count
50        );
51
52        self.tables.insert(name, table);
53        Ok(())
54    }
55
56    /// Store a temporary table, replacing if it already exists
57    pub fn insert_or_replace(&mut self, name: String, table: Arc<DataTable>) -> Result<()> {
58        if !name.starts_with('#') {
59            return Err(anyhow!(
60                "Temporary table name must start with #, got: {}",
61                name
62            ));
63        }
64
65        let row_count = table.row_count();
66        let col_count = table.column_count();
67
68        if self.tables.contains_key(&name) {
69            info!(
70                "Replacing temporary table {} ({} rows, {} columns)",
71                name, row_count, col_count
72            );
73        } else {
74            info!(
75                "Storing temporary table {} ({} rows, {} columns)",
76                name, row_count, col_count
77            );
78        }
79
80        self.tables.insert(name, table);
81        Ok(())
82    }
83
84    /// Get a temporary table by name
85    /// Returns None if table doesn't exist
86    pub fn get(&self, name: &str) -> Option<Arc<DataTable>> {
87        debug!("Looking up temporary table: {}", name);
88        self.tables.get(name).cloned()
89    }
90
91    /// Check if a temporary table exists
92    pub fn contains(&self, name: &str) -> bool {
93        self.tables.contains_key(name)
94    }
95
96    /// Drop a temporary table
97    /// Returns true if table existed and was dropped
98    pub fn drop(&mut self, name: &str) -> bool {
99        if let Some(_table) = self.tables.remove(name) {
100            info!("Dropped temporary table: {}", name);
101            true
102        } else {
103            debug!("Temporary table {} does not exist, cannot drop", name);
104            false
105        }
106    }
107
108    /// Get count of temp tables
109    pub fn count(&self) -> usize {
110        self.tables.len()
111    }
112
113    /// Get all temp table names
114    pub fn list_tables(&self) -> Vec<String> {
115        self.tables.keys().cloned().collect()
116    }
117
118    /// Clear all temp tables
119    /// Used when script execution completes
120    pub fn clear(&mut self) {
121        let count = self.tables.len();
122        self.tables.clear();
123        if count > 0 {
124            info!("Cleared {} temporary tables", count);
125        }
126    }
127
128    /// Check if registry is empty
129    pub fn is_empty(&self) -> bool {
130        self.tables.is_empty()
131    }
132}
133
134impl Default for TempTableRegistry {
135    fn default() -> Self {
136        Self::new()
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::data::datatable::{DataColumn, DataRow, DataValue};
144
145    fn create_test_table(rows: usize) -> DataTable {
146        let mut table = DataTable::new("test_table");
147
148        table.add_column(DataColumn::new("id".to_string()));
149        table.add_column(DataColumn::new("name".to_string()));
150
151        for i in 0..rows {
152            table.add_row(DataRow {
153                values: vec![
154                    DataValue::Integer(i as i64),
155                    DataValue::String(format!("row_{}", i)),
156                ],
157            });
158        }
159
160        table
161    }
162
163    #[test]
164    fn test_new_registry() {
165        let registry = TempTableRegistry::new();
166        assert_eq!(registry.count(), 0);
167        assert!(registry.is_empty());
168    }
169
170    #[test]
171    fn test_insert_and_get() {
172        let mut registry = TempTableRegistry::new();
173        let table = create_test_table(10);
174
175        let result = registry.insert("#test".to_string(), Arc::new(table));
176        assert!(result.is_ok());
177        assert_eq!(registry.count(), 1);
178
179        let retrieved = registry.get("#test");
180        assert!(retrieved.is_some());
181        assert_eq!(retrieved.unwrap().row_count(), 10);
182    }
183
184    #[test]
185    fn test_insert_without_hash_fails() {
186        let mut registry = TempTableRegistry::new();
187        let table = create_test_table(5);
188
189        let result = registry.insert("test".to_string(), Arc::new(table));
190        assert!(result.is_err());
191        assert!(result
192            .unwrap_err()
193            .to_string()
194            .contains("must start with #"));
195    }
196
197    #[test]
198    fn test_insert_duplicate_fails() {
199        let mut registry = TempTableRegistry::new();
200        let table1 = create_test_table(5);
201        let table2 = create_test_table(10);
202
203        registry
204            .insert("#test".to_string(), Arc::new(table1))
205            .unwrap();
206        let result = registry.insert("#test".to_string(), Arc::new(table2));
207
208        assert!(result.is_err());
209        assert!(result.unwrap_err().to_string().contains("already exists"));
210    }
211
212    #[test]
213    fn test_insert_or_replace() {
214        let mut registry = TempTableRegistry::new();
215        let table1 = create_test_table(5);
216        let table2 = create_test_table(10);
217
218        registry
219            .insert_or_replace("#test".to_string(), Arc::new(table1))
220            .unwrap();
221        assert_eq!(registry.get("#test").unwrap().row_count(), 5);
222
223        registry
224            .insert_or_replace("#test".to_string(), Arc::new(table2))
225            .unwrap();
226        assert_eq!(registry.get("#test").unwrap().row_count(), 10);
227    }
228
229    #[test]
230    fn test_contains() {
231        let mut registry = TempTableRegistry::new();
232        let table = create_test_table(5);
233
234        assert!(!registry.contains("#test"));
235        registry
236            .insert("#test".to_string(), Arc::new(table))
237            .unwrap();
238        assert!(registry.contains("#test"));
239    }
240
241    #[test]
242    fn test_drop() {
243        let mut registry = TempTableRegistry::new();
244        let table = create_test_table(5);
245
246        registry
247            .insert("#test".to_string(), Arc::new(table))
248            .unwrap();
249        assert_eq!(registry.count(), 1);
250
251        let dropped = registry.drop("#test");
252        assert!(dropped);
253        assert_eq!(registry.count(), 0);
254
255        let dropped_again = registry.drop("#test");
256        assert!(!dropped_again);
257    }
258
259    #[test]
260    fn test_list_tables() {
261        let mut registry = TempTableRegistry::new();
262        let table1 = create_test_table(5);
263        let table2 = create_test_table(10);
264
265        registry
266            .insert("#table1".to_string(), Arc::new(table1))
267            .unwrap();
268        registry
269            .insert("#table2".to_string(), Arc::new(table2))
270            .unwrap();
271
272        let tables = registry.list_tables();
273        assert_eq!(tables.len(), 2);
274        assert!(tables.contains(&"#table1".to_string()));
275        assert!(tables.contains(&"#table2".to_string()));
276    }
277
278    #[test]
279    fn test_clear() {
280        let mut registry = TempTableRegistry::new();
281        let table1 = create_test_table(5);
282        let table2 = create_test_table(10);
283
284        registry
285            .insert("#table1".to_string(), Arc::new(table1))
286            .unwrap();
287        registry
288            .insert("#table2".to_string(), Arc::new(table2))
289            .unwrap();
290
291        assert_eq!(registry.count(), 2);
292
293        registry.clear();
294        assert_eq!(registry.count(), 0);
295        assert!(registry.is_empty());
296    }
297
298    #[test]
299    fn test_multiple_tables() {
300        let mut registry = TempTableRegistry::new();
301
302        for i in 0..5 {
303            let table = create_test_table(i * 10);
304            registry
305                .insert(format!("#table{}", i), Arc::new(table))
306                .unwrap();
307        }
308
309        assert_eq!(registry.count(), 5);
310
311        for i in 0..5 {
312            let table = registry.get(&format!("#table{}", i));
313            assert!(table.is_some());
314            assert_eq!(table.unwrap().row_count(), i * 10);
315        }
316    }
317}
sql_cli/data/temp_table_registry.rs

sql_cli/data/
temp_table_registry.rs
