use crate::proto::{GetProducerOffsetsResponse, PbBucketOffset, PbProducerTableOffsets};
use crate::{BucketId, PartitionId, TableId};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BucketOffset {
pub partition_id: Option<PartitionId>,
pub bucket_id: BucketId,
pub log_end_offset: Option<i64>,
}
impl BucketOffset {
pub fn to_pb(&self) -> PbBucketOffset {
PbBucketOffset {
partition_id: self.partition_id,
bucket_id: self.bucket_id,
log_end_offset: self.log_end_offset,
}
}
pub fn from_pb(pb: &PbBucketOffset) -> Self {
Self {
partition_id: pb.partition_id,
bucket_id: pb.bucket_id,
log_end_offset: pb.log_end_offset,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProducerTableOffsets {
pub table_id: TableId,
pub bucket_offsets: Vec<BucketOffset>,
}
impl ProducerTableOffsets {
pub fn to_pb(&self) -> PbProducerTableOffsets {
PbProducerTableOffsets {
table_id: self.table_id,
bucket_offsets: self
.bucket_offsets
.iter()
.map(BucketOffset::to_pb)
.collect(),
}
}
pub fn from_pb(pb: &PbProducerTableOffsets) -> Self {
Self {
table_id: pb.table_id,
bucket_offsets: pb
.bucket_offsets
.iter()
.map(BucketOffset::from_pb)
.collect(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProducerOffsets {
pub producer_id: Option<String>,
pub expiration_time: Option<i64>,
pub table_offsets: Vec<ProducerTableOffsets>,
}
impl ProducerOffsets {
pub fn from_pb(pb: &GetProducerOffsetsResponse) -> Self {
Self {
producer_id: pb.producer_id.clone(),
expiration_time: pb.expiration_time,
table_offsets: pb
.table_offsets
.iter()
.map(ProducerTableOffsets::from_pb)
.collect(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bucket_offset_roundtrip() {
for b in [
BucketOffset {
partition_id: None,
bucket_id: 0,
log_end_offset: None,
},
BucketOffset {
partition_id: Some(42),
bucket_id: 3,
log_end_offset: Some(1234),
},
] {
assert_eq!(BucketOffset::from_pb(&b.to_pb()), b);
}
}
#[test]
fn test_producer_table_offsets_roundtrip() {
let t = ProducerTableOffsets {
table_id: 5,
bucket_offsets: vec![
BucketOffset {
partition_id: None,
bucket_id: 0,
log_end_offset: Some(100),
},
BucketOffset {
partition_id: Some(7),
bucket_id: 1,
log_end_offset: None,
},
],
};
assert_eq!(ProducerTableOffsets::from_pb(&t.to_pb()), t);
}
}