nordisk/knut

crates/knut-flight/src/bin/knut-flight-server.rs

//! **knut-flight-server (N4)** — a plain Apache Arrow Flight server over knut's
//! rel2graph output: nodes, edges, and the source relational tables of the
//! 12-table ERP bench, streamed as Arrow RecordBatches.
//!
//! Config is via flags/env only — no interactive setup — so this drops straight
//! into a container (the unit N5 wraps as a Skidbladnir/k8s service):
//!
//! | flag | env | default |
//! |---|---|---|
//! | `--addr <host:port>` | `KNUT_FLIGHT_ADDR` | `0.0.0.0:50055` |
//! | `--scale <n>` | `KNUT_FLIGHT_SCALE` | `1` |
//! | `--parquet <dir>` | `KNUT_FLIGHT_PARQUET` | *(unset → bench data)* |
//! | `--mapping <file>` | `KNUT_FLIGHT_MAPPING` | *(unset → bench 12-table spec)* |
//!
//! With `--parquet <dir>` the server serves a REAL graph read from an Iceberg /
//! object-store parquet export (the live-data seam) instead of the synthetic
//! bench generator; `--mapping` supplies the rel→graph YAML for that data (it
//! defaults to the bench's 12-table spec, so a `rel2graph-export --parquet` dump
//! serves out of the box). Without `--parquet`, `--scale` selects the bench size.
//!
//! `--version` / `-V` prints `knut-flight-server v<version> (<git-hash>)` and exits.
//!
//! ```text
//! cargo run -p knut-flight --bin knut-flight-server -- --addr 0.0.0.0:50055 --scale 1
//! cargo run -p knut-flight --bin knut-flight-server -- --parquet ./graph-export
//! ```

use std::net::SocketAddr;
use std::path::PathBuf;

use arrow_flight::flight_service_server::FlightServiceServer;
use knut_flight::{Catalog, DEFAULT_PORT, GraphFlightService, VERSION_LINE};
use tonic::transport::Server;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut addr: Option<String> = std::env::var("KNUT_FLIGHT_ADDR").ok();
    let mut scale: usize = std::env::var("KNUT_FLIGHT_SCALE")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(1);
    let mut parquet: Option<PathBuf> = std::env::var("KNUT_FLIGHT_PARQUET").ok().map(PathBuf::from);
    let mut mapping: Option<PathBuf> = std::env::var("KNUT_FLIGHT_MAPPING").ok().map(PathBuf::from);

    let mut args = std::env::args().skip(1);
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--version" | "-V" => {
                println!("knut-flight-server {VERSION_LINE}");
                return Ok(());
            }
            "--help" | "-h" => {
                println!(
                    "knut-flight-server {VERSION_LINE}\n\n\
                     Arrow Flight server for knut's rel2graph output.\n\n\
                     USAGE:\n  knut-flight-server [--addr <host:port>] [--scale <n>] [--parquet <dir>] [--mapping <file>]\n\n\
                     OPTIONS:\n\
                     \x20 --addr <host:port>  bind address (env KNUT_FLIGHT_ADDR, default 0.0.0.0:{DEFAULT_PORT})\n\
                     \x20 --scale <n>         bench scale factor (env KNUT_FLIGHT_SCALE, default 1)\n\
                     \x20 --parquet <dir>     serve a REAL graph from a parquet export dir (env KNUT_FLIGHT_PARQUET)\n\
                     \x20 --mapping <file>    rel→graph YAML for --parquet data (env KNUT_FLIGHT_MAPPING, default: bench spec)\n\
                     \x20 -V, --version       print version and exit\n\
                     \x20 -h, --help          print this help and exit"
                );
                return Ok(());
            }
            "--addr" => addr = Some(args.next().ok_or("--addr needs a value")?),
            "--scale" => {
                scale = args
                    .next()
                    .ok_or("--scale needs a value")?
                    .parse()
                    .map_err(|_| "--scale must be a positive integer")?;
            }
            "--parquet" => {
                parquet = Some(PathBuf::from(args.next().ok_or("--parquet needs a value")?))
            }
            "--mapping" => {
                mapping = Some(PathBuf::from(args.next().ok_or("--mapping needs a value")?))
            }
            other => return Err(format!("unknown argument: {other}").into()),
        }
    }

    let addr: SocketAddr = addr
        .unwrap_or_else(|| format!("0.0.0.0:{DEFAULT_PORT}"))
        .parse()?;

    eprintln!("knut-flight-server {VERSION_LINE}");
    let catalog = match &parquet {
        Some(dir) => {
            // Live-data path: read a real parquet export and run rel2graph over it.
            let spec = match &mapping {
                Some(path) => {
                    let yaml = std::fs::read_to_string(path)?;
                    knut_bifrost::MappingSpec::from_yaml(&yaml)?
                }
                None => knut_rel2graph_bench::mapping_spec(),
            };
            eprintln!("building rel2graph catalog from parquet export at {dir:?}…");
            Catalog::from_parquet_dir(&spec, dir)?
        }
        None => {
            eprintln!("building rel2graph catalog (synthetic 12-table ERP, scale {scale}x)…");
            Catalog::from_bench(scale)
        }
    };
    eprintln!("catalog ready: {} datasets", catalog.len());
    for ds in catalog.datasets() {
        eprintln!("  {:<40} {}", ds.name, ds.description);
    }
    eprintln!("listening on {addr} (Arrow Flight)");

    let service = GraphFlightService::new(catalog);
    Server::builder()
        .add_service(FlightServiceServer::new(service))
        .serve(addr)
        .await?;

    Ok(())
}