nordisk/skade

skade-flight/tests/nested_namespaces.rs

//! Nested namespaces through the Flight SQL door: a one-, two- and three-level
//! namespace each served as ONE schema named `a.b`-style, listed by
//! `GetDbSchemas` / `GetTables` (Power BI's navigator) and SELECTed by quoted
//! schema name — over real TCP with the stock `FlightSqlServiceClient`.

use std::sync::Arc;
use std::time::Duration;

use arrow_array::{Array, Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use futures::TryStreamExt;
use skade_flight::arrow_flight::sql::client::FlightSqlServiceClient;
use skade_flight::arrow_flight::sql::{CommandGetDbSchemas, CommandGetTables};
use skade_flight::engine::EngineOptions;
use skade_flight::{Auth, Engine};
use tonic::transport::Channel;

fn ids(n: i64) -> (Schema, RecordBatch) {
    let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]);
    let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(Int64Array::from((1..=n).collect::<Vec<_>>()))]).unwrap();
    (schema, batch)
}

async fn collect(client: &mut FlightSqlServiceClient<Channel>, info: arrow_flight::FlightInfo) -> Vec<RecordBatch> {
    let ticket = info.endpoint[0].ticket.clone().unwrap();
    client.do_get(ticket).await.unwrap().try_collect().await.unwrap()
}

fn strings(batches: &[RecordBatch], column: &str) -> Vec<String> {
    batches
        .iter()
        .flat_map(|b| {
            let col = b.column_by_name(column).unwrap_or_else(|| panic!("no column {column}"));
            let col = col.as_any().downcast_ref::<StringArray>().unwrap();
            (0..col.len()).map(|i| col.value(i).to_string()).collect::<Vec<_>>()
        })
        .collect()
}

async fn count(client: &mut FlightSqlServiceClient<Channel>, sql: &str) -> i64 {
    let info = client.execute(sql.to_string(), None).await.unwrap_or_else(|e| panic!("{sql}: {e}"));
    let batches = collect(client, info).await;
    batches[0].column(0).as_any().downcast_ref::<Int64Array>().unwrap().value(0)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn one_two_and_three_level_namespaces_list_and_select_over_flight_sql() {
    let dir = tempfile::tempdir().unwrap();
    let wh = skade::Warehouse::open(dir.path().join("wh")).await.unwrap();
    for (name, n) in [("one.t1", 1), ("a.b.t2", 2), ("x.y.z.t3", 3)] {
        let (s, b) = ids(n);
        wh.create_table(name, &s).await.unwrap().append(&[b]).await.unwrap();
    }
    let opts = EngineOptions {
        default_namespace: "one".into(),
        target_partitions: 2,
        refresh_after: Duration::from_millis(50),
        ..Default::default()
    };
    let engine = Arc::new(Engine::with_engine_options(wh, opts).await.unwrap());

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let _server = {
        let engine = engine.clone();
        tokio::spawn(async move {
            skade_flight::serve_on(engine, Arc::new(Auth::open()), listener, Default::default(), futures::future::pending())
                .await
                .unwrap();
        })
    };
    let channel = Channel::from_shared(format!("http://{addr}")).unwrap().connect().await.unwrap();
    let mut c = FlightSqlServiceClient::new(channel);

    // GetDbSchemas: every depth is one schema, parents included.
    let info = c.get_db_schemas(CommandGetDbSchemas::default()).await.unwrap();
    let schemas = strings(&collect(&mut c, info).await, "db_schema_name");
    for expected in ["main", "one", "a", "a.b", "x", "x.y", "x.y.z"] {
        assert!(schemas.iter().any(|s| s == expected), "{expected} missing from {schemas:?}");
    }

    // GetTables: each table under its flattened schema.
    let info = c.get_tables(CommandGetTables::default()).await.unwrap();
    let batches = collect(&mut c, info).await;
    let pairs: Vec<(String, String)> =
        strings(&batches, "db_schema_name").into_iter().zip(strings(&batches, "table_name")).collect();
    for (schema, table) in [("one", "t1"), ("a.b", "t2"), ("x.y.z", "t3")] {
        assert!(pairs.iter().any(|(s, t)| s == schema && t == table), "{schema}.{table} missing from {pairs:?}");
    }

    // SELECT by quoted flattened schema name, at each depth.
    assert_eq!(count(&mut c, "SELECT count(*) FROM skade.one.t1").await, 1);
    assert_eq!(count(&mut c, r#"SELECT count(*) FROM skade."a.b".t2"#).await, 2);
    assert_eq!(count(&mut c, r#"SELECT count(*) FROM skade."x.y.z".t3"#).await, 3);
    assert_eq!(count(&mut c, r#"SELECT count(*) FROM "a.b".t2"#).await, 2);
    // The default schema is the one-level namespace.
    assert_eq!(count(&mut c, "SELECT count(*) FROM t1").await, 1);

    // GetPrimaryKeys maps the flattened name back to the elements (no identity → empty, not an error).
    assert!(engine.primary_keys("x.y.z", "t3").await.unwrap().is_empty());
}