#[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();
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() {
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() {
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, parse_select_from};
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));
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, &[]);
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);
}
#[test]
fn parse_select_from_extracts_identifier() {
assert_eq!(
parse_select_from("SELECT * FROM people").as_deref(),
Some("people")
);
assert_eq!(
parse_select_from("select * from \"people\"").as_deref(),
Some("people")
);
assert_eq!(
parse_select_from("SELECT * FROM `people`").as_deref(),
Some("people")
);
assert_eq!(
parse_select_from("SELECT a, b FROM people WHERE a > 1").as_deref(),
Some("people")
);
assert_eq!(
parse_select_from("SELECT * FROM people;").as_deref(),
Some("people")
);
}
#[test]
fn parse_select_from_rejects_non_select() {
assert!(parse_select_from("MATCH (n) RETURN n").is_none());
assert!(parse_select_from("SELECT 1").is_none());
}
#[test]
fn sql_info_static_builds() {
let batch = crate::service::sql_info_batch();
assert!(
batch.num_rows() >= 3,
"advertises at least name/version/arrow"
);
}
}
#[cfg(test)]
mod service_directive_tests {
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>"));
}
}