nordisk/edda

crates/nornir-flight/src/traits.rs

//! The schema-agnostic producer/consumer seams.
//!
//! The whole point of the crate: callers deal only in [`RecordBatch`] plus a
//! routing **key** (the "who/what" carried by the Flight descriptor / ticket).
//! Framing, backpressure and the wire live in [`crate::server`]; a use-case
//! plugs in by implementing one or both of these traits.

use arrow_array::RecordBatch;
use async_trait::async_trait;

use crate::error::Result;

/// A **producer target**: `do_put` decodes the incoming Flight stream into
/// `RecordBatch`es and hands them here under the key extracted from the Flight
/// descriptor. Schema-agnostic — the sink decides what the batches mean.
///
/// The build-info emitter (slice 2) targets a `FlightSink`; the in-memory
/// [`MemFlightStore`](crate::MemFlightStore) is the default reference impl.
#[async_trait]
pub trait FlightSink: Send + Sync + 'static {
    /// Accept all `RecordBatch`es pushed under `key`. Called once per `do_put`
    /// (after the whole client stream is decoded). Returning `Err` maps onto a
    /// `tonic::Status` back to the producer.
    async fn put(&self, key: String, batches: Vec<RecordBatch>) -> Result<()>;
}

/// A **consumer source**: `do_get`/`do_exchange` resolve a ticket key to a set
/// of `RecordBatch`es, which the seam encodes back onto the wire. Schema-
/// agnostic — the source owns storage/replay/subscription policy.
///
/// The sealer and the viz live-tail (slice 5) target a `FlightSource`.
#[async_trait]
pub trait FlightSource: Send + Sync + 'static {
    /// Produce the `RecordBatch`es for ticket `key`. An unknown key should
    /// return [`Error::Route`](crate::Error::Route) (→ `NOT_FOUND`).
    async fn get(&self, key: String) -> Result<Vec<RecordBatch>>;
}