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