what_core/database/

sqlite.rs

//! SQLite database backend for the What framework.
//!
//! Each collection is stored as a table with:
//! - `id` INTEGER PRIMARY KEY AUTOINCREMENT
//! - `data` TEXT (JSON object of all fields)
//!
//! Key-value pairs use a special `_kv_store` table.
//!
//! Uses r2d2 connection pool + spawn_blocking to avoid blocking the async runtime.

use r2d2::Pool;
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite::{Connection, params};
use serde_json::{Map, Value, json};
use std::collections::HashMap;
use std::path::Path;

use super::CollectionQuery;
use crate::Result;

/// SQLite-backed database for collections and key-value storage.
/// Uses a connection pool for concurrent reads and spawn_blocking
/// to keep the async runtime responsive.
#[derive(Clone)]
pub struct SqliteDatabase {
    pool: Pool<SqliteConnectionManager>,
}

/// Connection customizer that sets WAL mode and busy timeout on each new connection.
#[derive(Debug)]
struct WhatCustomizer;

impl r2d2::CustomizeConnection<Connection, rusqlite::Error> for WhatCustomizer {
    fn on_acquire(&self, conn: &mut Connection) -> std::result::Result<(), rusqlite::Error> {
        conn.execute_batch("PRAGMA busy_timeout=5000; PRAGMA synchronous=NORMAL;")?;
        Ok(())
    }
}

impl SqliteDatabase {
    /// Open or create a SQLite database at the given path.
    /// Creates a connection pool with WAL mode for concurrent reads.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let manager = SqliteConnectionManager::file(path);
        let pool = Pool::builder()
            .max_size(8)
            .connection_customizer(Box::new(WhatCustomizer))
            .build(manager)
            .map_err(|e| crate::Error::Data(format!("Pool creation failed: {}", e)))?;
        let db = Self { pool };
        db.init()?;
        Ok(db)
    }

    /// Create an in-memory database (for testing).
    /// Uses pool_size=1 since in-memory DBs are per-connection.
    pub fn in_memory() -> Result<Self> {
        let manager = SqliteConnectionManager::memory();
        let pool = Pool::builder()
            .max_size(1)
            .connection_customizer(Box::new(WhatCustomizer))
            .build(manager)
            .map_err(|e| crate::Error::Data(format!("Pool creation failed: {}", e)))?;
        let db = Self { pool };
        db.init()?;
        Ok(db)
    }

    fn init(&self) -> Result<()> {
        let conn = self
            .pool
            .get()
            .map_err(|e| crate::Error::Data(format!("Pool get failed: {}", e)))?;
        // Set WAL mode once (requires exclusive lock, so do it before pool fills)
        conn.execute_batch("PRAGMA journal_mode=WAL;")?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS _kv_store (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            )",
            [],
        )?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS _collections (
                name TEXT PRIMARY KEY
            )",
            [],
        )?;
        Ok(())
    }

    /// Synchronous import of JSON items into a collection (for store.json migration).
    /// Only inserts if the collection table is currently empty.
    pub fn import_json_collection(&self, name: &str, items: &[Value]) {
        let conn = match self.pool.get() {
            Ok(c) => c,
            Err(_) => return,
        };
        let safe_name = sanitize_table_name(name);
        if ensure_table_sync(&conn, name).is_err() {
            return;
        }

        let count: i64 = conn
            .query_row(
                &format!("SELECT COUNT(*) FROM \"{}\"", safe_name),
                [],
                |r| r.get(0),
            )
            .unwrap_or(0);
        if count > 0 {
            return;
        }

        for item in items {
            let mut data_map = match item {
                Value::Object(map) => map.clone(),
                _ => continue,
            };
            data_map.remove("id");
            let data_str = serde_json::to_string(&Value::Object(data_map)).unwrap_or_default();
            conn.execute(
                &format!("INSERT INTO \"{}\" (data) VALUES (?1)", safe_name),
                params![data_str],
            )
            .ok();
        }
        tracing::info!(
            "Imported {} items from store.json into '{}'",
            items.len(),
            name
        );
    }

    /// Get all items from a collection
    pub async fn get_collection(&self, name: &str) -> Result<Vec<Value>> {
        let pool = self.pool.clone();
        let name = name.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            ensure_table_sync(&conn, &name)?;
            let safe_name = sanitize_table_name(&name);
            let mut stmt = conn.prepare(&format!(
                "SELECT id, data FROM \"{}\" ORDER BY id",
                safe_name
            ))?;
            let rows = stmt.query_map([], |row| {
                let id: i64 = row.get(0)?;
                let data_str: String = row.get(1)?;
                Ok((id, data_str))
            })?;
            let mut items = Vec::new();
            for row in rows {
                let (id, data_str) = row?;
                let mut item: Value = serde_json::from_str(&data_str).unwrap_or(json!({}));
                if let Value::Object(ref mut map) = item {
                    map.insert("id".to_string(), json!(id));
                }
                items.push(item);
            }
            Ok(items)
        })
        .await
        .unwrap()
    }

    /// Query a collection with sort, filter, search, limit, offset
    pub async fn query_collection(
        &self,
        name: &str,
        query: &CollectionQuery,
    ) -> Result<Vec<Value>> {
        let pool = self.pool.clone();
        let name = name.to_string();
        let query = query.clone();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            ensure_table_sync(&conn, &name)?;
            let safe_name = sanitize_table_name(&name);

            let mut sql = format!("SELECT id, data FROM \"{}\"", safe_name);
            let mut where_clauses = Vec::new();
            let mut bind_values: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();

            if let Some(ref filter_expr) = query.filter {
                let (clause, values) = build_filter_sql(filter_expr);
                if !clause.is_empty() {
                    where_clauses.push(clause);
                    bind_values.extend(values);
                }
            }

            // Policy-forced scope filters — AND-ed in, cannot be widened by the user.
            for forced in &query.forced_filters {
                let (clause, values) = build_filter_sql(forced);
                if !clause.is_empty() {
                    where_clauses.push(clause);
                    bind_values.extend(values);
                }
            }

            if let Some(ref search_term) = query.search {
                if !search_term.is_empty() {
                    let fields: Vec<&str> = query
                        .search_fields
                        .as_deref()
                        .map(|s| s.split(',').map(|f| f.trim()).collect())
                        .unwrap_or_default();

                    if fields.is_empty() {
                        where_clauses.push("data LIKE ?".to_string());
                        bind_values.push(Box::new(format!("%{}%", search_term)));
                    } else {
                        let field_conditions: Vec<String> = fields
                            .iter()
                            .map(|f| {
                                bind_values.push(Box::new(format!("%{}%", search_term)));
                                format!("json_extract(data, '$.{}') LIKE ?", sanitize_field_name(f))
                            })
                            .collect();
                        where_clauses.push(format!("({})", field_conditions.join(" OR ")));
                    }
                }
            }

            if !where_clauses.is_empty() {
                sql.push_str(" WHERE ");
                sql.push_str(&where_clauses.join(" AND "));
            }

            if let Some(ref sort_expr) = query.sort {
                let (field, desc) = parse_sort(sort_expr);
                let safe_field = sanitize_field_name(&field);
                let dir = if desc { "DESC" } else { "ASC" };
                sql.push_str(&format!(
                    " ORDER BY json_extract(data, '$.{}') {}",
                    safe_field, dir
                ));
            } else {
                sql.push_str(" ORDER BY id");
            }

            if let Some(limit) = query.limit {
                sql.push_str(&format!(" LIMIT {}", limit));
            }
            if let Some(offset) = query.offset {
                sql.push_str(&format!(" OFFSET {}", offset));
            }

            let mut stmt = conn.prepare(&sql)?;
            let params_ref: Vec<&dyn rusqlite::types::ToSql> =
                bind_values.iter().map(|b| b.as_ref()).collect();
            let rows = stmt.query_map(params_ref.as_slice(), |row| {
                let id: i64 = row.get(0)?;
                let data_str: String = row.get(1)?;
                Ok((id, data_str))
            })?;

            let mut items = Vec::new();
            for row in rows {
                let (id, data_str) = row?;
                let mut item: Value = serde_json::from_str(&data_str).unwrap_or(json!({}));
                if let Value::Object(ref mut map) = item {
                    map.insert("id".to_string(), json!(id));
                }
                items.push(item);
            }
            Ok(items)
        })
        .await
        .unwrap()
    }

    /// Find items by a field value
    pub async fn find_by(
        &self,
        collection: &str,
        field: &str,
        value: &Value,
    ) -> Result<Vec<Value>> {
        let pool = self.pool.clone();
        let collection = collection.to_string();
        let field = field.to_string();
        let value = value.clone();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            ensure_table_sync(&conn, &collection)?;
            let safe_name = sanitize_table_name(&collection);
            let safe_field = sanitize_field_name(&field);

            let value_str = match &value {
                Value::String(s) => s.clone(),
                Value::Number(n) => n.to_string(),
                Value::Bool(b) => b.to_string(),
                _ => serde_json::to_string(&value).unwrap_or_default(),
            };

            let sql = if field == "id" {
                format!("SELECT id, data FROM \"{}\" WHERE id = ?1", safe_name)
            } else {
                format!(
                    "SELECT id, data FROM \"{}\" WHERE json_extract(data, '$.{}') = ?1",
                    safe_name, safe_field
                )
            };

            let mut stmt = conn.prepare(&sql)?;
            let rows = stmt.query_map(params![value_str], |row| {
                let id: i64 = row.get(0)?;
                let data_str: String = row.get(1)?;
                Ok((id, data_str))
            })?;

            let mut items = Vec::new();
            for row in rows {
                let (id, data_str) = row?;
                let mut item: Value = serde_json::from_str(&data_str).unwrap_or(json!({}));
                if let Value::Object(ref mut map) = item {
                    map.insert("id".to_string(), json!(id));
                }
                items.push(item);
            }
            Ok(items)
        })
        .await
        .unwrap()
    }

    /// Find a single item by field value
    pub async fn find_one_by(
        &self,
        collection: &str,
        field: &str,
        value: &Value,
    ) -> Result<Option<Value>> {
        let items = self.find_by(collection, field, value).await?;
        Ok(items.into_iter().next())
    }

    /// Create an item in a collection
    pub async fn create(&self, collection: &str, item: Value) -> Result<Value> {
        let pool = self.pool.clone();
        let collection = collection.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            ensure_table_sync(&conn, &collection)?;
            let safe_name = sanitize_table_name(&collection);

            let mut data_map = match item {
                Value::Object(map) => map,
                _ => Map::new(),
            };
            data_map.remove("id");
            // Auto-stamp a sortable creation timestamp so `sort=created_at:desc`
            // works out of the box (millisecond precision avoids ties).
            if !data_map.contains_key("created_at") {
                data_map.insert(
                    "created_at".to_string(),
                    json!(chrono::Local::now()
                        .format("%Y-%m-%dT%H:%M:%S%.3f")
                        .to_string()),
                );
            }
            let data_str = serde_json::to_string(&Value::Object(data_map.clone()))?;

            conn.execute(
                &format!("INSERT INTO \"{}\" (data) VALUES (?1)", safe_name),
                params![data_str],
            )?;

            let id = conn.last_insert_rowid();
            data_map.insert("id".to_string(), json!(id));
            Ok(Value::Object(data_map))
        })
        .await
        .unwrap()
    }

    /// Update an item by ID
    pub async fn update(
        &self,
        collection: &str,
        id: &Value,
        updates: Value,
    ) -> Result<Option<Value>> {
        let pool = self.pool.clone();
        let collection = collection.to_string();
        let id = id.clone();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            ensure_table_sync(&conn, &collection)?;
            let safe_name = sanitize_table_name(&collection);

            let id_num = match &id {
                Value::Number(n) => n.as_i64().unwrap_or(0),
                Value::String(s) => s.parse::<i64>().unwrap_or(0),
                _ => 0,
            };

            let current: Option<String> = conn
                .query_row(
                    &format!("SELECT data FROM \"{}\" WHERE id = ?1", safe_name),
                    params![id_num],
                    |row| row.get(0),
                )
                .ok();

            let Some(current_str) = current else {
                return Ok(None);
            };

            let mut current_data: Map<String, Value> =
                serde_json::from_str(&current_str).unwrap_or_default();

            if let Value::Object(update_map) = updates {
                for (k, v) in update_map {
                    if k != "id" {
                        current_data.insert(k, v);
                    }
                }
            }

            let updated_str = serde_json::to_string(&Value::Object(current_data.clone()))?;
            conn.execute(
                &format!("UPDATE \"{}\" SET data = ?1 WHERE id = ?2", safe_name),
                params![updated_str, id_num],
            )?;

            current_data.insert("id".to_string(), json!(id_num));
            Ok(Some(Value::Object(current_data)))
        })
        .await
        .unwrap()
    }

    /// Delete an item by ID
    pub async fn delete(&self, collection: &str, id: &Value) -> Result<bool> {
        let pool = self.pool.clone();
        let collection = collection.to_string();
        let id = id.clone();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            ensure_table_sync(&conn, &collection)?;
            let safe_name = sanitize_table_name(&collection);

            let id_num = match &id {
                Value::Number(n) => n.as_i64().unwrap_or(0),
                Value::String(s) => s.parse::<i64>().unwrap_or(0),
                _ => 0,
            };

            let rows = conn.execute(
                &format!("DELETE FROM \"{}\" WHERE id = ?1", safe_name),
                params![id_num],
            )?;

            Ok(rows > 0)
        })
        .await
        .unwrap()
    }

    /// Set a key-value pair
    pub async fn set(&self, key: &str, value: Value) -> Result<()> {
        let pool = self.pool.clone();
        let key = key.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            let value_str = serde_json::to_string(&value)?;
            conn.execute(
                "INSERT OR REPLACE INTO _kv_store (key, value) VALUES (?1, ?2)",
                params![key, value_str],
            )?;
            Ok(())
        })
        .await
        .unwrap()
    }

    /// Get a value by key
    pub async fn get(&self, key: &str) -> Result<Option<Value>> {
        let pool = self.pool.clone();
        let key = key.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            let result: Option<String> = conn
                .query_row(
                    "SELECT value FROM _kv_store WHERE key = ?1",
                    params![key],
                    |row| row.get(0),
                )
                .ok();
            match result {
                Some(s) => Ok(Some(serde_json::from_str(&s)?)),
                None => Ok(None),
            }
        })
        .await
        .unwrap()
    }

    /// Get all data as template context
    pub async fn as_context(&self) -> Result<HashMap<String, Value>> {
        let pool = self.pool.clone();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            let mut context = HashMap::new();

            let mut stmt = conn.prepare("SELECT name FROM _collections")?;
            let names: Vec<String> = stmt
                .query_map([], |row| row.get(0))?
                .filter_map(|r| r.ok())
                .collect();
            drop(stmt);

            for name in names {
                let safe_name = sanitize_table_name(&name);
                let mut stmt = conn.prepare(&format!(
                    "SELECT id, data FROM \"{}\" ORDER BY id",
                    safe_name
                ))?;
                let items: Vec<Value> = stmt
                    .query_map([], |row| {
                        let id: i64 = row.get(0)?;
                        let data_str: String = row.get(1)?;
                        Ok((id, data_str))
                    })?
                    .filter_map(|r| r.ok())
                    .map(|(id, data_str)| {
                        let mut item: Value = serde_json::from_str(&data_str).unwrap_or(json!({}));
                        if let Value::Object(ref mut map) = item {
                            map.insert("id".to_string(), json!(id));
                        }
                        item
                    })
                    .collect();
                drop(stmt);
                context.insert(name, Value::Array(items));
            }

            let mut stmt = conn.prepare("SELECT key, value FROM _kv_store")?;
            let kvs: Vec<(String, String)> = stmt
                .query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
                .filter_map(|r| r.ok())
                .collect();
            drop(stmt);

            for (key, value_str) in kvs {
                if let Ok(value) = serde_json::from_str(&value_str) {
                    context.insert(key, value);
                }
            }

            Ok(context)
        })
        .await
        .unwrap()
    }

    /// Replace an entire collection
    pub async fn set_collection(&self, name: &str, items: Vec<Value>) -> Result<()> {
        let pool = self.pool.clone();
        let name = name.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            ensure_table_sync(&conn, &name)?;
            let safe_name = sanitize_table_name(&name);

            conn.execute(&format!("DELETE FROM \"{}\"", safe_name), [])?;

            let mut stmt =
                conn.prepare(&format!("INSERT INTO \"{}\" (data) VALUES (?1)", safe_name))?;
            for item in items {
                let mut data = match item {
                    Value::Object(map) => map,
                    _ => Map::new(),
                };
                data.remove("id");
                let data_str = serde_json::to_string(&Value::Object(data))?;
                stmt.execute(params![data_str])?;
            }

            Ok(())
        })
        .await
        .unwrap()
    }

    /// Atomically modify a key-value pair
    pub async fn atomic_modify<F>(&self, key: &str, f: F) -> Result<Value>
    where
        F: FnOnce(Option<&Value>) -> Value + Send + 'static,
    {
        let pool = self.pool.clone();
        let key = key.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            let current: Option<Value> = conn
                .query_row(
                    "SELECT value FROM _kv_store WHERE key = ?1",
                    params![key],
                    |row| {
                        let s: String = row.get(0)?;
                        Ok(serde_json::from_str(&s).ok())
                    },
                )
                .ok()
                .flatten();

            let new_value = f(current.as_ref());
            let value_str = serde_json::to_string(&new_value)?;
            conn.execute(
                "INSERT OR REPLACE INTO _kv_store (key, value) VALUES (?1, ?2)",
                params![key, value_str],
            )?;
            Ok(new_value)
        })
        .await
        .unwrap()
    }

    /// Delete a key-value pair
    pub async fn remove(&self, key: &str) -> Result<Option<Value>> {
        let pool = self.pool.clone();
        let key = key.to_string();
        tokio::task::spawn_blocking(move || {
            let conn = pool.get().map_err(|e| crate::Error::Data(e.to_string()))?;
            let current: Option<String> = conn
                .query_row(
                    "SELECT value FROM _kv_store WHERE key = ?1",
                    params![key],
                    |row| row.get(0),
                )
                .ok();
            let result = current.and_then(|s| serde_json::from_str(&s).ok());
            conn.execute("DELETE FROM _kv_store WHERE key = ?1", params![key])?;
            Ok(result)
        })
        .await
        .unwrap()
    }
}

// ---------------------------------------------------------------------------
// SQL Helpers
// ---------------------------------------------------------------------------

/// Ensure a collection table exists (called inside spawn_blocking with a pooled connection)
fn ensure_table_sync(conn: &Connection, name: &str) -> Result<()> {
    let safe_name = sanitize_table_name(name);
    conn.execute(
        &format!(
            "CREATE TABLE IF NOT EXISTS \"{}\" (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                data TEXT NOT NULL DEFAULT '{{}}'
            )",
            safe_name
        ),
        [],
    )?;
    conn.execute(
        "INSERT OR IGNORE INTO _collections (name) VALUES (?1)",
        params![name],
    )?;
    Ok(())
}

fn sanitize_table_name(name: &str) -> String {
    name.chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '_')
        .collect()
}

fn sanitize_field_name(name: &str) -> String {
    name.chars()
        .filter(|c| c.is_ascii_alphanumeric() || *c == '_')
        .collect()
}

fn build_filter_sql(filter_expr: &str) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
    let mut clauses = Vec::new();
    let mut values: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();

    let or_groups: Vec<&str> = filter_expr.split(',').collect();
    let mut or_parts = Vec::new();

    for group in or_groups {
        let and_conditions: Vec<&str> = group.split('&').collect();
        let mut and_parts = Vec::new();

        for cond in and_conditions {
            let cond = cond.trim();
            if let Some((field, val)) = cond.split_once(">=") {
                let field = sanitize_field_name(field.trim());
                values.push(Box::new(val.trim().to_string()));
                and_parts.push(format!("json_extract(data, '$.{}') >= ?", field));
            } else if let Some((field, val)) = cond.split_once("<=") {
                let field = sanitize_field_name(field.trim());
                values.push(Box::new(val.trim().to_string()));
                and_parts.push(format!("json_extract(data, '$.{}') <= ?", field));
            } else if let Some((field, val)) = cond.split_once('>') {
                let field = sanitize_field_name(field.trim());
                values.push(Box::new(val.trim().to_string()));
                and_parts.push(format!("json_extract(data, '$.{}') > ?", field));
            } else if let Some((field, val)) = cond.split_once('<') {
                let field = sanitize_field_name(field.trim());
                values.push(Box::new(val.trim().to_string()));
                and_parts.push(format!("json_extract(data, '$.{}') < ?", field));
            } else if let Some((field, val)) = cond.split_once('=') {
                let field = sanitize_field_name(field.trim());
                values.push(Box::new(val.trim().to_string()));
                and_parts.push(format!("json_extract(data, '$.{}') = ?", field));
            }
        }

        if !and_parts.is_empty() {
            or_parts.push(format!("({})", and_parts.join(" AND ")));
        }
    }

    if !or_parts.is_empty() {
        clauses.push(format!("({})", or_parts.join(" OR ")));
    }

    (clauses.join(" AND "), values)
}

fn parse_sort(expr: &str) -> (String, bool) {
    if let Some((field, dir)) = expr.rsplit_once(':') {
        (field.to_string(), dir.eq_ignore_ascii_case("desc"))
    } else {
        (expr.to_string(), false)
    }
}

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

    #[tokio::test]
    async fn test_create_and_get() {
        let db = SqliteDatabase::in_memory().unwrap();
        let item = json!({"title": "Hello", "content": "World"});
        let created = db.create("posts", item).await.unwrap();
        assert_eq!(created.get("id"), Some(&json!(1)));
        assert_eq!(created.get("title"), Some(&json!("Hello")));

        let items = db.get_collection("posts").await.unwrap();
        assert_eq!(items.len(), 1);
        assert_eq!(items[0].get("title"), Some(&json!("Hello")));
    }

    #[tokio::test]
    async fn test_find_by() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.create("users", json!({"name": "Alice", "role": "admin"}))
            .await
            .unwrap();
        db.create("users", json!({"name": "Bob", "role": "user"}))
            .await
            .unwrap();
        db.create("users", json!({"name": "Charlie", "role": "admin"}))
            .await
            .unwrap();

        let admins = db.find_by("users", "role", &json!("admin")).await.unwrap();
        assert_eq!(admins.len(), 2);
    }

    #[tokio::test]
    async fn test_find_by_id() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.create("posts", json!({"title": "First"})).await.unwrap();
        db.create("posts", json!({"title": "Second"}))
            .await
            .unwrap();

        let found = db.find_one_by("posts", "id", &json!(2)).await.unwrap();
        assert!(found.is_some());
        assert_eq!(found.unwrap().get("title"), Some(&json!("Second")));
    }

    #[tokio::test]
    async fn test_update() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.create("posts", json!({"title": "Draft"})).await.unwrap();

        let updated = db
            .update(
                "posts",
                &json!(1),
                json!({"title": "Published", "status": "live"}),
            )
            .await
            .unwrap();
        assert!(updated.is_some());
        let u = updated.unwrap();
        assert_eq!(u.get("title"), Some(&json!("Published")));
        assert_eq!(u.get("status"), Some(&json!("live")));
    }

    #[tokio::test]
    async fn test_delete() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.create("posts", json!({"title": "To Delete"}))
            .await
            .unwrap();

        let deleted = db.delete("posts", &json!(1)).await.unwrap();
        assert!(deleted);

        let items = db.get_collection("posts").await.unwrap();
        assert!(items.is_empty());
    }

    #[tokio::test]
    async fn test_kv_store() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.set("counter", json!(42)).await.unwrap();

        let val = db.get("counter").await.unwrap();
        assert_eq!(val, Some(json!(42)));

        let removed = db.remove("counter").await.unwrap();
        assert_eq!(removed, Some(json!(42)));

        let val = db.get("counter").await.unwrap();
        assert_eq!(val, None);
    }

    #[tokio::test]
    async fn test_query_with_filter() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.create("posts", json!({"title": "A", "status": "published"}))
            .await
            .unwrap();
        db.create("posts", json!({"title": "B", "status": "draft"}))
            .await
            .unwrap();
        db.create("posts", json!({"title": "C", "status": "published"}))
            .await
            .unwrap();

        let query = CollectionQuery {
            filter: Some("status=published".to_string()),
            ..Default::default()
        };
        let items = db.query_collection("posts", &query).await.unwrap();
        assert_eq!(items.len(), 2);
    }

    #[tokio::test]
    async fn test_query_with_sort_and_limit() {
        let db = SqliteDatabase::in_memory().unwrap();
        for i in 1..=5 {
            db.create("items", json!({"n": i})).await.unwrap();
        }

        let query = CollectionQuery {
            sort: Some("n:desc".to_string()),
            limit: Some(3),
            ..Default::default()
        };
        let items = db.query_collection("items", &query).await.unwrap();
        assert_eq!(items.len(), 3);
        assert_eq!(items[0]["n"], 5);
        assert_eq!(items[1]["n"], 4);
        assert_eq!(items[2]["n"], 3);
    }

    #[tokio::test]
    async fn test_query_with_search() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.create(
            "posts",
            json!({"title": "Rust Programming", "content": "Learn Rust"}),
        )
        .await
        .unwrap();
        db.create(
            "posts",
            json!({"title": "Python Basics", "content": "Learn Python"}),
        )
        .await
        .unwrap();

        let query = CollectionQuery {
            search: Some("rust".to_string()),
            search_fields: Some("title".to_string()),
            ..Default::default()
        };
        let items = db.query_collection("posts", &query).await.unwrap();
        assert_eq!(items.len(), 1);
        assert_eq!(items[0]["title"], "Rust Programming");
    }

    #[tokio::test]
    async fn test_as_context() {
        let db = SqliteDatabase::in_memory().unwrap();
        db.create("posts", json!({"title": "Hello"})).await.unwrap();
        db.set("site_name", json!("My Site")).await.unwrap();

        let ctx = db.as_context().await.unwrap();
        assert!(ctx.contains_key("posts"));
        assert!(ctx.contains_key("site_name"));
        assert_eq!(ctx["site_name"], json!("My Site"));
    }

    #[tokio::test]
    async fn test_set_collection() {
        let db = SqliteDatabase::in_memory().unwrap();
        let items = vec![json!({"name": "A"}), json!({"name": "B"})];
        db.set_collection("letters", items).await.unwrap();

        let result = db.get_collection("letters").await.unwrap();
        assert_eq!(result.len(), 2);
    }
}