Skip to main content

ruststream_rdkafka/
context.rs

1//! Per-delivery context fields exposed to handlers.
2//!
3//! [`KafkaContext`] carries the Kafka delivery metadata that is not part of the payload or the
4//! headers. Request it in a handler by typing the context parameter as
5//! `Context<'_, KafkaContext>` and read individual fields with the zero-sized keys in [`keys`].
6
7use bytes::Bytes;
8use ruststream::{BuildContext, Field};
9
10use crate::message::KafkaMessage;
11
12/// Native Kafka delivery metadata, built once per delivery.
13#[derive(Debug, Clone, Default, PartialEq, Eq)]
14pub struct KafkaContext {
15    topic: String,
16    partition: i32,
17    offset: i64,
18    timestamp_millis: Option<i64>,
19    key: Option<Bytes>,
20}
21
22impl KafkaContext {
23    /// The topic the record was consumed from.
24    #[must_use]
25    pub fn topic(&self) -> &str {
26        &self.topic
27    }
28
29    /// The partition the record was consumed from.
30    #[must_use]
31    pub fn partition(&self) -> i32 {
32        self.partition
33    }
34
35    /// The record's offset within its partition.
36    #[must_use]
37    pub fn offset(&self) -> i64 {
38        self.offset
39    }
40
41    /// The record's timestamp in milliseconds since the epoch, when the broker provided one.
42    #[must_use]
43    pub fn timestamp_millis(&self) -> Option<i64> {
44        self.timestamp_millis
45    }
46
47    /// The record key, or `None` for keyless records.
48    #[must_use]
49    pub fn key(&self) -> Option<&[u8]> {
50        self.key.as_deref()
51    }
52}
53
54impl BuildContext<KafkaMessage> for KafkaContext {
55    fn build(msg: &KafkaMessage) -> Self {
56        Self {
57            topic: msg.topic().to_owned(),
58            partition: msg.partition(),
59            offset: msg.offset(),
60            timestamp_millis: msg.timestamp_millis(),
61            key: msg.key().map(Bytes::copy_from_slice),
62        }
63    }
64}
65
66/// Zero-sized [`Field`] keys reading one [`KafkaContext`] field each.
67pub mod keys {
68    use ruststream::ContextField;
69
70    use super::{Field, KafkaContext};
71
72    use crate::eos::SourceOffset;
73
74    /// Reads the source topic name.
75    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
76    pub struct Topic;
77
78    impl Field<KafkaContext> for Topic {
79        type Value<'a> = &'a str;
80
81        fn get(self, src: &KafkaContext) -> &str {
82            src.topic()
83        }
84    }
85
86    impl ContextField for Topic {
87        type Context = KafkaContext;
88        type Value = String;
89        fn read(self, src: &KafkaContext) -> String {
90            src.topic().to_owned()
91        }
92    }
93
94    /// Reads the source partition.
95    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
96    pub struct Partition;
97
98    impl Field<KafkaContext> for Partition {
99        type Value<'a> = i32;
100
101        fn get(self, src: &KafkaContext) -> i32 {
102            src.partition()
103        }
104    }
105
106    impl ContextField for Partition {
107        type Context = KafkaContext;
108        type Value = i32;
109        fn read(self, src: &KafkaContext) -> i32 {
110            src.partition()
111        }
112    }
113
114    /// Reads the record offset.
115    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
116    pub struct Offset;
117
118    impl Field<KafkaContext> for Offset {
119        type Value<'a> = i64;
120
121        fn get(self, src: &KafkaContext) -> i64 {
122            src.offset()
123        }
124    }
125
126    impl ContextField for Offset {
127        type Context = KafkaContext;
128        type Value = i64;
129        fn read(self, src: &KafkaContext) -> i64 {
130            src.offset()
131        }
132    }
133
134    /// Reads the record timestamp in milliseconds since the epoch, when present.
135    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
136    pub struct TimestampMillis;
137
138    impl Field<KafkaContext> for TimestampMillis {
139        type Value<'a> = Option<i64>;
140
141        fn get(self, src: &KafkaContext) -> Option<i64> {
142            src.timestamp_millis()
143        }
144    }
145
146    impl ContextField for TimestampMillis {
147        type Context = KafkaContext;
148        type Value = Option<i64>;
149        fn read(self, src: &KafkaContext) -> Option<i64> {
150            src.timestamp_millis()
151        }
152    }
153
154    /// Reads the record key, when present.
155    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
156    pub struct Key;
157
158    impl Field<KafkaContext> for Key {
159        type Value<'a> = Option<&'a [u8]>;
160
161        fn get(self, src: &KafkaContext) -> Option<&[u8]> {
162            src.key()
163        }
164    }
165
166    impl ContextField for Key {
167        type Context = KafkaContext;
168        type Value = Option<Vec<u8>>;
169        fn read(self, src: &KafkaContext) -> Option<Vec<u8>> {
170            src.key().map(<[u8]>::to_vec)
171        }
172    }
173
174    /// Reads the delivery's source coordinates as one value, the form
175    /// [`EosPipeline::publish`](crate::EosPipeline::publish) takes.
176    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
177    pub struct Source;
178
179    impl Field<KafkaContext> for Source {
180        type Value<'a> = SourceOffset;
181
182        fn get(self, src: &KafkaContext) -> SourceOffset {
183            SourceOffset::new(src.topic(), src.partition(), src.offset())
184        }
185    }
186
187    impl ContextField for Source {
188        type Context = KafkaContext;
189        type Value = SourceOffset;
190        fn read(self, src: &KafkaContext) -> SourceOffset {
191            SourceOffset::new(src.topic(), src.partition(), src.offset())
192        }
193    }
194}