knut-bifrost 0.3.0

//! Flight transform plugins (#6): a [`Transform`] whose logic lives in an
//! **external Arrow-Flight service**, reached over the bidirectional
//! [`DoExchange`] stream — **Arrow record batches in, Arrow record batches
//! out** — instead of in-process Rust or a sandboxed wasm module. The service
//! can be written in any language with an Arrow Flight server (Python, Go,
//! C++, Java, …); bifröst just proxies the transform to it and keeps the same
//! [`GraphMutation`] contract.
//!
//! This is the sibling of [`crate::wasm::WasmTransform`]: same shape (rows →
//! transformed rows → one [`VertexSet`] per output batch under a configured
//! label + id column), but **out-of-process and polyglot** rather than an
//! in-process wasm interpreter.
//!
//! ## Wire protocol
//!
//! `apply` opens a `DoExchange` call to the configured endpoint, encodes the
//! input batches as a Flight data stream ([`FlightDataEncoderBuilder`]), and
//! reads the server's response stream back as record batches
//! ([`FlightRecordBatchStream`], via the client). The server is free to
//! filter / derive / reshape the rows; an identity server (echo the input)
//! yields the input rows as vertices.
//!
//! [`DoExchange`]: arrow_flight::flight_service_client::FlightServiceClient::do_exchange
//!
//! ```no_run
//! # #[cfg(feature = "flight")]
//! # async fn demo() -> knut_bifrost::Result<()> {
//! use knut_bifrost::flight::FlightTransform;
//! let t = FlightTransform::connect("http://transform-svc:50055", "Person", "_gar_id").await?;
//! # let _ = t; Ok(()) }
//! ```

use std::sync::Arc;

use arrow_array::RecordBatch;
use arrow_flight::{
    FlightClient, decode::FlightRecordBatchStream, encode::FlightDataEncoderBuilder,
};
use futures::TryStreamExt;
use tokio::sync::Mutex;
use tonic::transport::Channel;

use crate::error::{BifrostError, Result};
use crate::mutation::{GraphMutation, VertexSet};
use crate::transform::{TableContext, Transform};

/// A [`Transform`] that proxies to an external Arrow-Flight `DoExchange`
/// service (Arrow record batches in/out).
///
/// The connected [`FlightClient`] is shared behind an `Arc<Mutex<…>>` so the
/// transform is `Send + Sync` and a single TCP channel is reused across
/// `apply` calls (a `DoExchange` call borrows the client mutably).
pub struct FlightTransform {
    client: Arc<Mutex<FlightClient>>,
    label: String,
    id_column: String,
}

impl FlightTransform {
    /// Connect to an Arrow-Flight endpoint (e.g. `http://host:50055`) and bind
    /// the transform's output to `label` / `id_column`.
    ///
    /// The connection is established once here; each `apply` opens a fresh
    /// `DoExchange` call over the same channel.
    pub async fn connect(
        endpoint: impl Into<String>,
        label: impl Into<String>,
        id_column: impl Into<String>,
    ) -> Result<Self> {
        let endpoint = endpoint.into();
        let channel = Channel::from_shared(endpoint.clone())
            .map_err(|e| flight_err(format!("invalid endpoint {endpoint:?}: {e}")))?
            .connect()
            .await
            .map_err(|e| flight_err(format!("connect {endpoint:?}: {e}")))?;
        Ok(Self::from_client(
            FlightClient::new(channel),
            label,
            id_column,
        ))
    }

    /// Build from an already-connected [`FlightClient`] (e.g. over a test
    /// channel). `connect` is the usual entry point.
    pub fn from_client(
        client: FlightClient,
        label: impl Into<String>,
        id_column: impl Into<String>,
    ) -> Self {
        Self {
            client: Arc::new(Mutex::new(client)),
            label: label.into(),
            id_column: id_column.into(),
        }
    }

    /// Run one `DoExchange` round: send `batches`, collect the transformed
    /// batches the service streams back.
    async fn exchange(&self, batches: &[RecordBatch]) -> Result<Vec<RecordBatch>> {
        // Encode the input batches as a Flight data stream. The encoder emits
        // the schema message followed by one dictionary/record-batch message
        // set per batch — exactly what the server's `FlightRecordBatchStream`
        // decoder expects.
        let owned: Vec<RecordBatch> = batches.to_vec();
        let input =
            FlightDataEncoderBuilder::new().build(futures::stream::iter(owned.into_iter().map(Ok)));

        let mut client = self.client.lock().await;
        let response: FlightRecordBatchStream = client
            .do_exchange(input)
            .await
            .map_err(|e| flight_err(format!("do_exchange: {e}")))?;

        response
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| flight_err(format!("decode response: {e}")))
    }
}

impl Transform for FlightTransform {
    fn apply(&self, _ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
        if batches.is_empty() {
            return Ok(GraphMutation::default());
        }

        // `Transform::apply` is sync; bridge to the async Flight client the
        // same way the rest of the workspace does — drive the future on the
        // ambient Tokio runtime if there is one (the listener runs inside
        // `#[tokio::main]`), else spin up a throwaway current-thread runtime.
        let out_batches = block_on(self.exchange(batches))?;

        let vertices = out_batches
            .into_iter()
            .map(|batch| VertexSet {
                label: self.label.clone(),
                id_column: self.id_column.clone(),
                batch,
            })
            .collect();
        Ok(GraphMutation {
            edge_deletes: Vec::new(),
            vertices,
            ..Default::default()
        })
    }
}

/// Map a transport/Flight error string into a [`BifrostError`], prefixed so
/// it's recognisable without touching `error.rs`.
fn flight_err(msg: String) -> BifrostError {
    BifrostError::Other(format!("flight: {msg}"))
}

/// Drive an async transform to completion from the sync [`Transform::apply`].
///
/// `apply` may be called from inside a running Tokio runtime (the listener's
/// `#[tokio::main]`) or from a plain thread. Blocking a worker thread on a
/// future would deadlock the current-thread flavor, so when a runtime handle
/// is present we hand the future to a **dedicated blocking thread** with its
/// own current-thread runtime (`block_in_place` needs the multi-thread
/// flavor, which we can't assume). With no ambient runtime we just build a
/// throwaway current-thread runtime here.
fn block_on<F>(fut: F) -> F::Output
where
    F: std::future::Future + Send,
    F::Output: Send,
{
    match tokio::runtime::Handle::try_current() {
        Ok(_) => std::thread::scope(|scope| {
            scope
                .spawn(|| {
                    tokio::runtime::Builder::new_current_thread()
                        .enable_all()
                        .build()
                        .expect("build current-thread runtime")
                        .block_on(fut)
                })
                .join()
                .expect("flight transform worker thread panicked")
        }),
        Err(_) => tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("build current-thread runtime")
            .block_on(fut),
    }
}

#[cfg(test)]
mod tests {
    // Shadow the crate's single-arg `Result<T>` alias (pulled in via `super::*`)
    // with the std two-arg `Result` the `FlightService` trait methods need.
    use std::result::Result;

    use super::{
        Channel, FlightClient, FlightDataEncoderBuilder, FlightRecordBatchStream, FlightTransform,
        RecordBatch, TableContext, Transform,
    };
    use futures::TryStreamExt;
    use std::net::SocketAddr;
    use std::sync::Arc;

    use arrow_array::{Int64Array, StringArray};
    use arrow_flight::{
        FlightData, PutResult, SchemaResult, Ticket,
        error::FlightError,
        flight_service_server::{FlightService, FlightServiceServer},
    };
    use arrow_schema::{DataType, Field, Schema};
    use futures::stream::BoxStream;
    use tonic::transport::server::TcpIncoming;
    use tonic::{Request, Response, Status, Streaming};

    /// A minimal in-process `DoExchange` server: it decodes the incoming Flight
    /// stream into record batches and re-encodes them back unchanged (an
    /// **echo** transform). This exercises the full client round-trip — encode
    /// → DoExchange → decode → VertexSet — without a separate process. A real
    /// service would filter/derive/reshape between decode and re-encode.
    #[derive(Clone)]
    struct EchoExchange;

    type ExchangeStream = BoxStream<'static, Result<FlightData, Status>>;

    #[tonic::async_trait]
    impl FlightService for EchoExchange {
        type HandshakeStream = BoxStream<'static, Result<arrow_flight::HandshakeResponse, Status>>;
        type ListFlightsStream = BoxStream<'static, Result<arrow_flight::FlightInfo, Status>>;
        type DoGetStream = ExchangeStream;
        type DoPutStream = BoxStream<'static, Result<PutResult, Status>>;
        type DoExchangeStream = ExchangeStream;
        type DoActionStream = BoxStream<'static, Result<arrow_flight::Result, Status>>;
        type ListActionsStream = BoxStream<'static, Result<arrow_flight::ActionType, Status>>;

        async fn do_exchange(
            &self,
            request: Request<Streaming<FlightData>>,
        ) -> Result<Response<Self::DoExchangeStream>, Status> {
            // Decode the client's input stream into RecordBatches …
            let inbound = request.into_inner();
            let batches = FlightRecordBatchStream::new_from_flight_data(
                inbound.map_err(|s| FlightError::Tonic(Box::new(s))),
            )
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| Status::internal(e.to_string()))?;

            // … and echo them straight back out as a Flight data stream.
            let out: ExchangeStream = Box::pin(
                FlightDataEncoderBuilder::new()
                    .build(futures::stream::iter(batches.into_iter().map(Ok)))
                    .map_err(|e: FlightError| Status::internal(e.to_string())),
            );
            Ok(Response::new(out))
        }

        // ── Unused surface — this test server only speaks DoExchange ──────────
        async fn handshake(
            &self,
            _r: Request<Streaming<arrow_flight::HandshakeRequest>>,
        ) -> Result<Response<Self::HandshakeStream>, Status> {
            Err(Status::unimplemented("handshake"))
        }
        async fn list_flights(
            &self,
            _r: Request<arrow_flight::Criteria>,
        ) -> Result<Response<Self::ListFlightsStream>, Status> {
            Err(Status::unimplemented("list_flights"))
        }
        async fn get_flight_info(
            &self,
            _r: Request<arrow_flight::FlightDescriptor>,
        ) -> Result<Response<arrow_flight::FlightInfo>, Status> {
            Err(Status::unimplemented("get_flight_info"))
        }
        async fn poll_flight_info(
            &self,
            _r: Request<arrow_flight::FlightDescriptor>,
        ) -> Result<Response<arrow_flight::PollInfo>, Status> {
            Err(Status::unimplemented("poll_flight_info"))
        }
        async fn get_schema(
            &self,
            _r: Request<arrow_flight::FlightDescriptor>,
        ) -> Result<Response<SchemaResult>, Status> {
            Err(Status::unimplemented("get_schema"))
        }
        async fn do_get(&self, _r: Request<Ticket>) -> Result<Response<Self::DoGetStream>, Status> {
            Err(Status::unimplemented("do_get"))
        }
        async fn do_put(
            &self,
            _r: Request<Streaming<FlightData>>,
        ) -> Result<Response<Self::DoPutStream>, Status> {
            Err(Status::unimplemented("do_put"))
        }
        async fn do_action(
            &self,
            _r: Request<arrow_flight::Action>,
        ) -> Result<Response<Self::DoActionStream>, Status> {
            Err(Status::unimplemented("do_action"))
        }
        async fn list_actions(
            &self,
            _r: Request<arrow_flight::Empty>,
        ) -> Result<Response<Self::ListActionsStream>, Status> {
            Err(Status::unimplemented("list_actions"))
        }
    }

    fn sample() -> RecordBatch {
        RecordBatch::try_new(
            Arc::new(Schema::new(vec![
                Field::new("_gar_id", DataType::Int64, false),
                Field::new("name", DataType::Utf8, false),
            ])),
            vec![
                Arc::new(Int64Array::from(vec![1, 2, 3])),
                Arc::new(StringArray::from(vec!["a", "b", "c"])),
            ],
        )
        .unwrap()
    }

    /// Stand up the echo `DoExchange` server on an ephemeral port (background
    /// task), point a `FlightTransform` at it, run `apply` over a sample batch,
    /// and assert the resulting `GraphMutation.vertices` carry the rows back.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn flight_exchange_round_trips_to_vertices() {
        // Bind on an ephemeral port first so the address is known before the
        // server task starts (no fixed-port race). `TcpIncoming` lets the
        // server consume the already-bound listener.
        let incoming = TcpIncoming::bind("127.0.0.1:0".parse::<SocketAddr>().unwrap())
            .expect("bind ephemeral port");
        let addr: SocketAddr = incoming.local_addr().unwrap();

        tokio::spawn(async move {
            tonic::transport::Server::builder()
                .add_service(FlightServiceServer::new(EchoExchange))
                .serve_with_incoming(incoming)
                .await
                .unwrap();
        });

        let endpoint = format!("http://{addr}");
        let transform = FlightTransform::connect(&endpoint, "Person", "_gar_id")
            .await
            .expect("connect to echo flight service");

        let batch = sample();
        // `apply` is sync and bridges to async internally; run it off the
        // async test thread so the inner runtime isn't nested on this one.
        let mutation = tokio::task::spawn_blocking(move || {
            transform.apply(&TableContext { table: "t".into() }, &[batch])
        })
        .await
        .unwrap()
        .expect("apply succeeds");

        assert_eq!(
            mutation.vertices.len(),
            1,
            "one output batch → one vertex set"
        );
        let vs = &mutation.vertices[0];
        assert_eq!(vs.label, "Person");
        assert_eq!(vs.id_column, "_gar_id");
        assert_eq!(vs.batch.num_rows(), 3);
        assert_eq!(vs.batch, sample(), "echo service round-trips the rows");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn empty_input_is_a_no_op() {
        // No server needed: an empty batch slice short-circuits before any
        // network call.
        let client = FlightClient::new(Channel::from_static("http://127.0.0.1:1").connect_lazy());
        let transform = FlightTransform::from_client(client, "X", "_gar_id");
        let mutation = tokio::task::spawn_blocking(move || {
            transform.apply(&TableContext { table: "t".into() }, &[])
        })
        .await
        .unwrap()
        .unwrap();
        assert!(mutation.is_empty());
    }
}