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: Vec<u8>,
}
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,
}
impl LoadFormat {
const fn header_value(self) -> &'static str {
match self {
LoadFormat::Json => "json",
}
}
}
#[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 const DEFAULT_MAX_CHUNK_BYTES: usize = 64 * 1024 * 1024;
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,
})
}
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,
}
}
#[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_be_url(mut self, url: impl Into<String>) -> Self {
self.be_url = Some(url.into());
self
}
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: Vec::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 put_chunk(
&self,
label: &str,
merge_type: &str,
body: Vec<u8>,
) -> Result<(usize, bool)> {
let base = self.cfg.be_url.as_deref().unwrap_or(&self.cfg.fe_url);
let mut url = self.cfg.stream_load_url(base);
let mut headers = vec![
("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(),
),
("strip_outer_array".to_string(), "true".to_string()),
("merge_type".to_string(), merge_type.to_string()),
];
headers.push(("max_filter_ratio".to_string(), "0".to_string()));
let mut hops = 0_u8;
let resp = 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() {
break resp;
}
hops += 1;
if hops > 4 {
return Err(BifrostError::Other(format!(
"doris: more than 4 redirects starting at {}",
self.cfg.stream_load_url(base)
)));
}
url = resp.location.clone().expect("is_redirect checked location");
tracing::debug!(
target: "bifrost::doris",
event = "stream_load_redirect",
to = %url,
"FE redirected the load to a backend",
);
};
if resp.status != 200 {
return Err(BifrostError::Other(format!(
"doris: stream load {label} returned HTTP {} — {}",
resp.status, resp.body
)));
}
parse_load_result(label, &resp.body)
}
async fn ship(
&self,
window: &TableWindow<'_>,
batches: &[RecordBatch],
merge_type: &str,
chunk0: usize,
labels: &mut Vec<String>,
) -> Result<(usize, usize)> {
let bodies = encode_chunks(batches, self.cfg.max_chunk_bytes)?;
let mut loaded = 0;
let mut duplicate = 0;
for (i, body) in bodies.into_iter().enumerate() {
let label = stream_load_label(
&self.cfg.db,
&self.cfg.table,
window.snapshots.from,
window.snapshots.to,
chunk0 + i,
);
let (rows, was_dup) = self.put_chunk(&label, merge_type, body).await?;
if was_dup {
duplicate += rows;
} else {
loaded += rows;
}
labels.push(label);
}
Ok((loaded, duplicate))
}
}
#[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 mut labels = Vec::new();
let (del_loaded, del_dup) = if window.delete_rows() > 0 {
self.ship(window, window.deletes, "DELETE", 0, &mut labels)
.await?
} else {
(0, 0)
};
let chunk0 = labels.len();
let (up_loaded, up_dup) = if window.upsert_rows() > 0 {
self.ship(window, window.upserts, "APPEND", chunk0, &mut labels)
.await?
} else {
(0, 0)
};
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,
chunks = labels.len(),
loaded = up_loaded,
deleted = del_loaded,
duplicate = up_dup + del_dup,
"stream load complete",
);
Ok(token.ack(window.snapshots.to, labels).with_counts(
up_loaded,
del_loaded,
up_dup + del_dup,
))
}
}
pub fn parse_load_result(label: &str, body: &str) -> Result<(usize, bool)> {
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}"
))
})?;
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) -> Result<Vec<Vec<u8>>> {
let mut out: Vec<Vec<u8>> = Vec::new();
let mut pending: Vec<RecordBatch> = Vec::new();
let mut pending_bytes = 0usize;
for batch in batches {
if batch.num_rows() == 0 {
continue;
}
let encoded = encode_json(std::slice::from_ref(batch))?;
if encoded.len() > budget {
if !pending.is_empty() {
out.push(encode_json(&pending)?);
pending.clear();
pending_bytes = 0;
}
split_encode(batch, budget, &mut out)?;
continue;
}
if pending_bytes + encoded.len() > budget && !pending.is_empty() {
out.push(encode_json(&pending)?);
pending.clear();
pending_bytes = 0;
}
pending_bytes += encoded.len();
pending.push(batch.clone());
}
if !pending.is_empty() {
out.push(encode_json(&pending)?);
}
Ok(out)
}
fn split_encode(batch: &RecordBatch, budget: usize, out: &mut Vec<Vec<u8>>) -> Result<()> {
let encoded = encode_json(std::slice::from_ref(batch))?;
if encoded.len() <= budget || batch.num_rows() <= 1 {
out.push(encoded);
return Ok(());
}
let half = batch.num_rows() / 2;
split_encode(&batch.slice(0, half), budget, out)?;
split_encode(&batch.slice(half, batch.num_rows() - half), budget, out)
}
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 = 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,
})
}
}