ruststream_rdkafka/
context.rs1use bytes::Bytes;
8use ruststream::{BuildContext, Field};
9
10use crate::message::KafkaMessage;
11
12#[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 #[must_use]
25 pub fn topic(&self) -> &str {
26 &self.topic
27 }
28
29 #[must_use]
31 pub fn partition(&self) -> i32 {
32 self.partition
33 }
34
35 #[must_use]
37 pub fn offset(&self) -> i64 {
38 self.offset
39 }
40
41 #[must_use]
43 pub fn timestamp_millis(&self) -> Option<i64> {
44 self.timestamp_millis
45 }
46
47 #[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
66pub mod keys {
68 use ruststream::ContextField;
69
70 use super::{Field, KafkaContext};
71
72 use crate::eos::SourceOffset;
73
74 #[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 #[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 #[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 #[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 #[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 #[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}