crates/graphar-flight/tests/integration_flight_sql.rs
//! Power BI / ADBC path, end to end against a live FalkorDB.
//!
//! Starts the Cypher Flight SQL server with mutual-exclusive header auth
//! (Basic + bearer) and drives it with the stock `FlightSqlServiceClient` —
//! the same client the Arrow Flight SQL ODBC driver and Power BI's ADBC
//! connector are built on. Exercises:
//!
//! - the **Basic handshake** → bearer-token issuance flow,
//! - the metadata surface: `GetCatalogs`, `GetDbSchemas`, `GetTables`, `GetSqlInfo`,
//! - a registered, *named* query executed as `SELECT * FROM people`.
//!
//! ## Prerequisites
//!
//! - Podman socket: `systemctl --user start podman.socket`
//! - FalkorDB image built locally:
//! ```bash
//! podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .
//! ```
//!
//! ## Run
//!
//! ```bash
//! cargo test -p graphar-flight --test integration_flight_sql -- --ignored --nocapture
//! ```
use std::sync::Arc;
use std::time::Duration;
use arrow_array::{Int64Array, StringArray};
use arrow_flight::sql::CommandGetTables;
use arrow_flight::sql::client::FlightSqlServiceClient;
use arrow_schema::{DataType, Field, Schema};
use futures::TryStreamExt;
use graphar_flight::{SchemaRegistry, ServerOptions, run_server_with};
use testcontainers::{GenericImage, ImageExt, core::ContainerPort, runners::AsyncRunner};
const PEOPLE_QUERY: &str = "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name AS name";
fn ensure_docker_host() {
if std::env::var("DOCKER_HOST").is_err() {
let uid = std::process::Command::new("id")
.arg("-u")
.output()
.ok()
.and_then(|o| String::from_utf8(o.stdout).ok())
.and_then(|s| s.trim().parse::<u32>().ok())
.unwrap_or(1000);
let sock = format!("/run/user/{uid}/podman/podman.sock");
if std::path::Path::new(&sock).exists() {
unsafe { std::env::set_var("DOCKER_HOST", format!("unix://{sock}")) };
}
}
}
/// Grab a free localhost port by binding and immediately dropping the listener.
fn free_port() -> u16 {
std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port()
}
async fn collect_rows(
client: &mut FlightSqlServiceClient<tonic::transport::Channel>,
info: arrow_flight::FlightInfo,
) -> Vec<arrow_array::RecordBatch> {
let ticket = info.endpoint[0]
.ticket
.clone()
.expect("endpoint has a ticket");
client
.do_get(ticket)
.await
.expect("do_get")
.try_collect::<Vec<_>>()
.await
.expect("collect batches")
}
#[tokio::test]
#[ignore = "needs podman + locally built FalkorDB image; run with --ignored"]
async fn power_bi_path_handshake_metadata_and_query() {
ensure_docker_host();
// ── FalkorDB container ───────────────────────────────────────────────────
let falkor_image = std::env::var("KNUT_FALKORDB_IMAGE")
.unwrap_or_else(|_| "localhost/knut-falkordb-valkey9:latest".into());
let (fname, ftag) = falkor_image.rsplit_once(':').unwrap();
let falkor = GenericImage::new(fname, ftag)
.with_exposed_port(ContainerPort::Tcp(6379))
.with_startup_timeout(Duration::from_secs(120))
.start()
.await
.expect("start falkordb (build it first: podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .)");
let falkor_port = falkor.get_host_port_ipv4(6379).await.unwrap();
let redis_url = format!("redis://127.0.0.1:{falkor_port}");
let graph = "bi";
// ── seed a tiny graph ────────────────────────────────────────────────────
let client = redis::Client::open(redis_url.clone()).unwrap();
let mut conn = client.get_multiplexed_async_connection().await.unwrap();
let _: redis::Value = redis::cmd("GRAPH.QUERY")
.arg(graph)
.arg("CREATE (:Person {_gar_id: 0, name: 'alice'}), (:Person {_gar_id: 1, name: 'bob'})")
.query_async(&mut conn)
.await
.unwrap();
// ── start the Flight SQL server (Basic + bearer auth) ────────────────────
let registry = Arc::new(SchemaRegistry::new());
registry.register_named(
"people",
PEOPLE_QUERY,
Arc::new(Schema::new(vec![
Field::new("_gar_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
])),
);
let port = free_port();
let addr = format!("127.0.0.1:{port}").parse().unwrap();
let opts = ServerOptions {
bearer_token: Some("issued-token".into()),
basic_auth: Some(("analyst".into(), "p@ss".into())),
..ServerOptions::default()
};
let server_redis = redis_url.clone();
tokio::spawn(async move {
run_server_with(&server_redis, graph, registry, addr, opts)
.await
.unwrap();
});
// Give tonic a moment to bind.
tokio::time::sleep(Duration::from_millis(500)).await;
// ── connect the stock Flight SQL client ──────────────────────────────────
let channel = tonic::transport::Channel::from_shared(format!("http://127.0.0.1:{port}"))
.unwrap()
.connect()
.await
.expect("connect flight channel");
let mut fsql = FlightSqlServiceClient::new(channel);
// A call without credentials is rejected.
assert!(
fsql.get_catalogs().await.is_err(),
"unauthenticated call must fail"
);
// Basic handshake → bearer token stored on the client FROM THE HEADER. The
// server sends no HandshakeResponse (the Dremio/arrow-cpp driver hangs on
// one), so the returned payload is empty by design; what proves the token
// arrived is that every authenticated call below succeeds.
let token = fsql.handshake("analyst", "p@ss").await.expect("handshake");
assert!(token.is_empty(), "no HandshakeResponse payload: the bearer rides the header");
// ── metadata surface ─────────────────────────────────────────────────────
let catalogs_info = fsql.get_catalogs().await.unwrap();
let catalogs = collect_rows(&mut fsql, catalogs_info).await;
let cat_names = string_col(&catalogs, "catalog_name");
assert!(
cat_names.contains(&graph.to_string()),
"catalog = graph name"
);
let tables_info = fsql
.get_tables(CommandGetTables {
include_schema: false,
..Default::default()
})
.await
.unwrap();
let tables = collect_rows(&mut fsql, tables_info).await;
let table_names = string_col(&tables, "table_name");
assert!(
table_names.contains(&"people".to_string()),
"named query surfaces as a table"
);
let sql_info_info = fsql.get_sql_info(vec![]).await.unwrap();
let sql_info = collect_rows(&mut fsql, sql_info_info).await;
assert!(
sql_info.iter().map(|b| b.num_rows()).sum::<usize>() >= 3,
"advertises server info"
);
// ── execute the named query as a SQL statement ───────────────────────────
let info = fsql
.execute("SELECT * FROM people".into(), None)
.await
.unwrap();
let rows = collect_rows(&mut fsql, info).await;
let total: usize = rows.iter().map(|b| b.num_rows()).sum();
assert_eq!(total, 2, "two Person rows");
// _gar_id came back typed as Int64 (registered schema, not auto-string).
let ids = rows[0].column_by_name("_gar_id").unwrap();
assert!(
ids.as_any().downcast_ref::<Int64Array>().is_some(),
"_gar_id is Int64"
);
}
/// Collect a Utf8 column's values across batches by name.
fn string_col(batches: &[arrow_array::RecordBatch], name: &str) -> Vec<String> {
use arrow_array::Array;
let mut out = Vec::new();
for b in batches {
if let Some(col) = b.column_by_name(name)
&& let Some(arr) = col.as_any().downcast_ref::<StringArray>()
{
for i in 0..arr.len() {
if !arr.is_null(i) {
out.push(arr.value(i).to_string());
}
}
}
}
out
}