Skip to main content

kafrust_protocol/api/
add_partitions_to_txn.rs

1use crate::codec::{Decoder, Encoder};
2use crate::error::Result;
3use crate::header::RequestHeader;
4
5pub const API_KEY: i16 = 24;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct AddPartitionsToTxnRequestV0 {
9    pub correlation_id: i32,
10    pub client_id: Option<String>,
11    pub transactional_id: String,
12    pub producer_id: i64,
13    pub producer_epoch: i16,
14    pub topics: Vec<AddPartitionsToTxnTopic>,
15}
16
17impl AddPartitionsToTxnRequestV0 {
18    pub fn encode(&self) -> Result<Vec<u8>> {
19        let mut encoder = Encoder::new();
20        RequestHeader {
21            api_key: API_KEY,
22            api_version: 0,
23            correlation_id: self.correlation_id,
24            client_id: self.client_id.clone(),
25        }
26        .encode_v1(&mut encoder)?;
27        encoder.write_string(&self.transactional_id)?;
28        encoder.write_i64(self.producer_id);
29        encoder.write_i16(self.producer_epoch);
30        encoder.write_array(Some(&self.topics), |encoder, topic| topic.encode(encoder))?;
31        Ok(encoder.into_bytes())
32    }
33}
34
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct AddPartitionsToTxnTopic {
37    pub name: String,
38    pub partitions: Vec<i32>,
39}
40
41impl AddPartitionsToTxnTopic {
42    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
43        encoder.write_string(&self.name)?;
44        encoder.write_array(Some(&self.partitions), |encoder, partition| {
45            encoder.write_i32(*partition);
46            Ok(())
47        })
48    }
49
50    fn encode_flexible(&self, encoder: &mut Encoder) -> Result<()> {
51        encoder.write_compact_string(&self.name)?;
52        encoder.write_compact_array(Some(&self.partitions), |encoder, partition| {
53            encoder.write_i32(*partition);
54            Ok(())
55        })?;
56        encoder.write_empty_tagged_fields();
57        Ok(())
58    }
59}
60
61#[derive(Debug, Clone, PartialEq, Eq)]
62pub struct AddPartitionsToTxnResponseV0 {
63    pub throttle_time_ms: i32,
64    pub errors: Vec<AddPartitionsToTxnTopicResult>,
65}
66
67impl AddPartitionsToTxnResponseV0 {
68    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
69        Ok(Self {
70            throttle_time_ms: decoder.read_i32()?,
71            errors: decoder
72                .read_array(
73                    "add partitions to transaction topic results",
74                    AddPartitionsToTxnTopicResult::decode,
75                )?
76                .unwrap_or_default(),
77        })
78    }
79}
80
81#[derive(Debug, Clone, PartialEq, Eq)]
82pub struct AddPartitionsToTxnTopicResult {
83    pub name: String,
84    pub partitions: Vec<AddPartitionsToTxnPartitionResult>,
85}
86
87impl AddPartitionsToTxnTopicResult {
88    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
89        Ok(Self {
90            name: decoder.read_string()?,
91            partitions: decoder
92                .read_array(
93                    "add partitions to transaction partition results",
94                    AddPartitionsToTxnPartitionResult::decode,
95                )?
96                .unwrap_or_default(),
97        })
98    }
99}
100
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct AddPartitionsToTxnPartitionResult {
103    pub partition_index: i32,
104    pub error_code: i16,
105}
106
107impl AddPartitionsToTxnPartitionResult {
108    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
109        Ok(Self {
110            partition_index: decoder.read_i32()?,
111            error_code: decoder.read_i16()?,
112        })
113    }
114
115    fn decode_flexible(decoder: &mut Decoder<'_>) -> Result<Self> {
116        let result = Self {
117            partition_index: decoder.read_i32()?,
118            error_code: decoder.read_i16()?,
119        };
120        decoder.read_tagged_fields()?;
121        Ok(result)
122    }
123}
124
125/// AddPartitionsToTxn v3, the flexible form used by current Kafka brokers.
126#[derive(Debug, Clone, PartialEq, Eq)]
127pub struct AddPartitionsToTxnRequestV3 {
128    pub correlation_id: i32,
129    pub client_id: Option<String>,
130    pub transactional_id: String,
131    pub producer_id: i64,
132    pub producer_epoch: i16,
133    pub topics: Vec<AddPartitionsToTxnTopic>,
134}
135
136impl AddPartitionsToTxnRequestV3 {
137    pub fn encode(&self) -> Result<Vec<u8>> {
138        let mut encoder = Encoder::new();
139        RequestHeader {
140            api_key: API_KEY,
141            api_version: 3,
142            correlation_id: self.correlation_id,
143            client_id: self.client_id.clone(),
144        }
145        .encode_v2(&mut encoder)?;
146        encoder.write_compact_string(&self.transactional_id)?;
147        encoder.write_i64(self.producer_id);
148        encoder.write_i16(self.producer_epoch);
149        encoder.write_compact_array(Some(&self.topics), |encoder, topic| {
150            topic.encode_flexible(encoder)
151        })?;
152        encoder.write_empty_tagged_fields();
153        Ok(encoder.into_bytes())
154    }
155}
156
157/// AddPartitionsToTxn v3 response with flexible tagged fields.
158#[derive(Debug, Clone, PartialEq, Eq)]
159pub struct AddPartitionsToTxnResponseV3 {
160    pub throttle_time_ms: i32,
161    pub errors: Vec<AddPartitionsToTxnTopicResult>,
162}
163
164impl AddPartitionsToTxnResponseV3 {
165    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
166        let throttle_time_ms = decoder.read_i32()?;
167        let errors = decoder
168            .read_compact_array("add partitions to transaction topic results", |decoder| {
169                let name = decoder.read_compact_string()?;
170                let partitions = decoder
171                    .read_compact_array(
172                        "add partitions to transaction partition results",
173                        AddPartitionsToTxnPartitionResult::decode_flexible,
174                    )?
175                    .unwrap_or_default();
176                decoder.read_tagged_fields()?;
177                Ok(AddPartitionsToTxnTopicResult { name, partitions })
178            })?
179            .unwrap_or_default();
180        decoder.read_tagged_fields()?;
181        Ok(Self {
182            throttle_time_ms,
183            errors,
184        })
185    }
186}
187
188#[cfg(test)]
189#[allow(clippy::unwrap_used)]
190mod tests {
191    use super::{
192        AddPartitionsToTxnRequestV0, AddPartitionsToTxnRequestV3, AddPartitionsToTxnResponseV0,
193        AddPartitionsToTxnResponseV3, AddPartitionsToTxnTopic, API_KEY,
194    };
195    use crate::codec::{Decoder, Encoder};
196
197    #[test]
198    fn encodes_add_partitions_to_txn_v0_request() {
199        let request = AddPartitionsToTxnRequestV0 {
200            correlation_id: 41,
201            client_id: Some("kafrust".to_owned()),
202            transactional_id: "orders-tx".to_owned(),
203            producer_id: 42,
204            producer_epoch: 3,
205            topics: vec![AddPartitionsToTxnTopic {
206                name: "orders".to_owned(),
207                partitions: vec![0, 2],
208            }],
209        };
210        let encoded = request.encode().unwrap();
211
212        assert_eq!(&encoded[0..8], &[0, 24, 0, 0, 0, 0, 0, 41]);
213        assert_eq!(
214            &encoded[encoded.len() - 12..],
215            &[0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 2]
216        );
217        assert_eq!(API_KEY, 24);
218    }
219
220    #[test]
221    fn decodes_add_partitions_to_txn_v0_response() {
222        let bytes = [
223            0, 0, 0, 7, // throttle time
224            0, 0, 0, 1, // topic count
225            0, 6, b'o', b'r', b'd', b'e', b'r', b's', // topic
226            0, 0, 0, 2, // partition count
227            0, 0, 0, 0, 0, 0, // partition 0, success
228            0, 0, 0, 2, 0, 47, // partition 2, invalid producer epoch
229        ];
230        let mut decoder = Decoder::new(&bytes);
231        let response = AddPartitionsToTxnResponseV0::decode_body(&mut decoder).unwrap();
232
233        assert_eq!(response.throttle_time_ms, 7);
234        assert_eq!(response.errors[0].name, "orders");
235        assert_eq!(response.errors[0].partitions[0].error_code, 0);
236        assert_eq!(response.errors[0].partitions[1].partition_index, 2);
237        assert_eq!(response.errors[0].partitions[1].error_code, 47);
238        assert!(decoder.is_empty());
239    }
240
241    #[test]
242    fn encodes_add_partitions_to_txn_v3_request_with_flexible_fields() {
243        let request = AddPartitionsToTxnRequestV3 {
244            correlation_id: 42,
245            client_id: Some("kafrust".to_owned()),
246            transactional_id: "orders-tx".to_owned(),
247            producer_id: 42,
248            producer_epoch: 3,
249            topics: vec![AddPartitionsToTxnTopic {
250                name: "orders".to_owned(),
251                partitions: vec![0, 2],
252            }],
253        };
254        let encoded = request.encode().unwrap();
255
256        assert_eq!(&encoded[0..8], &[0, 24, 0, 3, 0, 0, 0, 42]);
257        assert!(encoded
258            .windows(b"orders".len())
259            .any(|window| window == b"orders"));
260        assert_eq!(encoded.last(), Some(&0));
261    }
262
263    #[test]
264    fn decodes_add_partitions_to_txn_v3_response_with_tagged_fields() {
265        let mut bytes = Encoder::new();
266        bytes.write_i32(7);
267        bytes
268            .write_compact_array(Some(&[()]), |encoder, ()| {
269                encoder.write_compact_string("orders")?;
270                encoder.write_compact_array(Some(&[()]), |encoder, ()| {
271                    encoder.write_i32(2);
272                    encoder.write_i16(47);
273                    encoder.write_empty_tagged_fields();
274                    Ok(())
275                })?;
276                encoder.write_empty_tagged_fields();
277                Ok(())
278            })
279            .unwrap();
280        bytes.write_empty_tagged_fields();
281        let encoded = bytes.into_bytes();
282        let mut decoder = Decoder::new(&encoded);
283
284        let response = AddPartitionsToTxnResponseV3::decode_body(&mut decoder).unwrap();
285
286        assert_eq!(response.throttle_time_ms, 7);
287        assert_eq!(response.errors[0].name, "orders");
288        assert_eq!(response.errors[0].partitions[0].partition_index, 2);
289        assert_eq!(response.errors[0].partitions[0].error_code, 47);
290        assert!(decoder.is_empty());
291    }
292}