datafusion-flight-sql-server 0.4.18

use std::sync::Arc;

use arrow_flight::{sql::client::FlightSqlServiceClient, FlightInfo};
use datafusion::arrow::{
    array::{Int32Array, RecordBatch, StringArray},
    datatypes::{DataType, Field, Schema},
};
use datafusion::{
    datasource::MemTable,
    execution::context::{SessionContext, SessionState},
    parquet::arrow::ArrowWriter,
    prelude::ParquetReadOptions,
};
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");
}

/// The query used by the schema-consistency tests below.
///
/// Aggregates over statistics-backed sources (e.g. Parquet) are a case where
/// the logical plan schema (advertised in `FlightInfo`) and the physical stream
/// schema disagree on field nullability: the `AggregateStatistics` optimizer
/// rewrites `MIN()` into a non-nullable literal taken from the file statistics,
/// while the logical schema keeps `MIN()` nullable. Strict clients (e.g. the
/// ADBC Flight SQL driver) reject a `DoGet` stream whose schema does not
/// exactly match the one advertised in `FlightInfo`.
const AGGREGATE_OVER_PARQUET: &str = "SELECT MIN(amount) AS lo, COUNT(*) AS n FROM orders_pq";

/// Registers a single-column Parquet table named `orders_pq`.
///
/// The file lives in the per-test-binary temp directory that Cargo manages, so
/// a panicking test cannot leave stray files behind in the system temp dir.
async fn create_parquet_session(file_name: &str) -> SessionState {
    let schema = Arc::new(Schema::new(vec![Field::new(
        "amount",
        DataType::Int32,
        false,
    )]));
    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![Arc::new(Int32Array::from(vec![50, 75, 100, 25]))],
    )
    .unwrap();

    let path = std::path::Path::new(env!("CARGO_TARGET_TMPDIR")).join(file_name);
    let file = std::fs::File::create(&path).expect("create parquet file");
    let mut writer = ArrowWriter::try_new(file, schema, None).expect("create writer");
    writer.write(&batch).expect("write batch");
    writer.close().expect("close writer");

    let ctx = SessionContext::new();
    ctx.register_parquet(
        "orders_pq",
        path.to_str().expect("utf-8 path"),
        ParquetReadOptions::default(),
    )
    .await
    .expect("register parquet");

    ctx.state()
}

/// Fetches the endpoint's `DoGet` stream, drains it, and returns the schema the
/// stream declared on the wire.
async fn do_get_stream_schema(
    client: &mut FlightSqlServiceClient<Channel>,
    flight_info: &FlightInfo,
) -> Schema {
    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");
    // Drain the stream: decoding the batches also validates them against the
    // declared schema, so a schema that does not fit the data fails here.
    while stream
        .try_next()
        .await
        .expect("Stream should work")
        .is_some()
    {}

    stream
        .schema()
        .expect("DoGet stream should declare a schema")
        .as_ref()
        .clone()
}

/// Asserts that `MIN(amount)` is still advertised as nullable, i.e. that the
/// query really does exercise the logical/physical schema divergence and the
/// test has not silently become vacuous.
fn assert_exercises_nullability_divergence(advertised: &Schema) {
    let lo = advertised.field_with_name("lo").expect("lo field");
    assert!(
        lo.is_nullable(),
        "MIN() must be advertised as nullable for this test to be meaningful, got {lo:?}"
    );
}

#[tokio::test]
async fn test_do_get_schema_matches_advertised_flight_info_schema() {
    let addr = "0.0.0.0:50069";
    let state = create_parquet_session("statement_query.parquet").await;
    start_test_server(addr.to_string(), state).await;

    let mut client = create_test_client(&format!("http://{}", addr)).await;

    let flight_info = client
        .execute(AGGREGATE_OVER_PARQUET.to_string(), None)
        .await
        .expect("Query should succeed");

    let advertised = flight_info
        .clone()
        .try_decode_schema()
        .expect("FlightInfo should carry a schema");
    assert_exercises_nullability_divergence(&advertised);

    let actual = do_get_stream_schema(&mut client, &flight_info).await;

    assert_eq!(
        advertised, actual,
        "DoGet stream schema must match the schema advertised in FlightInfo"
    );
}

#[tokio::test]
async fn test_prepared_do_get_schema_matches_advertised_flight_info_schema() {
    let addr = "0.0.0.0:50070";
    let state = create_parquet_session("prepared_statement_query.parquet").await;
    start_test_server(addr.to_string(), state).await;

    let mut client = create_test_client(&format!("http://{}", addr)).await;

    let mut prepared = client
        .prepare(AGGREGATE_OVER_PARQUET.to_string(), None)
        .await
        .expect("Prepare should succeed");

    let flight_info = prepared.execute().await.expect("Query should succeed");
    let advertised = flight_info
        .clone()
        .try_decode_schema()
        .expect("FlightInfo should carry a schema");
    assert_exercises_nullability_divergence(&advertised);

    let actual = do_get_stream_schema(&mut client, &flight_info).await;

    assert_eq!(
        advertised, actual,
        "DoGet stream schema must match the schema advertised in FlightInfo"
    );
}