xz-memory-engine 0.2.0

use async_trait::async_trait;
use sqlx::Row;
use sqlx::SqlitePool;

use xz_memory_core::StoreError;
use xz_memory_core::traits::store::EntryStore;
use xz_memory_core::types::entry::*;

/// SQLite-backed entry store implementing [`EntryStore`].
///
/// Stores entries in a SQLite database table with columns:
/// `id TEXT PRIMARY KEY`, `partition TEXT NOT NULL`,
/// `body TEXT NOT NULL`, `recorded_at INTEGER NOT NULL`.
///
/// # Examples
///
/// ```
/// use xz_memory_engine::backends::sqlite::SqliteEntryStore;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let store = SqliteEntryStore::new("sqlite::memory:").await?;
/// # Ok(())
/// # }
/// ```
pub struct SqliteEntryStore {
    pub pool: SqlitePool,
}

impl SqliteEntryStore {
    /// Create a new [`SqliteEntryStore`] backed by the SQLite database at `path`.
    ///
    /// Use `"sqlite::memory:"` for an in-memory database (useful for testing).
    /// The necessary table and index are created automatically if they do not exist.
    pub async fn new(path: &str) -> Result<Self, StoreError> {
        let pool =
            SqlitePool::connect(path).await.map_err(|e| StoreError::Backend(e.to_string()))?;
        let store = SqliteEntryStore { pool };
        store.migrate().await?;
        Ok(store)
    }

    /// Build a [`SqliteEntryStore`] over an existing pool.
    ///
    /// Runs the necessary migrations on the pool so the `entries` table and
    /// index are ready.  This is useful when the caller already manages the
    /// database connection and wants to avoid a second pool.
    pub async fn from_pool(pool: SqlitePool) -> Result<Self, StoreError> {
        let store = SqliteEntryStore { pool };
        store.migrate().await?;
        Ok(store)
    }

    pub async fn migrate(&self) -> Result<(), StoreError> {
        sqlx::query(
            "CREATE TABLE IF NOT EXISTS entries (
                id TEXT PRIMARY KEY,
                partition TEXT NOT NULL,
                body TEXT NOT NULL,
                recorded_at INTEGER NOT NULL
            )",
        )
        .execute(&self.pool)
        .await
        .map_err(|e| StoreError::Backend(e.to_string()))?;

        sqlx::query("CREATE INDEX IF NOT EXISTS idx_entries_partition ON entries(partition)")
            .execute(&self.pool)
            .await
            .map_err(|e| StoreError::Backend(e.to_string()))?;

        Ok(())
    }
}

#[async_trait]
impl EntryStore for SqliteEntryStore {
    async fn append(&self, entry: Entry) -> Result<(), StoreError> {
        sqlx::query(
            "INSERT OR REPLACE INTO entries (id, partition, body, recorded_at) VALUES (?, ?, ?, ?)",
        )
        .bind(&entry.id)
        .bind(&entry.partition)
        .bind(&entry.body)
        .bind(entry.recorded_at as i64)
        .execute(&self.pool)
        .await
        .map_err(|e| StoreError::Backend(e.to_string()))?;
        Ok(())
    }

    async fn query(
        &self,
        partition: &str,
        range: &TimeRange,
        opts: &QueryOptions,
    ) -> Result<Vec<Entry>, StoreError> {
        let order = match opts.sort {
            SortOrder::Ascending => "ASC",
            SortOrder::Descending => "DESC",
        };

        let mut conditions = String::new();
        if range.start.is_some() {
            conditions.push_str(" AND recorded_at >= ?");
        }
        if range.end.is_some() {
            conditions.push_str(" AND recorded_at <= ?");
        }

        let sql = format!(
            "SELECT id, partition, body, recorded_at FROM entries WHERE partition = ?{} ORDER BY recorded_at {} LIMIT ?",
            conditions, order,
        );

        let mut query = sqlx::query(&sql).bind(partition);
        if let Some(start) = range.start {
            query = query.bind(start as i64);
        }
        if let Some(end) = range.end {
            query = query.bind(end as i64);
        }
        query = query.bind(opts.limit as i64);

        let rows =
            query.fetch_all(&self.pool).await.map_err(|e| StoreError::Backend(e.to_string()))?;

        let entries: Vec<Entry> = rows
            .iter()
            .map(|row| Entry {
                id: row.get("id"),
                partition: row.get("partition"),
                body: row.get("body"),
                recorded_at: row.get::<i64, _>("recorded_at") as u64,
            })
            .collect();

        Ok(entries)
    }

    async fn evict(&self, partition: &str, keep: usize) -> Result<usize, StoreError> {
        let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM entries WHERE partition = ?")
            .bind(partition)
            .fetch_one(&self.pool)
            .await
            .map_err(|e| StoreError::Backend(e.to_string()))?;

        let total = count as usize;
        if total <= keep {
            return Ok(0);
        }

        let remove_count = total - keep;
        let result = sqlx::query(
            "DELETE FROM entries WHERE id IN (SELECT id FROM entries WHERE partition = ? ORDER BY recorded_at ASC LIMIT ?)",
        )
        .bind(partition)
        .bind(remove_count as i64)
        .execute(&self.pool)
        .await
        .map_err(|e| StoreError::Backend(e.to_string()))?;

        Ok(result.rows_affected() as usize)
    }

    async fn delete(&self, id: &str) -> Result<(), StoreError> {
        sqlx::query("DELETE FROM entries WHERE id = ?")
            .bind(id)
            .execute(&self.pool)
            .await
            .map_err(|e| StoreError::Backend(e.to_string()))?;
        Ok(())
    }

    async fn clear_partition(&self, partition: &str) -> Result<(), StoreError> {
        sqlx::query("DELETE FROM entries WHERE partition = ?")
            .bind(partition)
            .execute(&self.pool)
            .await
            .map_err(|e| StoreError::Backend(e.to_string()))?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn entry(id: &str, partition: &str, body: &str, recorded_at: u64) -> Entry {
        Entry { id: id.into(), partition: partition.into(), body: body.into(), recorded_at }
    }

    fn opts(limit: usize, sort: SortOrder) -> QueryOptions {
        QueryOptions { limit, sort }
    }

    fn range(start: Option<u64>, end: Option<u64>) -> TimeRange {
        TimeRange { start, end }
    }

    async fn new_store() -> Result<SqliteEntryStore, StoreError> {
        SqliteEntryStore::new("sqlite::memory:").await
    }

    #[tokio::test]
    async fn append_persists_entry() {
        let store = new_store().await.unwrap();
        store.append(entry("a", "p1", "hello", 100)).await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "a");
        assert_eq!(results[0].body, "hello");
        assert_eq!(results[0].recorded_at, 100);
    }

    #[tokio::test]
    async fn append_replace_same_id() {
        let store = new_store().await.unwrap();
        store.append(entry("a", "p1", "first", 100)).await.unwrap();
        store.append(entry("a", "p1", "second", 200)).await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].body, "second");
        assert_eq!(results[0].recorded_at, 200);
    }

    #[tokio::test]
    async fn query_returns_multiple_entries() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();
        store.append(entry("3", "p1", "c", 300)).await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(results.len(), 3);
    }

    #[tokio::test]
    async fn query_respects_time_range_start() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();
        store.append(entry("3", "p1", "c", 300)).await.unwrap();

        let results = store
            .query("p1", &range(Some(200), None), &opts(10, SortOrder::Ascending))
            .await
            .unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].id, "2");
        assert_eq!(results[1].id, "3");
    }

    #[tokio::test]
    async fn query_respects_time_range_end() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();
        store.append(entry("3", "p1", "c", 300)).await.unwrap();

        let results = store
            .query("p1", &range(None, Some(200)), &opts(10, SortOrder::Ascending))
            .await
            .unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].id, "1");
        assert_eq!(results[1].id, "2");
    }

    #[tokio::test]
    async fn query_respects_time_range_both() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();
        store.append(entry("3", "p1", "c", 300)).await.unwrap();

        let results = store
            .query("p1", &range(Some(200), Some(300)), &opts(10, SortOrder::Ascending))
            .await
            .unwrap();
        assert_eq!(results.len(), 2);
    }

    #[tokio::test]
    async fn query_respects_sort_ascending() {
        let store = new_store().await.unwrap();
        store.append(entry("3", "p1", "c", 300)).await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].recorded_at, 100);
        assert_eq!(results[1].recorded_at, 200);
        assert_eq!(results[2].recorded_at, 300);
    }

    #[tokio::test]
    async fn query_respects_sort_descending() {
        let store = new_store().await.unwrap();
        store.append(entry("3", "p1", "c", 300)).await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Descending)).await.unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].recorded_at, 300);
        assert_eq!(results[1].recorded_at, 200);
        assert_eq!(results[2].recorded_at, 100);
    }

    #[tokio::test]
    async fn query_respects_limit() {
        let store = new_store().await.unwrap();
        for i in 0..5 {
            store.append(entry(&format!("{i}"), "p1", "x", i * 100)).await.unwrap();
        }

        let results =
            store.query("p1", &range(None, None), &opts(3, SortOrder::Ascending)).await.unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].recorded_at, 0);
        assert_eq!(results[2].recorded_at, 200);
    }

    #[tokio::test]
    async fn query_empty_partition_returns_empty() {
        let store = new_store().await.unwrap();
        let results = store
            .query("nonexistent", &range(None, None), &opts(10, SortOrder::Ascending))
            .await
            .unwrap();
        assert!(results.is_empty());
    }

    #[tokio::test]
    async fn query_partition_isolation() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p2", "b", 200)).await.unwrap();

        let r1 =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        let r2 =
            store.query("p2", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();

        assert_eq!(r1.len(), 1);
        assert_eq!(r1[0].id, "1");
        assert_eq!(r2.len(), 1);
        assert_eq!(r2[0].id, "2");
    }

    #[tokio::test]
    async fn evict_removes_oldest_entries() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "oldest", 100)).await.unwrap();
        store.append(entry("2", "p1", "middle", 200)).await.unwrap();
        store.append(entry("3", "p1", "newest", 300)).await.unwrap();

        let removed = store.evict("p1", 1).await.unwrap();
        assert_eq!(removed, 2);

        let remaining =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(remaining.len(), 1);
        assert_eq!(remaining[0].id, "3");
    }

    #[tokio::test]
    async fn evict_keep_all_when_keep_exceeds_count() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();

        let removed = store.evict("p1", 5).await.unwrap();
        assert_eq!(removed, 0);

        let remaining =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(remaining.len(), 2);
    }

    #[tokio::test]
    async fn evict_returns_correct_count() {
        let store = new_store().await.unwrap();
        for i in 0..10 {
            store.append(entry(&format!("{i}"), "p1", "x", i * 100)).await.unwrap();
        }

        let removed = store.evict("p1", 3).await.unwrap();
        assert_eq!(removed, 7);

        let remaining =
            store.query("p1", &range(None, None), &opts(20, SortOrder::Ascending)).await.unwrap();
        assert_eq!(remaining.len(), 3);
        assert_eq!(remaining[0].recorded_at, 700);
        assert_eq!(remaining[2].recorded_at, 900);
    }

    #[tokio::test]
    async fn evict_empty_partition_returns_zero() {
        let store = new_store().await.unwrap();
        let removed = store.evict("empty", 1).await.unwrap();
        assert_eq!(removed, 0);
    }

    #[tokio::test]
    async fn evict_keep_zero_clears_partition() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();

        let removed = store.evict("p1", 0).await.unwrap();
        assert_eq!(removed, 2);

        let remaining =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert!(remaining.is_empty());
    }

    #[tokio::test]
    async fn evict_partition_isolation() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();
        store.append(entry("3", "p2", "c", 300)).await.unwrap();

        store.evict("p1", 1).await.unwrap();

        let p1_results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        let p2_results =
            store.query("p2", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();

        assert_eq!(p1_results.len(), 1);
        assert_eq!(p2_results.len(), 1);
    }

    #[tokio::test]
    async fn delete_removes_entry_by_id() {
        let store = new_store().await.unwrap();
        store.append(entry("a", "p1", "hello", 100)).await.unwrap();
        store.append(entry("b", "p1", "world", 200)).await.unwrap();

        store.delete("a").await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "b");
    }

    #[tokio::test]
    async fn delete_nonexistent_id_is_noop() {
        let store = new_store().await.unwrap();
        store.append(entry("a", "p1", "hello", 100)).await.unwrap();

        store.delete("nonexistent").await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert_eq!(results.len(), 1);
    }

    #[tokio::test]
    async fn clear_partition_removes_all_entries() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p1", "b", 200)).await.unwrap();
        store.append(entry("3", "p1", "c", 300)).await.unwrap();

        store.clear_partition("p1").await.unwrap();

        let results =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        assert!(results.is_empty());
    }

    #[tokio::test]
    async fn clear_partition_isolation() {
        let store = new_store().await.unwrap();
        store.append(entry("1", "p1", "a", 100)).await.unwrap();
        store.append(entry("2", "p2", "b", 200)).await.unwrap();

        store.clear_partition("p1").await.unwrap();

        let r1 =
            store.query("p1", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();
        let r2 =
            store.query("p2", &range(None, None), &opts(10, SortOrder::Ascending)).await.unwrap();

        assert!(r1.is_empty());
        assert_eq!(r2.len(), 1);
    }

    #[tokio::test]
    async fn clear_partition_nonexistent_is_noop() {
        let store = new_store().await.unwrap();
        store.clear_partition("nonexistent").await.unwrap();
    }
}