crates/graphar-flight/src/tests.rs
#[cfg(test)]
mod auto_schema_tests {
use crate::convert::response_to_batch_auto;
use redis::Value;
fn make_response(headers: Vec<&str>, rows: Vec<Vec<Value>>) -> Value {
Value::Array(vec![
Value::Array(
headers
.iter()
.map(|h| Value::BulkString(h.as_bytes().to_vec()))
.collect(),
),
Value::Array(rows.into_iter().map(Value::Array).collect()),
Value::Array(vec![]),
])
}
#[test]
fn infers_column_names_from_header() {
let response = make_response(
vec!["_gar_id", "name"],
vec![vec![Value::Int(0), Value::BulkString(b"Alice".to_vec())]],
);
let batch = response_to_batch_auto(response).unwrap();
assert_eq!(batch.schema().field(0).name(), "_gar_id");
assert_eq!(batch.schema().field(1).name(), "name");
}
#[test]
fn types_inferred_from_first_non_null_value() {
use arrow_schema::DataType;
let response = make_response(
vec!["id", "score", "active", "name"],
vec![vec![
Value::Int(1),
Value::Double(2.5),
Value::Boolean(true),
Value::BulkString(b"Alice".to_vec()),
]],
);
let batch = response_to_batch_auto(response).unwrap();
let schema = batch.schema();
assert_eq!(schema.field(0).data_type(), &DataType::Int64);
assert_eq!(schema.field(1).data_type(), &DataType::Float64);
assert_eq!(schema.field(2).data_type(), &DataType::Boolean);
assert_eq!(schema.field(3).data_type(), &DataType::Utf8);
}
#[test]
fn int_column_returns_int64() {
let response = make_response(vec!["id"], vec![vec![Value::Int(42)]]);
let batch = response_to_batch_auto(response).unwrap();
let col = batch
.column(0)
.as_any()
.downcast_ref::<arrow_array::Int64Array>()
.unwrap();
assert_eq!(col.value(0), 42i64);
}
#[test]
fn null_column_falls_back_to_utf8_and_is_null() {
let response = make_response(vec!["name"], vec![vec![Value::Nil]]);
let batch = response_to_batch_auto(response).unwrap();
// all-null → Utf8 fallback
assert_eq!(
batch.schema().field(0).data_type(),
&arrow_schema::DataType::Utf8
);
assert!(arrow_array::Array::is_null(batch.column(0).as_ref(), 0));
}
#[test]
fn nulls_skipped_during_sniff_type_from_later_row() {
// First row is null, second row has an Int — should still infer Int64.
let response = make_response(vec!["id"], vec![vec![Value::Nil], vec![Value::Int(7)]]);
let batch = response_to_batch_auto(response).unwrap();
assert_eq!(
batch.schema().field(0).data_type(),
&arrow_schema::DataType::Int64
);
let col = batch
.column(0)
.as_any()
.downcast_ref::<arrow_array::Int64Array>()
.unwrap();
assert!(arrow_array::Array::is_null(col, 0));
assert_eq!(col.value(1), 7i64);
}
#[test]
fn empty_result_returns_empty_schema() {
let batch = response_to_batch_auto(Value::Array(vec![])).unwrap();
assert_eq!(batch.num_rows(), 0);
assert_eq!(batch.num_columns(), 0);
}
#[test]
fn two_column_edge_schema_is_int64() {
// src + dst come back as Int64 — no registration needed.
let response = make_response(
vec!["src", "dst"],
vec![
vec![Value::Int(0), Value::Int(1)],
vec![Value::Int(0), Value::Int(2)],
],
);
let batch = response_to_batch_auto(response).unwrap();
assert_eq!(batch.num_columns(), 2);
assert_eq!(batch.num_rows(), 2);
let src = batch
.column(0)
.as_any()
.downcast_ref::<arrow_array::Int64Array>()
.unwrap();
assert_eq!(src.value(0), 0i64);
assert_eq!(src.value(1), 0i64);
}
}
#[cfg(test)]
mod convert_tests {
use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema};
use redis::Value;
use crate::convert::response_to_batch;
fn schema_i64_str() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
]))
}
fn make_response(rows: Vec<Vec<Value>>) -> Value {
Value::Array(vec![
Value::Array(vec![
Value::BulkString(b"id".to_vec()),
Value::BulkString(b"name".to_vec()),
]),
Value::Array(rows.into_iter().map(Value::Array).collect()),
Value::Array(vec![Value::BulkString(
b"Query internal execution time: 0.1 msec".to_vec(),
)]),
])
}
#[test]
fn empty_result_returns_zero_rows() {
let schema = schema_i64_str();
let response = Value::Array(vec![]);
let batch = response_to_batch(response, &schema).unwrap();
assert_eq!(batch.num_rows(), 0);
}
#[test]
fn two_rows_parsed_correctly() {
let schema = schema_i64_str();
let response = make_response(vec![
vec![Value::Int(1), Value::BulkString(b"Alice".to_vec())],
vec![Value::Int(2), Value::BulkString(b"Bob".to_vec())],
]);
let batch = response_to_batch(response, &schema).unwrap();
assert_eq!(batch.num_rows(), 2);
assert_eq!(batch.num_columns(), 2);
let ids = batch
.column(0)
.as_any()
.downcast_ref::<arrow_array::Int64Array>()
.unwrap();
assert_eq!(ids.value(0), 1);
assert_eq!(ids.value(1), 2);
let names = batch
.column(1)
.as_any()
.downcast_ref::<arrow_array::StringArray>()
.unwrap();
assert_eq!(names.value(0), "Alice");
assert_eq!(names.value(1), "Bob");
}
#[test]
fn null_value_becomes_none() {
let schema = schema_i64_str();
let response = make_response(vec![vec![Value::Int(42), Value::Nil]]);
let batch = response_to_batch(response, &schema).unwrap();
let names = batch
.column(1)
.as_any()
.downcast_ref::<arrow_array::StringArray>()
.unwrap();
assert!(arrow_array::Array::is_null(names, 0));
}
#[test]
fn float_column_parsed() {
let schema = Arc::new(Schema::new(vec![Field::new(
"score",
DataType::Float64,
true,
)]));
let response = Value::Array(vec![
Value::Array(vec![Value::BulkString(b"score".to_vec())]),
Value::Array(vec![
Value::Array(vec![Value::Double(2.5)]),
Value::Array(vec![Value::Int(2)]),
]),
Value::Array(vec![]),
]);
let batch = response_to_batch(response, &schema).unwrap();
let scores = batch
.column(0)
.as_any()
.downcast_ref::<arrow_array::Float64Array>()
.unwrap();
assert!((scores.value(0) - 2.5).abs() < 1e-9);
assert_eq!(scores.value(1), 2.0);
}
}
#[cfg(test)]
mod registry_tests {
use crate::registry::SchemaRegistry;
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
#[test]
fn register_and_get() {
let r = SchemaRegistry::new();
let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
r.register("my_query", schema.clone());
assert!(r.get("my_query").is_some());
assert!(r.get("other").is_none());
}
#[test]
fn overwrite_existing() {
let r = SchemaRegistry::new();
let s1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
let s2 = Arc::new(Schema::new(vec![Field::new("b", DataType::Utf8, false)]));
r.register("q", s1);
r.register("q", s2);
let got = r.get("q").unwrap();
assert_eq!(got.field(0).name(), "b");
}
#[test]
fn register_vertex_query_first_col_is_gar_id_int64() {
let r = SchemaRegistry::new();
r.register_vertex_query(
"MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name",
&["name"],
);
let schema = r
.get("MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name")
.unwrap();
let f0 = schema.field(0);
assert_eq!(f0.name(), "_gar_id");
assert_eq!(f0.data_type(), &DataType::Int64);
assert!(!f0.is_nullable());
let f1 = schema.field(1);
assert_eq!(f1.name(), "name");
assert_eq!(f1.data_type(), &DataType::Utf8);
assert!(f1.is_nullable());
}
#[test]
fn register_vertex_query_no_extra_cols() {
let r = SchemaRegistry::new();
r.register_vertex_query("MATCH (n) RETURN n._gar_id AS _gar_id", &[]);
let schema = r.get("MATCH (n) RETURN n._gar_id AS _gar_id").unwrap();
assert_eq!(schema.fields().len(), 1);
assert_eq!(schema.field(0).name(), "_gar_id");
}
#[test]
fn register_edge_query_first_two_cols_are_int64() {
let r = SchemaRegistry::new();
r.register_edge_query(
"MATCH (a)-[r:KNOWS]->(b) RETURN a._gar_id AS src, b._gar_id AS dst, r.since",
&["since"],
);
let schema = r
.get("MATCH (a)-[r:KNOWS]->(b) RETURN a._gar_id AS src, b._gar_id AS dst, r.since")
.unwrap();
assert_eq!(schema.field(0).name(), "src");
assert_eq!(schema.field(0).data_type(), &DataType::Int64);
assert_eq!(schema.field(1).name(), "dst");
assert_eq!(schema.field(1).data_type(), &DataType::Int64);
assert_eq!(schema.field(2).name(), "since");
assert_eq!(schema.field(2).data_type(), &DataType::Utf8);
}
#[test]
fn register_edge_query_no_extra_cols() {
let r = SchemaRegistry::new();
r.register_edge_query(
"MATCH (a)-[:FOLLOWS]->(b) RETURN a._gar_id AS src, b._gar_id AS dst",
&[],
);
let schema = r
.get("MATCH (a)-[:FOLLOWS]->(b) RETURN a._gar_id AS src, b._gar_id AS dst")
.unwrap();
assert_eq!(schema.fields().len(), 2);
}
#[test]
fn unregistered_query_returns_none() {
let r = SchemaRegistry::new();
assert!(r.get("MATCH (n) RETURN n").is_none());
}
}
#[cfg(test)]
mod named_table_tests {
use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema};
use crate::registry::{SchemaRegistry, SqlShape, classify_sql, table_candidates};
fn person_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("_gar_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
]))
}
#[test]
fn register_named_resolves_alias_to_query() {
let r = SchemaRegistry::new();
let q = "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name";
r.register_named("people", q, person_schema());
assert_eq!(r.resolve_table("people").as_deref(), Some(q));
// The underlying schema is reachable by the Cypher key too.
assert!(r.get(q).is_some());
}
#[test]
fn resolve_table_falls_back_to_raw_query_key() {
let r = SchemaRegistry::new();
let q = "MATCH (n) RETURN n._gar_id AS _gar_id";
r.register_vertex_query(q, &[]);
// Not a named alias, but a registered query — resolves to itself.
assert_eq!(r.resolve_table(q).as_deref(), Some(q));
}
#[test]
fn resolve_table_unknown_is_none() {
let r = SchemaRegistry::new();
assert!(r.resolve_table("nope").is_none());
}
#[test]
fn tables_lists_alias_for_named_and_query_for_unnamed() {
let r = SchemaRegistry::new();
r.register_named(
"people",
"MATCH (n:Person) RETURN n._gar_id AS _gar_id",
person_schema(),
);
r.register_vertex_query("MATCH (c:Company) RETURN c._gar_id AS _gar_id", &[]);
let mut names: Vec<String> = r.tables().into_iter().map(|(n, _)| n).collect();
names.sort();
assert_eq!(
names,
vec![
"MATCH (c:Company) RETURN c._gar_id AS _gar_id".to_string(),
"people".to_string(),
],
);
}
#[test]
fn tables_reports_registered_schema() {
let r = SchemaRegistry::new();
r.register_named(
"people",
"MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name",
person_schema(),
);
let (_, schema) = r.tables().into_iter().find(|(n, _)| n == "people").unwrap();
assert_eq!(schema.field(0).name(), "_gar_id");
assert_eq!(schema.field(0).data_type(), &DataType::Int64);
}
/// **The one honoured shape, in every spelling a driver emits.** Quotes,
/// backticks, a trailing `;`, lowercase, an `[AS] alias` — none of these
/// changes the result, so all of them are `Honoured`.
#[test]
fn select_star_from_one_table_is_the_honoured_shape() {
for sql in [
"SELECT * FROM people",
"select * from \"people\"",
"SELECT * FROM `people`",
"SELECT * FROM people;",
"SELECT * FROM people p",
"SELECT * FROM people AS \"t0\"",
" SELECT *\n FROM people\n",
] {
assert_eq!(
classify_sql(sql),
SqlShape::Honoured { table: "people".into() },
"{sql:?}"
);
}
}
/// **SQL the rewrite cannot honour is refused, naming what it would have
/// dropped.** This is the defect: the previous parser returned `people` for
/// `SELECT a, b FROM people WHERE a > 1` and the server sent the whole
/// table — projection gone, predicate gone, wrong count, no error. Power
/// BI folds by default, so a dashboard card showed the count of everything.
///
/// RED direction: put the old "first token after FROM" grab back and every
/// arm here is `Honoured` — the first assertion fails on the WHERE case.
#[test]
fn sql_the_engine_would_truncate_is_unhonoured_and_says_what_it_drops() {
let SqlShape::Unhonoured { table, dropped } =
classify_sql("SELECT a, b FROM people WHERE a > 1")
else {
panic!("a projection + predicate must never be served as the whole table");
};
assert_eq!(table, "people");
assert_eq!(
dropped,
vec![
"projection `a, b` (only `*` is served)".to_string(),
"WHERE".to_string()
],
"both clauses named, in source order"
);
// Each clause alone is enough to refuse.
let dropped_of = |sql: &str| match classify_sql(sql) {
SqlShape::Unhonoured { dropped, .. } => dropped,
other => panic!("{sql:?} must be unhonoured, got {other:?}"),
};
assert_eq!(dropped_of("SELECT count(*) FROM people"), vec!["projection `count(*)` (only `*` is served)"]);
assert_eq!(dropped_of("SELECT * FROM people WHERE name = 'x'"), vec!["WHERE"]);
assert_eq!(dropped_of("SELECT * FROM people LIMIT 10"), vec!["LIMIT"]);
assert_eq!(dropped_of("SELECT * FROM people ORDER BY name"), vec!["ORDER BY"]);
assert_eq!(dropped_of("SELECT * FROM people p JOIN cities c ON p.city = c.id"), vec!["JOIN"]);
assert_eq!(dropped_of("SELECT * FROM people, cities"), vec!["a second FROM table (comma join)"]);
assert_eq!(dropped_of("SELECT DISTINCT * FROM people"), vec!["DISTINCT"]);
assert_eq!(
dropped_of("SELECT * FROM people GROUP BY city HAVING count(*) > 1"),
vec!["GROUP BY", "HAVING"]
);
assert_eq!(dropped_of("WITH x AS (SELECT * FROM people) SELECT * FROM x"), vec!["WITH (common table expression)"]);
assert_eq!(dropped_of("SELECT * FROM people something odd"), vec!["trailing `something odd`"]);
}
/// Cypher, and SQL with no FROM, are not SQL for this engine: served verbatim.
#[test]
fn cypher_and_from_less_sql_are_not_sql_here() {
assert_eq!(classify_sql("MATCH (n) RETURN n"), SqlShape::NotSql);
assert_eq!(classify_sql("SELECT 1"), SqlShape::NotSql);
// A Cypher query that happens to contain the word FROM is still Cypher.
assert_eq!(classify_sql("MATCH (n:Person) WHERE n.from = 1 RETURN n"), SqlShape::NotSql);
}
/// **A qualified name resolves to the bare registered table.** Drivers
/// qualify what `GetTables` gave them; the registry holds `people`.
///
/// RED direction: drop the last-segment candidate and `public.people` finds
/// nothing.
#[test]
fn qualified_names_resolve_to_the_registered_table() {
for sql in [
"SELECT * FROM public.people",
"SELECT * FROM \"public\".\"people\"",
"SELECT * FROM `public`.`people`",
"SELECT * FROM \"public.people\"",
] {
let SqlShape::Honoured { table } = classify_sql(sql) else {
panic!("{sql:?} is the honoured shape");
};
assert_eq!(table, "public.people", "{sql:?}: quotes stripped per segment, dots kept");
assert_eq!(
table_candidates(&table),
vec!["public.people".to_string(), "people".to_string()],
"full name first, bare table second"
);
}
assert_eq!(table_candidates("people"), vec!["people".to_string()]);
}
/// **The handshake sends NO `HandshakeResponse` and ALWAYS an
/// `authorization: Bearer …` header** — the Dremio ODBC / arrow-cpp shape
/// (see `handshake_response`). RED direction: queue one HandshakeResponse
/// (the old reply) and the item count is 1; drop the header on an open
/// server and the second half fails.
#[tokio::test]
async fn handshake_sends_no_payload_and_always_a_bearer_header() {
use futures::StreamExt;
let resp = crate::service::handshake_response(Some("issued-token")).unwrap();
assert_eq!(
resp.metadata().get("authorization").unwrap().to_str().unwrap(),
"Bearer issued-token"
);
let items = resp.into_inner().collect::<Vec<_>>().await;
assert_eq!(items.len(), 0, "a queued HandshakeResponse hangs the C++ driver");
// Open server: a placeholder bearer, so the driver never sends an EMPTY one.
let open = crate::service::handshake_response(None).unwrap();
assert_eq!(
open.metadata().get("authorization").unwrap().to_str().unwrap(),
"Bearer open"
);
assert_eq!(open.into_inner().collect::<Vec<_>>().await.len(), 0);
}
/// **A statement's FlightInfo carries the request's FlightDescriptor.**
/// Both statement handlers omitted it while every metadata handler had it;
/// a generic C++ client refused every query result and accepted every
/// GetTables, and the difference was this field. RED direction: drop
/// `.with_descriptor` from `statement_info` and the descriptor is `None`.
#[test]
fn a_statement_flight_info_echoes_its_descriptor() {
use arrow_flight::FlightDescriptor;
let schema = person_schema();
let descriptor = FlightDescriptor::new_cmd(b"SELECT * FROM people".to_vec());
let info = crate::service::statement_info(
"MATCH (n:Person) RETURN n".into(),
schema.clone(),
descriptor.clone(),
)
.unwrap()
.into_inner();
assert_eq!(info.flight_descriptor, Some(descriptor), "the descriptor comes back");
assert_eq!(info.endpoint.len(), 1, "one endpoint");
assert!(info.endpoint[0].ticket.is_some(), "with a ticket");
assert_eq!(
info.try_decode_schema().unwrap(),
*schema,
"and the query's schema"
);
}
#[test]
fn sql_info_static_builds() {
// Exercises the lazily-built GetSqlInfo payload (server name/version).
let batch = crate::service::sql_info_batch();
assert!(
batch.num_rows() >= 3,
"advertises at least name/version/arrow"
);
}
}
#[cfg(test)]
mod service_directive_tests {
// parse_put_directive is private; test via the error type's Display.
use crate::error::FlightSqlError;
#[test]
fn unknown_directive_produces_correct_error() {
let err = FlightSqlError::UnknownPutCommand("whatever".into());
let msg = err.to_string();
assert!(msg.contains("VERTICES:<label>"));
assert!(msg.contains("EDGES:<src>:<type>:<dst>"));
}
}
/// **The zero-row Utf8 fallback contract** — `i64_column` is the inverse of
/// `sniff_column_types`, and these tests are what keeps the two in step. They
/// live WITH the implementation so a change to the coercion table cannot land
/// without facing them (the bug's whole cost was that the fix existed twice).
#[cfg(test)]
mod i64_column_tests {
use crate::convert::{i64_column, response_to_batch_auto};
use arrow_array::{Int32Array, RecordBatch, StringArray, UInt16Array};
use arrow_schema::{DataType, Field, Schema};
use redis::Value;
use std::sync::Arc;
/// **RED WHEN BROKEN — the graph bug, frozen.** A ZERO-ROW result has no
/// non-null cell for the sniff to inspect, so every column falls back to
/// `Utf8`. A hard `Int64Array` downcast returns `None` here, and the caller
/// turns that into the misleading `query must RETURN an Int64 src` error for
/// a query that was perfectly correct. Assert the column is PRESENT (Some)
/// and EMPTY — not absent.
///
/// Delete any integer arm or the `Utf8` arm of the coercion table and this
/// fails.
#[test]
fn zero_row_utf8_fallback_reads_as_present_and_empty() {
// Exactly what the sniff produces for an empty graph: every column Utf8.
let response = Value::Array(vec![
Value::Array(vec![
Value::BulkString(b"src".to_vec()),
Value::BulkString(b"dst".to_vec()),
]),
Value::Array(vec![]),
Value::Array(vec![]),
]);
let batch = response_to_batch_auto(response).unwrap();
assert_eq!(
batch.schema().field(0).data_type(),
&DataType::Utf8,
"the sniff really does fall a zero-row column back to Utf8 — \
this is the precondition the coercion exists for",
);
let src = i64_column(&batch, "src")
.expect("zero-row Utf8 `src` must read as an EMPTY column, not as absent");
let dst = i64_column(&batch, "dst")
.expect("zero-row Utf8 `dst` must read as an EMPTY column, not as absent");
assert!(
src.is_empty() && dst.is_empty(),
"no rows in, no values out"
);
}
/// The whole integer family coerces to i64, nulls survive as `None`, and a
/// non-integer column is rejected so a REAL shape mistake still errors.
#[test]
fn coerces_the_integer_family_and_rejects_non_integers() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("i32", DataType::Int32, true),
Field::new("u16", DataType::UInt16, true),
Field::new("strs", DataType::Utf8, true),
Field::new("floats", DataType::Float64, true),
])),
vec![
Arc::new(Int32Array::from(vec![Some(10), None, Some(-3)])) as _,
Arc::new(UInt16Array::from(vec![Some(7), Some(8), None])) as _,
// Utf8 with real content (a driver in string mode) — and whitespace,
// which the coercion trims.
Arc::new(StringArray::from(vec![Some(" 42 "), Some("nope"), None])) as _,
Arc::new(arrow_array::Float64Array::from(vec![Some(1.5), None, None])) as _,
],
)
.unwrap();
assert_eq!(
i64_column(&batch, "i32").unwrap(),
vec![Some(10), None, Some(-3)]
);
assert_eq!(
i64_column(&batch, "u16").unwrap(),
vec![Some(7), Some(8), None]
);
// "nope" is unparseable → None, not an error for the whole column.
assert_eq!(
i64_column(&batch, "strs").unwrap(),
vec![Some(42), None, None]
);
// RED WHEN BROKEN: widening the table to Float would silently accept a
// genuinely wrong query shape. Floats must stay rejected.
assert!(
i64_column(&batch, "floats").is_none(),
"a Float column is NOT an integer column"
);
assert!(
i64_column(&batch, "absent").is_none(),
"a missing column reads as absent"
);
}
/// A populated graph really does sniff `id(a)` to `Int64` — the arm the
/// zero-row case is the exception to.
#[test]
fn populated_int_column_still_sniffs_int64() {
let response = Value::Array(vec![
Value::Array(vec![Value::BulkString(b"src".to_vec())]),
Value::Array(vec![
Value::Array(vec![Value::Int(0)]),
Value::Array(vec![Value::Int(9)]),
]),
Value::Array(vec![]),
]);
let batch = response_to_batch_auto(response).unwrap();
assert_eq!(batch.schema().field(0).data_type(), &DataType::Int64);
assert_eq!(i64_column(&batch, "src").unwrap(), vec![Some(0), Some(9)]);
}
}