nornir 0.5.2

//! SQL over the warehouse — `nornir warehouse query "SELECT …"`.
//!
//! Executes arbitrary SQL over the Iceberg warehouse tables with
//! [DataFusion](https://datafusion.apache.org/). Each warehouse table is
//! exposed to DataFusion through the **native** `iceberg-datafusion`
//! [`IcebergStaticTableProvider`] — a read-only, catalog-backed provider that
//! plans an Iceberg table scan with **projection + filter pushdown** all the
//! way to the Parquet files. No table is copied into memory; the Parquet
//! footers' column/row-group bounds do the pruning.
//!
//! ## Why the native provider (not a MemTable copy)
//!
//! The alternative fallback — [`Warehouse::scan_arrow`](super::Warehouse) each
//! table into `RecordBatch`es and register a DataFusion `MemTable` — always
//! materialises the *whole* table into RAM before the planner can prune it, and
//! throws away Iceberg's manifest/row-group statistics. The native provider
//! keeps pushdown, so `SELECT … WHERE repo = 'x'` reads only the matching data
//! files. Both paths need DataFusion as the engine; the native provider is the
//! same dependency universe (arrow-58 / `iceberg-arrow58`) already in use by
//! [`super::iceberg`], so the `Table`/`Catalog` types unify with zero glue.
//!
//! ## Laziness & robustness
//!
//! Tables are registered through a lazy [`SchemaProvider`]: only the tables a
//! query actually references are loaded from the redb catalog, so a query never
//! pays to open all ~50 warehouse tables, and a single torn/blank table's
//! metadata cannot break an unrelated query. Bare table names resolve against
//! the warehouse namespace (`SELECT * FROM bench_runs`).
//!
//! ## Read-only & the single-writer lock
//!
//! This module never mutates. Callers open the warehouse with
//! [`super::open_read_only`] / [`super::iceberg::IcebergWarehouse::open_read_only`],
//! which — when a live `nornir-server` holds the exclusive redb lock — reads a
//! consistent copied-aside snapshot of `catalog.redb` instead of failing. The
//! table data still resolves to the live warehouse (absolute `file://` paths),
//! so queries see the catalog as of "now" without touching the writer's lock.

use std::sync::Arc;

use anyhow::{Context, Result};
use arrow::array::RecordBatch;
use async_trait::async_trait;
use datafusion::catalog::{SchemaProvider, TableProvider};
use datafusion::error::{DataFusionError, Result as DFResult};
use datafusion::prelude::SessionContext;
use iceberg::{Catalog, NamespaceIdent, TableIdent};
use iceberg_datafusion::IcebergStaticTableProvider;
use skade_katalog::RedbCatalog;

use super::iceberg::IcebergWarehouse;

/// A lazy DataFusion schema backed by the warehouse's Iceberg catalog. Holds
/// the (already-known) table-name list for planning; loads and wraps a table in
/// a read-only [`IcebergStaticTableProvider`] only when the planner asks for it.
#[derive(Debug)]
struct WarehouseSchema {
    catalog: Arc<RedbCatalog>,
    namespace: NamespaceIdent,
    table_names: Vec<String>,
}

#[async_trait]
impl SchemaProvider for WarehouseSchema {
    fn table_names(&self) -> Vec<String> {
        self.table_names.clone()
    }

    fn table_exist(&self, name: &str) -> bool {
        self.table_names.iter().any(|t| t == name)
    }

    async fn table(&self, name: &str) -> DFResult<Option<Arc<dyn TableProvider>>> {
        if !self.table_exist(name) {
            return Ok(None);
        }
        let ident = TableIdent::new(self.namespace.clone(), name.to_string());
        let table = self
            .catalog
            .load_table(&ident)
            .await
            .map_err(|e| DataFusionError::External(Box::new(e)))?;
        let provider = IcebergStaticTableProvider::try_new_from_table(table)
            .await
            .map_err(|e| DataFusionError::External(Box::new(e)))?;
        Ok(Some(Arc::new(provider)))
    }
}

/// Run `sql` over `wh`'s warehouse tables and return the result batches.
///
/// Every warehouse table is queryable by its warehouse name. Supports the full
/// DataFusion SQL surface — `SELECT … WHERE`, `JOIN` across warehouse tables,
/// `GROUP BY` / aggregates, `ORDER BY`, `LIMIT`, … Read-only: `wh` should be
/// opened via [`IcebergWarehouse::open_read_only`] so a live server's lock is
/// tolerated.
pub fn query(wh: &IcebergWarehouse, sql: &str) -> Result<Vec<RecordBatch>> {
    let schema = Arc::new(WarehouseSchema {
        catalog: wh.catalog().clone(),
        namespace: wh.namespace().clone(),
        table_names: wh.table_names()?,
    });

    // Drive DataFusion on the warehouse's own tokio runtime — the lazy
    // provider's `load_table`/scan futures run there, and `df.collect()` needs
    // a runtime context for its internal `tokio::spawn`s.
    wh.block_on(async move {
        let ctx = SessionContext::new();
        // Replace the default `datafusion.public` schema with the warehouse
        // schema so bare table names (`FROM bench_runs`) resolve to Iceberg
        // tables.
        let catalog = ctx
            .catalog("datafusion")
            .context("datafusion default catalog missing")?;
        catalog
            .register_schema("public", schema)
            .map_err(|e| anyhow::anyhow!("register warehouse schema: {e}"))?;

        let df = ctx
            .sql(sql)
            .await
            .with_context(|| format!("plan SQL: {sql}"))?;
        let batches = df.collect().await.context("execute SQL")?;
        Ok(batches)
    })
}

/// Render `batches` as a bordered text table (the default CLI output).
pub fn format_table(batches: &[RecordBatch]) -> Result<String> {
    Ok(arrow::util::pretty::pretty_format_batches(batches)
        .context("format result table")?
        .to_string())
}

/// Render `batches` as a JSON array of row objects (`--format json`).
pub fn format_json(batches: &[RecordBatch]) -> Result<String> {
    let mut buf = Vec::new();
    {
        let mut writer = arrow::json::ArrayWriter::new(&mut buf);
        for b in batches {
            writer.write(b).context("write JSON batch")?;
        }
        writer.finish().context("finish JSON")?;
    }
    if buf.is_empty() {
        // ArrayWriter emits nothing when no batch was written; normalise to an
        // empty JSON array so `--format json` is always valid JSON.
        return Ok("[]".to_string());
    }
    String::from_utf8(buf).context("JSON output was not valid UTF-8")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::knowledge::symbols::{CallEdgeRow, SymbolRow, SymbolScan};
    use arrow::array::{Array, Int64Array, StringArray};

    /// A `symbol_facts` row (build-free syn marker) for `crate_name`.
    fn sym(crate_name: &str, item: &str) -> SymbolRow {
        SymbolRow {
            crate_name: crate_name.into(),
            module_path: format!("{crate_name}::api"),
            item_kind: "fn".into(),
            item_name: item.into(),
            visibility: "pub".into(),
            file: format!("{crate_name}/src/api.rs"),
            line: 1,
            doc_lines: 0,
            signature: None,
        }
    }

    /// A `call_edges` row for `crate_name`: `caller` → `callee`.
    fn call(crate_name: &str, caller: &str, callee: &str) -> CallEdgeRow {
        CallEdgeRow {
            crate_name: crate_name.into(),
            caller_path: caller.into(),
            callee_ident: callee.into(),
            call_kind: "call".into(),
            file: format!("{crate_name}/src/api.rs"),
            line: 1,
        }
    }

    /// Seed a temp warehouse with two repos' symbol_facts + call_edges:
    ///   korp: symbols {draw, redraw}; calls {draw→render, redraw→paint}
    ///   knut: symbols {paint};        calls {paint→render}
    fn seeded() -> (tempfile::TempDir, IcebergWarehouse) {
        let dir = tempfile::tempdir().unwrap();
        let wh = IcebergWarehouse::open(dir.path()).unwrap();
        wh.append_symbol_scan(&SymbolScan {
            snapshot_id: uuid::Uuid::new_v4(),
            ts: chrono::Utc::now(),
            repo: "korp".into(),
            symbols: vec![sym("korp", "draw"), sym("korp", "redraw")],
            calls: vec![
                call("korp", "korp::draw", "render"),
                call("korp", "korp::redraw", "paint"),
            ],
            features: vec![],
            tests: vec![],
        })
        .unwrap();
        wh.append_symbol_scan(&SymbolScan {
            snapshot_id: uuid::Uuid::new_v4(),
            ts: chrono::Utc::now(),
            repo: "knut".into(),
            symbols: vec![sym("knut", "paint")],
            calls: vec![call("knut", "knut::paint", "render")],
            features: vec![],
            tests: vec![],
        })
        .unwrap();
        (dir, wh)
    }

    fn str_col(b: &RecordBatch, col: usize) -> Vec<String> {
        let a = b
            .column(col)
            .as_any()
            .downcast_ref::<StringArray>()
            .expect("string column");
        (0..a.len()).map(|i| a.value(i).to_string()).collect()
    }

    #[test]
    fn select_where_over_seeded_warehouse() {
        let (_d, wh) = seeded();
        let batches = query(
            &wh,
            "SELECT item_name FROM symbol_facts WHERE repo = 'korp' ORDER BY item_name",
        )
        .unwrap();
        let total: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total, 2, "korp seeded two symbols");
        let names: Vec<String> = batches.iter().flat_map(|b| str_col(b, 0)).collect();
        assert_eq!(names, vec!["draw".to_string(), "redraw".to_string()]);
    }

    #[test]
    fn join_across_two_warehouse_tables() {
        let (_d, wh) = seeded();
        // symbol_facts ⋈ call_edges on crate_name, filtered to one symbol —
        // for crate `korp`, symbol `draw` joins to korp's two call edges.
        let batches = query(
            &wh,
            "SELECT c.callee_ident \
             FROM symbol_facts s \
             JOIN call_edges c ON s.crate_name = c.crate_name \
             WHERE s.item_name = 'draw' \
             ORDER BY c.callee_ident",
        )
        .unwrap();
        let callees: Vec<String> = batches.iter().flat_map(|b| str_col(b, 0)).collect();
        assert_eq!(callees, vec!["paint".to_string(), "render".to_string()]);
    }

    #[test]
    fn aggregate_group_by_over_call_edges() {
        let (_d, wh) = seeded();
        let batches = query(
            &wh,
            "SELECT repo, COUNT(*) AS n FROM call_edges GROUP BY repo ORDER BY repo",
        )
        .unwrap();
        // Collect (repo, n) pairs across batches.
        let mut pairs: Vec<(String, i64)> = Vec::new();
        for b in &batches {
            let repos = str_col(b, 0);
            let counts = b
                .column(1)
                .as_any()
                .downcast_ref::<Int64Array>()
                .expect("count column is i64");
            for (i, repo) in repos.iter().enumerate() {
                pairs.push((repo.clone(), counts.value(i)));
            }
        }
        assert_eq!(
            pairs,
            vec![("knut".to_string(), 1), ("korp".to_string(), 2)],
            "knut has 1 call edge, korp has 2"
        );
    }

    #[test]
    fn json_output_is_valid_json_array() {
        let (_d, wh) = seeded();
        let batches = query(
            &wh,
            "SELECT item_name FROM symbol_facts WHERE repo = 'knut'",
        )
        .unwrap();
        let json = format_json(&batches).unwrap();
        let v: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert!(v.is_array(), "json output is an array: {json}");
        assert_eq!(v.as_array().unwrap().len(), 1);
        assert_eq!(v[0]["item_name"], "paint");
    }
}