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