crates/knut-bifrost/src/bin/bifrost.rs
//! `bifrost` — drive and observe the file → graph lifecycle headlessly.
//!
//! This is the CLI front door to the [`Listener`](knut_bifrost::Listener): it
//! watches a directory of file drops, runs them through a mapping into a graph
//! sink, and emits **structured tracing events** at every stage — `file_detected`,
//! `files_read`, `table_synced`, `sink_applied`, `round_complete`. With
//! `--json-logs` those become one JSON object per line, so a test or bencher
//! can drive this binary and assert on the lifecycle with no human in the loop.
//!
//! `--dry-run` swaps the FalkorDB sink for a counting null sink, so the
//! detect→transform half of the pipeline can be exercised with no database.
//!
//! ```text
//! bifrost run --mapping m.yml --files ./drops --graph demo \
//! --falkor redis://127.0.0.1:6379 --once --json-logs
//! bifrost run --mapping m.yml --files ./drops --dry-run --once # no DB
//! ```
use std::path::PathBuf;
use std::process::ExitCode;
use std::time::Duration;
use async_trait::async_trait;
use clap::{Parser, Subcommand};
use knut_bifrost::watch_cli::{LakeArgs, MappingArgs, TlsArgs, WatchArgs};
use knut_bifrost::{
Backoff, GraphMutation, Listener, Result as BifrostResult, SinkReport, SyncState,
TransformRegistry,
sink::{FalkorDbSink, GraphSink},
source::{FileSource, TableSource},
};
#[derive(Parser)]
#[command(
name = "bifrost",
about = "Drive the file → graph lifecycle with structured logs",
// `v<pkg> (<git-sha>)`, stamped pure-fs by build.rs on every build → clap
// prints `bifrost v<ver> (<hash>)` for `--version`.
version = knut_bifrost::VERSION_LINE
)]
struct Cli {
/// Emit one JSON object per log line (machine-readable) instead of pretty text.
#[arg(long, global = true)]
json_logs: bool,
#[command(subcommand)]
command: Command,
}
#[derive(Subcommand)]
// `Run` carries the whole CLI surface (mapping/files/falkor/backoff/metrics/…)
// and is constructed exactly once at startup, so the size gap to `GenSample`
// is intentional — boxing a parsed-once clap variant buys nothing.
#[allow(clippy::large_enum_variant)]
enum Command {
/// Watch a directory of file drops and materialise them into a graph.
Run {
/// Table names map to `<files>/<table>/`.
#[command(flatten)]
mapping: MappingArgs,
/// Root directory of file drops to watch.
#[arg(long)]
files: PathBuf,
/// FalkorDB / Redis URL (ignored with --dry-run).
#[arg(long, default_value = "redis://127.0.0.1:6379")]
falkor: String,
/// Mutual TLS to a FalkorDB that takes nothing else (`--falkor rediss://…`).
#[command(flatten)]
tls: TlsArgs,
/// Target graph name.
#[arg(long, default_value = "bifrost")]
graph: String,
/// Sync-state ledger path (default: <files>/.bifrost-state.json).
#[arg(long)]
state: Option<PathBuf>,
/// Use a **Valkey**-backed shared ledger at this URL instead of a file
/// (so multiple instances share progress). Valkey 9+ / Redis-API.
#[arg(long, value_name = "URL")]
state_valkey: Option<String>,
/// Valkey hash key for the ledger (with --state-valkey).
#[arg(long, default_value = "bifrost:state")]
state_key: String,
/// Run exactly one sync round and exit (else poll forever).
#[arg(long)]
once: bool,
/// Poll interval in seconds when looping.
#[arg(long, default_value_t = 5)]
interval: u64,
/// Cap (seconds) for the failure backoff. A fully-failed round doubles
/// the inter-round sleep up to this, resetting on the first round that
/// makes progress. Defaults to 10× `--interval`.
#[arg(long)]
max_backoff: Option<u64>,
/// On shutdown, write the cumulative run metrics as Prometheus text to
/// this file (a `/metrics` endpoint can serve the same later).
#[arg(long, value_name = "PATH")]
metrics_file: Option<PathBuf>,
/// Serve the cumulative run metrics live over HTTP at `/metrics` on this
/// `ip:port` (e.g. `127.0.0.1:9090`) so a Prometheus scraper can pull
/// them while bifröst runs — published after every round. Absent → no
/// server (today's behavior). Requires building with the `metrics-http`
/// feature.
#[arg(long, value_name = "ADDR")]
metrics_addr: Option<String>,
/// Skip the database: count the graph mutation instead of writing it.
#[arg(long)]
dry_run: bool,
/// Airflow API base URL — after each round, emit an AIP-82 asset event
/// per ingested table to `{url}/api/v2/assets/events`.
#[arg(long, value_name = "URL")]
emit_airflow: Option<String>,
/// Asset URI prefix for emitted events (event uri = `{prefix}{table}`).
#[arg(long, default_value = "graph://")]
asset_prefix: String,
/// Bearer token for the Airflow API, if required.
#[arg(long, value_name = "TOKEN")]
airflow_token: Option<String>,
},
/// Write a sample parquet drop (`id`, `name`, `score`) under
/// `<files>/<table>/<chunk>.parquet` — handy for trying `run` with no
/// other tooling.
GenSample {
/// Root drops directory.
#[arg(long)]
files: PathBuf,
/// Table (subdirectory) name.
#[arg(long, default_value = "people")]
table: String,
/// Number of rows.
#[arg(long, default_value_t = 1000)]
rows: usize,
/// Chunk file name.
#[arg(long, default_value = "chunk0.parquet")]
chunk: String,
},
/// Run the **knut ▸ Spark 4.1 + Iceberg 1.11.0 LOAD demo** against a live
/// Spark Connect endpoint: drop → create namespace → create table → insert,
/// then verify the landed row count AND that an Iceberg snapshot committed.
/// The target is `lake.knut_demo.people` (see `knut_bifrost::spark_demo`).
///
/// Needs `--features spark41`. Built WITHOUT it the subcommand still exists
/// and exits non-zero naming the missing feature — a load that cannot run
/// must say so, not vanish from `--help`.
IcebergDemo {
/// Spark Connect endpoint of the 4.1 + Iceberg server.
#[arg(long, default_value = "sc://localhost:15012")]
url: String,
/// Iceberg catalog (the blessed image bakes `lake`).
#[arg(long, default_value = knut_bifrost::spark_demo::DEMO_CATALOG)]
catalog: String,
/// Namespace under the catalog.
#[arg(long, default_value = knut_bifrost::spark_demo::DEMO_NAMESPACE)]
namespace: String,
/// Table under the namespace.
#[arg(long, default_value = knut_bifrost::spark_demo::DEMO_TABLE)]
table: String,
/// Rows the INSERT writes.
#[arg(long, default_value_t = knut_bifrost::spark_demo::DEMO_ROWS)]
rows: usize,
/// Print the SQL program and exit WITHOUT touching the network — the
/// container-free way to see exactly what the demo would run.
#[arg(long)]
plan_only: bool,
},
/// **Sip an Iceberg warehouse into a graph, detached** — the streaming half
/// of `run`, sourced from a **skade Iceberg warehouse** instead of file
/// drops. This is the "submit a Rust job that sits on Iceberg CDC" front
/// door (§g5): point it at a warehouse dir, hand it the same mapping, and it
/// polls snapshots forever, MERGE-ing only the rows each new snapshot added
/// (`mode: incremental`/`cdc` in the mapping) into FalkorDB — JSON logs,
/// graceful shutdown, a durable per-table snapshot cursor, and the same
/// `/metrics` endpoint as `run`. Pure Rust, no Spark, no JVM.
///
/// Needs `--features skade`. Built WITHOUT it the subcommand still exists and
/// exits non-zero naming the missing feature — a source that cannot run must
/// say so, not vanish from `--help` (mirrors `iceberg-demo`).
///
/// The lake is one of four (see `LakeArgs`): a skade warehouse directory,
/// Iceberg tables on S3/MinIO read with no catalog (`--s3-root`, feature
/// `skade-s3`), an Iceberg REST catalog (`--rest-uri`, feature `skade-rest`;
/// `--catalog aistor|polaris|unity|gravitino` with feature `catalog-<vendor>`),
/// or a remote skade-katalog (`--grpc-uri`, feature `skade-grpc`). S3 keys
/// come from `MINIO_ACCESS_KEY` / `MINIO_SECRET_KEY`, catalog secrets from
/// `BIFROST_CATALOG_CREDENTIAL` / `BIFROST_CATALOG_TOKEN`, never a flag.
/// `bifrost falkordb install|k8s` runs this as a service.
///
/// ```text
/// bifrost watch --warehouse /var/lib/lake --mapping m.yml \
/// --graph social --falkor redis://127.0.0.1:6379 --json-logs
/// bifrost watch --warehouse /var/lib/lake --mapping m.yml --once --dry-run
/// bifrost watch --rest-uri https://minio:9000/_iceberg --rest-warehouse analytics \
/// --sigv4 s3tables --s3-endpoint https://minio:9000 --mapping m.yml --state st.json
/// bifrost watch --s3-root s3://lake/warehouse --s3-endpoint https://minio:9000 \
/// --mapping m.yml --state st.json
/// BIFROST_CATALOG_CREDENTIAL=id:secret bifrost watch --catalog polaris \
/// --rest-uri https://polaris:8181 --rest-warehouse lake --mapping m.yml
/// bifrost watch --grpc-uri https://katalog:8182 --grpc-ca ca.pem --mapping m.yml
/// ```
Watch(WatchArgs),
/// **An Iceberg table's commits into Apache Doris** — native Doris tables
/// (indexes, the BE's own storage), fed one commit per Stream Load window.
/// The table is found in a skade warehouse, on S3 with no catalog
/// (`--s3-table`), or behind a REST catalog (`--rest-uri`, AIStor with
/// `--sigv4 s3tables`). `install`/`k8s` package the same run as a systemd
/// unit or a Kubernetes StatefulSet through skidbladnir.
///
/// Needs `--features service`. Built without it the subcommand still exists
/// and exits non-zero naming the missing feature.
#[cfg(feature = "service")]
Doris {
#[command(subcommand)]
cmd: knut_bifrost::doris_cli::DorisCommand,
},
/// Iceberg → Doris delivery (needs `--features service`; this build lacks it).
#[cfg(not(feature = "service"))]
Doris {
#[arg(trailing_var_arg = true, allow_hyphen_values = true, hide = true)]
rest: Vec<String>,
},
/// **`watch` as a service** — the Iceberg lake → FalkorDB stream installed
/// through skidbladnir: `install`/`uninstall`/`unit`/`start`/`stop`/`status`
/// (systemd) and `k8s` (StatefulSet + Secret + volume), all from the same
/// flags `watch` takes. Secrets come from the environment only.
///
/// Needs `--features falkordb-service` (or `service`). Built without it the
/// subcommand still exists and exits non-zero naming the missing feature.
#[cfg(feature = "falkordb-service")]
Falkordb {
#[command(subcommand)]
cmd: knut_bifrost::falkordb_cli::FalkordbCommand,
},
/// `watch` as a systemd / Kubernetes service (needs `--features falkordb-service`; this build lacks it).
#[cfg(not(feature = "falkordb-service"))]
Falkordb {
#[arg(trailing_var_arg = true, allow_hyphen_values = true, hide = true)]
rest: Vec<String>,
},
/// Headless self-test: run the core transform→mutation path over a known
/// in-memory table and report each invariant. Built `--features testmatrix`
/// it also emits `functional` rows into the nornir test matrix. Exits
/// non-zero if any check fails (so CI / the matrix sees a broken core).
SelfTest,
}
fn main() -> ExitCode {
let cli = Cli::parse();
// Shared knut bootstrap (one copy for bifröst + knutpipe; see `knut-tracing`).
knut_tracing::init_tracing(cli.json_logs);
match run(cli.command) {
Ok(()) => ExitCode::SUCCESS,
Err(e) => {
tracing::error!(target: "bifrost", error = %e, "fatal");
ExitCode::FAILURE
}
}
}
fn run(cmd: Command) -> Result<(), Box<dyn std::error::Error>> {
match cmd {
#[cfg(feature = "service")]
Command::Doris { cmd } => {
let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build()?;
Ok(rt.block_on(knut_bifrost::doris_cli::dispatch(cmd))?)
}
#[cfg(not(feature = "service"))]
Command::Doris { .. } => Err("`bifrost doris` needs a build with `--features service`".into()),
#[cfg(feature = "falkordb-service")]
Command::Falkordb { cmd } => Ok(knut_bifrost::falkordb_cli::dispatch(cmd)?),
#[cfg(not(feature = "falkordb-service"))]
Command::Falkordb { .. } => Err("`bifrost falkordb` needs a build with `--features falkordb-service`".into()),
Command::SelfTest => {
let ok = knut_bifrost::testmatrix::self_test();
tracing::info!(target: "bifrost", ok, "self-test complete");
if ok {
Ok(())
} else {
Err("bifröst self-test failed (a core invariant did not hold)".into())
}
}
Command::GenSample {
files,
table,
rows,
chunk,
} => gen_sample(files, table, rows, chunk),
Command::IcebergDemo {
url,
catalog,
namespace,
table,
rows,
plan_only,
} => iceberg_demo(url, catalog, namespace, table, rows, plan_only),
Command::Watch(w) => {
let falkor = w.falkor_url();
let tls = w.tls.identity()?;
watch_pipeline(
w.lake,
w.mapping,
falkor,
w.graph,
w.state,
w.once,
w.interval,
w.max_backoff,
w.metrics_file,
w.metrics_addr,
w.dry_run,
tls,
)
}
Command::Run {
mapping,
files,
falkor,
tls,
graph,
state,
state_valkey,
state_key,
once,
interval,
max_backoff,
metrics_file,
metrics_addr,
dry_run,
emit_airflow,
asset_prefix,
airflow_token,
} => run_pipeline(
mapping,
files,
falkor,
graph,
state,
state_valkey,
state_key,
once,
interval,
max_backoff,
metrics_file,
metrics_addr,
dry_run,
emit_airflow,
asset_prefix,
airflow_token,
tls.identity()?,
),
}
}
/// The `iceberg-demo` subcommand: build the plan (pure), optionally print it, else
/// EXECUTE it through `knut_bifrost::spark_demo::run_iceberg_demo`.
///
/// One call site for both builds. With `spark41` this is the real load; without it
/// `run_iceberg_demo` is the fail-loud half and this exits non-zero with the feature
/// name — the CLI never has to know which half it linked.
fn iceberg_demo(
url: String,
catalog: String,
namespace: String,
table: String,
rows: usize,
plan_only: bool,
) -> Result<(), Box<dyn std::error::Error>> {
use knut_bifrost::spark_demo;
let plan = spark_demo::iceberg_demo_plan_for(&url, &catalog, &namespace, &table, rows);
tracing::info!(
target: "bifrost",
event = "iceberg_demo_plan",
url = %plan.url,
table = %plan.table,
rows = plan.rows,
spark_client_linked = spark_demo::spark_client_linked(),
"knut -> Spark 4.1 + Iceberg demo plan",
);
for (i, stmt) in plan.program().iter().enumerate() {
tracing::info!(target: "bifrost", event = "iceberg_demo_sql", step = i, sql = %stmt);
}
if plan_only {
return Ok(());
}
let landed = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?
.block_on(spark_demo::run_iceberg_demo(&plan))?;
tracing::info!(
target: "bifrost",
event = "iceberg_demo_complete",
table = %plan.table,
rows = landed,
"knut landed {landed} rows in {} (Iceberg 1.11.0 on Spark 4.1)",
plan.table,
);
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn run_pipeline(
mapping: MappingArgs,
files: PathBuf,
falkor: String,
graph: String,
state: Option<PathBuf>,
state_valkey: Option<String>,
state_key: String,
once: bool,
interval: u64,
max_backoff: Option<u64>,
metrics_file: Option<PathBuf>,
metrics_addr: Option<String>,
dry_run: bool,
emit_airflow: Option<String>,
asset_prefix: String,
airflow_token: Option<String>,
tls: Option<graphar_falkordb::FalkorTls>,
) -> Result<(), Box<dyn std::error::Error>> {
let spec = mapping.load()?;
let state_path = state.unwrap_or_else(|| files.join(".bifrost-state.json"));
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?;
rt.block_on(async move {
let source = FileSource::new(&files);
// Valkey-backed shared ledger, or the default atomic file ledger.
let sync_state = match &state_valkey {
Some(url) => SyncState::open_valkey(url, &state_key).await?,
None => SyncState::open(&state_path)?,
};
// AIP-82 emitter (closes the loop: ingested → asset event → DAG).
let emitter: Option<(std::sync::Arc<dyn knut_bifrost::AssetEmitter>, String)> =
emit_airflow.map(|url| {
let e = knut_bifrost::airflow::AirflowEmitter::new(url, airflow_token);
(
std::sync::Arc::new(e) as std::sync::Arc<dyn knut_bifrost::AssetEmitter>,
asset_prefix,
)
});
tracing::info!(
target: "bifrost",
event = "start",
files = %files.display(),
graph = %graph,
dry_run,
once,
emit_airflow = emitter.is_some(),
metrics_http = metrics_addr.is_some(),
state = if state_valkey.is_some() { "valkey" } else { "file" },
"bifrost starting",
);
// Start the live `/metrics` HTTP server if `--metrics-addr` was given.
// Held for the whole run so the endpoint stays up; the poll loop
// publishes the cumulative metrics into its shared snapshot each round.
// The `MetricsServer` is dropped (stopped + joined) when this scope ends.
let metrics_server = start_metrics_server(metrics_addr.as_deref())?;
if dry_run {
let sink = NullSink;
let mut l = Listener::new(spec, TransformRegistry::new(), source, sink, sync_state);
if let Some((e, prefix)) = emitter {
l = l.with_asset_emitter(e, prefix);
}
drive(
l,
once,
interval,
max_backoff,
metrics_file,
&metrics_server,
)
.await
} else {
let sink = FalkorDbSink::connect_with_identity(&falkor, graph.clone(), tls.clone()).await?;
let mut l = Listener::new(spec, TransformRegistry::new(), source, sink, sync_state);
if let Some((e, prefix)) = emitter {
l = l.with_asset_emitter(e, prefix);
}
drive(
l,
once,
interval,
max_backoff,
metrics_file,
&metrics_server,
)
.await
}
})
}
/// The `watch` subcommand: stream a **skade Iceberg warehouse** into a graph,
/// reusing the exact `run` loop ([`drive`]) — only the source differs
/// ([`IcebergSource`](knut_bifrost::source::IcebergSource) instead of
/// [`FileSource`]). One call site for both builds: with `skade` this is the real
/// streamer; without it, the fail-loud stub below.
#[cfg(feature = "skade")]
#[allow(clippy::too_many_arguments)]
fn watch_pipeline(
lake: LakeArgs,
mapping: MappingArgs,
falkor: String,
graph: String,
state: Option<PathBuf>,
once: bool,
interval: u64,
max_backoff: Option<u64>,
metrics_file: Option<PathBuf>,
metrics_addr: Option<String>,
dry_run: bool,
tls: Option<graphar_falkordb::FalkorTls>,
) -> Result<(), Box<dyn std::error::Error>> {
use knut_bifrost::source::IcebergSource;
let spec = mapping.load()?;
let state_path = state.unwrap_or_else(|| lake.default_state());
let state_dir = state_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(std::path::Path::new("."));
let lake = lake.lake(state_dir)?;
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?;
rt.block_on(async move {
let lake_desc = lake.describe();
let source = IcebergSource::from_lake(lake);
let sync_state = SyncState::open(&state_path)?;
tracing::info!(
target: "bifrost",
event = "start",
lake = %lake_desc,
graph = %graph,
dry_run,
once,
source = "iceberg",
metrics_http = metrics_addr.is_some(),
"bifrost watch starting (Iceberg lake → graph)",
);
// Same live `/metrics` endpoint as `run`; held for the whole run.
let metrics_server = start_metrics_server(metrics_addr.as_deref())?;
if dry_run {
let sink = NullSink;
let l = Listener::new(spec, TransformRegistry::new(), source, sink, sync_state);
drive(
l,
once,
interval,
max_backoff,
metrics_file,
&metrics_server,
)
.await
} else {
let sink = FalkorDbSink::connect_with_identity(&falkor, graph.clone(), tls.clone()).await?;
let l = Listener::new(spec, TransformRegistry::new(), source, sink, sync_state);
drive(
l,
once,
interval,
max_backoff,
metrics_file,
&metrics_server,
)
.await
}
})
}
/// Without `skade` there is no `IcebergSource` to construct — the subcommand
/// still parses (so it never vanishes from `--help`) but fails loud naming the
/// feature, mirroring `iceberg-demo`'s feature-off half.
#[cfg(not(feature = "skade"))]
#[allow(clippy::too_many_arguments)]
fn watch_pipeline(
_lake: LakeArgs,
_mapping: MappingArgs,
_falkor: String,
_graph: String,
_state: Option<PathBuf>,
_once: bool,
_interval: u64,
_max_backoff: Option<u64>,
_metrics_file: Option<PathBuf>,
_metrics_addr: Option<String>,
_dry_run: bool,
_tls: Option<graphar_falkordb::FalkorTls>,
) -> Result<(), Box<dyn std::error::Error>> {
Err(
"bifrost watch needs the `skade` feature (the Iceberg warehouse source); \
this build linked no IcebergSource, so it cannot stream a warehouse — \
rebuild with --features skade"
.into(),
)
}
// ── Live `/metrics` HTTP endpoint (feature `metrics-http`) ───────────────────
//
// The server start + per-round publish are abstracted behind `OptMetricsServer`
// and two helpers so the `drive` loop is identical whether or not the feature is
// compiled in. With the feature on, `OptMetricsServer` is an
// `Option<MetricsServer>` and `publish_metrics` forwards the cumulative metrics
// into the shared snapshot each round; without it, it's a zero-sized unit and
// both helpers are no-ops (and `--metrics-addr` warns that the build lacks the
// feature).
/// Carrier for the optional running metrics server, kept alive for the run.
#[cfg(feature = "metrics-http")]
type OptMetricsServer = Option<knut_bifrost::MetricsServer>;
#[cfg(not(feature = "metrics-http"))]
type OptMetricsServer = ();
/// Start the live `/metrics` HTTP server when `--metrics-addr` is given (and the
/// `metrics-http` feature is built in). A bind failure is fatal — the operator
/// asked for the endpoint, so failing loudly beats silently dropping it.
#[cfg(feature = "metrics-http")]
fn start_metrics_server(
addr: Option<&str>,
) -> Result<OptMetricsServer, Box<dyn std::error::Error>> {
match addr {
Some(addr) => {
let server = knut_bifrost::MetricsServer::start(addr)?;
let bound = server
.local_addr()
.map(|a| a.to_string())
.unwrap_or_else(|_| addr.into());
tracing::info!(
target: "metrics",
event = "metrics_http_started",
addr = %bound,
"serving Prometheus metrics at http://{bound}/metrics",
);
Ok(Some(server))
}
None => Ok(None),
}
}
/// Without the `metrics-http` feature there is no server; warn if the flag was
/// nonetheless supplied so the operator knows the endpoint isn't running.
#[cfg(not(feature = "metrics-http"))]
fn start_metrics_server(
addr: Option<&str>,
) -> Result<OptMetricsServer, Box<dyn std::error::Error>> {
if addr.is_some() {
tracing::warn!(
target: "metrics",
event = "metrics_http_unavailable",
"--metrics-addr was given but this build lacks the `metrics-http` feature; \
no /metrics server started (rebuild with --features metrics-http)",
);
}
Ok(())
}
/// Publish the listener's cumulative metrics into the live server's shared
/// snapshot so the next scrape renders them. No-op without the feature / server.
#[cfg(feature = "metrics-http")]
fn publish_metrics<Src: TableSource, Snk: GraphSink>(
server: &OptMetricsServer,
listener: &Listener<Src, Snk>,
) {
if let Some(server) = server {
server.handle().publish(listener.metrics());
}
}
#[cfg(not(feature = "metrics-http"))]
fn publish_metrics<Src: TableSource, Snk: GraphSink>(
_server: &OptMetricsServer,
_listener: &Listener<Src, Snk>,
) {
}
/// Write a sample parquet drop so `bifrost run` can be tried with no other
/// tooling. Columns: `id: Int64`, `name: Utf8`, `score: Float64`.
fn gen_sample(
files: PathBuf,
table: String,
rows: usize,
chunk: String,
) -> Result<(), Box<dyn std::error::Error>> {
use arrow_array::{Float64Array, Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
let ids: Vec<i64> = (0..rows as i64).collect();
let names: Vec<String> = (0..rows).map(|i| format!("person_{i}")).collect();
let scores: Vec<f64> = (0..rows).map(|i| i as f64 * 0.5).collect();
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("score", DataType::Float64, false),
])),
vec![
Arc::new(Int64Array::from(ids)),
Arc::new(StringArray::from(names)),
Arc::new(Float64Array::from(scores)),
],
)?;
let dir = files.join(&table);
std::fs::create_dir_all(&dir)?;
let path = dir.join(&chunk);
graphar::io::write_chunk(
&path,
std::slice::from_ref(&batch),
&graphar::FileType::Parquet,
)?;
println!("wrote {rows} rows → {}", path.display());
Ok(())
}
/// Run one round, or poll until **Ctrl-C** — finishing the current round
/// cleanly before exiting (graceful shutdown). The library's tracing events
/// report each round (including the cumulative `metrics` event). Generic over
/// the sink so `--dry-run` and the real sink share one loop.
///
/// Between rounds the sleep is the base `interval`, but a fully-failed
/// (no-progress) round backs the next sleep off exponentially up to
/// `max_backoff` (default 10× `interval`), resetting on the first round that
/// makes progress. On shutdown, the cumulative metrics are optionally written
/// to `metrics_file` as Prometheus text. When a live `/metrics` server is
/// running (`metrics_server`), the cumulative metrics are published into its
/// shared snapshot after every round so each scrape sees the current counters.
async fn drive<Src: TableSource, Snk: GraphSink>(
mut listener: Listener<Src, Snk>,
once: bool,
interval: u64,
max_backoff: Option<u64>,
metrics_file: Option<PathBuf>,
metrics_server: &OptMetricsServer,
) -> Result<(), Box<dyn std::error::Error>> {
if once {
listener.sync_once().await?;
publish_metrics(metrics_server, &listener);
dump_metrics(&listener, &metrics_file)?;
return Ok(());
}
let interval = Duration::from_secs(interval);
let max = max_backoff
.map(Duration::from_secs)
.unwrap_or_else(|| interval.saturating_mul(10));
let mut backoff = Backoff::new(interval, max);
loop {
let report = listener.sync_once().await?;
// Publish the round's cumulative metrics so a live scrape sees them.
publish_metrics(metrics_server, &listener);
let sleep = backoff.next_sleep(report.is_failure());
if sleep > backoff.interval() {
tracing::warn!(
target: "bifrost::round",
event = "backoff",
sleep_secs = sleep.as_secs(),
interval_secs = backoff.interval().as_secs(),
"round made no progress — backing off before retry",
);
}
// Sleep until the next poll OR a shutdown signal — Ctrl-C during the
// wait exits immediately, but only between rounds (the current round
// always completes and its state is persisted first).
tokio::select! {
_ = tokio::signal::ctrl_c() => {
tracing::info!(target: "bifrost", event = "shutdown", "graceful shutdown");
dump_metrics(&listener, &metrics_file)?;
return Ok(());
}
_ = tokio::time::sleep(sleep) => {}
}
}
}
/// Write the listener's cumulative metrics as Prometheus text to `path` (if
/// any). A `/metrics` HTTP endpoint can serve the same `render_prometheus`
/// output later without an HTTP dependency in the library.
fn dump_metrics<Src: TableSource, Snk: GraphSink>(
listener: &Listener<Src, Snk>,
path: &Option<PathBuf>,
) -> Result<(), Box<dyn std::error::Error>> {
if let Some(path) = path {
std::fs::write(path, knut_bifrost::render_prometheus(listener.metrics()))?;
tracing::info!(
target: "metrics",
event = "metrics_written",
path = %path.display(),
"wrote Prometheus metrics",
);
}
Ok(())
}
/// A graph sink that writes nothing — it just counts what *would* be written,
/// so `--dry-run` can exercise detect→transform with no database. The library's
/// `sink_applied` event still fires via the listener's per-table log.
#[derive(Default)]
struct NullSink;
#[async_trait]
impl GraphSink for NullSink {
async fn apply(&mut self, mutation: &GraphMutation) -> BifrostResult<SinkReport> {
let vertices: usize = mutation.vertices.iter().map(|v| v.batch.num_rows()).sum();
let edges: usize = mutation.edges.iter().map(|e| e.batch.num_rows()).sum();
let deleted: usize = mutation.deletes.iter().map(|d| d.batch.num_rows()).sum();
tracing::info!(
target: "bifrost::sink",
event = "dry_run_mutation",
vertices,
edges,
deleted,
"dry-run: counted graph mutation (not written)",
);
Ok(SinkReport {
vertices_written: vertices,
edges_written: edges,
deleted,
})
}
}