graphar-flight 0.1.2

//! Power BI / ADBC path, end to end against a live FalkorDB.
//!
//! Starts the Cypher Flight SQL server with mutual-exclusive header auth
//! (Basic + bearer) and drives it with the stock `FlightSqlServiceClient` —
//! the same client the Arrow Flight SQL ODBC driver and Power BI's ADBC
//! connector are built on. Exercises:
//!
//! - the **Basic handshake** → bearer-token issuance flow,
//! - the metadata surface: `GetCatalogs`, `GetDbSchemas`, `GetTables`, `GetSqlInfo`,
//! - a registered, *named* query executed as `SELECT * FROM people`.
//!
//! ## Prerequisites
//!
//! - Podman socket: `systemctl --user start podman.socket`
//! - FalkorDB image built locally:
//!   ```bash
//!   podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .
//!   ```
//!
//! ## Run
//!
//! ```bash
//! cargo test -p graphar-flight --test integration_flight_sql -- --ignored --nocapture
//! ```

use std::sync::Arc;
use std::time::Duration;

use arrow_array::{Int64Array, StringArray};
use arrow_flight::sql::CommandGetTables;
use arrow_flight::sql::client::FlightSqlServiceClient;
use arrow_schema::{DataType, Field, Schema};
use futures::TryStreamExt;
use graphar_flight::{SchemaRegistry, ServerOptions, run_server_with};
use testcontainers::{GenericImage, ImageExt, core::ContainerPort, runners::AsyncRunner};

const PEOPLE_QUERY: &str = "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name AS name";

fn ensure_docker_host() {
    if std::env::var("DOCKER_HOST").is_err() {
        let uid = std::process::Command::new("id")
            .arg("-u")
            .output()
            .ok()
            .and_then(|o| String::from_utf8(o.stdout).ok())
            .and_then(|s| s.trim().parse::<u32>().ok())
            .unwrap_or(1000);
        let sock = format!("/run/user/{uid}/podman/podman.sock");
        if std::path::Path::new(&sock).exists() {
            unsafe { std::env::set_var("DOCKER_HOST", format!("unix://{sock}")) };
        }
    }
}

/// Grab a free localhost port by binding and immediately dropping the listener.
fn free_port() -> u16 {
    std::net::TcpListener::bind("127.0.0.1:0")
        .unwrap()
        .local_addr()
        .unwrap()
        .port()
}

async fn collect_rows(
    client: &mut FlightSqlServiceClient<tonic::transport::Channel>,
    info: arrow_flight::FlightInfo,
) -> Vec<arrow_array::RecordBatch> {
    let ticket = info.endpoint[0]
        .ticket
        .clone()
        .expect("endpoint has a ticket");
    client
        .do_get(ticket)
        .await
        .expect("do_get")
        .try_collect::<Vec<_>>()
        .await
        .expect("collect batches")
}

#[tokio::test]
#[ignore = "needs podman + locally built FalkorDB image; run with --ignored"]
async fn power_bi_path_handshake_metadata_and_query() {
    ensure_docker_host();

    // ── FalkorDB container ───────────────────────────────────────────────────
    let falkor_image = std::env::var("KNUT_FALKORDB_IMAGE")
        .unwrap_or_else(|_| "localhost/knut-falkordb-valkey9:latest".into());
    let (fname, ftag) = falkor_image.rsplit_once(':').unwrap();
    let falkor = GenericImage::new(fname, ftag)
        .with_exposed_port(ContainerPort::Tcp(6379))
        .with_startup_timeout(Duration::from_secs(120))
        .start()
        .await
        .expect("start falkordb (build it first: podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .)");
    let falkor_port = falkor.get_host_port_ipv4(6379).await.unwrap();
    let redis_url = format!("redis://127.0.0.1:{falkor_port}");
    let graph = "bi";

    // ── seed a tiny graph ────────────────────────────────────────────────────
    let client = redis::Client::open(redis_url.clone()).unwrap();
    let mut conn = client.get_multiplexed_async_connection().await.unwrap();
    let _: redis::Value = redis::cmd("GRAPH.QUERY")
        .arg(graph)
        .arg("CREATE (:Person {_gar_id: 0, name: 'alice'}), (:Person {_gar_id: 1, name: 'bob'})")
        .query_async(&mut conn)
        .await
        .unwrap();

    // ── start the Flight SQL server (Basic + bearer auth) ────────────────────
    let registry = Arc::new(SchemaRegistry::new());
    registry.register_named(
        "people",
        PEOPLE_QUERY,
        Arc::new(Schema::new(vec![
            Field::new("_gar_id", DataType::Int64, false),
            Field::new("name", DataType::Utf8, true),
        ])),
    );
    let port = free_port();
    let addr = format!("127.0.0.1:{port}").parse().unwrap();
    let opts = ServerOptions {
        bearer_token: Some("issued-token".into()),
        basic_auth: Some(("analyst".into(), "p@ss".into())),
        ..ServerOptions::default()
    };
    let server_redis = redis_url.clone();
    tokio::spawn(async move {
        run_server_with(&server_redis, graph, registry, addr, opts)
            .await
            .unwrap();
    });
    // Give tonic a moment to bind.
    tokio::time::sleep(Duration::from_millis(500)).await;

    // ── connect the stock Flight SQL client ──────────────────────────────────
    let channel = tonic::transport::Channel::from_shared(format!("http://127.0.0.1:{port}"))
        .unwrap()
        .connect()
        .await
        .expect("connect flight channel");
    let mut fsql = FlightSqlServiceClient::new(channel);

    // A call without credentials is rejected.
    assert!(
        fsql.get_catalogs().await.is_err(),
        "unauthenticated call must fail"
    );

    // Basic handshake → bearer token stored on the client.
    let token = fsql.handshake("analyst", "p@ss").await.expect("handshake");
    assert_eq!(token.as_ref(), b"issued-token");

    // ── metadata surface ─────────────────────────────────────────────────────
    let catalogs_info = fsql.get_catalogs().await.unwrap();
    let catalogs = collect_rows(&mut fsql, catalogs_info).await;
    let cat_names = string_col(&catalogs, "catalog_name");
    assert!(
        cat_names.contains(&graph.to_string()),
        "catalog = graph name"
    );

    let tables_info = fsql
        .get_tables(CommandGetTables {
            include_schema: false,
            ..Default::default()
        })
        .await
        .unwrap();
    let tables = collect_rows(&mut fsql, tables_info).await;
    let table_names = string_col(&tables, "table_name");
    assert!(
        table_names.contains(&"people".to_string()),
        "named query surfaces as a table"
    );

    let sql_info_info = fsql.get_sql_info(vec![]).await.unwrap();
    let sql_info = collect_rows(&mut fsql, sql_info_info).await;
    assert!(
        sql_info.iter().map(|b| b.num_rows()).sum::<usize>() >= 3,
        "advertises server info"
    );

    // ── execute the named query as a SQL statement ───────────────────────────
    let info = fsql
        .execute("SELECT * FROM people".into(), None)
        .await
        .unwrap();
    let rows = collect_rows(&mut fsql, info).await;
    let total: usize = rows.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total, 2, "two Person rows");

    // _gar_id came back typed as Int64 (registered schema, not auto-string).
    let ids = rows[0].column_by_name("_gar_id").unwrap();
    assert!(
        ids.as_any().downcast_ref::<Int64Array>().is_some(),
        "_gar_id is Int64"
    );
}

/// Collect a Utf8 column's values across batches by name.
fn string_col(batches: &[arrow_array::RecordBatch], name: &str) -> Vec<String> {
    use arrow_array::Array;
    let mut out = Vec::new();
    for b in batches {
        if let Some(col) = b.column_by_name(name)
            && let Some(arr) = col.as_any().downcast_ref::<StringArray>()
        {
            for i in 0..arr.len() {
                if !arr.is_null(i) {
                    out.push(arr.value(i).to_string());
                }
            }
        }
    }
    out
}