datafusion-flight-sql-server: add a few tests

Author: hozan23

datafusion-flight-sql-server/tests/integration_test.rs

@@ -0,0 +1,328 @@
+use std::sync::Arc;
+
+use arrow::{
+    array::{Int32Array, RecordBatch, StringArray},
+    datatypes::{DataType, Field, Schema},
+};
+use arrow_flight::sql::client::FlightSqlServiceClient;
+use datafusion::{
+    datasource::MemTable,
+    execution::context::{SessionContext, SessionState},
+};
+use datafusion_flight_sql_server::service::FlightSqlService;
+use futures::TryStreamExt;
+use tokio::time::{sleep, Duration};
+use tonic::transport::{Channel, Endpoint};
+
+fn create_test_session() -> SessionState {
+    let ctx = SessionContext::new();
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Int32, false),
+        Field::new("name", DataType::Utf8, false),
+    ]));
+
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3])),
+            Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie"])),
+        ],
+    )
+    .unwrap();
+
+    let table = MemTable::try_new(schema, vec![vec![batch]]).unwrap();
+    ctx.register_table("users", Arc::new(table)).unwrap();
+
+    let orders_schema = Arc::new(Schema::new(vec![
+        Field::new("order_id", DataType::Int32, false),
+        Field::new("user_id", DataType::Int32, false),
+        Field::new("amount", DataType::Int32, false),
+    ]));
+
+    let orders_batch = RecordBatch::try_new(
+        orders_schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![100, 101, 102, 103])),
+            Arc::new(Int32Array::from(vec![1, 2, 1, 3])),
+            Arc::new(Int32Array::from(vec![50, 75, 100, 25])),
+        ],
+    )
+    .unwrap();
+
+    let orders_table = MemTable::try_new(orders_schema, vec![vec![orders_batch]]).unwrap();
+    ctx.register_table("orders", Arc::new(orders_table))
+        .unwrap();
+
+    ctx.state()
+}
+
+async fn start_test_server(addr: String, state: SessionState) {
+    tokio::spawn(async move {
+        FlightSqlService::new(state)
+            .serve(addr)
+            .await
+            .expect("Server should start successfully");
+    });
+
+    sleep(Duration::from_millis(500)).await;
+}
+
+async fn create_test_client(addr: &str) -> FlightSqlServiceClient<Channel> {
+    let endpoint = Endpoint::new(addr.to_string()).expect("Valid endpoint");
+    let channel = endpoint.connect().await.expect("Connection successful");
+    FlightSqlServiceClient::new(channel)
+}
+
+#[tokio::test]
+async fn test_basic_query_execution() {
+    let addr = "0.0.0.0:50061";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let flight_info = client
+        .execute("SELECT * FROM users".to_string(), None)
+        .await
+        .expect("Query should succeed");
+
+    let ticket = flight_info
+        .endpoint
+        .first()
+        .expect("Should have endpoint")
+        .ticket
+        .clone()
+        .expect("Should have ticket");
+
+    let mut stream = client.do_get(ticket).await.expect("do_get should succeed");
+
+    let mut batches = Vec::new();
+    while let Some(batch) = stream.try_next().await.expect("Stream should work") {
+        batches.push(batch);
+    }
+
+    assert!(!batches.is_empty(), "Should have result batches");
+
+    let first_batch = &batches[0];
+    assert_eq!(first_batch.num_columns(), 2);
+    assert_eq!(first_batch.schema().field(0).name(), "id");
+    assert_eq!(first_batch.schema().field(1).name(), "name");
+
+    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(total_rows, 3);
+}
+
+#[tokio::test]
+async fn test_query_with_filter() {
+    let addr = "0.0.0.0:50062";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let flight_info = client
+        .execute("SELECT name FROM users WHERE id > 1".to_string(), None)
+        .await
+        .expect("Query should succeed");
+
+    let ticket = flight_info
+        .endpoint
+        .first()
+        .expect("Should have endpoint")
+        .ticket
+        .clone()
+        .expect("Should have ticket");
+
+    let mut stream = client.do_get(ticket).await.expect("do_get should succeed");
+
+    let mut batches = Vec::new();
+    while let Some(batch) = stream.try_next().await.expect("Stream should work") {
+        batches.push(batch);
+    }
+
+    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(total_rows, 2, "Should have 2 rows after filter");
+}
+
+#[tokio::test]
+async fn test_prepared_statement_creation() {
+    let addr = "0.0.0.0:50063";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let query = "SELECT * FROM users WHERE id = $1";
+    let prepared = client
+        .prepare(query.to_string(), None)
+        .await
+        .expect("Prepare should succeed");
+
+    let dataset_schema = prepared
+        .dataset_schema()
+        .expect("Should have dataset schema");
+    assert_eq!(dataset_schema.fields().len(), 2);
+
+    let parameter_schema = prepared
+        .parameter_schema()
+        .expect("Should have parameter schema");
+    assert_eq!(parameter_schema.fields().len(), 1);
+}
+
+#[tokio::test]
+async fn test_get_schemas() {
+    let addr = "0.0.0.0:50064";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let flight_info = client
+        .get_db_schemas(arrow_flight::sql::CommandGetDbSchemas {
+            catalog: Some("datafusion".to_string()),
+            db_schema_filter_pattern: None,
+        })
+        .await
+        .expect("GetDbSchemas should succeed");
+
+    let ticket = flight_info
+        .endpoint
+        .first()
+        .expect("Should have endpoint")
+        .ticket
+        .clone()
+        .expect("Should have ticket");
+
+    let mut stream = client.do_get(ticket).await.expect("do_get should succeed");
+
+    let mut batches = Vec::new();
+    while let Some(batch) = stream.try_next().await.expect("Stream should work") {
+        batches.push(batch);
+    }
+
+    assert!(!batches.is_empty(), "Should have schema results");
+}
+
+#[tokio::test]
+async fn test_get_tables() {
+    let addr = "0.0.0.0:50065";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let flight_info = client
+        .get_tables(arrow_flight::sql::CommandGetTables {
+            catalog: Some("datafusion".to_string()),
+            db_schema_filter_pattern: None,
+            table_name_filter_pattern: None,
+            table_types: vec![],
+            include_schema: true,
+        })
+        .await
+        .expect("GetTables should succeed");
+
+    let ticket = flight_info
+        .endpoint
+        .first()
+        .expect("Should have endpoint")
+        .ticket
+        .clone()
+        .expect("Should have ticket");
+
+    let mut stream = client.do_get(ticket).await.expect("do_get should succeed");
+
+    let mut batches = Vec::new();
+    while let Some(batch) = stream.try_next().await.expect("Stream should work") {
+        batches.push(batch);
+    }
+
+    assert!(!batches.is_empty(), "Should have table results");
+
+    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert!(total_rows > 0, "Should have at least one table");
+}
+
+#[tokio::test]
+async fn test_invalid_query() {
+    let addr = "0.0.0.0:50066";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let result = client
+        .execute("SELECT * FROM nonexistent_table".to_string(), None)
+        .await;
+
+    assert!(result.is_err(), "Query should fail for nonexistent table");
+}
+
+#[tokio::test]
+async fn test_query_with_aggregation() {
+    let addr = "0.0.0.0:50067";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let flight_info = client
+        .execute("SELECT COUNT(*) as count FROM users".to_string(), None)
+        .await
+        .expect("Query should succeed");
+
+    let ticket = flight_info
+        .endpoint
+        .first()
+        .expect("Should have endpoint")
+        .ticket
+        .clone()
+        .expect("Should have ticket");
+
+    let mut stream = client.do_get(ticket).await.expect("do_get should succeed");
+
+    let mut batches = Vec::new();
+    while let Some(batch) = stream.try_next().await.expect("Stream should work") {
+        batches.push(batch);
+    }
+
+    assert!(!batches.is_empty(), "Should have result batches");
+
+    let first_batch = &batches[0];
+    assert_eq!(first_batch.num_columns(), 1);
+    assert_eq!(first_batch.schema().field(0).name(), "count");
+}
+
+#[tokio::test]
+async fn test_query_with_join() {
+    let addr = "0.0.0.0:50068";
+    let state = create_test_session();
+    start_test_server(addr.to_string(), state).await;
+
+    let mut client = create_test_client(&format!("http://{}", addr)).await;
+
+    let flight_info = client
+        .execute(
+            r#"
+            SELECT u.id, u.name, o.order_id 
+                FROM users u 
+                JOIN orders o 
+                    ON u.id = o.user_id "#
+                .to_string(),
+            None,
+        )
+        .await
+        .expect("Join query should succeed");
+
+    let ticket = flight_info.endpoint[0].ticket.clone().unwrap();
+    let mut stream = client.do_get(ticket).await.expect("do_get should succeed");
+
+    let mut batches = Vec::new();
+    while let Some(batch) = stream.try_next().await.expect("Stream should work") {
+        batches.push(batch);
+    }
+
+    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(total_rows, 4, "Should have 4 rows from join");
+}

datafusion-flight-sql-server/tests/state_test.rs

@@ -0,0 +1,37 @@
+use datafusion_flight_sql_server::state::QueryHandle;
+
+#[test]
+fn test_query_handle_with_complex_sql() {
+    let query = r#"
+        SELECT
+            a.id,
+            a.name,
+            COUNT(b.order_id) as order_count
+        FROM customers a
+        LEFT JOIN orders b ON a.id = b.customer_id
+        WHERE a.created_at > $1 AND a.status = $2
+        GROUP BY a.id, a.name
+        HAVING COUNT(b.order_id) > $3
+        ORDER BY order_count DESC
+        LIMIT 100
+    "#
+    .to_string();
+
+    let handle = QueryHandle::new(query.clone(), None);
+
+    let encoded = handle.clone().encode();
+    let decoded = QueryHandle::try_decode(encoded).expect("Should decode");
+
+    assert_eq!(decoded.query(), query);
+}
+
+#[test]
+fn test_query_handle_empty_query() {
+    let query = String::new();
+    let handle = QueryHandle::new(query.clone(), None);
+
+    let encoded = handle.clone().encode();
+    let decoded = QueryHandle::try_decode(encoded).expect("Should decode");
+
+    assert_eq!(decoded.query(), "");
+}