nordisk/edda

crates/nornir-flight/src/client.rs

//! A thin client over `arrow_flight::FlightClient` for the three verbs.
//!
//! Consumers (a remote build-info producer, the viz live-tail) use this to talk
//! to a [`FlightServer`](crate::FlightServer) without re-deriving the framing
//! and auth-header plumbing. Keys map onto the descriptor path (`put`/`exchange`)
//! and the ticket bytes (`get`), matching the server's key extraction.

use arrow_array::RecordBatch;
use arrow_flight::{
    FlightClient, FlightDescriptor, Ticket, encode::FlightDataEncoderBuilder, error::FlightError,
};
use futures::{TryStreamExt, stream};
use tonic::transport::Channel;

use crate::error::{Error, Result};

/// A connected Flight client bound to one endpoint, optionally carrying a bearer
/// token on every call (to satisfy the shared auth interceptor).
pub struct FlightConn {
    client: FlightClient,
}

impl FlightConn {
    /// Connect to `endpoint` (e.g. `http://127.0.0.1:50051`).
    pub async fn connect(endpoint: impl Into<String>) -> Result<Self> {
        let channel = Channel::from_shared(endpoint.into())
            .map_err(|e| Error::Route(format!("bad endpoint: {e}")))?
            .connect()
            .await?;
        Ok(Self::from_channel(channel))
    }

    /// Connect to a TLS endpoint (e.g. `https://host:50051`) trusting the given
    /// [`ClientTlsConfig`](tonic::transport::ClientTlsConfig) — the client sibling
    /// of [`FlightServer::serve_with_tls`](crate::FlightServer::serve_with_tls).
    /// Gated behind the `tls` cargo feature so plaintext consumers never pull in
    /// tonic's TLS stack.
    #[cfg(feature = "tls")]
    pub async fn connect_tls(
        endpoint: impl Into<String>,
        tls: tonic::transport::ClientTlsConfig,
    ) -> Result<Self> {
        let channel = Channel::from_shared(endpoint.into())
            .map_err(|e| Error::Route(format!("bad endpoint: {e}")))?
            .tls_config(tls)
            .map_err(|e| Error::Route(format!("bad client tls: {e}")))?
            .connect()
            .await?;
        Ok(Self::from_channel(channel))
    }

    /// Wrap an already-built tonic [`Channel`].
    pub fn from_channel(channel: Channel) -> Self {
        Self {
            client: FlightClient::new(channel),
        }
    }

    /// Attach `authorization: Bearer <token>` to every subsequent call. Chainable.
    pub fn with_bearer(mut self, token: &str) -> Result<Self> {
        self.client
            .add_header("authorization", &format!("Bearer {token}"))
            .map_err(|e| Error::Route(format!("bad token header: {e}")))?;
        Ok(self)
    }

    /// `do_put`: push `batches` under routing `key` (the descriptor path).
    pub async fn put(&mut self, key: &str, batches: Vec<RecordBatch>) -> Result<()> {
        let descriptor = FlightDescriptor::new_path(vec![key.to_string()]);
        let input = FlightDataEncoderBuilder::new()
            .with_flight_descriptor(Some(descriptor))
            .build(stream::iter(batches.into_iter().map(Ok::<_, FlightError>)));
        // Drain the PutResult acks so the whole stream is flushed before return.
        let _acks: Vec<_> = self.client.do_put(input).await?.try_collect().await?;
        Ok(())
    }

    /// `do_get`: pull the `RecordBatch`es for ticket `key`.
    pub async fn get(&mut self, key: &str) -> Result<Vec<RecordBatch>> {
        let ticket = Ticket::new(key.as_bytes().to_vec());
        let batches = self.client.do_get(ticket).await?.try_collect().await?;
        Ok(batches)
    }

    /// `do_exchange`: send `batches` under `key` and receive the echoed batches.
    pub async fn exchange(&mut self, key: &str, batches: Vec<RecordBatch>) -> Result<Vec<RecordBatch>> {
        let descriptor = FlightDescriptor::new_path(vec![key.to_string()]);
        let input = FlightDataEncoderBuilder::new()
            .with_flight_descriptor(Some(descriptor))
            .build(stream::iter(batches.into_iter().map(Ok::<_, FlightError>)));
        let out = self.client.do_exchange(input).await?.try_collect().await?;
        Ok(out)
    }
}