limbo/tests/integration/query_processing/test_expr_index.rs

use std::sync::Arc;

use crate::common::{limbo_exec_rows, TempDatabase};
use turso_core::StepResult;

fn explain_plans(conn: &Arc<turso_core::Connection>, sql: &str) -> anyhow::Result<Vec<String>> {
    let mut stmt = conn.prepare(format!("EXPLAIN QUERY PLAN {sql}"))?;
    let mut plans = Vec::new();
    loop {
        match stmt.step()? {
            StepResult::Row => {
                let row = stmt.row().unwrap();
                plans.push(row.get::<String>(3)?);
            }
            StepResult::IO => {
                stmt.run_once()?;
            }
            StepResult::Done => break,
            _ => unreachable!(),
        }
    }
    Ok(plans)
}

#[test]
fn expression_index_used_for_where() -> anyhow::Result<()> {
    let _ = env_logger::try_init();
    let tmp_db = TempDatabase::new_with_rusqlite("CREATE TABLE t (a INT, b INT, c INT);");
    let conn = tmp_db.connect_limbo();

    conn.execute("INSERT INTO t VALUES (1, 2, 0)")?;
    conn.execute("INSERT INTO t VALUES (3, 4, 0)")?;
    conn.execute("INSERT INTO t VALUES (5, 6, 0)")?;

    conn.execute("CREATE INDEX idx_expr ON t(a + b)")?;

    let plans = explain_plans(&conn, "SELECT * FROM t WHERE a + b = 7")?;
    assert!(
        plans.iter().any(|p| p.contains("idx_expr")),
        "expected query plan to mention idx_expr, got {plans:?}"
    );

    let rows = limbo_exec_rows(&tmp_db, &conn, "SELECT a, b FROM t WHERE a + b = 7");
    assert_eq!(rows, vec![vec![3.into(), 4.into()]]);
    Ok(())
}

#[test]
fn expression_index_used_for_order_by() -> anyhow::Result<()> {
    let _ = env_logger::try_init();
    let tmp_db = TempDatabase::new_with_rusqlite("CREATE TABLE t (a INT, b INT);");
    let conn = tmp_db.connect_limbo();

    conn.execute("INSERT INTO t VALUES (2, 2)")?;
    conn.execute("INSERT INTO t VALUES (1, 3)")?;
    conn.execute("INSERT INTO t VALUES (0, 5)")?;

    conn.execute("CREATE INDEX idx_expr_order ON t(a + b)")?;

    let plans = explain_plans(
        &conn,
        "SELECT a, b FROM t WHERE a + b > 0 ORDER BY a + b DESC LIMIT 1 OFFSET 0",
    )?;
    assert!(
        plans.iter().any(|p| p.contains("idx_expr_order")),
        "expected query plan to mention idx_expr_order, got {plans:?}"
    );

    let rows = limbo_exec_rows(
        &tmp_db,
        &conn,
        "SELECT a, b FROM t WHERE a + b > 0 ORDER BY a + b DESC LIMIT 1",
    );
    assert_eq!(rows, vec![vec![0.into(), 5.into()]]);
    Ok(())
}

#[test]
fn expression_index_covering_scan() -> anyhow::Result<()> {
    let _ = env_logger::try_init();
    let tmp_db = TempDatabase::new_with_rusqlite("CREATE TABLE t (a INT, b INT);");
    let conn = tmp_db.connect_limbo();

    conn.execute("INSERT INTO t VALUES (1, 2)")?;
    conn.execute("INSERT INTO t VALUES (3, 4)")?;
    conn.execute("INSERT INTO t VALUES (5, 6)")?;

    conn.execute("CREATE INDEX idx_expr_proj ON t(a + b)")?;

    let plans = explain_plans(&conn, "SELECT a + b FROM t")?;
    assert!(
        plans
            .iter()
            .any(|p| p.contains("USING COVERING INDEX idx_expr_proj")),
        "expected covering index usage, got {plans:?}"
    );

    let rows = limbo_exec_rows(&tmp_db, &conn, "SELECT a + b FROM t ORDER BY a + b");
    assert_eq!(rows, vec![vec![3.into()], vec![7.into()], vec![11.into()]]);
    Ok(())
}