unifier-cli 0.5.0

//! Managed SQLite databases under `sqlite/<name>.sqlite`.
//!
//! The reserved [`TRIPLES_DB`] database holds RDF-style subject/predicate/object
//! rows and is created automatically on first use.

use std::fs;
use std::path::PathBuf;

use rusqlite::{params, types::ValueRef, Connection, Statement};

use crate::constants::{SQLITE, TRIPLES_DB};
use crate::error::{Error, Result};
use crate::home::UnifierHome;
use crate::scope::validate_chroot_name;

const SQLITE_SUFFIX: &str = ".sqlite";

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColumnInfo {
    pub name: String,
    pub decl_type: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TableInfo {
    pub name: String,
    pub columns: Vec<ColumnInfo>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueryResult {
    pub columns: Vec<String>,
    pub rows: Vec<Vec<String>>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExecResult {
    pub rows_affected: usize,
}

pub fn sqlite_dir(home: &UnifierHome) -> PathBuf {
    home.path().join(SQLITE)
}

pub fn validate_db_name(name: &str) -> Result<()> {
    validate_chroot_name(name).map_err(|_| Error::msg(format!("invalid database name: {name}")))
}

pub fn db_path(home: &UnifierHome, name: &str) -> Result<PathBuf> {
    validate_db_name(name)?;
    Ok(sqlite_dir(home).join(format!("{name}{SQLITE_SUFFIX}")))
}

/// Create an empty managed database (or ensure `triples` schema).
pub fn create_db(home: &UnifierHome, name: &str) -> Result<PathBuf> {
    let path = db_path(home, name)?;
    let _conn = open_connection(home, name, true)?;
    Ok(path)
}

pub fn list_databases(home: &UnifierHome) -> Result<Vec<String>> {
    let dir = sqlite_dir(home);
    if !dir.is_dir() {
        return Ok(Vec::new());
    }
    let mut names = Vec::new();
    for entry in fs::read_dir(&dir)? {
        let entry = entry?;
        if !entry.file_type()?.is_file() {
            continue;
        }
        let file_name = entry.file_name();
        let Some(name) = file_name.to_str() else {
            continue;
        };
        if let Some(stem) = name.strip_suffix(SQLITE_SUFFIX) {
            if validate_db_name(stem).is_ok() {
                names.push(stem.to_string());
            }
        }
    }
    names.sort();
    Ok(names)
}

pub fn list_tables(home: &UnifierHome, name: &str) -> Result<Vec<TableInfo>> {
    let conn = open_connection(home, name, false)?;
    let mut stmt = conn.prepare(
        "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name",
    )?;
    let names = stmt
        .query_map([], |row| row.get::<_, String>(0))?
        .collect::<std::result::Result<Vec<_>, _>>()?;

    let mut tables = Vec::with_capacity(names.len());
    for table_name in names {
        let columns = table_columns(&conn, &table_name)?;
        tables.push(TableInfo {
            name: table_name,
            columns,
        });
    }
    Ok(tables)
}

pub fn exec_sql(home: &UnifierHome, name: &str, sql: &str) -> Result<SqlOutcome> {
    let sql = sql.trim();
    if sql.is_empty() {
        return Err(Error::msg("sql must not be empty"));
    }
    let conn = open_connection(home, name, true)?;
    if looks_like_query(sql) {
        return Ok(SqlOutcome::Query(run_query(&conn, sql)?));
    }
    let trimmed = sql.trim_end_matches(';').trim();
    if trimmed.contains(';') {
        conn.execute_batch(sql)?;
        Ok(SqlOutcome::Exec(ExecResult { rows_affected: 0 }))
    } else {
        let rows_affected = conn.execute(sql, [])?;
        Ok(SqlOutcome::Exec(ExecResult { rows_affected }))
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SqlOutcome {
    Query(QueryResult),
    Exec(ExecResult),
}

pub fn insert_triple(
    home: &UnifierHome,
    subject: &str,
    predicate: &str,
    object: &str,
) -> Result<()> {
    if subject.is_empty() || predicate.is_empty() || object.is_empty() {
        return Err(Error::msg(
            "triple subject, predicate, and object must be non-empty",
        ));
    }
    let conn = open_connection(home, TRIPLES_DB, true)?;
    conn.execute(
        "INSERT OR IGNORE INTO triples (subject, predicate, object) VALUES (?1, ?2, ?3)",
        params![subject, predicate, object],
    )?;
    Ok(())
}

pub fn query_triples(
    home: &UnifierHome,
    subject: Option<&str>,
    predicate: Option<&str>,
    object: Option<&str>,
) -> Result<Vec<(String, String, String)>> {
    let conn = open_connection(home, TRIPLES_DB, true)?;
    let mut sql = String::from("SELECT subject, predicate, object FROM triples WHERE 1 = 1");
    let mut values: Vec<String> = Vec::new();
    if let Some(s) = subject {
        sql.push_str(" AND subject = ?");
        values.push(s.to_string());
    }
    if let Some(p) = predicate {
        sql.push_str(" AND predicate = ?");
        values.push(p.to_string());
    }
    if let Some(o) = object {
        sql.push_str(" AND object = ?");
        values.push(o.to_string());
    }
    sql.push_str(" ORDER BY subject, predicate, object");

    let mut stmt = conn.prepare(&sql)?;
    let params: Vec<&dyn rusqlite::types::ToSql> = values
        .iter()
        .map(|v| v as &dyn rusqlite::types::ToSql)
        .collect();
    let rows = stmt
        .query_map(params.as_slice(), |row| {
            Ok((row.get(0)?, row.get(1)?, row.get(2)?))
        })?
        .collect::<std::result::Result<Vec<_>, _>>()?;
    Ok(rows)
}

fn open_connection(home: &UnifierHome, name: &str, create: bool) -> Result<Connection> {
    let path = db_path(home, name)?;
    if !create && !path.is_file() {
        return Err(Error::msg(format!("database not found: {name}")));
    }
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    let conn = Connection::open(&path).map_err(sql_err)?;
    if name == TRIPLES_DB {
        ensure_triples_schema(&conn)?;
    }
    Ok(conn)
}

fn ensure_triples_schema(conn: &Connection) -> Result<()> {
    conn.execute_batch(
        "CREATE TABLE IF NOT EXISTS triples (
            subject TEXT NOT NULL,
            predicate TEXT NOT NULL,
            object TEXT NOT NULL,
            created_at TEXT NOT NULL DEFAULT (datetime('now')),
            PRIMARY KEY (subject, predicate, object)
        );
        CREATE INDEX IF NOT EXISTS idx_triples_predicate ON triples(predicate);
        CREATE INDEX IF NOT EXISTS idx_triples_object ON triples(object);",
    )
    .map_err(sql_err)?;
    Ok(())
}

fn table_columns(conn: &Connection, table: &str) -> Result<Vec<ColumnInfo>> {
    // table name comes from sqlite_master; still quote carefully
    let pragma = format!("PRAGMA table_info({})", quote_ident(table));
    let mut stmt = conn.prepare(&pragma)?;
    let cols = stmt
        .query_map([], |row| {
            Ok(ColumnInfo {
                name: row.get(1)?,
                decl_type: row.get::<_, Option<String>>(2)?.unwrap_or_default(),
            })
        })?
        .collect::<std::result::Result<Vec<_>, _>>()?;
    Ok(cols)
}

fn quote_ident(name: &str) -> String {
    format!("\"{}\"", name.replace('"', "\"\""))
}

fn looks_like_query(sql: &str) -> bool {
    let head = sql
        .trim_start()
        .trim_start_matches('(')
        .to_ascii_lowercase();
    head.starts_with("select")
        || head.starts_with("with")
        || head.starts_with("pragma")
        || head.starts_with("explain")
        || head.starts_with("values")
}

fn run_query(conn: &Connection, sql: &str) -> Result<QueryResult> {
    let mut stmt = conn.prepare(sql)?;
    let columns = stmt
        .column_names()
        .iter()
        .map(|s| (*s).to_string())
        .collect::<Vec<_>>();
    let rows = collect_rows(&mut stmt)?;
    Ok(QueryResult { columns, rows })
}

fn collect_rows(stmt: &mut Statement<'_>) -> Result<Vec<Vec<String>>> {
    let column_count = stmt.column_count();
    let mut rows = Vec::new();
    let mut query = stmt.query([])?;
    while let Some(row) = query.next()? {
        let mut values = Vec::with_capacity(column_count);
        for i in 0..column_count {
            values.push(value_to_string(row.get_ref(i)?));
        }
        rows.push(values);
    }
    Ok(rows)
}

fn value_to_string(value: ValueRef<'_>) -> String {
    match value {
        ValueRef::Null => String::new(),
        ValueRef::Integer(v) => v.to_string(),
        ValueRef::Real(v) => v.to_string(),
        ValueRef::Text(v) => String::from_utf8_lossy(v).into_owned(),
        ValueRef::Blob(v) => format!("\\x{}", hex::encode_fallback(v)),
    }
}

fn sql_err(e: rusqlite::Error) -> Error {
    Error::msg(format!("sqlite: {e}"))
}

impl From<rusqlite::Error> for Error {
    fn from(e: rusqlite::Error) -> Self {
        sql_err(e)
    }
}

mod hex {
    pub fn encode_fallback(bytes: &[u8]) -> String {
        const HEX: &[u8; 16] = b"0123456789abcdef";
        let mut out = String::with_capacity(bytes.len() * 2);
        for b in bytes {
            out.push(HEX[(b >> 4) as usize] as char);
            out.push(HEX[(b & 0xf) as usize] as char);
        }
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::home::UnifierHome;
    use tempfile::tempdir;

    fn home() -> (tempfile::TempDir, UnifierHome) {
        let tmp = tempdir().unwrap();
        let home = UnifierHome::resolve(Some(tmp.path().to_path_buf()), None).unwrap();
        home.ensure().unwrap();
        (tmp, home)
    }

    #[test]
    fn create_list_exec_and_tables() {
        let (_tmp, home) = home();
        create_db(&home, "app").unwrap();
        assert_eq!(list_databases(&home).unwrap(), vec!["app".to_string()]);

        match exec_sql(
            &home,
            "app",
            "CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT); INSERT INTO items(name) VALUES ('a'), ('b');",
        )
        .unwrap()
        {
            SqlOutcome::Exec(_) => {}
            other => panic!("expected exec, got {other:?}"),
        }

        let tables = list_tables(&home, "app").unwrap();
        assert_eq!(tables.len(), 1);
        assert_eq!(tables[0].name, "items");
        assert_eq!(tables[0].columns[0].name, "id");
        assert_eq!(tables[0].columns[1].name, "name");

        match exec_sql(&home, "app", "SELECT name FROM items ORDER BY id").unwrap() {
            SqlOutcome::Query(q) => {
                assert_eq!(q.columns, vec!["name"]);
                assert_eq!(q.rows, vec![vec!["a".to_string()], vec!["b".to_string()]]);
            }
            other => panic!("expected query, got {other:?}"),
        }
    }

    #[test]
    fn triples_insert_and_query() {
        let (_tmp, home) = home();
        insert_triple(&home, "alice", "knows", "bob").unwrap();
        insert_triple(&home, "alice", "knows", "bob").unwrap(); // ignore dup
        insert_triple(&home, "alice", "likes", "tea").unwrap();

        let all = query_triples(&home, None, None, None).unwrap();
        assert_eq!(all.len(), 2);
        let knows = query_triples(&home, Some("alice"), Some("knows"), None).unwrap();
        assert_eq!(
            knows,
            vec![("alice".to_string(), "knows".to_string(), "bob".to_string())]
        );

        assert!(list_databases(&home)
            .unwrap()
            .contains(&"triples".to_string()));
        let tables = list_tables(&home, "triples").unwrap();
        assert!(tables.iter().any(|t| t.name == "triples"));
    }

    #[test]
    fn rejects_bad_db_name() {
        let (_tmp, home) = home();
        assert!(create_db(&home, "../x").is_err());
        assert!(create_db(&home, "a/b").is_err());
    }
}