use std::collections::BTreeSet;
use arrow_array::RecordBatch;
use arrow_schema::Schema;
use async_trait::async_trait;
use crate::{
error::{BifrostError, Result},
sink::table::{AckToken, LoadAck, TableSink, TableWindow, stream_load_label},
};
#[derive(Debug, Clone)]
pub struct StreamLoadRequest {
pub method: &'static str,
pub url: String,
pub headers: Vec<(String, String)>,
pub body: bytes::Bytes,
}
impl StreamLoadRequest {
pub fn header(&self, name: &str) -> Option<&str> {
self.headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case(name))
.map(|(_, v)| v.as_str())
}
}
#[derive(Debug, Clone)]
pub struct StreamLoadResponse {
pub status: u16,
pub location: Option<String>,
pub body: String,
}
impl StreamLoadResponse {
pub fn ok(body: impl Into<String>) -> Self {
Self {
status: 200,
location: None,
body: body.into(),
}
}
pub fn redirect(location: impl Into<String>) -> Self {
Self {
status: 307,
location: Some(location.into()),
body: String::new(),
}
}
pub fn is_redirect(&self) -> bool {
matches!(self.status, 301 | 302 | 303 | 307 | 308) && self.location.is_some()
}
}
#[async_trait]
pub trait StreamLoadTransport: Send + Sync {
async fn send(&self, req: &StreamLoadRequest) -> Result<StreamLoadResponse>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LoadMode {
Append,
Upsert,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum LoadFormat {
#[default]
Json,
Arrow,
}
impl LoadFormat {
const fn header_value(self) -> &'static str {
match self {
LoadFormat::Json => "json",
LoadFormat::Arrow => "arrow",
}
}
}
#[derive(Debug, Clone)]
pub struct DorisConfig {
pub fe_url: String,
pub db: String,
pub table: String,
pub user: String,
pub password: String,
pub mode: LoadMode,
pub format: LoadFormat,
pub max_chunk_bytes: usize,
pub connect_timeout: std::time::Duration,
pub request_timeout: std::time::Duration,
pub be_url: Option<String>,
pub max_parallel_loads: usize,
pub two_phase_commit: bool,
pub sequence: bool,
pub key_columns: Vec<String>,
}
pub const DEFAULT_MAX_CHUNK_BYTES: usize = 64 * 1024 * 1024;
pub const DEFAULT_MAX_PARALLEL_LOADS: usize = 4;
pub const DEFAULT_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
pub const DEFAULT_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
impl DorisConfig {
pub fn from_env(
fe_url: impl Into<String>,
db: impl Into<String>,
table: impl Into<String>,
) -> Result<Self> {
let user = std::env::var("DORIS_USER")
.map_err(|_| BifrostError::Other("DORIS_USER is not set".into()))?;
let password = std::env::var("DORIS_PASSWORD")
.map_err(|_| BifrostError::Other("DORIS_PASSWORD is not set".into()))?;
Ok(Self {
fe_url: fe_url.into(),
db: db.into(),
table: table.into(),
user,
password,
mode: LoadMode::Append,
format: LoadFormat::default(),
max_chunk_bytes: DEFAULT_MAX_CHUNK_BYTES,
connect_timeout: DEFAULT_CONNECT_TIMEOUT,
request_timeout: DEFAULT_REQUEST_TIMEOUT,
be_url: None,
max_parallel_loads: DEFAULT_MAX_PARALLEL_LOADS,
two_phase_commit: true,
sequence: true,
key_columns: Vec::new(),
})
}
pub fn new(
fe_url: impl Into<String>,
db: impl Into<String>,
table: impl Into<String>,
user: impl Into<String>,
password: impl Into<String>,
) -> Self {
Self {
fe_url: fe_url.into(),
db: db.into(),
table: table.into(),
user: user.into(),
password: password.into(),
mode: LoadMode::Append,
format: LoadFormat::default(),
max_chunk_bytes: DEFAULT_MAX_CHUNK_BYTES,
connect_timeout: DEFAULT_CONNECT_TIMEOUT,
request_timeout: DEFAULT_REQUEST_TIMEOUT,
be_url: None,
max_parallel_loads: DEFAULT_MAX_PARALLEL_LOADS,
two_phase_commit: true,
sequence: true,
key_columns: Vec::new(),
}
}
#[must_use]
pub fn with_timeouts(
mut self,
connect: std::time::Duration,
request: std::time::Duration,
) -> Self {
self.connect_timeout = connect;
self.request_timeout = request;
self
}
#[must_use]
pub fn with_mode(mut self, mode: LoadMode) -> Self {
self.mode = mode;
self
}
#[must_use]
pub fn with_max_chunk_bytes(mut self, n: usize) -> Self {
self.max_chunk_bytes = n.max(1);
self
}
#[must_use]
pub fn with_max_parallel_loads(mut self, n: usize) -> Self {
self.max_parallel_loads = n.max(1);
self
}
#[must_use]
pub fn with_be_url(mut self, url: impl Into<String>) -> Self {
self.be_url = Some(url.into());
self
}
#[must_use]
pub fn with_two_phase_commit(mut self, on: bool) -> Self {
self.two_phase_commit = on;
self
}
#[must_use]
pub fn with_sequence(mut self, on: bool) -> Self {
self.sequence = on;
self
}
#[must_use]
pub fn with_key_columns(mut self, keys: Vec<String>) -> Self {
self.key_columns = keys;
self
}
fn two_phase_url(&self, base: &str) -> String {
format!(
"{}/api/{}/{}/_stream_load_2pc",
base.trim_end_matches('/'),
self.db,
self.table
)
}
fn stream_load_url(&self, base: &str) -> String {
format!(
"{}/api/{}/{}/_stream_load",
base.trim_end_matches('/'),
self.db,
self.table
)
}
fn schema_url(&self) -> String {
format!(
"{}/api/{}/{}/_schema",
self.fe_url.trim_end_matches('/'),
self.db,
self.table
)
}
fn basic_auth(&self) -> String {
format!(
"Basic {}",
base64_std(format!("{}:{}", self.user, self.password).as_bytes())
)
}
}
pub struct DorisSink<T: StreamLoadTransport> {
cfg: DorisConfig,
transport: T,
probed: bool,
}
impl<T: StreamLoadTransport> DorisSink<T> {
pub fn new(cfg: DorisConfig, transport: T) -> Self {
Self {
cfg,
transport,
probed: false,
}
}
pub fn config(&self) -> &DorisConfig {
&self.cfg
}
async fn remote_columns(&self) -> Result<BTreeSet<String>> {
let req = StreamLoadRequest {
method: "GET",
url: self.cfg.schema_url(),
headers: vec![("Authorization".into(), self.cfg.basic_auth())],
body: bytes::Bytes::new(),
};
let resp = self.transport.send(&req).await?;
if resp.status != 200 {
return Err(BifrostError::Other(format!(
"doris: schema probe of {}.{} returned HTTP {} — {}",
self.cfg.db, self.cfg.table, resp.status, resp.body
)));
}
let v: serde_json::Value = serde_json::from_str(&resp.body).map_err(|e| {
BifrostError::Other(format!("doris: schema probe returned non-JSON: {e}"))
})?;
let props = v
.get("data")
.and_then(|d| d.get("properties"))
.and_then(serde_json::Value::as_array)
.ok_or_else(|| {
BifrostError::Other(format!(
"doris: schema probe of {}.{} has no data.properties — {}",
self.cfg.db, self.cfg.table, resp.body
))
})?;
Ok(props
.iter()
.filter_map(|p| p.get("name").and_then(serde_json::Value::as_str))
.map(str::to_owned)
.collect())
}
async fn exchange(&self, url: String, headers: Vec<(String, String)>, body: bytes::Bytes) -> Result<StreamLoadResponse> {
let start = url.clone();
let mut url = url;
let mut hops = 0_u8;
loop {
let req = StreamLoadRequest {
method: "PUT",
url: url.clone(),
headers: headers.clone(),
body: body.clone(),
};
let resp = self.transport.send(&req).await.map_err(|e| {
if hops > 0 {
BifrostError::Other(format!(
"doris: the FE redirected this load to {url} and the request did not \
complete — this is a NETWORK REACHABILITY failure. Measured \
behaviour: an unreachable BE does not refuse the connection, it \
HANGS until the timeout ({}s connect / {}s request), while the FE \
keeps answering /api/health with 200. The loader must reach the FE \
({}) AND every BE it can be redirected to on :8040. Underlying \
error: {e}",
self.cfg.connect_timeout.as_secs(),
self.cfg.request_timeout.as_secs(),
self.cfg.fe_url
))
} else {
e
}
})?;
if !resp.is_redirect() {
return Ok(resp);
}
hops += 1;
if hops > 4 {
return Err(BifrostError::Other(format!("doris: more than 4 redirects starting at {start}")));
}
url = resp.location.clone().expect("is_redirect checked location");
tracing::debug!(
target: "bifrost::doris",
event = "stream_load_redirect",
to = %url,
"FE redirected the request to a backend",
);
}
}
async fn put_chunk(&self, label: &str, load: &PlannedLoad, body: bytes::Bytes) -> Result<LoadOutcome> {
let base = self.cfg.be_url.as_deref().unwrap_or(&self.cfg.fe_url);
let mut headers = vec![
("columns".to_string(), load.columns.clone()),
("Authorization".to_string(), self.cfg.basic_auth()),
("Expect".to_string(), "100-continue".to_string()),
("label".to_string(), label.to_string()),
("format".to_string(), self.cfg.format.header_value().to_string()),
("merge_type".to_string(), load.merge_type.to_string()),
];
if let Some(cond) = load.delete_condition {
headers.push(("delete".to_string(), cond.to_string()));
}
if load.sequenced {
headers.push(("function_column.sequence_col".to_string(), SEQ_COLUMN.to_string()));
}
if self.cfg.two_phase_commit {
headers.push(("two_phase_commit".to_string(), "true".to_string()));
}
if self.cfg.format == LoadFormat::Json {
headers.push(("strip_outer_array".to_string(), "true".to_string()));
}
headers.push(("max_filter_ratio".to_string(), "0".to_string()));
let resp = match self.exchange(self.cfg.stream_load_url(base), headers, body).await {
Ok(resp) => resp,
Err(send_err) => {
return match self.label_state(label).await {
Ok(state) => match state.as_str() {
"PRECOMMITTED" => Ok(LoadOutcome::AlreadyPending),
"COMMITTED" | "VISIBLE" => Ok(LoadOutcome::AlreadyFinished),
_ => Err(send_err),
},
Err(_) => Err(send_err),
};
}
};
if resp.status != 200 {
return Err(BifrostError::Other(format!(
"doris: stream load {label} returned HTTP {} — {}",
resp.status, resp.body
)));
}
parse_load_outcome(label, &resp.body)
}
async fn label_state(&self, label: &str) -> Result<String> {
let encoded: String = label
.bytes()
.map(|b| match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => (b as char).to_string(),
_ => format!("%{b:02X}"),
})
.collect();
let req = StreamLoadRequest {
method: "GET",
url: format!(
"{}/api/{}/get_load_state?label={encoded}",
self.cfg.fe_url.trim_end_matches('/'),
self.cfg.db
),
headers: vec![("Authorization".to_string(), self.cfg.basic_auth())],
body: bytes::Bytes::new(),
};
let resp = self.transport.send(&req).await?;
let v: serde_json::Value = serde_json::from_str(&resp.body).map_err(|e| {
BifrostError::Other(format!("doris: get_load_state for {label} returned non-JSON: {e} — {}", resp.body))
})?;
match (v.get("code").and_then(serde_json::Value::as_i64), v.get("data").and_then(serde_json::Value::as_str)) {
(Some(0), Some(state)) => Ok(state.to_ascii_uppercase()),
_ => Err(BifrostError::Other(format!("doris: get_load_state for {label} failed: {}", resp.body))),
}
}
async fn txn_op(&self, label: &str, op: &str) -> Result<()> {
let base = self.cfg.be_url.as_deref().unwrap_or(&self.cfg.fe_url);
let headers = vec![
("Authorization".to_string(), self.cfg.basic_auth()),
("Expect".to_string(), "100-continue".to_string()),
("label".to_string(), label.to_string()),
("txn_operation".to_string(), op.to_string()),
];
let resp = self.exchange(self.cfg.two_phase_url(base), headers, bytes::Bytes::new()).await;
let result = match resp {
Ok(resp) if resp.status != 200 => Err(BifrostError::Other(format!(
"doris: 2pc {op} of {label} returned HTTP {} — {}",
resp.status, resp.body
))),
Ok(resp) => parse_txn_result(label, op, &resp.body),
Err(e) => Err(e),
};
match result {
Ok(()) => Ok(()),
Err(e) => match self.label_state(label).await.as_deref() {
Ok("COMMITTED" | "VISIBLE") if op == "commit" => {
tracing::info!(target: "bifrost::doris", event = "commit_replayed", label = %label, "already committed");
Ok(())
}
Ok("ABORTED" | "UNKNOWN") if op == "abort" => Ok(()),
_ => Err(e),
},
}
}
fn plan_loads(&self, window: &TableWindow<'_>) -> Result<Vec<PlannedLoad>> {
if self.cfg.mode == LoadMode::Append {
return Ok(if window.upsert_rows() > 0 {
vec![PlannedLoad::new("APPEND", None, false, window.upserts.to_vec())]
} else {
Vec::new()
});
}
let keys: &[String] = if self.cfg.key_columns.is_empty() { window.key_columns } else { &self.cfg.key_columns };
let seq = if self.cfg.sequence {
Some(window.sequence.ok_or_else(|| {
BifrostError::Other(format!(
"doris: {}: the window carries no sequence number but the sink versions rows by \
it (DorisConfig::sequence). Deliver one snapshot per window, or turn the \
sequence off for a table without a sequence column.",
window.table
))
})?)
} else {
None
};
let deletes: Vec<RecordBatch> = if window.delete_rows() > 0 && window.upsert_rows() > 0 {
if keys.is_empty() {
return Err(BifrostError::Other(format!(
"doris: {}: this window deletes and upserts rows but no key columns are known, \
so an update (delete + insert of one key) cannot be told from two changes. \
Name them (--key / DorisConfig::with_key_columns) or give the Iceberg table \
identifier fields.",
window.table
)));
}
drop_upserted_keys(window.deletes, window.upserts, keys)?
} else {
window.deletes.to_vec()
};
let (whole, identity) = split_identity_deletes(deletes, keys)?;
let mut loads = Vec::new();
if !identity.is_empty() {
let projected = identity
.iter()
.map(|b| project_named(b, keys))
.collect::<Result<Vec<_>>>()?;
let tagged = projected.iter().map(|b| tag(b, None, seq)).collect::<Result<Vec<_>>>()?;
loads.push(PlannedLoad::new("DELETE", None, seq.is_some(), tagged));
}
if window.upsert_rows() > 0 || !whole.is_empty() {
let mut body = Vec::new();
for b in window.upserts {
body.push(tag(b, Some(0), seq)?);
}
for b in &whole {
body.push(tag(b, Some(1), seq)?);
}
let body = unify(body)?;
loads.push(PlannedLoad::new("MERGE", Some(DELETE_CONDITION), seq.is_some(), body));
}
Ok(loads)
}
}
pub const OP_COLUMN: &str = "__bifrost_op";
pub const SEQ_COLUMN: &str = "__bifrost_seq";
const DELETE_CONDITION: &str = "__bifrost_op=1";
struct PlannedLoad {
merge_type: &'static str,
delete_condition: Option<&'static str>,
sequenced: bool,
columns: String,
batches: Vec<RecordBatch>,
}
impl PlannedLoad {
fn new(merge_type: &'static str, delete_condition: Option<&'static str>, sequenced: bool, batches: Vec<RecordBatch>) -> Self {
let columns = batches
.first()
.map(|b| b.schema().fields().iter().map(|f| f.name().as_str()).collect::<Vec<_>>().join(","))
.unwrap_or_default();
Self { merge_type, delete_condition, sequenced, columns, batches }
}
fn rows(&self) -> usize {
self.batches.iter().map(RecordBatch::num_rows).sum()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LoadOutcome {
Loaded { rows: usize },
AlreadyFinished,
AlreadyPending,
}
fn tag(b: &RecordBatch, op: Option<i8>, seq: Option<i64>) -> Result<RecordBatch> {
use arrow_array::{ArrayRef, Int8Array, Int64Array};
use std::sync::Arc;
let mut fields: Vec<arrow_schema::FieldRef> = b.schema().fields().iter().cloned().collect();
let mut cols: Vec<ArrayRef> = b.columns().to_vec();
if let Some(op) = op {
fields.push(Arc::new(arrow_schema::Field::new(OP_COLUMN, arrow_schema::DataType::Int8, false)));
cols.push(Arc::new(Int8Array::from(vec![op; b.num_rows()])));
}
if let Some(seq) = seq {
fields.push(Arc::new(arrow_schema::Field::new(SEQ_COLUMN, arrow_schema::DataType::Int64, false)));
cols.push(Arc::new(Int64Array::from(vec![seq; b.num_rows()])));
}
Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), cols)?)
}
fn unify(batches: Vec<RecordBatch>) -> Result<Vec<RecordBatch>> {
use std::sync::Arc;
let Some(first) = batches.first() else { return Ok(batches) };
let fields: Vec<arrow_schema::Field> =
first.schema().fields().iter().map(|f| f.as_ref().clone().with_nullable(true)).collect();
let schema = Arc::new(Schema::new(fields));
batches
.iter()
.map(|b| {
let cols = schema
.fields()
.iter()
.map(|f| {
let c = b.column_by_name(f.name()).ok_or_else(|| {
BifrostError::Other(format!("doris: a delete row has no column {:?} the upserts carry", f.name()))
})?;
if c.data_type() == f.data_type() {
Ok(Arc::clone(c))
} else {
arrow::compute::cast(c, f.data_type()).map_err(BifrostError::from)
}
})
.collect::<Result<Vec<_>>>()?;
Ok(RecordBatch::try_new(Arc::clone(&schema), cols)?)
})
.collect()
}
fn project_named(b: &RecordBatch, names: &[String]) -> Result<RecordBatch> {
let idx = names
.iter()
.map(|k| {
b.schema()
.index_of(k)
.map_err(|_| BifrostError::Source(format!("delete rows have no identity column {k:?}")))
})
.collect::<Result<Vec<_>>>()?;
b.project(&idx).map_err(|e| BifrostError::Source(format!("projecting delete keys: {e}")))
}
fn split_identity_deletes(deletes: Vec<RecordBatch>, keys: &[String]) -> Result<(Vec<RecordBatch>, Vec<RecordBatch>)> {
use arrow_array::Array as _;
if keys.is_empty() {
return Ok((deletes, Vec::new()));
}
let mut whole = Vec::new();
let mut identity = Vec::new();
for b in deletes {
if b.num_rows() == 0 {
continue;
}
let schema = b.schema();
let partial = schema
.fields()
.iter()
.enumerate()
.any(|(i, f)| !keys.iter().any(|k| k == f.name()) && b.column(i).null_count() > 0);
if partial { identity.push(b) } else { whole.push(b) }
}
Ok((whole, identity))
}
fn drop_upserted_keys(deletes: &[RecordBatch], upserts: &[RecordBatch], keys: &[String]) -> Result<Vec<RecordBatch>> {
use arrow::row::{RowConverter, SortField};
let Some(first) = upserts.first().or_else(|| deletes.first()) else { return Ok(Vec::new()) };
let fields = keys
.iter()
.map(|k| {
first
.schema()
.field_with_name(k)
.map(|f| SortField::new(f.data_type().clone()))
.map_err(|_| BifrostError::Source(format!("no key column {k:?} in the window")))
})
.collect::<Result<Vec<_>>>()?;
let conv = RowConverter::new(fields)?;
let key_cols = |b: &RecordBatch| -> Result<Vec<arrow_array::ArrayRef>> {
Ok(project_named(b, keys)?.columns().to_vec())
};
let mut upserted = std::collections::HashSet::new();
for b in upserts {
let rows = conv.convert_columns(&key_cols(b)?)?;
for r in rows.iter() {
upserted.insert(r.as_ref().to_vec());
}
}
let mut out = Vec::new();
for b in deletes {
let rows = conv.convert_columns(&key_cols(b)?)?;
let keep: arrow_array::BooleanArray = rows.iter().map(|r| Some(!upserted.contains(r.as_ref()))).collect();
let f = arrow::compute::filter_record_batch(b, &keep)?;
if f.num_rows() > 0 {
out.push(f);
}
}
Ok(out)
}
#[async_trait]
impl<T: StreamLoadTransport> TableSink for DorisSink<T> {
async fn probe(&mut self, schema: &Schema) -> Result<()> {
let remote = self.remote_columns().await?;
let missing: Vec<&str> = schema
.fields()
.iter()
.map(|f| f.name().as_str())
.filter(|n| !remote.contains(*n))
.collect();
let ok = missing.is_empty();
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"doris_sink::schema_probe",
ok,
&format!(
"{}.{}: {} source column(s), {} missing remotely",
self.cfg.db,
self.cfg.table,
schema.fields().len(),
missing.len()
),
);
if !ok {
return Err(BifrostError::Other(format!(
"doris: {}.{} is missing column(s) the source sends: {} — refusing to \
start. Discovering this mid-stream, after a partial load, is worse.",
self.cfg.db,
self.cfg.table,
missing.join(", ")
)));
}
self.probed = true;
Ok(())
}
async fn apply(&mut self, window: &TableWindow<'_>, token: AckToken) -> Result<LoadAck> {
if self.cfg.mode == LoadMode::Append && window.delete_rows() > 0 {
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"doris_sink::append_mode_refuses_deletes",
false,
&format!(
"{}: {} delete row(s) in Append mode",
window.table,
window.delete_rows()
),
);
return Err(BifrostError::Other(format!(
"doris: {} delete row(s) arrived for {}.{} in Append mode — the source is \
not append-only after all. Refusing rather than silently diverging from \
the lake. Use a unique-key table with LoadMode::Upsert, or fix the writer.",
window.delete_rows(),
self.cfg.db,
self.cfg.table
)));
}
let loads = self.plan_loads(window)?;
let merged_deletes: usize = loads
.iter()
.filter(|l| l.merge_type == "MERGE")
.map(PlannedLoad::rows)
.sum::<usize>()
.saturating_sub(window.upsert_rows());
let mut jobs_in: Vec<(String, usize, bytes::Bytes)> = Vec::new();
let mut chunk_rows: Vec<usize> = Vec::new();
for (li, load) in loads.iter().enumerate() {
for (body, rows) in encode_chunks(&load.batches, self.cfg.max_chunk_bytes, self.cfg.format)? {
chunk_rows.push(rows);
let label = stream_load_label(
&self.cfg.db,
&self.cfg.table,
window.snapshots.from,
window.snapshots.to,
jobs_in.len(),
);
jobs_in.push((label, li, bytes::Bytes::from(body)));
}
}
let loads_ref = &loads;
let this: &Self = &*self;
let jobs = jobs_in.into_iter().map(move |(label, li, body)| {
let load = &loads_ref[li];
async move {
let r = this.put_chunk(&label, load, body).await;
Ok::<_, anyhow::Error>((label, li, r))
}
});
let done = znippy_zoomies::gatling::io::run_ordered(jobs, self.cfg.max_parallel_loads.max(1), |x| x)
.await
.map_err(|e| BifrostError::Other(format!("{e:#}")))?;
let mut labels = Vec::with_capacity(done.len());
let mut to_commit = Vec::new();
let mut fresh_precommitted = Vec::new();
let mut first_err = None;
let (mut loaded, mut deleted, mut duplicate) = (0usize, 0usize, 0usize);
for (ci, (label, li, r)) in done.into_iter().enumerate() {
match r {
Ok(LoadOutcome::Loaded { rows }) => {
if loads[li].merge_type == "DELETE" {
deleted += rows;
} else {
loaded += rows;
}
if self.cfg.two_phase_commit {
to_commit.push(label.clone());
fresh_precommitted.push(label.clone());
}
}
Ok(LoadOutcome::AlreadyFinished) => {
tracing::info!(
target: "bifrost::doris",
event = "duplicate_label_rejected",
label = %label,
"replay rejected by label — chunk already loaded",
);
duplicate += chunk_rows[ci];
}
Ok(LoadOutcome::AlreadyPending) => {
if !self.cfg.two_phase_commit {
first_err.get_or_insert_with(|| BifrostError::Other(format!(
"doris: label {label} is still RUNNING from an earlier attempt — retry later"
)));
} else {
to_commit.push(label.clone());
}
}
Err(e) => {
first_err.get_or_insert(e);
}
}
labels.push(label);
}
if let Some(e) = first_err {
for label in &fresh_precommitted {
if let Err(abort) = self.txn_op(label, "abort").await {
tracing::warn!(target: "bifrost::doris", event = "abort_failed", label = %label, error = %abort,
"could not abort a pre-committed chunk; the replay will commit it by label");
}
}
return Err(e);
}
for label in &to_commit {
self.txn_op(label, "commit").await?;
}
tracing::info!(
target: "bifrost::doris",
event = "stream_load_window",
table = %window.table,
dest = %format!("{}.{}", self.cfg.db, self.cfg.table),
from = ?window.snapshots.from,
to = window.snapshots.to,
sequence = ?window.sequence,
chunks = labels.len(),
committed = to_commit.len(),
loaded,
deleted,
duplicate,
"stream load complete",
);
let moved = merged_deletes.min(loaded);
loaded -= moved;
deleted += moved;
Ok(token.ack(window.snapshots.to, labels).with_counts(loaded, deleted, duplicate))
}
}
pub fn parse_load_result(label: &str, body: &str) -> Result<(usize, bool)> {
Ok(match parse_load_outcome(label, body)? {
LoadOutcome::Loaded { rows } => (rows, false),
LoadOutcome::AlreadyFinished | LoadOutcome::AlreadyPending => (0, true),
})
}
pub fn parse_load_outcome(label: &str, body: &str) -> Result<LoadOutcome> {
let v: serde_json::Value = serde_json::from_str(body).map_err(|e| {
BifrostError::Other(format!("doris: stream load {label} returned non-JSON: {e} — {body}"))
})?;
if v.get("Status").and_then(serde_json::Value::as_str) == Some("Label Already Exists") {
let existing = v.get("ExistingJobStatus").and_then(serde_json::Value::as_str).unwrap_or("FINISHED");
return Ok(if existing.eq_ignore_ascii_case("FINISHED") {
LoadOutcome::AlreadyFinished
} else {
LoadOutcome::AlreadyPending
});
}
let (rows, _) = parse_fresh_load(label, &v)?;
Ok(LoadOutcome::Loaded { rows })
}
pub fn parse_txn_result(label: &str, op: &str, body: &str) -> Result<()> {
let v: serde_json::Value = serde_json::from_str(body).map_err(|e| {
BifrostError::Other(format!("doris: 2pc {op} of {label} returned non-JSON: {e} — {body}"))
})?;
let field = |k: &str| {
v.as_object()
.and_then(|o| o.iter().find(|(name, _)| name.eq_ignore_ascii_case(k)))
.and_then(|(_, x)| x.as_str())
.unwrap_or("")
.to_string()
};
let status = field("status");
let msg = { let m = field("msg"); if m.is_empty() { field("message") } else { m } };
if status.eq_ignore_ascii_case("success") {
return Ok(());
}
let m = msg.to_ascii_lowercase();
if op == "commit" && m.contains("already") && (m.contains("visible") || m.contains("committed")) {
tracing::info!(target: "bifrost::doris", event = "commit_replayed", label = %label, "already committed");
return Ok(());
}
if op == "abort" && m.contains("already") && m.contains("abort") {
return Ok(());
}
Err(BifrostError::Other(format!("doris: 2pc {op} of {label} failed: {status} — {msg}")))
}
fn parse_fresh_load(label: &str, v: &serde_json::Value) -> Result<(usize, bool)> {
let status = v
.get("Status")
.and_then(serde_json::Value::as_str)
.unwrap_or("");
let msg = v
.get("Message")
.and_then(serde_json::Value::as_str)
.unwrap_or("");
let loaded = v
.get("NumberLoadedRows")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0) as usize;
match status {
"Success" | "Publish Timeout" => {
let filtered = v
.get("NumberFilteredRows")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0);
if filtered > 0 {
return Err(BifrostError::Other(format!(
"doris: stream load {label} filtered {filtered} row(s) — refusing. A \
dropped row is the worst failure a change feed can have. {}",
v.get("ErrorURL")
.and_then(serde_json::Value::as_str)
.unwrap_or("")
)));
}
Ok((loaded, false))
}
"Label Already Exists" => {
tracing::info!(
target: "bifrost::doris",
event = "duplicate_label_rejected",
label = %label,
"replay rejected by label — window already loaded",
);
Ok((0, true))
}
other => Err(BifrostError::Other(format!(
"doris: stream load {label} failed with status `{other}` — {msg}. {}",
v.get("ErrorURL")
.and_then(serde_json::Value::as_str)
.unwrap_or("")
))),
}
}
fn encode_chunks(batches: &[RecordBatch], budget: usize, format: LoadFormat) -> Result<Vec<(Vec<u8>, usize)>> {
let mut out: Vec<(Vec<u8>, usize)> = Vec::new();
let mut pending: Vec<RecordBatch> = Vec::new();
let mut pending_bytes = 0usize;
let rows = |p: &[RecordBatch]| p.iter().map(RecordBatch::num_rows).sum::<usize>();
for batch in batches {
if batch.num_rows() == 0 {
continue;
}
let encoded = encode(format, std::slice::from_ref(batch))?;
if encoded.len() > budget {
if !pending.is_empty() {
out.push((encode(format, &pending)?, rows(&pending)));
pending.clear();
pending_bytes = 0;
}
split_encode(batch, budget, format, &mut out)?;
continue;
}
if pending_bytes + encoded.len() > budget && !pending.is_empty() {
out.push((encode(format, &pending)?, rows(&pending)));
pending.clear();
pending_bytes = 0;
}
pending_bytes += encoded.len();
pending.push(batch.clone());
}
if !pending.is_empty() {
out.push((encode(format, &pending)?, rows(&pending)));
}
Ok(out)
}
fn split_encode(batch: &RecordBatch, budget: usize, format: LoadFormat, out: &mut Vec<(Vec<u8>, usize)>) -> Result<()> {
let encoded = encode(format, std::slice::from_ref(batch))?;
if encoded.len() <= budget || batch.num_rows() <= 1 {
out.push((encoded, batch.num_rows()));
return Ok(());
}
let half = batch.num_rows() / 2;
split_encode(&batch.slice(0, half), budget, format, out)?;
split_encode(&batch.slice(half, batch.num_rows() - half), budget, format, out)
}
fn encode(format: LoadFormat, batches: &[RecordBatch]) -> Result<Vec<u8>> {
match format {
LoadFormat::Json => encode_json(batches),
LoadFormat::Arrow => encode_arrow(batches),
}
}
fn encode_arrow(batches: &[RecordBatch]) -> Result<Vec<u8>> {
let Some(first) = batches.first() else { return Ok(Vec::new()) };
let bytes: usize = batches.iter().map(RecordBatch::get_array_memory_size).sum();
let mut writer = arrow::ipc::writer::StreamWriter::try_new(Vec::with_capacity(bytes), &first.schema())?;
for b in batches {
writer.write(b)?;
}
writer.finish()?;
Ok(writer.into_inner()?)
}
fn encode_json(batches: &[RecordBatch]) -> Result<Vec<u8>> {
let mut writer = arrow::json::ArrayWriter::new(Vec::new());
for b in batches {
writer.write(b)?;
}
writer.finish()?;
Ok(writer.into_inner())
}
fn base64_std(input: &[u8]) -> String {
const T: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
for c in input.chunks(3) {
let b = [c[0], *c.get(1).unwrap_or(&0), *c.get(2).unwrap_or(&0)];
let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]);
out.push(T[(n >> 18 & 63) as usize] as char);
out.push(T[(n >> 12 & 63) as usize] as char);
out.push(if c.len() > 1 {
T[(n >> 6 & 63) as usize] as char
} else {
'='
});
out.push(if c.len() > 2 {
T[(n & 63) as usize] as char
} else {
'='
});
}
out
}
pub struct ReqwestTransport {
client: reqwest::Client,
}
impl ReqwestTransport {
pub fn new(cfg: &DorisConfig) -> Result<Self> {
let client = tls_roots(reqwest::Client::builder())
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(cfg.connect_timeout)
.timeout(cfg.request_timeout)
.build()
.map_err(|e| BifrostError::Other(format!("doris: http client: {e}")))?;
Ok(Self { client })
}
}
#[async_trait]
impl StreamLoadTransport for ReqwestTransport {
async fn send(&self, req: &StreamLoadRequest) -> Result<StreamLoadResponse> {
let method = reqwest::Method::from_bytes(req.method.as_bytes())
.map_err(|e| BifrostError::Other(format!("doris: bad method: {e}")))?;
let mut rb = self.client.request(method, &req.url);
for (k, v) in &req.headers {
rb = rb.header(k.as_str(), v.as_str());
}
if !req.body.is_empty() {
rb = rb.body(req.body.clone()); }
let resp = rb
.send()
.await
.map_err(|e| BifrostError::Other(format!("doris: {} {}: {e}", req.method, req.url)))?;
let status = resp.status().as_u16();
let location = resp
.headers()
.get(reqwest::header::LOCATION)
.and_then(|v| v.to_str().ok())
.map(str::to_owned);
let body = resp
.text()
.await
.map_err(|e| BifrostError::Other(format!("doris: reading response body: {e}")))?;
Ok(StreamLoadResponse {
status,
location,
body,
})
}
}
#[allow(dead_code)]
fn tls_roots(b: reqwest::ClientBuilder) -> reqwest::ClientBuilder {
let _ = rustls::crypto::ring::default_provider().install_default();
b.tls_certs_only(
webpki_root_certs::TLS_SERVER_ROOT_CERTS
.iter()
.filter_map(|c| reqwest::Certificate::from_der(c).ok()),
)
}