use std::time::Duration;
use barnabas_core::{IsolationLevel, Partitioner};
use kafka_protocol::records::Compression;
use crate::{Consumer, Credentials, Producer, Result, Transport, EARLIEST, LATEST};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StartOffset {
Earliest,
Latest,
At(i64),
}
impl StartOffset {
fn as_i64(self) -> i64 {
match self {
Self::Earliest => EARLIEST,
Self::Latest => LATEST,
Self::At(offset) => offset,
}
}
}
pub struct ConsumerBuilder<T> {
transport: T,
}
impl<T: Transport> ConsumerBuilder<T> {
pub(crate) fn new(transport: T) -> Self {
Self { transport }
}
#[must_use]
pub fn bootstrap<I, S>(self, addrs: I) -> ConsumerNeedsClientId<T>
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
ConsumerNeedsClientId {
transport: self.transport,
bootstrap: addrs.into_iter().map(Into::into).collect(),
}
}
}
pub struct ConsumerNeedsClientId<T> {
transport: T,
bootstrap: Vec<String>,
}
impl<T: Transport> ConsumerNeedsClientId<T> {
#[must_use]
pub fn client_id(self, client_id: impl Into<String>) -> ConsumerReady<T> {
ConsumerReady {
transport: self.transport,
bootstrap: self.bootstrap,
client_id: client_id.into(),
isolation: IsolationLevel::ReadCommitted,
credentials: None,
max_wait: None,
prefetch: None,
incremental: None,
assignments: Vec::new(),
every_partition: Vec::new(),
}
}
}
pub struct ConsumerReady<T> {
transport: T,
bootstrap: Vec<String>,
client_id: String,
isolation: IsolationLevel,
credentials: Option<Credentials>,
max_wait: Option<Duration>,
prefetch: Option<bool>,
incremental: Option<bool>,
assignments: Vec<(String, i32, StartOffset)>,
every_partition: Vec<(String, StartOffset)>,
}
impl<T: Transport> ConsumerReady<T> {
#[must_use]
pub fn isolation(mut self, isolation: IsolationLevel) -> Self {
self.isolation = isolation;
self
}
#[must_use]
pub fn credentials(mut self, credentials: Credentials) -> Self {
self.credentials = Some(credentials);
self
}
#[must_use]
pub fn assign(mut self, topic: impl Into<String>, partition: i32, start: StartOffset) -> Self {
self.assignments.push((topic.into(), partition, start));
self
}
#[must_use]
pub fn assign_range(
mut self,
topic: impl Into<String>,
partitions: impl IntoIterator<Item = i32>,
start: StartOffset,
) -> Self {
let topic = topic.into();
for partition in partitions {
self.assignments.push((topic.clone(), partition, start));
}
self
}
#[must_use]
pub fn assign_all(mut self, topic: impl Into<String>, start: StartOffset) -> Self {
self.every_partition.push((topic.into(), start));
self
}
#[must_use]
pub fn max_wait(mut self, max_wait: Duration) -> Self {
self.max_wait = Some(max_wait);
self
}
#[must_use]
pub fn prefetch(mut self, prefetch: bool) -> Self {
self.prefetch = Some(prefetch);
self
}
#[must_use]
pub fn incremental_fetch(mut self, incremental: bool) -> Self {
self.incremental = Some(incremental);
self
}
pub async fn build(self) -> Result<Consumer<T>> {
let mut consumer = if self.credentials.is_some() {
let mut cluster =
crate::Cluster::connect(self.transport, &self.bootstrap, &self.client_id).await?;
if let Some(credentials) = self.credentials {
cluster.set_credentials(credentials);
}
Consumer::from_cluster(cluster, self.isolation)
} else {
Consumer::new(
self.transport,
&self.bootstrap,
&self.client_id,
self.isolation,
)
.await?
};
if let Some(max_wait) = self.max_wait {
consumer.set_max_wait(max_wait);
}
if let Some(prefetch) = self.prefetch {
consumer.set_prefetch(prefetch);
}
if let Some(incremental) = self.incremental {
consumer.set_incremental_fetch(incremental);
}
for (topic, start) in self.every_partition {
let count = consumer.partition_count(&topic).await?;
for partition in 0..count {
consumer.assign(&topic, partition, start.as_i64()).await?;
}
}
for (topic, partition, start) in self.assignments {
consumer.assign(&topic, partition, start.as_i64()).await?;
}
Ok(consumer)
}
}
pub struct ProducerBuilder<T> {
transport: T,
}
impl<T: Transport> ProducerBuilder<T> {
pub(crate) fn new(transport: T) -> Self {
Self { transport }
}
#[must_use]
pub fn bootstrap<I, S>(self, addrs: I) -> ProducerNeedsClientId<T>
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
ProducerNeedsClientId {
transport: self.transport,
bootstrap: addrs.into_iter().map(Into::into).collect(),
}
}
}
pub struct ProducerNeedsClientId<T> {
transport: T,
bootstrap: Vec<String>,
}
impl<T: Transport> ProducerNeedsClientId<T> {
#[must_use]
pub fn client_id(self, client_id: impl Into<String>) -> ProducerReady<T> {
ProducerReady {
transport: self.transport,
bootstrap: self.bootstrap,
client_id: client_id.into(),
transactional_id: None,
compression: None,
partitioner: None,
max_in_flight: None,
}
}
}
pub struct ProducerReady<T> {
transport: T,
bootstrap: Vec<String>,
client_id: String,
transactional_id: Option<String>,
compression: Option<Compression>,
partitioner: Option<Partitioner>,
max_in_flight: Option<usize>,
}
impl<T: Transport> ProducerReady<T> {
#[must_use]
pub fn transactional_id(mut self, id: impl Into<String>) -> Self {
self.transactional_id = Some(id.into());
self
}
#[must_use]
pub fn compression(mut self, compression: Compression) -> Self {
self.compression = Some(compression);
self
}
#[must_use]
pub fn partitioner(mut self, partitioner: Partitioner) -> Self {
self.partitioner = partitioner.into();
self
}
#[must_use]
pub fn max_in_flight(mut self, max: usize) -> Self {
self.max_in_flight = Some(max);
self
}
pub async fn build(self) -> Result<Producer<T>> {
let mut producer = match &self.transactional_id {
Some(id) => {
Producer::transactional(self.transport, &self.bootstrap, &self.client_id, id)
.await?
}
None => Producer::idempotent(self.transport, &self.bootstrap, &self.client_id).await?,
};
if let Some(compression) = self.compression {
producer.set_compression(compression);
}
if let Some(partitioner) = self.partitioner {
producer.set_partitioner(partitioner);
}
if let Some(max) = self.max_in_flight {
producer.set_max_in_flight(max);
}
Ok(producer)
}
}