mi6_core/model/

storage.rs

use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use chrono::{DateTime, Utc};

use crate::context::GitBranchInfo;
use crate::input::transcript::FilePosition;

use super::error::StorageError;
use super::{Event, Session, SessionQuery};

/// Aggregated statistics across all sessions from the storage layer.
///
/// Provides a summary of session counts, token usage, costs, and API requests
/// computed at the database level for efficiency.
///
/// Note: This is distinct from `mi6_client::GlobalStats` which includes
/// runtime OS metrics (CPU, memory) for TUI display.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct StorageStats {
    /// Total number of sessions matching the query
    pub session_count: u32,
    /// Number of active sessions (ended_at IS NULL)
    pub active_session_count: u32,
    /// Total tokens across all sessions (input + output + cache_read + cache_write)
    pub total_tokens: i64,
    /// Total cost in USD across all sessions
    pub total_cost_usd: f64,
    /// Total API requests across all sessions
    pub total_api_requests: u32,
}

/// Query builder for storage statistics.
///
/// `StorageStatsQuery` provides a fluent API for filtering which sessions
/// are included in the aggregation.
#[derive(Debug, Clone, Default)]
pub struct StorageStatsQuery {
    /// When true, only include active sessions (ended_at IS NULL)
    pub active_only: bool,
    /// Filter by framework (claude, cursor, gemini, codex)
    pub framework: Option<String>,
}

impl StorageStatsQuery {
    /// Create a new empty query (includes all sessions)
    pub fn new() -> Self {
        Self::default()
    }

    /// Filter to only active sessions
    pub fn active_only(mut self) -> Self {
        self.active_only = true;
        self
    }

    /// Filter by framework
    pub fn with_framework(mut self, framework: impl Into<String>) -> Self {
        self.framework = Some(framework.into());
        self
    }
}

/// Sort order for query results
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Order {
    /// Ascending order (oldest first for timestamps, lowest first for IDs)
    Asc,
    /// Descending order (newest first for timestamps, highest first for IDs)
    #[default]
    Desc,
}

/// Column to order event query results by.
///
/// This enum makes the ordering column explicit, eliminating the previous
/// implicit behavior where ordering was determined by whether `after_id` was set.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum EventOrder {
    /// Order by event timestamp (default: newest first)
    #[default]
    Timestamp,
    /// Order by event ID (default: ascending for watching new events)
    Id,
}

/// Composable query builder for event retrieval.
///
/// `EventQuery` provides a fluent API for constructing event queries with
/// filters, ordering, and pagination.
///
/// # Ordering Behavior
///
/// The query provides **explicit ordering** via the `order_by` field:
/// - [`EventOrder::Timestamp`]: Order by event timestamp (default)
/// - [`EventOrder::Id`]: Order by event ID
///
/// Use `order_by_timestamp()` or `order_by_id()` to set the ordering column,
/// and `with_direction()` to control ascending/descending order.
///
/// # Use Cases
///
/// - **Recent events**: Use `order_by_timestamp()` with `Order::Desc` (default)
/// - **Watching new events**: Use `order_by_id()` with `after_id` for pagination
///
/// # Defaults
///
/// - `order_by`: [`EventOrder::Timestamp`] - orders by timestamp
/// - `direction`: [`Order::Desc`] for timestamp, [`Order::Asc`] for ID
///
/// # Example
/// ```ignore
/// // Get 50 most recent events for a session (orders by timestamp DESC)
/// let query = EventQuery::new()
///     .order_by_timestamp()
///     .with_session(Some("session-123".to_string()))
///     .with_limit(50);
///
/// // Watch for new events after ID 100 (orders by ID ASC)
/// let query = EventQuery::new()
///     .order_by_id()
///     .with_after_id(100)
///     .with_limit(10);
/// ```
#[derive(Default, Clone)]
pub struct EventQuery {
    /// Filter by session ID (mutually exclusive with `session_ids`)
    pub session_id: Option<String>,
    /// Filter by multiple session IDs (mutually exclusive with `session_id`)
    pub session_ids: Option<Vec<String>>,
    /// Filter by event type (case-insensitive)
    pub event_type: Option<String>,
    /// Filter by permission mode
    pub permission_mode: Option<String>,
    /// Filter by framework (claude, cursor, gemini)
    pub framework: Option<String>,
    /// Filter events after this timestamp (exclusive)
    pub after_ts: Option<DateTime<Utc>>,
    /// Filter events before this timestamp (exclusive)
    pub before_ts: Option<DateTime<Utc>>,
    /// Filter events after this ID (exclusive)
    pub after_id: Option<i64>,
    /// Maximum number of events to return
    pub limit: Option<usize>,
    /// Column to order results by (default: Timestamp)
    pub order_by: EventOrder,
    /// Sort direction (default: Desc for timestamp, Asc for ID)
    pub direction: Option<Order>,
    /// When true, only return API request events (events with token data)
    pub api_requests_only: bool,
    /// When true, exclude API request events
    pub exclude_api_requests: bool,
}

impl EventQuery {
    /// Create a new empty query.
    ///
    /// The query defaults to ordering by timestamp (descending).
    /// Use `order_by_id()` for ID-based ordering (e.g., when watching for new events).
    pub fn new() -> Self {
        Self::default()
    }

    /// Order results by timestamp.
    ///
    /// This is the default ordering and is suitable for viewing recent events.
    /// Default direction is [`Order::Desc`] (newest first).
    ///
    /// # Example
    /// ```ignore
    /// let query = EventQuery::new()
    ///     .order_by_timestamp()
    ///     .with_limit(50);
    /// ```
    pub fn order_by_timestamp(mut self) -> Self {
        self.order_by = EventOrder::Timestamp;
        self
    }

    /// Order results by ID.
    ///
    /// Use this when watching for new events with `after_id` pagination.
    /// Default direction is [`Order::Asc`] (oldest first, for forward pagination).
    ///
    /// # Example
    /// ```ignore
    /// let query = EventQuery::new()
    ///     .order_by_id()
    ///     .with_after_id(100)
    ///     .with_limit(10);
    /// ```
    pub fn order_by_id(mut self) -> Self {
        self.order_by = EventOrder::Id;
        self
    }

    /// Filter events after the given timestamp (exclusive)
    pub fn with_after_ts(mut self, ts: DateTime<Utc>) -> Self {
        self.after_ts = Some(ts);
        self
    }

    /// Filter events before the given timestamp (exclusive)
    pub fn with_before_ts(mut self, ts: DateTime<Utc>) -> Self {
        self.before_ts = Some(ts);
        self
    }

    /// Filter events between two timestamps (exclusive on both ends).
    ///
    /// Note: If `start >= end`, the filter will return no results.
    /// The caller is responsible for ensuring `start < end` for meaningful queries.
    pub fn with_between_ts(mut self, start: DateTime<Utc>, end: DateTime<Utc>) -> Self {
        self.after_ts = Some(start);
        self.before_ts = Some(end);
        self
    }

    /// Filter events after the given ID (exclusive).
    ///
    /// Use with `order_by_id()` for consistent pagination when watching for new events.
    ///
    /// # Example
    /// ```ignore
    /// let query = EventQuery::new()
    ///     .order_by_id()
    ///     .with_after_id(last_seen_id)
    ///     .with_limit(10);
    /// ```
    pub fn with_after_id(mut self, id: i64) -> Self {
        self.after_id = Some(id);
        self
    }

    /// Set the maximum number of events to return
    pub fn with_limit(mut self, limit: usize) -> Self {
        self.limit = Some(limit);
        self
    }

    /// Set the sort direction (Asc/Desc).
    ///
    /// This controls whether results are in ascending or descending order.
    /// If not set, defaults to [`Order::Desc`] for timestamp ordering
    /// and [`Order::Asc`] for ID ordering.
    pub fn with_direction(mut self, direction: Order) -> Self {
        self.direction = Some(direction);
        self
    }

    /// Set session filter if value is Some.
    ///
    /// Note: This clears any previously set `session_ids` filter.
    pub fn with_session(mut self, session_id: Option<String>) -> Self {
        self.session_id = session_id;
        self.session_ids = None;
        self
    }

    /// Filter events by multiple session IDs.
    ///
    /// Use this to efficiently query events for multiple sessions in a single
    /// database query, avoiding the N+1 query problem.
    ///
    /// Note: This clears any previously set `session_id` filter.
    /// Passing an empty slice will result in no events being returned.
    ///
    /// # Example
    /// ```ignore
    /// // Get events for multiple sessions at once
    /// let events = storage.query(
    ///     &EventQuery::new()
    ///         .with_sessions(&["session-1", "session-2", "session-3"])
    ///         .with_event_type(Some("UserPromptSubmit".to_string()))
    ///         .with_limit(100)
    /// )?;
    /// ```
    pub fn with_sessions(mut self, session_ids: &[impl AsRef<str>]) -> Self {
        // Store the session_ids - even empty vec is stored to distinguish from None
        self.session_ids = Some(session_ids.iter().map(|s| s.as_ref().to_string()).collect());
        self.session_id = None;
        self
    }

    /// Set event type filter if value is Some
    pub fn with_event_type(mut self, event_type: Option<String>) -> Self {
        self.event_type = event_type;
        self
    }

    /// Set permission mode filter if value is Some
    pub fn with_permission_mode(mut self, permission_mode: Option<String>) -> Self {
        self.permission_mode = permission_mode;
        self
    }

    /// Filter by framework
    pub fn with_framework(mut self, framework: impl Into<String>) -> Self {
        self.framework = Some(framework.into());
        self
    }

    /// Only return API request events (events with token data)
    pub fn api_requests_only(mut self) -> Self {
        self.api_requests_only = true;
        self.exclude_api_requests = false;
        self
    }

    /// Exclude API request events from results
    pub fn exclude_api_requests(mut self) -> Self {
        self.exclude_api_requests = true;
        self.api_requests_only = false;
        self
    }

    /// Returns the effective sort direction for this query.
    ///
    /// Storage implementations should use this method to determine the sort
    /// direction when executing queries.
    ///
    /// # Returns
    /// - If `direction` is explicitly set, returns that value.
    /// - For [`EventOrder::Id`], defaults to [`Order::Asc`] (for watching new events).
    /// - For [`EventOrder::Timestamp`], defaults to [`Order::Desc`] (most recent first).
    ///
    /// # Example
    /// ```
    /// use mi6_core::{EventQuery, Order};
    ///
    /// // Default: Desc for timestamp ordering
    /// let query = EventQuery::new();
    /// assert_eq!(query.effective_direction(), Order::Desc);
    ///
    /// // With order_by_id: Asc for watching
    /// let query = EventQuery::new().order_by_id();
    /// assert_eq!(query.effective_direction(), Order::Asc);
    ///
    /// // Explicit direction overrides default
    /// let query = EventQuery::new().with_direction(Order::Asc);
    /// assert_eq!(query.effective_direction(), Order::Asc);
    /// ```
    pub fn effective_direction(&self) -> Order {
        self.direction.unwrap_or(match self.order_by {
            EventOrder::Id => Order::Asc,
            EventOrder::Timestamp => Order::Desc,
        })
    }

    /// Returns true if this query should order by ID instead of timestamp.
    ///
    /// Storage implementations should use this method to determine which column
    /// to use for ordering results.
    ///
    /// # Example
    /// ```
    /// use mi6_core::EventQuery;
    ///
    /// // Default: order by timestamp
    /// let query = EventQuery::new();
    /// assert!(!query.orders_by_id());
    ///
    /// // With order_by_id(): order by ID
    /// let query = EventQuery::new().order_by_id();
    /// assert!(query.orders_by_id());
    /// ```
    pub fn orders_by_id(&self) -> bool {
        matches!(self.order_by, EventOrder::Id)
    }
}

/// Storage backend trait for mi6 events
///
/// All events (including API requests) are stored in a unified events table.
/// API requests are distinguished by having the `ApiRequest` event type and
/// populated token fields.
pub trait Storage {
    /// Insert a new event, returns the event ID
    fn insert(&self, event: &Event) -> Result<i64, StorageError>;

    /// Query events using the composable EventQuery builder.
    ///
    /// This is the primary method for retrieving events. Use `EventQuery` to
    /// specify filters, ordering, and pagination.
    ///
    /// # Ordering
    /// - Use `order_by_timestamp()` for recent events (default, descending)
    /// - Use `order_by_id()` for watching new events (ascending by default)
    ///
    /// # API Requests
    /// Use `api_requests_only()` to filter to only API request events,
    /// or `exclude_api_requests()` to exclude them.
    ///
    /// # Examples
    /// ```ignore
    /// // Get 50 most recent events (including API requests)
    /// let events = storage.query(&EventQuery::new().order_by_timestamp().with_limit(50))?;
    ///
    /// // Get events for a specific session
    /// let events = storage.query(&EventQuery::new().with_session(Some("sess-123".to_string())).with_limit(100))?;
    ///
    /// // Watch for new events after a known ID
    /// let events = storage.query(&EventQuery::new().order_by_id().with_after_id(last_id).with_limit(10))?;
    ///
    /// // Get only API request events
    /// let api_events = storage.query(&EventQuery::new().api_requests_only().with_limit(50))?;
    /// ```
    fn query(&self, query: &EventQuery) -> Result<Vec<Event>, StorageError>;

    /// Delete events older than the retention period, returns count deleted.
    fn gc(&self, retention: Duration) -> Result<usize, StorageError>;

    /// Count events older than the retention period (for dry-run).
    fn count_expired(&self, retention: Duration) -> Result<usize, StorageError>;

    /// Count total events
    fn count(&self) -> Result<usize, StorageError>;

    /// List sessions matching the query criteria.
    ///
    /// Sessions contain denormalized metadata that is incrementally updated
    /// when events are inserted. This enables single-query monitoring by
    /// returning all session data in one call.
    ///
    /// # Examples
    /// ```ignore
    /// // Get all active sessions, most recent first
    /// let sessions = storage.list_sessions(&SessionQuery::new().active_only())?;
    ///
    /// // Get sessions for a specific framework
    /// let sessions = storage.list_sessions(&SessionQuery::new().with_framework("claude"))?;
    ///
    /// // Get 10 most recently active sessions
    /// let sessions = storage.list_sessions(
    ///     &SessionQuery::new()
    ///         .with_order_by(SessionOrder::LastActivity)
    ///         .with_limit(10)
    /// )?;
    /// ```
    fn list_sessions(&self, query: &SessionQuery) -> Result<Vec<Session>, StorageError>;

    /// Get a single session by ID.
    ///
    /// This queries by session_id alone. If multiple sessions exist with the same
    /// session_id (on different machines), returns the most recently active one.
    ///
    /// For exact lookups with a known machine_id, use [`get_session_by_key`].
    ///
    /// Returns `Ok(None)` if no session with the given ID exists.
    fn get_session(&self, session_id: &str) -> Result<Option<Session>, StorageError>;

    /// Get a single session by composite key (machine_id, session_id).
    ///
    /// This is the exact lookup using the composite primary key, which ensures
    /// uniqueness across multiple machines.
    ///
    /// Returns `Ok(None)` if no session with the given key exists.
    fn get_session_by_key(
        &self,
        machine_id: &str,
        session_id: &str,
    ) -> Result<Option<Session>, StorageError>;

    /// Get a session by its process ID.
    ///
    /// Returns the most recently active session associated with the given PID.
    /// This is useful for looking up sessions when only the PID is known
    /// (e.g., from process monitoring tools).
    ///
    /// Returns `Ok(None)` if no session with the given PID exists.
    fn get_session_by_pid(&self, pid: i32) -> Result<Option<Session>, StorageError>;

    /// Update git branch information for a session.
    ///
    /// This updates the session's git_branch, git_pr_number, and git_issue_number
    /// fields. It is typically called:
    /// - On SessionStart, to capture the initial branch
    /// - After detecting a git branch-changing command in PostToolUse
    ///
    /// Returns `Ok(false)` if the session doesn't exist.
    fn update_session_git_info(
        &self,
        session_id: &str,
        git_info: &GitBranchInfo,
    ) -> Result<bool, StorageError>;

    /// Update the GitHub repository for a session.
    ///
    /// This updates the session's github_repo field if it's currently empty.
    /// Uses COALESCE semantics to avoid overwriting existing values.
    ///
    /// It is typically called:
    /// - On SessionStart, to capture the initial repo from git remote
    /// - When a worktree is detected
    ///
    /// Returns `Ok(false)` if the session doesn't exist.
    fn update_session_github_repo(
        &self,
        session_id: &str,
        github_repo: &str,
    ) -> Result<bool, StorageError>;

    /// Upsert git directory path and GitHub repository for a session.
    ///
    /// This is called on every hook event BEFORE the event is inserted, to ensure
    /// git context is updated first. This allows branch parsing (which happens
    /// during event insert) to correctly set github_issue/github_pr for the
    /// current repo.
    ///
    /// When github_repo changes, github_issue and github_pr are cleared because
    /// issue/PR numbers are only meaningful in the context of a specific repository.
    ///
    /// Creates the session if it doesn't exist (UPSERT semantics).
    ///
    /// This is safe to call frequently (~50µs with direct file access).
    fn upsert_session_git_context(
        &self,
        session_id: &str,
        machine_id: &str,
        framework: &str,
        timestamp: i64,
        local_git_dir: Option<&str>,
        github_repo: Option<&str>,
    ) -> Result<(), StorageError>;

    /// Update the transcript path for a session.
    ///
    /// This updates the session's transcript_path field if it's currently empty.
    /// Uses COALESCE semantics to avoid overwriting existing values.
    ///
    /// This is primarily used for Codex sessions where the transcript path
    /// is not available from hooks but is known from session file scanning.
    ///
    /// Returns `Ok(false)` if the session doesn't exist.
    fn update_session_transcript_path(
        &self,
        machine_id: &str,
        session_id: &str,
        transcript_path: &str,
    ) -> Result<bool, StorageError>;

    /// Get aggregated statistics across all sessions.
    ///
    /// Returns session counts, token totals, costs, and API request counts
    /// computed at the database level via a single aggregation query. This is
    /// more efficient than iterating over individual sessions in application code.
    ///
    /// # Examples
    /// ```ignore
    /// // Get stats for all sessions
    /// let stats = storage.storage_stats(&StorageStatsQuery::new())?;
    /// println!("Total sessions: {}", stats.session_count);
    /// println!("Active sessions: {}", stats.active_session_count);
    /// println!("Total cost: ${:.2}", stats.total_cost_usd);
    ///
    /// // Get stats for only active sessions
    /// let stats = storage.storage_stats(&StorageStatsQuery::new().active_only())?;
    ///
    /// // Get stats for a specific framework
    /// let stats = storage.storage_stats(&StorageStatsQuery::new().with_framework("claude"))?;
    /// ```
    fn storage_stats(&self, query: &StorageStatsQuery) -> Result<StorageStats, StorageError>;

    // ========================================================================
    // Transcript scanning methods (optional, with default no-op implementations)
    // ========================================================================

    /// Get the last scanned position for a transcript file.
    ///
    /// Returns `Ok(None)` if the file has never been scanned.
    /// Default implementation returns `Ok(None)`.
    fn get_transcript_position(&self, _path: &Path) -> Result<Option<FilePosition>, StorageError> {
        Ok(None)
    }

    /// Set the scanned position for a transcript file.
    ///
    /// Default implementation does nothing.
    fn set_transcript_position(
        &self,
        _path: &Path,
        _position: &FilePosition,
    ) -> Result<(), StorageError> {
        Ok(())
    }

    /// Check if an event with the given UUID already exists for a session.
    ///
    /// Used for deduplication when scanning transcripts.
    /// Default implementation returns `Ok(false)`.
    fn event_exists_by_uuid(&self, _session_id: &str, _uuid: &str) -> Result<bool, StorageError> {
        Ok(false)
    }

    /// Query all transcript file positions.
    ///
    /// Returns a list of (path, position) tuples for all tracked transcript files.
    /// Default implementation returns an empty list.
    fn query_transcript_positions(&self) -> Result<Vec<(String, FilePosition)>, StorageError> {
        Ok(vec![])
    }
}

/// Blanket implementation of `Storage` for `Arc<T>` where `T: Storage`.
///
/// This enables shared ownership of storage instances. Wrap your storage in an
/// `Arc` and it can be used anywhere a `Storage` implementation is expected.
///
/// # Use Cases
///
/// - Sharing storage across multiple owned structs
/// - Single-threaded async runtimes (e.g., `tokio::task::LocalSet`)
/// - Multi-threaded access when `T: Send + Sync` (see note below)
///
/// # Thread Safety Note
///
/// For multi-threaded use, `T` must be `Send + Sync`. `SqliteStorage` is not
/// `Sync` (since `rusqlite::Connection` is not `Sync`), so `Arc<SqliteStorage>`
/// cannot be shared across threads directly. For multi-threaded scenarios with
/// `SqliteStorage`, use `Arc<Mutex<SqliteStorage>>` or a connection pool.
///
/// # Example
/// ```ignore
/// use mi6_storage_sqlite::SqliteStorage;
/// use mi6_core::{Storage, EventQuery};
/// use std::sync::Arc;
///
/// // Create shared storage
/// let storage = Arc::new(SqliteStorage::open(path)?);
///
/// // Clone for use in multiple places (single-threaded)
/// let storage_for_queries = Arc::clone(&storage);
/// let storage_for_inserts = Arc::clone(&storage);
///
/// // Both references can use Storage trait methods
/// let events = storage_for_queries.query(&EventQuery::new().with_limit(10))?;
/// storage_for_inserts.insert(&event)?;
/// ```
impl<T: Storage> Storage for Arc<T> {
    fn insert(&self, event: &Event) -> Result<i64, StorageError> {
        (**self).insert(event)
    }

    fn query(&self, query: &EventQuery) -> Result<Vec<Event>, StorageError> {
        (**self).query(query)
    }

    fn gc(&self, retention: Duration) -> Result<usize, StorageError> {
        (**self).gc(retention)
    }

    fn count_expired(&self, retention: Duration) -> Result<usize, StorageError> {
        (**self).count_expired(retention)
    }

    fn count(&self) -> Result<usize, StorageError> {
        (**self).count()
    }

    fn list_sessions(&self, query: &SessionQuery) -> Result<Vec<Session>, StorageError> {
        (**self).list_sessions(query)
    }

    fn get_session(&self, session_id: &str) -> Result<Option<Session>, StorageError> {
        (**self).get_session(session_id)
    }

    fn get_session_by_key(
        &self,
        machine_id: &str,
        session_id: &str,
    ) -> Result<Option<Session>, StorageError> {
        (**self).get_session_by_key(machine_id, session_id)
    }

    fn get_session_by_pid(&self, pid: i32) -> Result<Option<Session>, StorageError> {
        (**self).get_session_by_pid(pid)
    }

    fn update_session_git_info(
        &self,
        session_id: &str,
        git_info: &GitBranchInfo,
    ) -> Result<bool, StorageError> {
        (**self).update_session_git_info(session_id, git_info)
    }

    fn update_session_github_repo(
        &self,
        session_id: &str,
        github_repo: &str,
    ) -> Result<bool, StorageError> {
        (**self).update_session_github_repo(session_id, github_repo)
    }

    fn upsert_session_git_context(
        &self,
        session_id: &str,
        machine_id: &str,
        framework: &str,
        timestamp: i64,
        local_git_dir: Option<&str>,
        github_repo: Option<&str>,
    ) -> Result<(), StorageError> {
        (**self).upsert_session_git_context(
            session_id,
            machine_id,
            framework,
            timestamp,
            local_git_dir,
            github_repo,
        )
    }

    fn update_session_transcript_path(
        &self,
        machine_id: &str,
        session_id: &str,
        transcript_path: &str,
    ) -> Result<bool, StorageError> {
        (**self).update_session_transcript_path(machine_id, session_id, transcript_path)
    }

    fn storage_stats(&self, query: &StorageStatsQuery) -> Result<StorageStats, StorageError> {
        (**self).storage_stats(query)
    }

    fn get_transcript_position(&self, path: &Path) -> Result<Option<FilePosition>, StorageError> {
        (**self).get_transcript_position(path)
    }

    fn set_transcript_position(
        &self,
        path: &Path,
        position: &FilePosition,
    ) -> Result<(), StorageError> {
        (**self).set_transcript_position(path, position)
    }

    fn event_exists_by_uuid(&self, session_id: &str, uuid: &str) -> Result<bool, StorageError> {
        (**self).event_exists_by_uuid(session_id, uuid)
    }

    fn query_transcript_positions(&self) -> Result<Vec<(String, FilePosition)>, StorageError> {
        (**self).query_transcript_positions()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::StringError;
    use std::error::Error;

    #[test]
    fn test_error_source_preserved() {
        // Create an underlying IO error
        let io_err = std::io::Error::new(std::io::ErrorKind::NotFound, "file not found");
        let storage_err = StorageError::Connection(Box::new(io_err));

        // Verify the error chain is preserved
        let source = storage_err.source();
        assert!(source.is_some(), "StorageError should have a source");
        assert!(source.is_some_and(|s| s.to_string().contains("file not found")));
    }

    #[test]
    fn test_string_error_works() {
        let msg = "custom error message";
        let storage_err = StorageError::Query(Box::new(StringError(msg.to_string())));

        // Verify the error message is preserved
        assert!(storage_err.to_string().contains(msg));
        assert!(storage_err.source().is_some());
    }
}