skade-sql/src/lib.rs
//! DataFusion SQL over a [`skade`] Iceberg warehouse.
//!
//! This is the code that used to sit in `skade` behind its `sql` feature. It
//! moved out for one reason: it depends on `apache/datafusion` **git**, because
//! no published datafusion is on arrow 59 (54.1.0, the newest release, resolves
//! arrow 58). A git dependency with no version requirement makes the whole crate
//! unpublishable, and skade has to publish — otherwise knut and every other
//! consumer keeps resolving the crates.io skade, which is arrow 58, and lands a
//! second `arrow_array` major in its graph.
//!
//! The API is the same calls it always was, reached through two extension
//! traits instead of inherent methods, since Rust will not let one crate write
//! `impl` blocks for another crate's types:
//!
//! ```ignore
//! use skade_sql::{TableSql, WarehouseSql};
//!
//! let batches = warehouse.sql("SELECT … FROM t JOIN u …").await?;
//! let batches = table.sql("SELECT … FROM events").await?;
//! ```
//!
//! Move back into skade and delete this crate the day apache publishes a
//! datafusion on arrow 59.
use std::sync::Arc;
use arrow_array::RecordBatch;
use datafusion::catalog::CatalogProvider;
use datafusion::prelude::SessionContext;
use iceberg::Catalog;
use iceberg_datafusion::{IcebergCatalogProvider, IcebergStaticTableProvider};
use skade::{Result, SkadeError, Table, Warehouse};
/// DataFusion errors reach the caller as [`SkadeError::Other`].
///
/// `skade` used to carry a `SkadeError::DataFusion(#[from] DataFusionError)`
/// variant. It cannot any more — the variant would put `datafusion` back in
/// skade's manifest and un-publish it again, which is the entire thing this
/// crate exists to avoid. The message is preserved in full; only the variant a
/// caller would match on is different.
fn df(e: impl std::fmt::Display) -> SkadeError {
SkadeError::other(format!("datafusion: {e}"))
}
/// SQL over a single [`Table`].
#[async_trait::async_trait]
pub trait TableSql {
/// Run one SQL statement over **this** table, registered under its bare
/// name (`SELECT … FROM <name> …`), against the handle's current snapshot.
/// For multi-table SQL use [`WarehouseSql::sql`].
async fn sql(&self, query: &str) -> Result<Vec<RecordBatch>>;
}
#[async_trait::async_trait]
impl TableSql for Table {
async fn sql(&self, query: &str) -> Result<Vec<RecordBatch>> {
let ctx = SessionContext::new();
let provider = IcebergStaticTableProvider::try_new_from_table(self.inner().clone())
.await
.map_err(df)?;
ctx.register_table(self.ident().name(), Arc::new(provider))
.map_err(df)?;
Ok(ctx
.sql(query)
.await
.map_err(df)?
.collect()
.await
.map_err(df)?)
}
}
/// SQL over a whole [`Warehouse`].
#[async_trait::async_trait]
pub trait WarehouseSql {
/// A fresh DataFusion session over the warehouse's *current* state.
///
/// Every table is reachable as `skade.<ns>.<table>`; every table is
/// additionally registered under its bare name, so `SELECT … FROM t JOIN u
/// …` works unqualified.
///
/// The bare-name registration is the **catalog-backed** provider — the same
/// object `skade.<ns>.<table>` resolves to — not a snapshot-pinned static
/// one. That is what makes `MERGE INTO t USING … ` and `INSERT INTO t`
/// work through a bare name: a provider that cannot write is not a SQL
/// surface, it is a view. Use [`TableSql::sql`] on a [`Table`] handle when
/// you want a query pinned to one snapshot.
async fn session(&self) -> Result<SessionContext>;
/// Run one SQL statement over the warehouse (any number of tables) and
/// collect the result batches. Convenience for
/// `self.session().await?.sql(query).await?.collect()`.
async fn sql(&self, query: &str) -> Result<Vec<RecordBatch>>;
}
#[async_trait::async_trait]
impl WarehouseSql for Warehouse {
async fn session(&self) -> Result<SessionContext> {
let ctx = SessionContext::new();
let catalog = self.catalog();
let provider = IcebergCatalogProvider::try_new(catalog.clone() as Arc<dyn Catalog>)
.await
.map_err(df)?;
let provider: Arc<dyn CatalogProvider> = Arc::new(provider);
ctx.register_catalog("skade", provider.clone());
// Register every table (all namespaces) under its bare name so
// unqualified `SELECT … FROM <table>` resolves. Cross-namespace name
// collisions are last-wins (rare in an embedded warehouse); use the
// `skade.<ns>.<table>` catalog path to disambiguate.
//
// The provider comes OUT OF THE CATALOG rather than being built here as
// a static one, so the bare name and the `skade.<ns>.<table>` path are
// literally the same writable object. Before this they were not, and a
// `MERGE INTO`/`INSERT INTO` against a bare name hit the read-only
// provider and was refused.
for ns in catalog.list_namespaces(None).await? {
let Some(schema) = provider.schema(&ns.join(".")) else {
continue;
};
for ident in catalog.list_tables(&ns).await? {
// A table listed a moment ago can be removed by a concurrent
// compaction swap (or a drop) before we load it — skip it
// rather than failing the whole session. The query only needs
// the tables it actually references; a transiently-absent
// scratch table must not 404 an unrelated `SELECT … FROM events`.
let table = match schema.table(ident.name()).await {
Ok(Some(t)) => t,
Ok(None) => continue,
Err(e) => return Err(df(e)),
};
ctx.register_table(ident.name(), table).map_err(df)?;
}
}
Ok(ctx)
}
async fn sql(&self, query: &str) -> Result<Vec<RecordBatch>> {
let ctx = self.session().await?;
Ok(ctx
.sql(query)
.await
.map_err(df)?
.collect()
.await
.map_err(df)?)
}
}