intelli_shell/storage/mod.rs
1use std::{
2 path::Path,
3 sync::{Arc, atomic::AtomicBool},
4};
5
6use client::{SqliteClient, SqliteClientBuilder};
7use color_eyre::eyre::Context;
8use itertools::Itertools;
9use migrations::MIGRATIONS;
10use regex::Regex;
11use rusqlite::{OpenFlags, functions::FunctionFlags};
12
13use crate::{
14 errors::Result,
15 utils::{COMMAND_VARIABLE_REGEX_QUOTES, SplitCaptures, SplitItem},
16};
17
18mod client;
19mod migrations;
20mod queries;
21mod sea_query_rusqlite;
22
23mod command;
24mod completion;
25mod import_export;
26mod release;
27mod variable;
28mod version;
29
30pub(crate) use sea_query_rusqlite::RusqliteBinder;
31
32type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
33
34/// `SqliteStorage` provides an interface for interacting with a SQLite database to store and retrieve application data,
35/// primarily [`Command`] and [`VariableValue`] entities
36#[derive(Clone)]
37pub struct SqliteStorage {
38 /// Whether the workspace-level temp tables are created
39 workspace_tables_loaded: Arc<AtomicBool>,
40 /// The SQLite client used for database operations
41 client: Arc<SqliteClient>,
42}
43
44impl SqliteStorage {
45 /// Creates a new instance of [`SqliteStorage`] using a persistent database file.
46 ///
47 /// If INTELLI_STORAGE environment variable is set, it will use the specified path for the database file.
48 pub async fn new(data_dir: impl AsRef<Path>) -> Result<Self> {
49 let builder = if let Some(path) = std::env::var_os("INTELLI_STORAGE") {
50 // If INTELLI_STORAGE is set, use it as the database path
51 tracing::info!("Using INTELLI_STORAGE path: {}", path.to_string_lossy());
52 SqliteClientBuilder::new().path(path)
53 } else {
54 // Otherwise, use the provided data directory
55 let db_path = data_dir.as_ref().join("storage.db3");
56 tracing::info!("Using default storage path: {}", db_path.display());
57 SqliteClientBuilder::new().path(db_path)
58 };
59 Ok(Self {
60 workspace_tables_loaded: Arc::new(AtomicBool::new(false)),
61 client: Arc::new(Self::open_client(builder).await?),
62 })
63 }
64
65 /// Creates a new in-memory instance of [`SqliteStorage`].
66 ///
67 /// This is primarily intended for testing purposes, where a persistent database is not required.
68 #[cfg(test)]
69 pub async fn new_in_memory() -> Result<Self> {
70 let client = Self::open_client(SqliteClientBuilder::new()).await?;
71 Ok(Self {
72 workspace_tables_loaded: Arc::new(AtomicBool::new(false)),
73 client: Arc::new(client),
74 })
75 }
76
77 /// Opens and initializes an SQLite client.
78 ///
79 /// This internal helper function configures the client with necessary PRAGMA settings for optimal performance and
80 /// data integrity (WAL mode, normal sync, foreign keys) and applies all pending database migrations.
81 async fn open_client(builder: SqliteClientBuilder) -> Result<SqliteClient> {
82 // Build the client
83 let client = builder
84 .flags(OpenFlags::default())
85 .open()
86 .await
87 .wrap_err("Error initializing SQLite client")?;
88
89 // Use Write-Ahead Logging (WAL) mode for better concurrency and performance.
90 client
91 .conn(|conn| {
92 Ok(conn
93 .pragma_update(None, "journal_mode", "wal")
94 .wrap_err("Error applying journal mode pragma")?)
95 })
96 .await?;
97
98 // Set synchronous mode to NORMAL. This means SQLite will still sync at critical moments, but less frequently
99 // than FULL, offering a good balance between safety and performance.
100 client
101 .conn(|conn| {
102 Ok(conn
103 .pragma_update(None, "synchronous", "normal")
104 .wrap_err("Error applying synchronous pragma")?)
105 })
106 .await?;
107
108 // Enforce foreign key constraints to maintain data integrity.
109 // This has a slight performance cost but is crucial for relational data.
110 client
111 .conn(|conn| {
112 Ok(conn
113 .pragma_update(None, "foreign_keys", "on")
114 .wrap_err("Error applying foreign keys pragma")?)
115 })
116 .await?;
117
118 // Store temp schema in memory
119 client
120 .conn(|conn| {
121 Ok(conn
122 .pragma_update(None, "temp_store", "memory")
123 .wrap_err("Error applying temp store pragma")?)
124 })
125 .await?;
126
127 // Apply all defined database migrations to bring the schema to the latest version.
128 // This is done atomically within a transaction.
129 client
130 .conn_mut(|conn| Ok(MIGRATIONS.to_latest(conn).wrap_err("Error applying migrations")?))
131 .await?;
132
133 // Add a regexp function to the client
134 client
135 .conn(|conn| {
136 Ok(conn
137 .create_scalar_function(
138 "regexp",
139 2,
140 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
141 |ctx| {
142 assert_eq!(ctx.len(), 2, "regexp() called with unexpected number of arguments");
143
144 let text = ctx
145 .get_raw(1)
146 .as_str_or_null()
147 .map_err(|e| rusqlite::Error::UserFunctionError(e.into()))?;
148
149 let Some(text) = text else {
150 return Ok(false);
151 };
152
153 let cached_re: Arc<Regex> =
154 ctx.get_or_create_aux(0, |vr| Ok::<_, BoxError>(Regex::new(vr.as_str()?)?))?;
155
156 Ok(cached_re.is_match(text))
157 },
158 )
159 .wrap_err("Error adding regexp function")?)
160 })
161 .await?;
162
163 // Add a cmd-to-regex function
164 client
165 .conn(|conn| {
166 Ok(conn
167 .create_scalar_function(
168 "cmd_to_regex",
169 1,
170 FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
171 |ctx| {
172 assert_eq!(
173 ctx.len(),
174 1,
175 "cmd_to_regex() called with unexpected number of arguments"
176 );
177 let cmd_template = ctx.get::<String>(0)?;
178
179 // Use the SplitCaptures iterator to process both unmatched literals and captured variables
180 let regex_body = SplitCaptures::new(&COMMAND_VARIABLE_REGEX_QUOTES, &cmd_template)
181 .filter_map(|item| match item {
182 // For unmatched parts, trim them and escape any special regex chars
183 SplitItem::Unmatched(s) => {
184 let trimmed = s.trim();
185 if trimmed.is_empty() {
186 None
187 } else {
188 Some(regex::escape(trimmed))
189 }
190 }
191 // For captured parts (the variables), replace them with a capture group
192 SplitItem::Captured(caps) => {
193 // Check which capture group matched to see if the placeholder was quoted
194 let placeholder_regex = if caps.get(1).is_some() {
195 // Group 1 matched '{{...}}', so expect a single-quoted argument
196 r"('[^']*')"
197 } else if caps.get(2).is_some() {
198 // Group 2 matched "{{...}}", so expect a double-quoted argument
199 r#"("[^"]*")"#
200 } else {
201 // Group 3 matched {{...}}, so expect a generic argument
202 r#"('[^']*'|"[^"]*"|\S+)"#
203 };
204 Some(String::from(placeholder_regex))
205 }
206 })
207 // Join them by any number of whitespaces
208 .join(r"\s+");
209
210 // Build the final regex
211 Ok(format!("^{regex_body}$"))
212 },
213 )
214 .wrap_err("Error adding cmd-to-regex function")?)
215 })
216 .await?;
217
218 Ok(client)
219 }
220
221 #[cfg(debug_assertions)]
222 pub async fn query(&self, sql: String) -> Result<String> {
223 self.client
224 .conn(move |conn| {
225 use prettytable::{Cell, Row, Table};
226 use rusqlite::types::Value;
227
228 let mut stmt = conn.prepare(&sql)?;
229 let column_names = stmt
230 .column_names()
231 .into_iter()
232 .map(String::from)
233 .collect::<Vec<String>>();
234 let columns_len = column_names.len();
235 let mut table = Table::new();
236 table.add_row(Row::from(column_names));
237 let rows = stmt.query_map([], |row| {
238 let mut cells = Vec::new();
239 for i in 0..columns_len {
240 let value: Value = row.get(i)?;
241 let cell_value = match value {
242 Value::Null => "NULL".to_string(),
243 Value::Integer(i) => i.to_string(),
244 Value::Real(f) => f.to_string(),
245 Value::Text(t) => t,
246 Value::Blob(_) => "[BLOB]".to_string(),
247 };
248 cells.push(Cell::new(&cell_value));
249 }
250 Ok(Row::from(cells))
251 })?;
252 for row in rows {
253 table.add_row(row?);
254 }
255 Ok(table.to_string())
256 })
257 .await
258 }
259}