nornir 0.5.4

//! 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;
use super::Warehouse;

/// 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)
    })
}

/// One column of a warehouse table's schema — the discovery companion to
/// [`query`]. A consumer that knows the table + column names (from
/// [`list_tables`] / [`describe_table`]) can then write `SELECT`s against them.
/// The POD now lives in the `nornir-warehouse-trait` edda leaf; re-exported here
/// at its old path (`crate::warehouse::sql::ColumnInfo`) so call sites are
/// untouched.
pub use nornir_warehouse_trait::ColumnInfo;

/// List every queryable warehouse table, sorted for stable output. This is the
/// discovery surface for [`query`]: these names are exactly the ones usable in a
/// `FROM` clause. Cheap — reads only the catalog's table-name index, never a
/// Parquet file.
///
/// Routed through the [`Warehouse`] trait object (only [`Warehouse::table_names`]
/// is needed), so it works against any backend — the local Iceberg tree or a
/// [`super::remote::RemoteWarehouse`] (which answers `table_names` via the
/// `Warehouse.Tables` RPC).
pub fn list_tables(wh: &dyn Warehouse) -> Result<Vec<String>> {
    let mut names = wh.table_names()?;
    names.sort();
    Ok(names)
}

/// Describe one warehouse table's columns (name · Arrow type · nullability) — so
/// a consumer can see what a `SELECT` may reference before writing SQL. Loads
/// only that one table's catalog metadata (its current Iceberg schema mapped to
/// Arrow), not any data files. Errors with a clear message if `table` is not a
/// known warehouse table.
///
/// Routed through the [`Warehouse`] trait object via
/// [`Warehouse::describe_columns`], so it is backend-agnostic: the Iceberg tree
/// maps its stored schema to Arrow, while a [`super::remote::RemoteWarehouse`]
/// returns a typed `Unsupported` (column types are lost on the stringified scan
/// wire — the Arrow Flight follow-up).
pub fn describe_table(wh: &dyn Warehouse, table: &str) -> Result<Vec<ColumnInfo>> {
    wh.describe_columns(table)
}

/// 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 list_tables_surfaces_the_seeded_tables_sorted() {
        let (_d, wh) = seeded();
        let tables = list_tables(&wh).unwrap();
        // The discovery surface must name the tables `query` can SELECT from.
        assert!(
            tables.iter().any(|t| t == "symbol_facts"),
            "list_tables must include symbol_facts: {tables:?}"
        );
        assert!(
            tables.iter().any(|t| t == "call_edges"),
            "list_tables must include call_edges: {tables:?}"
        );
        // Sorted for stable output.
        let mut sorted = tables.clone();
        sorted.sort();
        assert_eq!(tables, sorted, "list_tables output is sorted");
    }

    #[test]
    fn describe_table_reports_columns_and_types() {
        let (_d, wh) = seeded();
        let cols = describe_table(&wh, "symbol_facts").unwrap();
        let names: Vec<&str> = cols.iter().map(|c| c.name.as_str()).collect();
        // The columns the SQL tests above SELECT/WHERE on must be discoverable.
        assert!(names.contains(&"item_name"), "columns: {names:?}");
        assert!(names.contains(&"repo"), "columns: {names:?}");
        // item_name is a non-null Utf8 string column — a red-when-broken shape check.
        let item = cols.iter().find(|c| c.name == "item_name").unwrap();
        assert_eq!(item.data_type, "Utf8", "item_name is Utf8: {item:?}");
    }

    #[test]
    fn describe_unknown_table_errors_with_the_known_list() {
        let (_d, wh) = seeded();
        let err = describe_table(&wh, "no_such_table").unwrap_err().to_string();
        assert!(err.contains("no_such_table"), "err names the bad table: {err}");
        assert!(
            err.contains("symbol_facts"),
            "err lists known tables to guide the consumer: {err}"
        );
    }

    /// **The detach proof.** Schema discovery is routed through the
    /// [`Warehouse`] TRAIT OBJECT: this drives `list_tables` / `describe_table`
    /// purely through a `Box<dyn Warehouse>` (the warehouse is moved into the
    /// box — no concrete handle survives). RED-when-broken two ways: (1) a revert
    /// of either signature back to `&IcebergWarehouse` fails to COMPILE here
    /// (`Box<dyn Warehouse>` no longer coerces), and (2) if `describe_columns`
    /// stopped delegating to the real Iceberg schema read, the column asserts
    /// fail. This is the seam the `RemoteWarehouse` also stands behind.
    #[test]
    fn discovery_is_routed_through_the_warehouse_trait_object() {
        let (_d, wh) = seeded();
        let boxed: Box<dyn Warehouse> = Box::new(wh);
        let tables = list_tables(boxed.as_ref()).unwrap();
        assert!(
            tables.iter().any(|t| t == "symbol_facts"),
            "table discovery over the trait object names symbol_facts: {tables:?}"
        );
        let cols = describe_table(boxed.as_ref(), "symbol_facts").unwrap();
        let item = cols
            .iter()
            .find(|c| c.name == "item_name")
            .expect("item_name discoverable through the trait object");
        assert_eq!(item.data_type, "Utf8", "typed describe survives the trait route: {item:?}");
    }

    #[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");
    }
}