use std::path::{Path, PathBuf};
use std::time::Duration;
use clap::{Args, Subcommand, ValueEnum};
use crate::delivery::{DeliveryOptions, IcebergDelivery, IcebergLocator, bucket_of, s3_ca, s3_config};
use crate::error::{BifrostError, Result};
use crate::sink::doris::{DorisConfig, DorisSink, LoadFormat, LoadMode, ReqwestTransport};
use crate::state::SyncState;
pub const SECRET_ENV: [&str; 5] = [
"MINIO_ACCESS_KEY",
"MINIO_SECRET_KEY",
"DORIS_USER",
"DORIS_PASSWORD",
"BIFROST_S3_CA_PEM",
];
pub const IMAGE_EXEC: &str = "/usr/local/bin/bifrost";
pub const K8S_STATE_DIR: &str = "/var/lib/bifrost";
#[derive(Subcommand, Debug, Clone)]
pub enum DorisCommand {
Run(RunArgs),
Unit(ServiceArgs),
Install(ServiceArgs),
Uninstall(LifecycleArgs),
Start(LifecycleArgs),
Stop(LifecycleArgs),
Status(LifecycleArgs),
K8s(K8sArgs),
}
#[derive(Args, Debug, Clone, Default, PartialEq, Eq)]
pub struct SourceArgs {
#[arg(long, value_name = "DIR", conflicts_with_all = ["s3_table", "rest_uri", "grpc_uri"])]
pub warehouse: Option<PathBuf>,
#[arg(long, value_name = "S3_URL", conflicts_with_all = ["rest_uri", "grpc_uri"])]
pub s3_table: Option<String>,
#[arg(long, value_name = "URL", conflicts_with = "grpc_uri")]
pub rest_uri: Option<String>,
#[arg(long, value_name = "URL")]
pub grpc_uri: Option<String>,
#[arg(long)]
pub rest_warehouse: Option<String>,
#[arg(long, value_name = "SERVICE")]
pub sigv4: Option<String>,
#[arg(long)]
pub table: Option<String>,
#[arg(long, value_name = "URL")]
pub s3_endpoint: Option<String>,
#[arg(long, default_value = "us-east-1")]
pub s3_region: String,
#[arg(long)]
pub s3_bucket: Option<String>,
#[arg(long, value_name = "PEM")]
pub s3_ca: Option<PathBuf>,
}
#[derive(Args, Debug, Clone, PartialEq, Eq)]
pub struct DorisArgs {
#[arg(long, value_name = "URL")]
pub doris_fe: String,
#[arg(long, value_name = "URL")]
pub doris_be: Option<String>,
#[arg(long)]
pub doris_db: String,
#[arg(long)]
pub doris_table: String,
#[arg(long, value_enum, default_value_t = ModeArg::Upsert)]
pub mode: ModeArg,
#[arg(long, value_enum, default_value_t = FormatArg::Json)]
pub format: FormatArg,
#[arg(long, default_value_t = 4)]
pub parallel_loads: usize,
#[arg(long, default_value_t = 64)]
pub chunk_mb: usize,
#[arg(long = "key", value_delimiter = ',', value_name = "COL")]
pub key: Vec<String>,
#[arg(long)]
pub no_sequence: bool,
#[arg(long)]
pub no_two_phase_commit: bool,
}
#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModeArg {
Upsert,
Append,
}
#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum FormatArg {
Json,
Arrow,
}
impl ModeArg {
fn as_str(self) -> &'static str {
match self {
ModeArg::Upsert => "upsert",
ModeArg::Append => "append",
}
}
}
impl FormatArg {
fn as_str(self) -> &'static str {
match self {
FormatArg::Json => "json",
FormatArg::Arrow => "arrow",
}
}
}
#[derive(Args, Debug, Clone, PartialEq, Eq)]
pub struct RunArgs {
#[command(flatten)]
pub source: SourceArgs,
#[command(flatten)]
pub doris: DorisArgs,
#[arg(long, value_name = "FILE", default_value = "bifrost-doris-cursor.json")]
pub state: PathBuf,
#[arg(long)]
pub once: bool,
#[arg(long, default_value_t = 5)]
pub interval: u64,
#[arg(long, default_value_t = 300)]
pub max_backoff: u64,
#[arg(long, default_value_t = 64)]
pub max_windows: usize,
}
#[derive(Args, Debug, Clone, PartialEq, Eq)]
pub struct ServiceArgs {
#[command(flatten)]
pub run: RunArgs,
#[command(flatten)]
pub life: LifecycleArgs,
}
#[derive(Args, Debug, Clone, PartialEq, Eq)]
pub struct LifecycleArgs {
#[arg(long, default_value = "bifrost-doris")]
pub name: String,
#[arg(long)]
pub system: bool,
#[arg(long)]
pub purge: bool,
}
#[derive(Args, Debug, Clone, PartialEq, Eq)]
pub struct K8sArgs {
#[command(flatten)]
pub run: RunArgs,
#[arg(long, default_value = "bifrost-doris")]
pub name: String,
#[arg(long)]
pub image: String,
#[arg(long, default_value = "default")]
pub namespace: String,
#[arg(long, default_value = "1Gi")]
pub state_size: String,
#[arg(long, value_name = "DIR")]
pub out: PathBuf,
}
fn env(name: &str) -> Option<String> {
std::env::var(name).ok().filter(|v| !v.is_empty())
}
impl SourceArgs {
fn s3(&self, state_dir: &Path) -> Result<skade::S3Config> {
let endpoint = self
.s3_endpoint
.clone()
.ok_or_else(|| BifrostError::Other("--s3-endpoint is required for --s3-table / --rest-uri".into()))?;
let bucket = self
.s3_bucket
.clone()
.or_else(|| self.s3_table.as_deref().and_then(bucket_of).map(str::to_string))
.or_else(|| self.rest_warehouse.clone())
.ok_or_else(|| BifrostError::Other("--s3-bucket is required (no bucket in --s3-table or --rest-warehouse)".into()))?;
let ca = s3_ca(self.s3_ca.as_deref(), state_dir)?;
s3_config(endpoint, self.s3_region.clone(), bucket, ca)
}
pub fn locator(&self, state_dir: &Path) -> Result<IcebergLocator> {
if let Some(dir) = &self.warehouse {
let table = self.table.clone().ok_or_else(|| BifrostError::Other("--warehouse needs --table".into()))?;
return Ok(IcebergLocator::Warehouse { dir: dir.clone(), table });
}
if let Some(loc) = &self.s3_table {
return Ok(IcebergLocator::S3(skade::connect::S3Table { s3: self.s3(state_dir)?, location: loc.clone() }));
}
if let Some(uri) = &self.rest_uri {
let table = self.table.clone().ok_or_else(|| BifrostError::Other("--rest-uri needs --table ns.name".into()))?;
let s3 = self.s3(state_dir)?;
let tls = s3.tls.clone();
let catalog = skade::connect::IcebergRest {
uri: uri.clone(),
warehouse: self.rest_warehouse.clone(),
prefix_or_ref: None,
token: None,
tls,
s3,
read_only: true,
sigv4: self.sigv4.clone().map(|service| skade::connect::SigV4 { service }),
};
return Ok(IcebergLocator::catalog(catalog, table));
}
if let Some(uri) = &self.grpc_uri {
let table = self.table.clone().ok_or_else(|| BifrostError::Other("--grpc-uri needs --table ns.name".into()))?;
let s3 = match &self.s3_endpoint {
Some(_) => Some(self.s3(state_dir)?),
None => None,
};
let catalog = skade::connect::SkadeGrpc { uri: uri.clone(), s3, read_only: true };
return Ok(IcebergLocator::catalog(catalog, table));
}
Err(BifrostError::Other(
"name the table: --warehouse DIR, --s3-table s3://…, --rest-uri URL, or --grpc-uri URL".into(),
))
}
fn to_args(&self) -> Vec<String> {
let mut a = Vec::new();
let mut push = |k: &str, v: Option<String>| {
if let Some(v) = v {
a.push(format!("--{k}"));
a.push(v);
}
};
push("warehouse", self.warehouse.as_ref().map(|p| p.display().to_string()));
push("s3-table", self.s3_table.clone());
push("rest-uri", self.rest_uri.clone());
push("rest-warehouse", self.rest_warehouse.clone());
push("grpc-uri", self.grpc_uri.clone());
push("sigv4", self.sigv4.clone());
push("table", self.table.clone());
push("s3-endpoint", self.s3_endpoint.clone());
push("s3-region", Some(self.s3_region.clone()));
push("s3-bucket", self.s3_bucket.clone());
push("s3-ca", self.s3_ca.as_ref().map(|p| p.display().to_string()));
a
}
}
impl DorisArgs {
fn config(&self) -> Result<DorisConfig> {
let mut cfg = DorisConfig::from_env(&self.doris_fe, &self.doris_db, &self.doris_table)?
.with_mode(match self.mode {
ModeArg::Upsert => LoadMode::Upsert,
ModeArg::Append => LoadMode::Append,
})
.with_max_chunk_bytes(self.chunk_mb.max(1) * 1024 * 1024)
.with_max_parallel_loads(self.parallel_loads)
.with_key_columns(self.key.clone())
.with_sequence(!self.no_sequence)
.with_two_phase_commit(!self.no_two_phase_commit);
cfg.format = match self.format {
FormatArg::Json => LoadFormat::Json,
FormatArg::Arrow => LoadFormat::Arrow,
};
if let Some(be) = &self.doris_be {
cfg = cfg.with_be_url(be);
}
Ok(cfg)
}
fn destination(&self) -> String {
format!("doris:{}.{}", self.doris_db, self.doris_table)
}
fn to_args(&self) -> Vec<String> {
let mut a = vec![
"--doris-fe".into(),
self.doris_fe.clone(),
"--doris-db".into(),
self.doris_db.clone(),
"--doris-table".into(),
self.doris_table.clone(),
"--mode".into(),
self.mode.as_str().into(),
"--format".into(),
self.format.as_str().into(),
"--chunk-mb".into(),
self.chunk_mb.to_string(),
"--parallel-loads".into(),
self.parallel_loads.to_string(),
];
if let Some(be) = &self.doris_be {
a.extend(["--doris-be".into(), be.clone()]);
}
if !self.key.is_empty() {
a.extend(["--key".into(), self.key.join(",")]);
}
if self.no_sequence {
a.push("--no-sequence".into());
}
if self.no_two_phase_commit {
a.push("--no-two-phase-commit".into());
}
a
}
}
impl RunArgs {
pub fn command_args(&self, state: &Path) -> Vec<String> {
let mut a = vec!["doris".to_string(), "run".to_string()];
a.extend(self.source.to_args());
a.extend(self.doris.to_args());
a.extend([
"--state".into(),
state.display().to_string(),
"--interval".into(),
self.interval.to_string(),
"--max-backoff".into(),
self.max_backoff.to_string(),
"--max-windows".into(),
self.max_windows.to_string(),
]);
a
}
fn state_dir(&self) -> PathBuf {
self.state.parent().filter(|p| !p.as_os_str().is_empty()).map(Path::to_path_buf).unwrap_or_else(|| PathBuf::from("."))
}
}
pub async fn run(args: &RunArgs) -> Result<()> {
let locator = args.source.locator(&args.state_dir())?;
let cfg = args.doris.config()?;
let transport = ReqwestTransport::new(&cfg)?;
let sink = DorisSink::new(cfg, transport);
let state = SyncState::open(&args.state)?;
let opts = DeliveryOptions { max_windows_per_round: args.max_windows.max(1), key_columns: args.doris.key.clone() };
let mut d = IcebergDelivery::new(locator, &args.doris.destination(), sink, state, opts);
tracing::info!(target: "bifrost::delivery", event = "start", table = %d.cursor_key(), cursor = ?d.cursor(), "delivering");
if args.once {
let o = d.round().await?;
tracing::info!(
target: "bifrost::delivery",
event = "round",
table = %d.cursor_key(),
windows = o.windows,
loaded = o.rows_loaded,
deleted = o.rows_deleted,
duplicate = o.rows_duplicate,
at = ?o.at_snapshot,
behind = o.behind,
read_ms = o.read.as_millis() as u64,
load_ms = o.load.as_millis() as u64,
"delivered",
);
return Ok(());
}
d.run(Duration::from_secs(args.interval.max(1)), Duration::from_secs(args.max_backoff.max(1)), shutdown_signal())
.await
}
async fn shutdown_signal() {
#[cfg(unix)]
{
use tokio::signal::unix::{SignalKind, signal};
match signal(SignalKind::terminate()) {
Ok(mut term) => {
tokio::select! {
_ = term.recv() => {}
_ = tokio::signal::ctrl_c() => {}
}
}
Err(_) => {
let _ = tokio::signal::ctrl_c().await;
}
}
}
#[cfg(not(unix))]
{
let _ = tokio::signal::ctrl_c().await;
}
}
fn home() -> PathBuf {
std::env::var_os("HOME").map(PathBuf::from).unwrap_or_else(|| PathBuf::from("/"))
}
pub fn installed_state_dir(name: &str, system: bool) -> PathBuf {
if system {
PathBuf::from("/var/lib").join(name)
} else {
home().join(".local/state").join(name)
}
}
pub fn service_spec(
name: &str,
exec: &str,
run: &RunArgs,
state: &Path,
system: bool,
) -> skidbladnir::service::ServiceSpec {
use skidbladnir::service::{AgentMode, ServiceSpec};
use skidbladnir::systemd::{RestartPolicy, Scope};
let mut spec = ServiceSpec::new(name, "bifrost: Iceberg commits into Apache Doris", exec)
.scope(if system { Scope::System } else { Scope::User })
.restart(RestartPolicy::OnFailure)
.agent(AgentMode::Fat)
.env("RUST_LOG", "info");
if let Some(dir) = state.parent() {
spec = spec.working_dir(dir.to_string_lossy());
}
if system {
spec = spec.user("bifrost");
}
for a in run.command_args(state) {
spec = spec.arg(a);
}
for k in SECRET_ENV {
if let Some(v) = env(k) {
spec = spec.secret_env(k, v);
}
}
spec
}
fn lifecycle(run: Option<&RunArgs>, life: &LifecycleArgs) -> Result<skidbladnir::service::Lifecycle> {
let exec = std::env::current_exe()
.map(|p| p.to_string_lossy().into_owned())
.unwrap_or_else(|_| "/usr/local/bin/bifrost".into());
let state_dir = installed_state_dir(&life.name, life.system);
let state = state_dir.join("cursor.json");
let mut spec = match run {
Some(r) => service_spec(&life.name, &exec, r, &state, life.system),
None => skidbladnir::service::ServiceSpec::new(&life.name, "bifrost: Iceberg commits into Apache Doris", &exec).scope(
if life.system { skidbladnir::systemd::Scope::System } else { skidbladnir::systemd::Scope::User },
),
};
if life.purge {
let root = skidbladnir::uninstall::DataRoot::from_env("BIFROST_DORIS_DATA", state_dir)
.map_err(|e| BifrostError::Other(format!("state dir: {e}")))?;
spec = spec.data_root(&root);
}
Ok(skidbladnir::service::Lifecycle::new(spec))
}
fn sk(e: anyhow::Error) -> BifrostError {
BifrostError::Other(format!("skidbladnir: {e:#}"))
}
pub fn k8s_files(args: &K8sArgs) -> Vec<skidbladnir::service::RenderedFile> {
use skidbladnir::service::{ServiceRenderer as _, Storage};
let state = Path::new(K8S_STATE_DIR).join("cursor.json");
let spec = service_spec(&args.name, IMAGE_EXEC, &args.run, &state, false)
.image(&args.image)
.namespace(&args.namespace)
.replicas(1)
.stateful(true)
.storage(Storage::new(&args.state_size, K8S_STATE_DIR));
skidbladnir::k8s::KubernetesRenderer.render(&spec).files
}
pub async fn dispatch(cmd: DorisCommand) -> Result<()> {
match cmd {
DorisCommand::Run(a) => run(&a).await,
DorisCommand::Unit(a) => {
print!("{}", lifecycle(Some(&a.run), &a.life)?.unit_text());
Ok(())
}
DorisCommand::Install(a) => {
let lc = lifecycle(Some(&a.run), &a.life)?;
lc.install().map_err(sk)?;
lc.enable().map_err(sk)?;
lc.start().map_err(sk)?;
println!("bifrost doris: INSTALLED and started ({})", lc.unit());
Ok(())
}
DorisCommand::Uninstall(l) => {
let lc = lifecycle(None, &l)?;
let r = lc.uninstall().map_err(sk)?;
println!("bifrost doris: unit removed = {}, erased {} path(s)", r.unit_removed, r.erased.len());
Ok(())
}
DorisCommand::Start(l) => lifecycle(None, &l)?.start().map_err(sk),
DorisCommand::Stop(l) => lifecycle(None, &l)?.stop().map_err(sk),
DorisCommand::Status(l) => {
println!("{}", if lifecycle(None, &l)?.status().map_err(sk)? { "active" } else { "inactive" });
Ok(())
}
DorisCommand::K8s(a) => {
std::fs::create_dir_all(&a.out)?;
for f in k8s_files(&a) {
let file = a.out.join(Path::new(&f.path).file_name().unwrap_or(f.path.as_ref()));
std::fs::write(&file, &f.contents)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&file, std::fs::Permissions::from_mode(f.mode))?;
}
println!("{}", file.display());
}
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use clap::Parser;
use super::*;
#[derive(Parser)]
struct T {
#[command(subcommand)]
cmd: DorisCommand,
}
fn parse(s: &str) -> DorisCommand {
T::try_parse_from(std::iter::once("t").chain(s.split_whitespace())).unwrap().cmd
}
const AISTOR: &str = "--rest-uri https://minio:9000/_iceberg --rest-warehouse red --sigv4 s3tables \
--table bench.people --s3-endpoint https://minio:9000 \
--doris-fe http://fe:8030 --doris-db lake --doris-table people";
#[test]
fn the_command_line_round_trips_through_its_own_rendering() {
let DorisCommand::Run(a) = parse(&format!("run {AISTOR} --state /s/c.json --interval 5")) else { panic!() };
let rendered = a.command_args(Path::new("/s/c.json"));
assert_eq!(&rendered[..2], ["doris", "run"]);
let DorisCommand::Run(b) = T::try_parse_from(std::iter::once("t".to_string()).chain(rendered.into_iter().skip(1)))
.unwrap()
.cmd
else {
panic!()
};
assert_eq!(a, b, "what the service runs parses back to what was asked for");
}
#[test]
fn two_sources_at_once_is_refused_by_clap() {
let r = T::try_parse_from(
"t run --warehouse /w --s3-table s3://b/x/y --doris-fe f --doris-db d --doris-table t".split_whitespace(),
);
assert!(r.is_err());
for other in ["--warehouse /w", "--s3-table s3://b/x/y", "--rest-uri http://r:8181"] {
let r = T::try_parse_from(
format!("t run --grpc-uri http://k:8182 {other} --table db.t --doris-fe f --doris-db d --doris-table t")
.split_whitespace(),
);
assert!(r.is_err(), "--grpc-uri with {other} must be refused");
}
}
const GRPC: &str = "--grpc-uri http://katalog:8182 --table db.people --doris-fe http://fe:8030 --doris-db lake --doris-table people";
#[test]
fn grpc_uri_selects_the_skade_katalog_grpc_catalog() {
let DorisCommand::Run(a) = parse(&format!("run {GRPC} --state /s/c.json")) else { panic!() };
let loc = a.source.locator(Path::new("/s")).unwrap();
assert_eq!(loc.source_id(), "grpc:http://katalog:8182#db.people");
let IcebergLocator::Catalog { connector, .. } = &loc else { panic!("{loc:?}") };
assert_eq!(connector.kind(), "skade-grpc");
assert!(connector.read_only(), "bifrost never writes the lake");
assert!(connector.describe().contains("files on this box"), "{}", connector.describe());
let rendered = a.command_args(Path::new("/s/c.json"));
let DorisCommand::Run(b) = T::try_parse_from(std::iter::once("t".to_string()).chain(rendered.into_iter().skip(1)))
.unwrap()
.cmd
else {
panic!()
};
assert_eq!(a, b, "what the service runs parses back to what was asked for");
let DorisCommand::Run(no_table) = parse("run --grpc-uri http://k:8182 --doris-fe f --doris-db d --doris-table t") else {
panic!()
};
let e = no_table.source.locator(Path::new("/s")).unwrap_err().to_string();
assert!(e.contains("--grpc-uri needs --table"), "{e}");
}
#[test]
fn the_rest_cursor_key_is_unchanged() {
unsafe {
std::env::set_var("MINIO_ACCESS_KEY", "ak-test");
std::env::set_var("MINIO_SECRET_KEY", "sk-test");
}
let DorisCommand::Run(a) = parse(&format!("run {AISTOR} --state /s/c.json")) else { panic!() };
let loc = a.source.locator(Path::new("/s")).unwrap();
assert_eq!(loc.source_id(), "rest:https://minio:9000/_iceberg#red#bench.people");
assert!(!format!("{loc:?}").contains("sk-test"), "no secret in the locator's Debug");
}
#[test]
fn systemd_and_k8s_run_the_same_command_and_no_secret_is_on_it() {
unsafe { std::env::set_var("DORIS_PASSWORD", "hunter2-not-on-argv") };
let DorisCommand::K8s(k) = parse(&format!("k8s {AISTOR} --image reg/bifrost:1 --out /tmp/x")) else { panic!() };
let state = Path::new(K8S_STATE_DIR).join("cursor.json");
let spec = service_spec("bifrost-doris", IMAGE_EXEC, &k.run, &state, false);
let unit = spec.render_systemd_unit();
let files = k8s_files(&k);
let workload = files.iter().find(|f| f.contents.contains("kind: StatefulSet")).expect("a StatefulSet");
for a in k.run.command_args(&state) {
assert!(workload.contents.contains(&a), "k8s args carry {a}");
assert!(unit.contains(&a), "the unit carries {a}");
}
assert!(!unit.contains("hunter2"), "a secret never lands in the 0644 unit");
assert!(!workload.contents.contains("hunter2"), "a secret never lands in the workload");
let secret = files.iter().find(|f| f.contents.contains("kind: Secret")).expect("a Secret for the credentials");
assert_eq!(secret.mode, 0o600);
assert!(workload.contents.contains(K8S_STATE_DIR), "the cursor lives on the volume");
}
#[test]
fn bucket_comes_from_the_table_url() {
assert_eq!(bucket_of("s3://red/warehouse/bench/t"), Some("red"));
assert_eq!(bucket_of("s3a://b/x"), Some("b"));
assert_eq!(bucket_of("/local/x"), None);
}
}