algeon-datafusion 0.1.2

use super::*;

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_get_catalogs_lists_default_catalog() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let flight_info = client.get_catalogs().await.unwrap();
    let batches = collect_metadata_batches(&mut client, &flight_info).await;

    server.abort();

    let names = collect_string_column(&batches, "catalog_name");
    assert!(names.contains(&"datafusion".to_string()), "expected default catalog, got {names:?}");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_get_db_schemas_lists_default_schemas() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let flight_info = client
        .get_db_schemas(CommandGetDbSchemas { catalog: Some("datafusion".into()), db_schema_filter_pattern: None })
        .await
        .unwrap();
    let batches = collect_metadata_batches(&mut client, &flight_info).await;

    server.abort();

    // `information_schema` is exposed by `SessionState` and is absent from the
    // catalog provider, so it does not surface here. The default `public` schema does.
    let schemas = collect_string_column(&batches, "db_schema_name");
    assert!(schemas.contains(&"public".to_string()), "missing public: {schemas:?}");
    let catalogs = collect_string_column(&batches, "catalog_name");
    assert!(catalogs.iter().all(|c| c == "datafusion"), "filter leaked: {catalogs:?}");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_get_tables_distinguishes_views_from_base_tables() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    // Seed both a real base table and a view. The assertions below check that
    // the metadata layer reports each with the correct Flight SQL table_type.
    let dir = tempfile::Builder::new().prefix("algeon-datafusion-flight-sql-metadata-").tempdir().unwrap();
    let parquet_path = dir.path().join("metadata_probe.parquet");
    write_update_edges_parquet(&parquet_path);
    let source_ddl = format!(
        "CREATE EXTERNAL TABLE metadata_probe_base STORED AS PARQUET LOCATION '{}'",
        parquet_path.to_string_lossy()
    );
    let _ = execute_batches(&mut client, &source_ddl).await;
    let _ =
        execute_batches(&mut client, "CREATE VIEW metadata_probe_view AS SELECT src AS id FROM metadata_probe_base")
            .await;

    let flight_info = client
        .get_tables(CommandGetTables {
            catalog: Some("datafusion".into()),
            db_schema_filter_pattern: Some("public".into()),
            table_name_filter_pattern: Some("metadata_probe%".into()),
            table_types: vec![],
            include_schema: false,
        })
        .await
        .unwrap();
    let batches = collect_metadata_batches(&mut client, &flight_info).await;

    server.abort();

    let pairs = collect_name_and_type(&batches);
    assert!(
        pairs.contains(&("metadata_probe_base".to_string(), "TABLE".to_string())),
        "missing seeded base table or wrong type: {pairs:?}"
    );
    assert!(
        pairs.contains(&("metadata_probe_view".to_string(), "VIEW".to_string())),
        "missing seeded view or wrong type: {pairs:?}"
    );
    let catalogs = collect_string_column(&batches, "catalog_name");
    assert!(catalogs.iter().all(|c| c == "datafusion"), "catalog filter leaked: {catalogs:?}");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_get_tables_respects_table_types_filter() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let _ = execute_batches(&mut client, "CREATE TABLE table_types_filter_base AS SELECT 1 AS x").await;
    let _ =
        execute_batches(&mut client, "CREATE VIEW table_types_filter_view AS SELECT x FROM table_types_filter_base")
            .await;

    let flight_info = client
        .get_tables(CommandGetTables {
            catalog: Some("datafusion".into()),
            db_schema_filter_pattern: Some("public".into()),
            table_name_filter_pattern: Some("table_types_filter%".into()),
            table_types: vec!["VIEW".into()],
            include_schema: false,
        })
        .await
        .unwrap();
    let batches = collect_metadata_batches(&mut client, &flight_info).await;

    server.abort();

    let pairs = collect_name_and_type(&batches);
    assert_eq!(pairs, vec![("table_types_filter_view".to_string(), "VIEW".to_string())]);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_get_tables_with_include_schema_emits_arrow_schema() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let _ = execute_batches(
        &mut client,
        "CREATE TABLE include_schema_probe AS \
         SELECT 1::INT AS id, 'a' AS label",
    )
    .await;

    let flight_info = client
        .get_tables(CommandGetTables {
            catalog: Some("datafusion".into()),
            db_schema_filter_pattern: Some("public".into()),
            table_name_filter_pattern: Some("include_schema_probe".into()),
            table_types: vec![],
            include_schema: true,
        })
        .await
        .unwrap();
    let batches = collect_metadata_batches(&mut client, &flight_info).await;

    server.abort();

    let pairs = collect_name_and_type(&batches);
    assert_eq!(pairs, vec![("include_schema_probe".to_string(), "TABLE".to_string())]);

    let schemas: Vec<Schema> = batches
        .iter()
        .filter_map(|batch| {
            let index = batch.schema().index_of("table_schema").ok()?;
            let bytes =
                batch.column(index).as_any().downcast_ref::<BinaryArray>().expect("table_schema column must be Binary");
            (0..bytes.len()).map(|i| try_schema_from_ipc_buffer(bytes.value(i)).unwrap()).next()
        })
        .collect();
    assert_eq!(schemas.len(), 1, "expected exactly one table_schema row");
    let field_names: Vec<String> = schemas[0].fields().iter().map(|f| f.name().clone()).collect();
    assert_eq!(field_names, vec!["id".to_string(), "label".to_string()]);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_get_table_types_returns_table_and_view() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let flight_info = client.get_table_types().await.unwrap();
    let batches = collect_metadata_batches(&mut client, &flight_info).await;

    server.abort();

    let types = collect_string_column(&batches, "table_type");
    assert_eq!(types, vec!["TABLE".to_string(), "VIEW".to_string()]);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_get_primary_keys_returns_empty_with_schema() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let cmd = CommandGetPrimaryKeys { catalog: None, db_schema: None, table: "any_table".to_string() };
    let flight_info = client.get_primary_keys(cmd).await.unwrap();
    let (schema, batches) = collect_metadata_schema_and_batches(&mut client, &flight_info).await;

    server.abort();

    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total_rows, 0, "primary keys must be empty for DataFusion catalog");
    let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
    assert_eq!(
        field_names,
        vec!["catalog_name", "db_schema_name", "table_name", "column_name", "key_name", "key_sequence"]
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flight_sql_foreign_keys_rpcs_return_empty_with_shared_schema() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let imported_info = client
        .get_imported_keys(CommandGetImportedKeys { catalog: None, db_schema: None, table: "t".to_string() })
        .await
        .unwrap();
    let exported_info = client
        .get_exported_keys(CommandGetExportedKeys { catalog: None, db_schema: None, table: "t".to_string() })
        .await
        .unwrap();
    let xref_info = client
        .get_cross_reference(CommandGetCrossReference {
            pk_catalog: None,
            pk_db_schema: None,
            pk_table: "p".to_string(),
            fk_catalog: None,
            fk_db_schema: None,
            fk_table: "f".to_string(),
        })
        .await
        .unwrap();

    let (imported_schema, imported) = collect_metadata_schema_and_batches(&mut client, &imported_info).await;
    let (exported_schema, exported) = collect_metadata_schema_and_batches(&mut client, &exported_info).await;
    let (xref_schema, xref) = collect_metadata_schema_and_batches(&mut client, &xref_info).await;

    server.abort();

    let expected_names = vec![
        "pk_catalog_name",
        "pk_db_schema_name",
        "pk_table_name",
        "pk_column_name",
        "fk_catalog_name",
        "fk_db_schema_name",
        "fk_table_name",
        "fk_column_name",
        "key_sequence",
        "fk_key_name",
        "pk_key_name",
        "update_rule",
        "delete_rule",
    ];
    for (label, schema, batches) in [
        ("imported", &imported_schema, &imported),
        ("exported", &exported_schema, &exported),
        ("cross_reference", &xref_schema, &xref),
    ] {
        let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(rows, 0, "{label} keys must be empty");
        let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(names, expected_names, "{label} keys schema mismatch");
        let update_rule = schema.field_with_name("update_rule").unwrap();
        assert_eq!(update_rule.data_type(), &DataType::UInt8, "{label} update_rule must be UInt8");
    }
    // Reference the UInt8Array import to keep it warning-free even when batches are empty.
    let _ = UInt8Array::from(Vec::<u8>::new());
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn strict_flight_sql_metadata_domain_uses_cpu_without_attempt_or_user_scan() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;
    config.native_optimizer_config =
        config.native_optimizer_config.clone().with_execution_mode(AlgeonGpuExecutionMode::NativeRequired);
    let backend = build_injected_test_backend(&mut config);
    config.native_optimizer_config = AlgeonNativeOptimizerConfig::default();
    assert_eq!(backend.optimizer_config().execution_mode(), AlgeonGpuExecutionMode::NativeRequired);
    let ctx = build_session_context_with_backend(&config, &backend).await.unwrap();
    assert_eq!(ctx.state().config_options().execution.target_partitions, 1);
    let source_dir = tempfile::Builder::new().prefix("algeon-datafusion-strict-metadata-").tempdir().unwrap();
    let source_path = source_dir.path().join("registered_source.parquet");
    write_update_edges_parquet(&source_path);
    let source_ddl = format!(
        "CREATE EXTERNAL TABLE metadata_registered_source STORED AS PARQUET LOCATION '{}'",
        source_path.to_string_lossy()
    );
    ctx.sql(&source_ddl).await.unwrap().collect().await.unwrap();
    let scan_probe = Arc::new(ScanFailingTableProvider::new(Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("label", DataType::Utf8, true),
    ]))));
    ctx.register_table("metadata_probe", Arc::clone(&scan_probe) as _).unwrap();
    ctx.register_table(
        "user_rows",
        Arc::new(EmptyTable::new(Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int64, false),
            Field::new("table_name", DataType::Utf8, true),
        ])))),
    )
    .unwrap();
    let validation_edges = Arc::new(ScanFailingTableProvider::new(Arc::new(Schema::new(vec![
        Field::new("src", DataType::Int64, false),
        Field::new("dst", DataType::Int64, false),
        Field::new("weight", DataType::Float64, false),
    ]))));
    ctx.register_table("validation_edges", Arc::clone(&validation_edges) as _).unwrap();
    let replacement = ctx.state().create_logical_plan("SELECT * FROM user_rows").await.unwrap();
    let controlled_rewrite = Arc::new(ControlledPlanReplacement::new(replacement));
    ctx.register_relation_planner(Arc::clone(&controlled_rewrite) as _).unwrap();
    let ctx = Arc::new(ctx);
    let shutdown = CancellationToken::new();
    let server = FlightSqlServer::builder()
        .config(config)
        .backend(backend.clone())
        .session_context(Arc::clone(&ctx))
        .listener(listener)
        .shutdown(ServerShutdown::Token(shutdown.clone()))
        .build()
        .await
        .unwrap();
    let server = tokio::spawn(async move { server.serve().await });
    let mut client = connect_with_retry(addr).await;
    let attempts_before = backend.attempt_counters_snapshot();
    let admission_before = backend.diagnostics().unwrap().admission;

    let get_tables = client
        .get_tables(CommandGetTables {
            catalog: Some("datafusion".into()),
            db_schema_filter_pattern: Some("public".into()),
            table_name_filter_pattern: Some("metadata_probe".into()),
            table_types: vec![],
            include_schema: false,
        })
        .await
        .unwrap();
    let get_tables_batches = collect_metadata_batches(&mut client, &get_tables).await;
    assert_eq!(collect_name_and_type(&get_tables_batches), vec![("metadata_probe".into(), "TABLE".into())]);

    let show_tables_info = client.execute("SHOW TABLES".to_string(), None).await.unwrap();
    let show_tables_query_id = query_id_from_statement_flight_info(&show_tables_info);
    let (show_tables_schema, show_tables) = collect_metadata_schema_and_batches(&mut client, &show_tables_info).await;
    let show_tables_schema_names: Vec<&str> =
        show_tables_schema.fields().iter().map(|field| field.name().as_str()).collect();
    assert_eq!(show_tables_schema_names, vec!["table_catalog", "table_schema", "table_name", "table_type"]);
    assert_eq!(
        collect_name_and_type(&show_tables).into_iter().find(|(name, _)| name == "metadata_probe"),
        Some(("metadata_probe".to_string(), "BASE TABLE".to_string()))
    );

    let show_columns_info = client.execute("SHOW COLUMNS FROM metadata_probe".to_string(), None).await.unwrap();
    let direct_columns = collect_metadata_batches(&mut client, &show_columns_info).await;
    assert_eq!(collect_string_column(&direct_columns, "column_name"), vec!["id", "label"]);

    let show_create_info =
        client.execute("SHOW CREATE TABLE metadata_registered_source".to_string(), None).await.unwrap();
    let show_create_query_id = query_id_from_statement_flight_info(&show_create_info);
    let show_create = collect_metadata_batches(&mut client, &show_create_info).await;
    assert_eq!(collect_string_column(&show_create, "table_name"), vec!["metadata_registered_source"]);
    assert_eq!(collect_string_column(&show_create, "definition"), vec![source_ddl.clone()]);
    let scan_free_create_info = client.execute("SHOW CREATE TABLE metadata_probe".to_string(), None).await.unwrap();
    let scan_free_create = collect_metadata_batches(&mut client, &scan_free_create_info).await;
    assert_eq!(collect_string_column(&scan_free_create, "table_name"), vec!["metadata_probe"]);

    let describe_info = client.execute("DESCRIBE metadata_probe".to_string(), None).await.unwrap();
    let described = collect_metadata_batches(&mut client, &describe_info).await;
    assert_eq!(collect_string_column(&described, "column_name"), vec!["id", "label"]);

    let information_schema_info = client
        .execute(
            "SELECT t.table_name, COUNT(c.column_name) AS column_count \
             FROM information_schema.tables t \
             JOIN information_schema.columns c \
               ON t.table_catalog = c.table_catalog \
              AND t.table_schema = c.table_schema \
              AND t.table_name = c.table_name \
             WHERE t.table_name = 'metadata_probe' \
             GROUP BY t.table_name"
                .to_string(),
            None,
        )
        .await
        .unwrap();
    let information_schema = collect_metadata_batches(&mut client, &information_schema_info).await;
    assert_eq!(collect_string_column(&information_schema, "table_name"), vec!["metadata_probe"]);
    let column_count = information_schema.iter().find_map(|batch| {
        batch.schema().index_of("column_count").ok().and_then(|index| int_value(batch.column(index)))
    });
    assert_eq!(column_count, Some(2));

    let gpu_list_info = client
        .execute(
            "SELECT function_name, provider, available, summary \
             FROM gpu_list_functions() \
             WHERE provider IN ('cugraph', 'cuvs') \
             ORDER BY function_name"
                .to_string(),
            None,
        )
        .await
        .unwrap();
    let gpu_list_query_id = query_id_from_statement_flight_info(&gpu_list_info);
    let gpu_list = collect_metadata_batches(&mut client, &gpu_list_info).await;
    let providers = collect_string_column(&gpu_list, "provider");
    assert!(providers.iter().any(|provider| provider == "cugraph"), "{providers:?}");
    assert!(providers.iter().any(|provider| provider == "cuvs"), "{providers:?}");

    let describe_info = client
        .execute("SELECT function_name FROM gpu_describe_function('cugraph_bfs')".to_string(), None)
        .await
        .unwrap();
    let described_function = collect_metadata_batches(&mut client, &describe_info).await;
    assert_eq!(collect_string_column(&described_function, "function_name"), vec!["cugraph_bfs"]);

    let validation_info = client
        .execute(
            r#"SELECT valid, relations_resolved
               FROM gpu_validate_call(
                 'cugraph_bfs',
                 '{"relations":{"edges":{"table":"validation_edges"}},"options":{"source_vertex":0,"depth_limit":2}}'
               )"#
            .to_string(),
            None,
        )
        .await
        .unwrap();
    let validation = collect_metadata_batches(&mut client, &validation_info).await;
    assert_eq!(validation.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
    assert!(!validation_edges.scan_called(), "metadata validation must inspect only the registered table schema");

    let mut prepared = client.prepare("SHOW COLUMNS FROM metadata_probe".to_string(), None).await.unwrap();
    let prepared_info = prepared.execute().await.unwrap();
    let prepared_columns = collect_metadata_batches(&mut client, &prepared_info).await;
    prepared.close().await.unwrap();
    assert_eq!(
        collect_string_column(&prepared_columns, "column_name"),
        collect_string_column(&direct_columns, "column_name")
    );

    let mut prepared_create =
        client.prepare("SHOW CREATE TABLE metadata_registered_source".to_string(), None).await.unwrap();
    let prepared_create_info = prepared_create.execute().await.unwrap();
    let prepared_create_rows = collect_metadata_batches(&mut client, &prepared_create_info).await;
    prepared_create.close().await.unwrap();
    assert_eq!(collect_string_column(&prepared_create_rows, "definition"), vec![source_ddl]);

    let cancelled_info = client.execute("SHOW TABLES".to_string(), None).await.unwrap();
    let cancelled_query_id = query_id_from_statement_flight_info(&cancelled_info);
    assert_eq!(cancel_query_raw(&mut client, &cancelled_info).await.result, CANCEL_RESULT_CANCELLING);

    assert!(!scan_probe.scan_called(), "metadata discovery must not call the user table provider's scan method");
    assert_eq!(
        backend.attempt_counters_snapshot(),
        attempts_before,
        "metadata execution must not begin Algeon attempts"
    );
    assert_eq!(
        backend.diagnostics().unwrap().admission,
        admission_before,
        "metadata execution must not alter GPU admission"
    );

    let metadata_stats = get_query_stats_raw(&mut client, &show_tables_query_id).await;
    assert!(metadata_stats.contains("# outcome=succeeded"), "{metadata_stats}");
    assert!(metadata_stats.contains("[physical_plan]"), "{metadata_stats}");
    assert!(!metadata_stats.contains("# attempt_id="), "{metadata_stats}");
    let show_create_stats = get_query_stats_raw(&mut client, &show_create_query_id).await;
    assert!(show_create_stats.contains("# outcome=succeeded"), "{show_create_stats}");
    assert!(!show_create_stats.contains("# attempt_id="), "{show_create_stats}");
    let gpu_list_stats = get_query_stats_raw(&mut client, &gpu_list_query_id).await;
    assert!(gpu_list_stats.contains("# outcome=succeeded"), "{gpu_list_stats}");
    assert!(!gpu_list_stats.contains("# attempt_id="), "{gpu_list_stats}");
    assert!(!gpu_list_stats.contains("[planning]"), "{gpu_list_stats}");
    let cancelled_stats = get_query_stats_raw(&mut client, &cancelled_query_id).await;
    assert!(cancelled_stats.contains("# outcome=cancelled"), "{cancelled_stats}");
    assert!(!cancelled_stats.contains("# attempt_id="), "{cancelled_stats}");

    let proof_rejection_info =
        client.execute("SELECT table_name FROM information_schema.tables".to_string(), None).await.unwrap();
    controlled_rewrite.set_enabled(true);
    let proof_rejection_ticket =
        proof_rejection_info.endpoint.first().and_then(|endpoint| endpoint.ticket.clone()).unwrap();
    let proof_rejection_status = match client.do_get(proof_rejection_ticket).await {
        Err(error) => unwrap_flight_status(error),
        Ok(stream) => unwrap_flight_status(stream.try_collect::<Vec<_>>().await.unwrap_err()),
    };
    controlled_rewrite.set_enabled(false);
    assert_error_metadata(&proof_rejection_status, "invalid_request", "invalid_input", "permanent");
    assert_eq!(
        backend.attempt_counters_snapshot(),
        attempts_before,
        "a retained metadata ticket must fail instead of retrying through the data domain"
    );

    for sql in [
        "SELECT id FROM user_rows WHERE CAST(id AS VARCHAR) ILIKE '1%'",
        "SELECT i.table_name FROM information_schema.tables i JOIN user_rows u ON true WHERE i.table_name ILIKE '%probe%'",
        "SELECT m.function_name FROM gpu_list_functions() m JOIN user_rows u ON true",
    ] {
        let flight_info = client.execute(sql.to_string(), None).await.unwrap();
        let ticket = flight_info.endpoint.first().and_then(|endpoint| endpoint.ticket.clone()).unwrap();
        let status = match client.do_get(ticket).await {
            Err(error) => unwrap_flight_status(error),
            Ok(stream) => unwrap_flight_status(stream.try_collect::<Vec<_>>().await.unwrap_err()),
        };
        assert_error_metadata(&status, "final_plan_requirement_unsatisfied", "planning", "permanent");
    }
    assert_eq!(
        backend.attempt_counters_snapshot().attempts_begun,
        attempts_before.attempts_begun + 3,
        "mixed and user-table statements must retain data-domain attempt ownership"
    );

    drop(client);
    shutdown.cancel();
    timeout(Duration::from_secs(10), server)
        .await
        .expect("token shutdown should drain promptly")
        .expect("server task should not panic")
        .expect("token shutdown should stop the server cleanly");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_sql_info_advertises_core_server_capabilities() {
    let (listener, addr) = free_local_addr();
    let mut config = test_server_config();
    declare_test_gpu_ownership(&mut config);
    config.bind_addr = addr;

    let server = tokio::spawn(async move { serve_with_listener(config, listener).await });
    let mut client = connect_with_retry(addr).await;

    let requested = vec![
        SqlInfo::FlightSqlServerName,
        SqlInfo::FlightSqlServerSql,
        SqlInfo::FlightSqlServerReadOnly,
        SqlInfo::FlightSqlServerTransaction,
        SqlInfo::FlightSqlServerCancel,
        SqlInfo::SqlIdentifierQuoteChar,
        SqlInfo::SqlNullOrdering,
        SqlInfo::SqlSupportedGroupBy,
        SqlInfo::SqlOuterJoinsSupportLevel,
        SqlInfo::SqlSchemaTerm,
        SqlInfo::SqlCatalogTerm,
        SqlInfo::SqlCatalogAtStart,
        SqlInfo::SqlSupportedUnions,
        SqlInfo::SqlTransactionsSupported,
        SqlInfo::SqlSupportedResultSetTypes,
        SqlInfo::SqlSupportedConcurrenciesForResultSetForwardOnly,
    ];
    let flight_info = client.get_sql_info(requested).await.unwrap();
    let ticket = flight_info.endpoint.first().and_then(|endpoint| endpoint.ticket.clone()).unwrap();
    let batches = client.do_get(ticket).await.unwrap().try_collect::<Vec<_>>().await.unwrap();

    server.abort();

    assert_eq!(batches.len(), 1);
    let batch = &batches[0];
    assert_eq!(batch.num_rows(), 16);
    assert_eq!(sql_info_string(batch, SqlInfo::FlightSqlServerName), "algeon-datafusion");
}