src/warehouse/sql.rs
//! 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;
#[cfg(feature = "sql")]
use async_trait::async_trait;
#[cfg(feature = "sql")]
use datafusion::catalog::{SchemaProvider, TableProvider};
#[cfg(feature = "sql")]
use datafusion::error::{DataFusionError, Result as DFResult};
#[cfg(feature = "sql")]
use datafusion::prelude::SessionContext;
#[cfg(feature = "sql")]
use iceberg::{Catalog, NamespaceIdent, TableIdent};
#[cfg(feature = "sql")]
use iceberg_datafusion::IcebergStaticTableProvider;
#[cfg(feature = "sql")]
use skade_katalog::RedbCatalog;
use super::Warehouse;
use super::iceberg::IcebergWarehouse;
#[cfg(feature = "sql")]
/// 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>,
}
#[cfg(feature = "sql")]
#[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)))
}
}
/// `sql` OFF — the DataFusion surface is not compiled in.
///
/// The API keeps its SHAPE rather than vanishing — the same pattern
/// `znippy-common` uses for its `openzl` feature — so every call site still
/// compiles and the failure is a named runtime error instead of a build break
/// in a downstream crate.
#[cfg(not(feature = "sql"))]
pub fn query(_wh: &IcebergWarehouse, _sql: &str) -> Result<Vec<RecordBatch>> {
anyhow::bail!(
"nornir was built WITHOUT the `sql` feature, so the DataFusion SQL surface is \
not compiled in. Rebuild with `--features sql`. It is off by default because \
datafusion 54 pins arrow ^58.3.0 and re-exports its own arrow, which pins the \
whole constellation to arrow 58."
)
}
#[cfg(feature = "sql")]
/// 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")
}
// ----- Mimir SQL query-templates (TASK #46) ----------------------------------
//
// The canned SQL templates (def-use, change-radius, call-chain,
// table-access-lineage) queried `symbol_facts` / `call_edges` /
// `warehouse_access_edges`; those tables have no producer since Mímir v2, so the
// templates were removed. Raw SQL over the warehouse ([`query`]) is unchanged.
#[cfg(test)]
mod tests {
use super::*;
use crate::knowledge::git_heat::{GitHeatRow, GitHeatScan};
use crate::warehouse::feature_dep_edges::{DepKind, EnableKind, FeatureDepEdgeRow};
use arrow::array::{Array, Int64Array, StringArray};
/// A `git_heat_facts` scan for `repo` over `files`.
fn heat(repo: &str, files: &[&str]) -> GitHeatScan {
let ts = chrono::Utc::now();
GitHeatScan {
snapshot_id: uuid::Uuid::new_v4(),
ts,
repo: repo.into(),
files: files
.iter()
.map(|f| GitHeatRow {
file: (*f).into(),
commits_total: 1,
commits_30d: 1,
commits_90d: 1,
authors_total: 1,
last_commit_ts: ts,
})
.collect(),
}
}
/// A `feature_dep_edges` row: `crate_name`'s feature `gui` enables `enables`.
fn edge(repo: &str, enables: &str) -> FeatureDepEdgeRow {
FeatureDepEdgeRow {
repo: repo.into(),
crate_name: repo.into(),
feature: "gui".into(),
enables: enables.into(),
kind: EnableKind::Feature,
optional: false,
target_cfg: None,
dep_kind: DepKind::Normal,
}
}
/// Seed a temp warehouse with two repos' git_heat_facts + feature_dep_edges:
/// korp: files {src/draw.rs, src/redraw.rs}; edges {gui→paint, gui→render}
/// knut: files {src/paint.rs}; edges {gui→render}
fn seeded() -> (tempfile::TempDir, IcebergWarehouse) {
let dir = tempfile::tempdir().unwrap();
let wh = IcebergWarehouse::open(dir.path()).unwrap();
wh.append_git_heat_scan(&heat("korp", &["src/redraw.rs", "src/draw.rs"]))
.unwrap();
wh.append_git_heat_scan(&heat("knut", &["src/paint.rs"]))
.unwrap();
wh.append_feature_dep_edges(
&uuid::Uuid::new_v4().to_string(),
&[edge("korp", "render"), edge("korp", "paint"), edge("knut", "render")],
)
.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 file FROM git_heat_facts WHERE repo = 'korp' ORDER BY file",
)
.unwrap();
let total: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total, 2, "korp seeded two files");
let names: Vec<String> = batches.iter().flat_map(|b| str_col(b, 0)).collect();
assert_eq!(names, vec!["src/draw.rs".to_string(), "src/redraw.rs".to_string()]);
}
#[test]
fn join_across_two_warehouse_tables() {
let (_d, wh) = seeded();
// git_heat_facts ⋈ feature_dep_edges on repo, filtered to one file —
// korp's `src/draw.rs` joins to korp's two feature edges.
let batches = query(
&wh,
"SELECT e.enables \
FROM git_heat_facts h \
JOIN feature_dep_edges e ON h.repo = e.repo \
WHERE h.file = 'src/draw.rs' \
ORDER BY e.enables",
)
.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_feature_dep_edges() {
let (_d, wh) = seeded();
let batches = query(
&wh,
"SELECT repo, COUNT(*) AS n FROM feature_dep_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 feature 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 == "git_heat_facts"),
"list_tables must include git_heat_facts: {tables:?}"
);
assert!(
tables.iter().any(|t| t == "feature_dep_edges"),
"list_tables must include feature_dep_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, "git_heat_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(&"file"), "columns: {names:?}");
assert!(names.contains(&"repo"), "columns: {names:?}");
// file is a Utf8 string column — a red-when-broken shape check.
let item = cols.iter().find(|c| c.name == "file").unwrap();
assert_eq!(item.data_type, "Utf8", "file 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("git_heat_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 == "git_heat_facts"),
"table discovery over the trait object names git_heat_facts: {tables:?}"
);
let cols = describe_table(boxed.as_ref(), "git_heat_facts").unwrap();
let item = cols
.iter()
.find(|c| c.name == "file")
.expect("file 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 file FROM git_heat_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]["file"], "src/paint.rs");
}
}