crates/nornir-flight/src/server.rs
//! The tonic Flight service + [`FlightServer::serve`] bootstrap.
//!
//! Mirrors `knut/crates/graphar-flight/src/server.rs` as a **precedent** for the
//! tonic bootstrap, but this is a *separate, generic* Flight use: raw
//! `do_put`/`do_get`/`do_exchange` over `RecordBatch`, no Flight-SQL. Auth is
//! **not re-rolled** — the server composes `nornir_service::auth_interceptor`
//! and drives `nornir_service::serve` (the one shared tonic serve loop).
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use arrow_flight::{
Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
HandshakeRequest, HandshakeResponse, PollInfo, PutResult, Result as FlightProtoResult,
SchemaResult, Ticket, decode::FlightRecordBatchStream, encode::FlightDataEncoderBuilder,
error::FlightError, flight_service_server::{FlightService, FlightServiceServer},
};
use futures::{Stream, TryStreamExt, stream};
use tonic::{Request, Response, Status, Streaming};
use crate::error::Result;
use crate::traits::{FlightSink, FlightSource};
/// A boxed server-streaming response of `FlightData` (the `do_get`/`do_exchange`
/// output).
type FlightStream = Pin<Box<dyn Stream<Item = std::result::Result<FlightData, Status>> + Send + 'static>>;
/// The producer/consumer wiring handed to [`FlightServer::serve`]: a sink for
/// `do_put` and a source for `do_get`. `do_exchange` uses both (store-then-echo).
/// Both are `Arc<dyn …>` so one store (e.g. [`MemFlightStore`](crate::MemFlightStore))
/// can back both.
#[derive(Clone)]
pub struct FlightRoutes {
/// Where `do_put` delivers decoded batches.
pub sink: Arc<dyn FlightSink>,
/// Where `do_get` resolves a ticket to batches.
pub source: Arc<dyn FlightSource>,
}
impl FlightRoutes {
/// Route both verbs at one object implementing both seams (the common case).
pub fn shared<T: FlightSink + FlightSource>(store: Arc<T>) -> Self {
Self { sink: store.clone(), source: store }
}
/// Route the two verbs at distinct sink/source impls.
pub fn new(sink: Arc<dyn FlightSink>, source: Arc<dyn FlightSource>) -> Self {
Self { sink, source }
}
}
/// The generic Arrow Flight transport service. Schema-agnostic: it frames
/// `RecordBatch`es on/off the wire and routes them through [`FlightRoutes`].
#[derive(Clone)]
pub struct NornirFlightService {
routes: FlightRoutes,
}
impl NornirFlightService {
/// Build the service over the given routes.
pub fn new(routes: FlightRoutes) -> Self {
Self { routes }
}
}
/// The routing key carried by a Flight descriptor: its `path` (joined by `/`),
/// else its `cmd` bytes as UTF-8, else empty. This is the "who/what" a caller
/// stamps on a `do_put`/`do_exchange`.
fn descriptor_key(d: Option<&FlightDescriptor>) -> String {
match d {
Some(fd) if !fd.path.is_empty() => fd.path.join("/"),
Some(fd) if !fd.cmd.is_empty() => String::from_utf8_lossy(&fd.cmd).into_owned(),
_ => String::new(),
}
}
/// The ticket key for `do_get`: the ticket bytes as UTF-8.
fn ticket_key(t: &Ticket) -> String {
String::from_utf8_lossy(&t.ticket).into_owned()
}
/// Encode owned batches into a boxed `FlightData` response stream.
fn encode_batches(batches: Vec<arrow_array::RecordBatch>) -> FlightStream {
Box::pin(
FlightDataEncoderBuilder::new()
.build(stream::iter(batches.into_iter().map(Ok::<_, FlightError>)))
.map_err(|e: FlightError| Status::internal(e.to_string())),
)
}
/// Decode a fully-buffered `FlightData` vec into `RecordBatch`es.
async fn decode_batches(raw: Vec<FlightData>) -> std::result::Result<Vec<arrow_array::RecordBatch>, Status> {
FlightRecordBatchStream::new_from_flight_data(stream::iter(
raw.into_iter().map(Ok::<_, FlightError>),
))
.try_collect()
.await
.map_err(|e| Status::internal(format!("flight decode: {e}")))
}
#[tonic::async_trait]
impl FlightService for NornirFlightService {
type HandshakeStream =
Pin<Box<dyn Stream<Item = std::result::Result<HandshakeResponse, Status>> + Send + 'static>>;
type ListFlightsStream =
Pin<Box<dyn Stream<Item = std::result::Result<FlightInfo, Status>> + Send + 'static>>;
type DoGetStream = FlightStream;
type DoPutStream =
Pin<Box<dyn Stream<Item = std::result::Result<PutResult, Status>> + Send + 'static>>;
type DoExchangeStream = FlightStream;
type DoActionStream =
Pin<Box<dyn Stream<Item = std::result::Result<FlightProtoResult, Status>> + Send + 'static>>;
type ListActionsStream =
Pin<Box<dyn Stream<Item = std::result::Result<ActionType, Status>> + Send + 'static>>;
/// A producer pushes `RecordBatch`es under the descriptor key → the sink.
async fn do_put(
&self,
request: Request<Streaming<FlightData>>,
) -> std::result::Result<Response<Self::DoPutStream>, Status> {
let raw: Vec<FlightData> = request
.into_inner()
.try_collect()
.await
.map_err(|e| Status::internal(format!("do_put recv: {e}")))?;
let key = descriptor_key(raw.first().and_then(|d| d.flight_descriptor.as_ref()));
let batches = decode_batches(raw).await?;
let n = batches.len();
self.routes
.sink
.put(key.clone(), batches)
.await
.map_err(Status::from)?;
crate::functional_status(
"nornir-flight",
"do_put",
true,
&format!("key={key:?} batches={n}"),
);
let out = stream::iter(vec![Ok(PutResult::default())]);
Ok(Response::new(Box::pin(out)))
}
/// A consumer resolves a ticket key → the source → a `RecordBatch` stream.
async fn do_get(
&self,
request: Request<Ticket>,
) -> std::result::Result<Response<Self::DoGetStream>, Status> {
let key = ticket_key(&request.into_inner());
let batches = self
.routes
.source
.get(key.clone())
.await
.map_err(Status::from)?;
let n = batches.len();
crate::functional_status(
"nornir-flight",
"do_get",
true,
&format!("key={key:?} batches={n}"),
);
Ok(Response::new(encode_batches(batches)))
}
/// Bidirectional: decode the incoming batches, deliver them to the sink,
/// then echo them back byte-identically (the round-trip / offload seam).
async fn do_exchange(
&self,
request: Request<Streaming<FlightData>>,
) -> std::result::Result<Response<Self::DoExchangeStream>, Status> {
let raw: Vec<FlightData> = request
.into_inner()
.try_collect()
.await
.map_err(|e| Status::internal(format!("do_exchange recv: {e}")))?;
let key = descriptor_key(raw.first().and_then(|d| d.flight_descriptor.as_ref()));
let batches = decode_batches(raw).await?;
let n = batches.len();
self.routes
.sink
.put(key.clone(), batches.clone())
.await
.map_err(Status::from)?;
crate::functional_status(
"nornir-flight",
"do_exchange",
true,
&format!("key={key:?} batches={n}"),
);
Ok(Response::new(encode_batches(batches)))
}
// --- unimplemented verbs: this is a raw record-batch transport, not a
// discovery/action service. Each returns UNIMPLEMENTED. ---
async fn handshake(
&self,
_request: Request<Streaming<HandshakeRequest>>,
) -> std::result::Result<Response<Self::HandshakeStream>, Status> {
// Auth rides the `authorization` header via the shared interceptor, so
// no Flight handshake round-trip is needed.
Err(Status::unimplemented("handshake: use the authorization header"))
}
async fn list_flights(
&self,
_request: Request<Criteria>,
) -> std::result::Result<Response<Self::ListFlightsStream>, Status> {
Err(Status::unimplemented("list_flights"))
}
async fn get_flight_info(
&self,
_request: Request<FlightDescriptor>,
) -> std::result::Result<Response<FlightInfo>, Status> {
Err(Status::unimplemented("get_flight_info"))
}
async fn poll_flight_info(
&self,
_request: Request<FlightDescriptor>,
) -> std::result::Result<Response<PollInfo>, Status> {
Err(Status::unimplemented("poll_flight_info"))
}
async fn get_schema(
&self,
_request: Request<FlightDescriptor>,
) -> std::result::Result<Response<SchemaResult>, Status> {
Err(Status::unimplemented("get_schema"))
}
async fn do_action(
&self,
_request: Request<Action>,
) -> std::result::Result<Response<Self::DoActionStream>, Status> {
Err(Status::unimplemented("do_action"))
}
async fn list_actions(
&self,
_request: Request<Empty>,
) -> std::result::Result<Response<Self::ListActionsStream>, Status> {
Err(Status::unimplemented("list_actions"))
}
}
/// The Flight transport bootstrap. Composes the shared `nornir-service` auth
/// interceptor and serve loop — the transport does NOT re-roll auth (LAW #5).
pub struct FlightServer;
impl FlightServer {
/// Serve the Flight transport over `routes` at `addr`, **open** (no auth).
/// Equivalent to [`serve_with_auth`](FlightServer::serve_with_auth) with an
/// empty token.
pub async fn serve(routes: FlightRoutes, addr: SocketAddr) -> Result<()> {
Self::serve_with_auth(routes, addr, Arc::new(String::new())).await
}
/// Serve with the shared bearer-auth interceptor.
///
/// * `token` empty ⇒ **open** (dev / single-box); the interceptor passes
/// every request through.
/// * `token` non-empty ⇒ every request MUST carry
/// `authorization: Bearer <token>` (constant-time compared by
/// `nornir_service::token_matches`), else `UNAUTHENTICATED`.
///
/// The interceptor is always installed; on an empty token it is a no-op, so
/// "empty = open" holds without a branch.
pub async fn serve_with_auth(
routes: FlightRoutes,
addr: SocketAddr,
token: Arc<String>,
) -> Result<()> {
let service = NornirFlightService::new(routes);
let interceptor = nornir_service::auth_interceptor(token);
let flight = FlightServiceServer::with_interceptor(service, interceptor);
let tonic_routes = tonic::service::Routes::new(flight);
crate::functional_status("nornir-flight", "serve", true, &format!("addr={addr}"));
nornir_service::serve(tonic_routes, addr).await?;
Ok(())
}
/// Serve the Flight transport over TLS, **open** (no bearer auth) — the
/// TLS-terminating sibling of [`serve`](FlightServer::serve). Gated behind
/// the `tls` cargo feature (which pulls tonic's TLS stack + turns on
/// `nornir-service`'s own `tls` feature); the plaintext consumers never
/// compile it in.
///
/// The transport does NOT re-roll the TLS serve tail (LAW #5): it hands the
/// caller's [`ServerTlsConfig`](tonic::transport::ServerTlsConfig) straight to
/// the shared `nornir_service::serve_with_tls` — the same one holger-server
/// consumes — so there is one TLS-aware serve loop across the constellation.
#[cfg(feature = "tls")]
pub async fn serve_with_tls(
routes: FlightRoutes,
addr: SocketAddr,
tls: tonic::transport::ServerTlsConfig,
) -> Result<()> {
Self::serve_with_tls_and_auth(routes, addr, Arc::new(String::new()), tls).await
}
/// Serve over TLS **and** the shared bearer-auth interceptor — the TLS
/// sibling of [`serve_with_auth`](FlightServer::serve_with_auth). Token
/// semantics are identical (empty ⇒ open, non-empty ⇒ every request must
/// carry `authorization: Bearer <token>`); TLS is layered *under* that auth
/// by the shared `nornir_service::serve_with_tls`.
#[cfg(feature = "tls")]
pub async fn serve_with_tls_and_auth(
routes: FlightRoutes,
addr: SocketAddr,
token: Arc<String>,
tls: tonic::transport::ServerTlsConfig,
) -> Result<()> {
let service = NornirFlightService::new(routes);
let interceptor = nornir_service::auth_interceptor(token);
let flight = FlightServiceServer::with_interceptor(service, interceptor);
let tonic_routes = tonic::service::Routes::new(flight);
crate::functional_status("nornir-flight", "serve_with_tls", true, &format!("addr={addr}"));
nornir_service::serve_with_tls(tonic_routes, addr, Some(tls)).await?;
Ok(())
}
}