use super::sink::{PageWrite, RowPolicy, SinkDriver};
use crate::check::{CheckContext, CheckReport};
use crate::dlq::BatchAtomicity;
use crate::error::FaucetError;
use crate::io::{CallContext, CallSlot};
use crate::traits::RowOutcome;
use async_trait::async_trait;
use futures::stream::{FuturesUnordered, StreamExt};
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum HttpBatching {
PerRecord {
concurrency: usize,
},
Array {
batch_size: usize,
},
}
#[async_trait]
pub trait HttpDriver: Send + Sync {
type Session: Send + Sync;
fn connector_name(&self) -> &'static str;
fn config_schema(&self) -> Value {
serde_json::json!({"type": "object", "properties": {}})
}
fn dataset_uri(&self) -> String {
format!("{}://unknown", self.connector_name())
}
fn call_context(&self) -> CallContext {
CallContext::new()
}
async fn check(
&self,
calls: &CallSlot,
ctx: &CheckContext,
) -> Result<CheckReport, FaucetError> {
let _ = (calls, ctx);
Ok(CheckReport::not_implemented())
}
fn batch_atomicity(&self) -> BatchAtomicity {
BatchAtomicity::BestEffort
}
fn batching(&self) -> HttpBatching;
async fn session(&self, calls: &CallSlot) -> Result<Self::Session, FaucetError>;
async fn send(
&self,
calls: &CallSlot,
session: &Self::Session,
body: &Value,
) -> Result<(), FaucetError>;
}
pub struct HttpSinkDriver<D> {
driver: D,
}
impl<D: HttpDriver> HttpSinkDriver<D> {
pub fn new(driver: D) -> Self {
Self { driver }
}
pub fn driver(&self) -> &D {
&self.driver
}
pub fn driver_mut(&mut self) -> &mut D {
&mut self.driver
}
async fn per_record(
&self,
calls: &CallSlot,
session: &D::Session,
rows: &[Value],
concurrency: usize,
policy: RowPolicy,
) -> Result<PageWrite, FaucetError> {
let concurrency = concurrency.max(1);
if policy == RowPolicy::WholePage {
let mut in_flight = FuturesUnordered::new();
let mut iter = rows.iter();
for record in iter.by_ref().take(concurrency) {
in_flight.push(self.driver.send(calls, session, record));
}
while let Some(result) = in_flight.next().await {
result?;
if let Some(record) = iter.next() {
in_flight.push(self.driver.send(calls, session, record));
}
}
return Ok(PageWrite::accepted(rows.len()));
}
let pending: Vec<_> = rows
.iter()
.enumerate()
.map(|(i, record)| async move { (i, self.driver.send(calls, session, record).await) })
.collect();
let mut indexed: Vec<(usize, RowOutcome)> = futures::stream::iter(pending)
.buffer_unordered(concurrency)
.collect()
.await;
indexed.sort_by_key(|(i, _)| *i);
Ok(PageWrite::from_outcomes(
indexed.into_iter().map(|(_, o)| o).collect(),
))
}
async fn arrays(
&self,
calls: &CallSlot,
session: &D::Session,
rows: &[Value],
batch_size: usize,
policy: RowPolicy,
) -> Result<PageWrite, FaucetError> {
let mut delivered = 0usize;
let mut failed = Vec::new();
let mut cause: Option<String> = None;
for chunk in crate::batching::chunks(rows, batch_size) {
let start = delivered + failed.len();
if let Some(msg) = &cause {
failed.extend((start..start + chunk.len()).map(|i| {
(
i,
FaucetError::Sink(format!(
"array-mode chunk not sent after earlier failure: {msg}"
)),
)
}));
continue;
}
match self
.driver
.send(calls, session, &Value::Array(chunk.to_vec()))
.await
{
Ok(()) => delivered += chunk.len(),
Err(e) if policy == RowPolicy::WholePage || delivered == 0 => return Err(e),
Err(e) => {
failed.extend((start..start + chunk.len()).map(|i| {
(
i,
FaucetError::Sink(
"array-mode chunk POST failed; rows not delivered".into(),
),
)
}));
cause = Some(e.to_string());
}
}
}
Ok(PageWrite::accepted(delivered).with_failures(failed))
}
}
#[async_trait]
impl<D: HttpDriver> SinkDriver for HttpSinkDriver<D> {
fn connector_name(&self) -> &'static str {
self.driver.connector_name()
}
fn config_schema(&self) -> Value {
self.driver.config_schema()
}
fn dataset_uri(&self) -> String {
self.driver.dataset_uri()
}
fn call_context(&self) -> CallContext {
self.driver.call_context()
}
async fn check(
&self,
calls: &CallSlot,
ctx: &CheckContext,
) -> Result<CheckReport, FaucetError> {
self.driver.check(calls, ctx).await
}
fn batch_atomicity(&self) -> BatchAtomicity {
self.driver.batch_atomicity()
}
async fn write(
&self,
calls: &CallSlot,
rows: &[Value],
policy: RowPolicy,
) -> Result<PageWrite, FaucetError> {
if rows.is_empty() {
return Ok(PageWrite::accepted(0));
}
let session = self.driver.session(calls).await?;
match self.driver.batching() {
HttpBatching::PerRecord { concurrency } => {
self.per_record(calls, &session, rows, concurrency, policy)
.await
}
HttpBatching::Array { batch_size } => {
self.arrays(calls, &session, rows, batch_size, policy).await
}
}
}
}