nordisk/skade

skade-flight/tests/with_catalog.rs

//! `Engine::with_catalog`: the engine over any `iceberg::Catalog` (here a
//! skade warehouse's own `RedbCatalog`, passed as `dyn Catalog` — the same
//! shape a Nessie REST catalog arrives in), a non-default namespace, and the
//! read-only refusal.

use std::sync::Arc;

use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use futures::TryStreamExt;
use skade_flight::engine::EngineOptions;
use skade_flight::{Engine, Error};

fn people() -> (Schema, RecordBatch) {
    let schema = Schema::new(vec![Field::new("id", DataType::Int64, false), Field::new("name", DataType::Utf8, false)]);
    let batch = RecordBatch::try_new(
        Arc::new(schema.clone()),
        vec![Arc::new(Int64Array::from(vec![1, 2, 3])), Arc::new(StringArray::from(vec!["ada", "bob", "cyd"]))],
    )
    .unwrap();
    (schema, batch)
}

async fn rows(engine: &Engine, sql: &str) -> usize {
    let stream = engine.execute_stream(sql).await.unwrap_or_else(|e| panic!("{sql}: {e}"));
    let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
    batches.iter().map(|b| b.num_rows()).sum()
}

/// A warehouse with `lake.people`, and its catalog as a plain `dyn Catalog`.
async fn lake() -> (tempfile::TempDir, skade::Warehouse, Arc<dyn iceberg::Catalog>) {
    let dir = tempfile::tempdir().unwrap();
    let wh = skade::Warehouse::open(dir.path().join("wh")).await.unwrap();
    let (s, b) = people();
    wh.create_table("lake.people", &s).await.unwrap().append(&[b]).await.unwrap();
    let catalog: Arc<dyn iceberg::Catalog> = wh.catalog();
    (dir, wh, catalog)
}

#[tokio::test]
async fn with_catalog_serves_select_in_its_default_namespace() {
    let (_dir, _wh, catalog) = lake().await;
    let opts = EngineOptions { default_namespace: "lake".into(), target_partitions: 2, ..Default::default() };
    let engine = Engine::with_catalog(catalog, opts).await.unwrap();

    assert!(engine.try_warehouse().is_none(), "a foreign catalog has no skade Warehouse");
    assert_eq!(rows(&engine, "SELECT id, name FROM people").await, 3, "unqualified resolves in `lake`");
    assert_eq!(rows(&engine, "SELECT id FROM skade.lake.people WHERE id > 1").await, 2);
    assert!(engine.tables().await.unwrap().iter().any(|t| t.schema == "lake" && t.name == "people"));
    // Not read-only by default: a session view is allowed.
    assert_eq!(engine.execute_update("CREATE VIEW v AS SELECT id FROM people").await.unwrap(), 0);
}

#[tokio::test]
async fn read_only_refuses_every_write_path_and_still_reads() {
    let (_dir, _wh, catalog) = lake().await;
    let opts = EngineOptions { default_namespace: "lake".into(), read_only: true, target_partitions: 2, ..Default::default() };
    let engine = Engine::with_catalog(catalog, opts).await.unwrap();

    for sql in ["CREATE VIEW v AS SELECT id FROM people", "DELETE FROM people", "SELECT 1"] {
        match engine.execute_update(sql).await {
            Err(Error::ReadOnly(m)) => assert!(m.contains("read-only"), "{m}"),
            other => panic!("{sql}: expected ReadOnly, got {other:?}"),
        }
    }
    // The statement path applies a session-view DDL on execution: refused too.
    match engine.execute_stream("CREATE VIEW v AS SELECT id FROM people").await {
        Err(Error::ReadOnly(_)) => {}
        Err(e) => panic!("expected ReadOnly, got {e}"),
        Ok(_) => panic!("expected ReadOnly, got a stream"),
    }
    let status = tonic::Status::from(Error::ReadOnly("x".into()));
    assert_eq!(status.code(), tonic::Code::PermissionDenied);

    assert_eq!(rows(&engine, "SELECT * FROM people").await, 3, "reads are untouched");
}

#[tokio::test]
async fn warehouse_constructors_still_delegate() {
    let dir = tempfile::tempdir().unwrap();
    let wh = skade::Warehouse::open(dir.path().join("wh")).await.unwrap();
    let (s, b) = people();
    wh.create_table("people", &s).await.unwrap().append(&[b]).await.unwrap();
    let engine = Engine::new(wh).await.unwrap();
    assert!(engine.try_warehouse().is_some());
    assert_eq!(engine.options().default_namespace, skade::DEFAULT_NAMESPACE);
    assert!(!engine.options().read_only);
    assert_eq!(rows(&engine, "SELECT id FROM people").await, 3);
}

/// A connector that says it is read-only, over a skade directory.
#[derive(Debug)]
struct ReadOnlyLocal(skade::connect::SkadeLocal);

#[async_trait::async_trait]
impl skade::connect::CatalogConnector for ReadOnlyLocal {
    fn kind(&self) -> &'static str {
        "test-read-only"
    }
    fn describe(&self) -> String {
        self.0.describe()
    }
    fn read_only(&self) -> bool {
        true
    }
    async fn open(&self) -> skade::Result<Arc<dyn iceberg::Catalog>> {
        self.0.open().await
    }
}

#[tokio::test]
async fn with_connector_opens_through_the_trait_and_a_read_only_connector_forces_read_only() {
    use skade::connect::{CatalogConnector, SkadeLocal};

    let dir = tempfile::tempdir().unwrap();
    {
        let wh = skade::Warehouse::open(dir.path().join("wh")).await.unwrap();
        let (s, b) = people();
        wh.create_table("lake.people", &s).await.unwrap().append(&[b]).await.unwrap();
    } // writer dropped: the connector takes the redb lock

    let local: Arc<dyn CatalogConnector> = Arc::new(SkadeLocal::new(dir.path().join("wh")));
    let opts = EngineOptions { default_namespace: "lake".into(), target_partitions: 2, ..Default::default() };
    let engine = Engine::with_connector(local, opts.clone()).await.unwrap();
    assert!(!engine.options().read_only, "skade-local is writable");
    assert_eq!(rows(&engine, "SELECT id FROM people").await, 3);
    drop(engine); // release the redb lock for the next open

    let ro: Arc<dyn CatalogConnector> = Arc::new(ReadOnlyLocal(SkadeLocal::new(dir.path().join("wh"))));
    let engine = Engine::with_connector(ro, opts).await.unwrap();
    assert!(engine.options().read_only, "a read-only connector forces read_only");
    assert!(matches!(engine.execute_update("CREATE VIEW v AS SELECT id FROM people").await, Err(Error::ReadOnly(_))));
    assert_eq!(rows(&engine, "SELECT id FROM people").await, 3);
}