intelli_shell/storage/

mod.rs

use std::{
    path::Path,
    sync::{Arc, atomic::AtomicBool},
};

use client::{SqliteClient, SqliteClientBuilder};
use color_eyre::{Result, eyre::Context};
use itertools::Itertools;
use migrations::MIGRATIONS;
use regex::Regex;
use rusqlite::{OpenFlags, functions::FunctionFlags};

use crate::utils::{COMMAND_VARIABLE_REGEX_QUOTES, SplitCaptures, SplitItem};

mod client;
mod migrations;
mod queries;

mod command;
mod variable;
mod version;

type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;

/// `SqliteStorage` provides an interface for interacting with a SQLite database to store and retrieve application data,
/// primarily [`Command`] and [`VariableValue`] entities
#[derive(Clone)]
pub struct SqliteStorage {
    /// Whether the workspace-level temp tables are created
    workspace_tables_loaded: Arc<AtomicBool>,
    /// The SQLite client used for database operations
    client: Arc<SqliteClient>,
}

impl SqliteStorage {
    /// Creates a new instance of [`SqliteStorage`] using a persistent database file.
    ///
    /// If INTELLI_STORAGE environment variable is set, it will use the specified path for the database file.
    pub async fn new(data_dir: impl AsRef<Path>) -> Result<Self> {
        let builder = if let Some(path) = std::env::var_os("INTELLI_STORAGE") {
            // If INTELLI_STORAGE is set, use it as the database path
            tracing::info!("Using INTELLI_STORAGE path: {}", path.to_string_lossy());
            SqliteClientBuilder::new().path(path)
        } else {
            // Otherwise, use the provided data directory
            tracing::info!("Using default storage path: {}", data_dir.as_ref().display());
            SqliteClientBuilder::new().path(data_dir.as_ref().join("storage.db3"))
        };
        Ok(Self {
            workspace_tables_loaded: Arc::new(AtomicBool::new(false)),
            client: Arc::new(Self::open_client(builder).await?),
        })
    }

    /// Creates a new in-memory instance of [`SqliteStorage`].
    ///
    /// This is primarily intended for testing purposes, where a persistent database is not required.
    #[cfg(test)]
    pub async fn new_in_memory() -> Result<Self> {
        let client = Self::open_client(SqliteClientBuilder::new()).await?;
        Ok(Self {
            workspace_tables_loaded: Arc::new(AtomicBool::new(false)),
            client: Arc::new(client),
        })
    }

    /// Opens and initializes an SQLite client.
    ///
    /// This internal helper function configures the client with necessary PRAGMA settings for optimal performance and
    /// data integrity (WAL mode, normal sync, foreign keys) and applies all pending database migrations.
    async fn open_client(builder: SqliteClientBuilder) -> Result<SqliteClient> {
        // Build the client
        let client = builder
            .flags(OpenFlags::default())
            .open()
            .await
            .wrap_err("Error initializing SQLite client")?;

        // Use Write-Ahead Logging (WAL) mode for better concurrency and performance.
        client
            .conn(|conn| {
                conn.pragma_update(None, "journal_mode", "wal")
                    .wrap_err("Error applying journal mode pragma")
            })
            .await?;

        // Set synchronous mode to NORMAL. This means SQLite will still sync at critical moments, but less frequently
        // than FULL, offering a good balance between safety and performance.
        client
            .conn(|conn| {
                conn.pragma_update(None, "synchronous", "normal")
                    .wrap_err("Error applying synchronous pragma")
            })
            .await?;

        // Enforce foreign key constraints to maintain data integrity.
        // This has a slight performance cost but is crucial for relational data.
        client
            .conn(|conn| {
                conn.pragma_update(None, "foreign_keys", "on")
                    .wrap_err("Error applying foreign keys pragma")
            })
            .await?;

        // Store temp schema in memory
        client
            .conn(|conn| {
                conn.pragma_update(None, "temp_store", "memory")
                    .wrap_err("Error applying temp store pragma")
            })
            .await?;

        // Apply all defined database migrations to bring the schema to the latest version.
        // This is done atomically within a transaction.
        client
            .conn_mut(|conn| MIGRATIONS.to_latest(conn).wrap_err("Error applying migrations"))
            .await?;

        // Add a regexp function to the client
        client
            .conn(|conn| {
                conn.create_scalar_function(
                    "regexp",
                    2,
                    FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
                    |ctx| {
                        assert_eq!(ctx.len(), 2, "regexp() called with unexpected number of arguments");

                        let text = ctx
                            .get_raw(1)
                            .as_str_or_null()
                            .map_err(|e| rusqlite::Error::UserFunctionError(e.into()))?;

                        let Some(text) = text else {
                            return Ok(false);
                        };

                        let cached_re: Arc<Regex> =
                            ctx.get_or_create_aux(0, |vr| Ok::<_, BoxError>(Regex::new(vr.as_str()?)?))?;

                        Ok(cached_re.is_match(text))
                    },
                )
                .wrap_err("Error adding regexp function")
            })
            .await?;

        // Add a cmd-to-regex function
        client
            .conn(|conn| {
                conn.create_scalar_function(
                    "cmd_to_regex",
                    1,
                    FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
                    |ctx| {
                        assert_eq!(
                            ctx.len(),
                            1,
                            "cmd_to_regex() called with unexpected number of arguments"
                        );
                        let cmd_template = ctx.get::<String>(0)?;

                        // Use the SplitCaptures iterator to process both unmatched literals and captured variables
                        let regex_body = SplitCaptures::new(&COMMAND_VARIABLE_REGEX_QUOTES, &cmd_template)
                            .filter_map(|item| match item {
                                // For unmatched parts, trim them and escape any special regex chars
                                SplitItem::Unmatched(s) => {
                                    let trimmed = s.trim();
                                    if trimmed.is_empty() {
                                        None
                                    } else {
                                        Some(regex::escape(trimmed))
                                    }
                                }
                                // For captured parts (the variables), replace them with a capture group
                                SplitItem::Captured(caps) => {
                                    // Check which capture group matched to see if the placeholder was quoted
                                    let placeholder_regex = if caps.get(1).is_some() {
                                        // Group 1 matched '{{...}}', so expect a single-quoted argument
                                        r"('[^']*')"
                                    } else if caps.get(2).is_some() {
                                        // Group 2 matched "{{...}}", so expect a double-quoted argument
                                        r#"("[^"]*")"#
                                    } else {
                                        // Group 3 matched {{...}}, so expect a generic argument
                                        r#"('[^']*'|"[^"]*"|\S+)"#
                                    };
                                    Some(String::from(placeholder_regex))
                                },
                            })
                            // Join them by any number of whitespaces
                            .join(r"\s+");

                        // Build the final regex
                        Ok(format!("^{regex_body}$"))
                    },
                )
                .wrap_err("Error adding cmd-to-regex function")
            })
            .await?;

        Ok(client)
    }

    #[cfg(debug_assertions)]
    pub async fn query(&self, sql: String) -> Result<String> {
        self.client
            .conn(move |conn| {
                use prettytable::{Cell, Row, Table};
                use rusqlite::types::Value;

                let mut stmt = conn.prepare(&sql)?;
                let column_names = stmt
                    .column_names()
                    .into_iter()
                    .map(String::from)
                    .collect::<Vec<String>>();
                let columns_len = column_names.len();
                let mut table = Table::new();
                table.add_row(Row::from(column_names));
                let rows = stmt.query_map([], |row| {
                    let mut cells = Vec::new();
                    for i in 0..columns_len {
                        let value: Value = row.get(i)?;
                        let cell_value = match value {
                            Value::Null => "NULL".to_string(),
                            Value::Integer(i) => i.to_string(),
                            Value::Real(f) => f.to_string(),
                            Value::Text(t) => t,
                            Value::Blob(_) => "[BLOB]".to_string(),
                        };
                        cells.push(Cell::new(&cell_value));
                    }
                    Ok(Row::from(cells))
                })?;
                for row in rows {
                    table.add_row(row?);
                }
                Ok(table.to_string())
            })
            .await
    }
}