nordisk/skade

skade-flight/src/namespaces.rs

//! **Nested namespaces, served as flat schemas.**
//!
//! Iceberg namespaces nest (`a` → `a.b` → `a.b.orders`); a SQL catalog has one
//! level of schemas. Every namespace, at any depth, is served as ONE DataFusion
//! schema named by its elements joined with `.` — so `skade."a.b".orders`
//! resolves, and Flight SQL `GetDbSchemas` / `GetTables` (what Power BI's
//! navigator shows) list `a.b` as a schema. The Iceberg keys are untouched:
//! the flattened name maps straight back to the namespace's elements.
//!
//! A namespace with an element that itself contains `.` (Nessie allows it)
//! would flatten ambiguously (`["a.b"]` and `["a","b"]` are both `a.b`), so it
//! is not served — refused BY NAME in a log line, never mis-mapped.
//!
//! (The stock `IcebergCatalogProvider` lists only top-level namespaces and turns
//! each element into its own schema, which is why it is not used here.)

use std::collections::{BTreeMap, HashSet, VecDeque};
use std::sync::Arc;

use datafusion::catalog::{CatalogProvider, SchemaProvider};
use futures::future::try_join_all;
use iceberg::{Catalog, NamespaceIdent};
use iceberg_datafusion::IcebergSchemaProvider;

use crate::error::{Error, Result};

/// What joins a namespace's elements into its schema name.
pub const SEPARATOR: char = '.';

/// The schema name a namespace is served under — or, for a namespace with an
/// element containing [`SEPARATOR`], the log line saying why it is not.
pub fn schema_name(ns: &NamespaceIdent) -> std::result::Result<String, String> {
    let elements: &[String] = ns.as_ref();
    if let Some(bad) = elements.iter().find(|e| e.contains(SEPARATOR)) {
        return Err(format!(
            "namespace {elements:?} not served: its element {bad:?} contains '{SEPARATOR}', so the flattened schema name would be ambiguous"
        ));
    }
    Ok(elements.join(&SEPARATOR.to_string()))
}

/// The namespace a flattened schema name stands for.
pub fn namespace_of(schema: &str) -> Result<NamespaceIdent> {
    Ok(NamespaceIdent::from_strs(schema.split(SEPARATOR))?)
}

/// Every namespace in `catalog`, parents before children, each once.
pub async fn walk(catalog: &dyn Catalog) -> Result<Vec<NamespaceIdent>> {
    let mut out = Vec::new();
    let mut seen = HashSet::new();
    let mut queue: VecDeque<NamespaceIdent> = catalog.list_namespaces(None).await?.into();
    while let Some(ns) = queue.pop_front() {
        if !seen.insert(ns.as_ref().clone()) {
            continue;
        }
        queue.extend(catalog.list_namespaces(Some(&ns)).await?);
        out.push(ns);
    }
    Ok(out)
}

/// The DataFusion catalog over an Iceberg catalog: one schema per namespace, at
/// any depth, named by [`schema_name`].
#[derive(Debug)]
pub struct FlatNamespaceCatalog {
    schemas: BTreeMap<String, Arc<dyn SchemaProvider>>,
}

impl FlatNamespaceCatalog {
    /// Walk `catalog` and build a schema provider per servable namespace.
    /// Returns the catalog and one log line per namespace not served.
    pub async fn try_new(catalog: Arc<dyn Catalog>) -> Result<(Self, Vec<String>)> {
        let mut served = Vec::new();
        let mut skipped = Vec::new();
        for ns in walk(catalog.as_ref()).await? {
            match schema_name(&ns) {
                Ok(name) => served.push((name, ns)),
                Err(why) => skipped.push(why),
            }
        }
        let providers = try_join_all(
            served.iter().map(|(_, ns)| IcebergSchemaProvider::try_new(Arc::clone(&catalog), ns.clone())),
        )
        .await
        .map_err(Error::from)?;
        let schemas = served
            .into_iter()
            .zip(providers)
            .map(|((name, _), provider)| (name, Arc::new(provider) as Arc<dyn SchemaProvider>))
            .collect();
        Ok((Self { schemas }, skipped))
    }
}

impl CatalogProvider for FlatNamespaceCatalog {
    fn schema_names(&self) -> Vec<String> {
        self.schemas.keys().cloned().collect()
    }

    fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
        self.schemas.get(name).cloned()
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use iceberg::CatalogBuilder;
    use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};

    use super::*;

    fn ns(elements: &[&str]) -> NamespaceIdent {
        NamespaceIdent::from_strs(elements.iter().copied()).unwrap()
    }

    #[test]
    fn a_namespace_flattens_and_maps_back() {
        assert_eq!(schema_name(&ns(&["main"])).unwrap(), "main");
        assert_eq!(schema_name(&ns(&["a", "b"])).unwrap(), "a.b");
        assert_eq!(schema_name(&ns(&["x", "y", "z"])).unwrap(), "x.y.z");
        assert_eq!(namespace_of("x.y.z").unwrap(), ns(&["x", "y", "z"]));
        assert_eq!(namespace_of("main").unwrap(), ns(&["main"]));
        let why = schema_name(&ns(&["a", "b.c"])).unwrap_err();
        assert!(why.contains("\"b.c\"") && why.contains("not served"), "{why}");
    }

    #[tokio::test]
    async fn every_depth_is_one_schema_and_a_dotted_element_is_skipped_by_name() {
        let tmp = tempfile::tempdir().unwrap();
        let catalog = MemoryCatalogBuilder::default()
            .with_storage_factory(Arc::new(iceberg::io::LocalFsStorageFactory))
            .load("memory", HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), format!("file://{}", tmp.path().display()))]))
            .await
            .unwrap();
        for path in [&["one"][..], &["a"], &["a", "b"], &["x"], &["x", "y"], &["x", "y", "z"], &["dot.ted"], &["a", "b.c"]] {
            catalog.create_namespace(&ns(path), HashMap::new()).await.unwrap();
        }

        let (flat, skipped) = FlatNamespaceCatalog::try_new(Arc::new(catalog)).await.unwrap();
        assert_eq!(flat.schema_names(), vec!["a", "a.b", "one", "x", "x.y", "x.y.z"]);
        assert!(flat.schema("x.y.z").is_some() && flat.schema("dot.ted").is_none());
        assert_eq!(skipped.len(), 2, "{skipped:?}");
        assert!(skipped.iter().any(|l| l.contains("\"dot.ted\"")), "{skipped:?}");
        assert!(skipped.iter().any(|l| l.contains("\"b.c\"")), "{skipped:?}");
    }
}