1use std::sync::Arc;
7
8use arrow_array::builder::{
9 BinaryBuilder, BooleanBuilder, Int64Builder, ListBuilder, StringBuilder, StructBuilder,
10 TimestampMillisecondBuilder, UInt8Builder,
11};
12use arrow_array::ArrayRef;
13use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef, TimeUnit};
14use lapin::message::Delivery;
15use lapin::types::{AMQPValue, ShortString};
16
17use crate::error::ConnectorError;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
21pub enum MetadataColumn {
22 Exchange,
24 RoutingKey,
26 Redelivered,
28 DeliveryCount,
30 MessageId,
32 CorrelationId,
34 Timestamp,
36 ContentType,
38 ContentEncoding,
40 DeliveryMode,
42 Priority,
44 Expiration,
46 ReplyTo,
48 Type,
50 UserId,
52 AppId,
54 Headers,
56}
57
58impl MetadataColumn {
59 const ALL: [Self; 17] = [
60 Self::Exchange,
61 Self::RoutingKey,
62 Self::Redelivered,
63 Self::DeliveryCount,
64 Self::MessageId,
65 Self::CorrelationId,
66 Self::Timestamp,
67 Self::ContentType,
68 Self::ContentEncoding,
69 Self::DeliveryMode,
70 Self::Priority,
71 Self::Expiration,
72 Self::ReplyTo,
73 Self::Type,
74 Self::UserId,
75 Self::AppId,
76 Self::Headers,
77 ];
78
79 #[must_use]
81 pub const fn option_name(self) -> &'static str {
82 match self {
83 Self::Exchange => "exchange",
84 Self::RoutingKey => "routing_key",
85 Self::Redelivered => "redelivered",
86 Self::DeliveryCount => "delivery_count",
87 Self::MessageId => "message_id",
88 Self::CorrelationId => "correlation_id",
89 Self::Timestamp => "timestamp",
90 Self::ContentType => "content_type",
91 Self::ContentEncoding => "content_encoding",
92 Self::DeliveryMode => "delivery_mode",
93 Self::Priority => "priority",
94 Self::Expiration => "expiration",
95 Self::ReplyTo => "reply_to",
96 Self::Type => "type",
97 Self::UserId => "user_id",
98 Self::AppId => "app_id",
99 Self::Headers => "headers",
100 }
101 }
102
103 fn field(self) -> Field {
104 let name = format!("_{}", self.option_name());
105 match self {
106 Self::Exchange | Self::RoutingKey => Field::new(name, DataType::Utf8, false),
107 Self::Redelivered => Field::new(name, DataType::Boolean, false),
108 Self::DeliveryCount => Field::new(name, DataType::Int64, true),
109 Self::Timestamp => {
110 Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true)
111 }
112 Self::DeliveryMode | Self::Priority => Field::new(name, DataType::UInt8, true),
113 Self::Headers => Field::new(
114 name,
115 DataType::List(Arc::new(Field::new_struct("item", header_fields(), false))),
116 false,
117 ),
118 Self::MessageId
119 | Self::CorrelationId
120 | Self::ContentType
121 | Self::ContentEncoding
122 | Self::Expiration
123 | Self::ReplyTo
124 | Self::Type
125 | Self::UserId
126 | Self::AppId => Field::new(name, DataType::Utf8, true),
127 }
128 }
129}
130
131fn header_fields() -> Fields {
132 Fields::from(vec![
133 Field::new("key", DataType::Utf8, false),
134 Field::new("value", DataType::Binary, true),
135 ])
136}
137
138#[derive(Debug, Clone, Default, PartialEq, Eq)]
140pub struct MetadataColumns(Vec<MetadataColumn>);
141
142impl MetadataColumns {
143 pub fn parse(list: Option<&str>) -> Result<Self, ConnectorError> {
148 let mut columns = Vec::new();
149 for name in list
150 .unwrap_or_default()
151 .split(',')
152 .map(str::trim)
153 .filter(|name| !name.is_empty())
154 {
155 let column = MetadataColumn::ALL
156 .into_iter()
157 .find(|column| column.option_name().eq_ignore_ascii_case(name))
158 .ok_or_else(|| {
159 let known: Vec<&str> = MetadataColumn::ALL
160 .iter()
161 .map(|c| c.option_name())
162 .collect();
163 ConnectorError::ConfigurationError(format!(
164 "unknown RabbitMQ metadata column '{name}' (expected one of {})",
165 known.join(", ")
166 ))
167 })?;
168 columns.push(column);
169 }
170 columns.sort_unstable();
171 columns.dedup();
172 Ok(Self(columns))
173 }
174
175 #[must_use]
177 pub fn is_empty(&self) -> bool {
178 self.0.is_empty()
179 }
180
181 pub fn validate_payload(&self, payload: &Schema) -> Result<(), ConnectorError> {
186 let collision = payload.fields().iter().find(|field| {
187 self.0.iter().any(|column| {
188 field
189 .name()
190 .eq_ignore_ascii_case(&format!("_{}", column.option_name()))
191 })
192 });
193 match collision {
194 Some(field) => Err(ConnectorError::SchemaMismatch(format!(
195 "RabbitMQ payload schema contains configured metadata column '{}'",
196 field.name()
197 ))),
198 None => Ok(()),
199 }
200 }
201
202 #[must_use]
204 pub fn output_schema(&self, payload: &SchemaRef) -> SchemaRef {
205 if self.0.is_empty() {
206 return Arc::clone(payload);
207 }
208 let mut fields = payload.fields().to_vec();
209 fields.extend(self.0.iter().map(|column| Arc::new(column.field())));
210 Arc::new(Schema::new_with_metadata(
211 fields,
212 payload.metadata().clone(),
213 ))
214 }
215
216 pub fn payload_schema(&self, output: &SchemaRef) -> Result<SchemaRef, ConnectorError> {
221 let payload_len = output
222 .fields()
223 .len()
224 .checked_sub(self.0.len())
225 .ok_or_else(|| {
226 ConnectorError::SchemaMismatch("source schema lacks its metadata columns".into())
227 })?;
228 let payload = Arc::new(
229 output
230 .project(&(0..payload_len).collect::<Vec<_>>())
231 .map_err(|error| ConnectorError::SchemaMismatch(error.to_string()))?,
232 );
233 if self.output_schema(&payload) != *output {
234 return Err(ConnectorError::SchemaMismatch(
235 "RabbitMQ source schema does not end with its configured metadata columns".into(),
236 ));
237 }
238 Ok(payload)
239 }
240
241 pub(super) fn arrays(&self, deliveries: &[&Delivery]) -> Result<Vec<ArrayRef>, ConnectorError> {
246 self.0
247 .iter()
248 .map(|column| column_array(*column, deliveries))
249 .collect()
250 }
251}
252
253impl std::fmt::Display for MetadataColumns {
254 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
255 let names: Vec<&str> = self.0.iter().map(|column| column.option_name()).collect();
256 f.write_str(&names.join(","))
257 }
258}
259
260fn column_array(
261 column: MetadataColumn,
262 deliveries: &[&Delivery],
263) -> Result<ArrayRef, ConnectorError> {
264 let text = |value: fn(&Delivery) -> Option<&ShortString>| -> ArrayRef {
265 let mut builder = StringBuilder::with_capacity(deliveries.len(), deliveries.len() * 16);
266 for delivery in deliveries {
267 builder.append_option(value(delivery).map(ShortString::as_str));
268 }
269 Arc::new(builder.finish())
270 };
271 let octet = |value: fn(&Delivery) -> Option<u8>| -> ArrayRef {
272 let mut builder = UInt8Builder::with_capacity(deliveries.len());
273 for delivery in deliveries {
274 builder.append_option(value(delivery));
275 }
276 Arc::new(builder.finish())
277 };
278 Ok(match column {
279 MetadataColumn::Exchange => text(|delivery| Some(&delivery.exchange)),
280 MetadataColumn::RoutingKey => text(|delivery| Some(&delivery.routing_key)),
281 MetadataColumn::Redelivered => {
282 let mut builder = BooleanBuilder::with_capacity(deliveries.len());
283 for delivery in deliveries {
284 builder.append_value(delivery.redelivered);
285 }
286 Arc::new(builder.finish())
287 }
288 MetadataColumn::DeliveryCount => {
289 let mut builder = Int64Builder::with_capacity(deliveries.len());
290 for delivery in deliveries {
291 builder.append_option(delivery_count(delivery)?);
292 }
293 Arc::new(builder.finish())
294 }
295 MetadataColumn::MessageId => text(|delivery| delivery.properties.message_id().as_ref()),
296 MetadataColumn::CorrelationId => {
297 text(|delivery| delivery.properties.correlation_id().as_ref())
298 }
299 MetadataColumn::Timestamp => {
300 let mut builder = TimestampMillisecondBuilder::with_capacity(deliveries.len());
301 for delivery in deliveries {
302 let millis = delivery
303 .properties
304 .timestamp()
305 .map(|seconds| {
306 i64::try_from(seconds)
307 .ok()
308 .and_then(|seconds| seconds.checked_mul(1_000))
309 .ok_or_else(|| {
310 ConnectorError::ReadError(format!(
311 "AMQP timestamp {seconds} exceeds the millisecond range"
312 ))
313 })
314 })
315 .transpose()?;
316 builder.append_option(millis);
317 }
318 Arc::new(builder.finish())
319 }
320 MetadataColumn::ContentType => text(|delivery| delivery.properties.content_type().as_ref()),
321 MetadataColumn::ContentEncoding => {
322 text(|delivery| delivery.properties.content_encoding().as_ref())
323 }
324 MetadataColumn::DeliveryMode => octet(|delivery| *delivery.properties.delivery_mode()),
325 MetadataColumn::Priority => octet(|delivery| *delivery.properties.priority()),
326 MetadataColumn::Expiration => text(|delivery| delivery.properties.expiration().as_ref()),
327 MetadataColumn::ReplyTo => text(|delivery| delivery.properties.reply_to().as_ref()),
328 MetadataColumn::Type => text(|delivery| delivery.properties.kind().as_ref()),
329 MetadataColumn::UserId => text(|delivery| delivery.properties.user_id().as_ref()),
330 MetadataColumn::AppId => text(|delivery| delivery.properties.app_id().as_ref()),
331 MetadataColumn::Headers => headers_array(deliveries),
332 })
333}
334
335fn delivery_count(delivery: &Delivery) -> Result<Option<i64>, ConnectorError> {
336 let Some(value) = delivery
337 .properties
338 .headers()
339 .as_ref()
340 .and_then(|headers| headers.inner().get(&ShortString::from("x-delivery-count")))
341 else {
342 return Ok(None);
343 };
344 integer(value).map(Some).ok_or_else(|| {
345 ConnectorError::ReadError("x-delivery-count header is not an integer".into())
346 })
347}
348
349fn integer(value: &AMQPValue) -> Option<i64> {
350 match value {
351 AMQPValue::ShortShortInt(v) => Some(i64::from(*v)),
352 AMQPValue::ShortShortUInt(v) => Some(i64::from(*v)),
353 AMQPValue::ShortInt(v) => Some(i64::from(*v)),
354 AMQPValue::ShortUInt(v) => Some(i64::from(*v)),
355 AMQPValue::LongInt(v) => Some(i64::from(*v)),
356 AMQPValue::LongUInt(v) => Some(i64::from(*v)),
357 AMQPValue::LongLongInt(v) => Some(*v),
358 _ => None,
359 }
360}
361
362fn headers_array(deliveries: &[&Delivery]) -> ArrayRef {
363 let entry = StructBuilder::new(
364 header_fields(),
365 vec![
366 Box::new(StringBuilder::new()),
367 Box::new(BinaryBuilder::new()),
368 ],
369 );
370 let mut builder = ListBuilder::new(entry).with_field(Arc::new(Field::new_struct(
371 "item",
372 header_fields(),
373 false,
374 )));
375 for delivery in deliveries {
376 if let Some(headers) = delivery.properties.headers() {
377 for (key, value) in headers {
378 let entry = builder.values();
379 entry
380 .field_builder::<StringBuilder>(0)
381 .expect("header key builder is Utf8")
382 .append_value(key.as_str());
383 entry
384 .field_builder::<BinaryBuilder>(1)
385 .expect("header value builder is Binary")
386 .append_option(render_header(value));
387 entry.append(true);
388 }
389 }
390 builder.append(true);
391 }
392 Arc::new(builder.finish())
393}
394
395pub(super) fn render_header(value: &AMQPValue) -> Option<Vec<u8>> {
399 match value {
400 AMQPValue::LongString(text) => Some(text.as_bytes().to_vec()),
401 AMQPValue::ShortString(text) => Some(text.as_str().as_bytes().to_vec()),
402 AMQPValue::ByteArray(bytes) => Some(bytes.as_slice().to_vec()),
403 AMQPValue::Void => None,
404 AMQPValue::FieldTable(_) | AMQPValue::FieldArray(_) => {
405 Some(json_value(value).to_string().into_bytes())
406 }
407 scalar => Some(scalar_text(scalar).into_bytes()),
408 }
409}
410
411fn scalar_text(value: &AMQPValue) -> String {
412 if let Some(integer) = integer(value) {
413 return integer.to_string();
414 }
415 match value {
416 AMQPValue::Boolean(flag) => flag.to_string(),
417 AMQPValue::Float(float) => float.to_string(),
418 AMQPValue::Double(double) => double.to_string(),
419 AMQPValue::Timestamp(seconds) => seconds.to_string(),
420 AMQPValue::DecimalValue(decimal) => format!("{}e-{}", decimal.value, decimal.scale),
421 _ => String::new(),
422 }
423}
424
425fn json_value(value: &AMQPValue) -> serde_json::Value {
426 match value {
427 AMQPValue::FieldTable(table) => serde_json::Value::Object(
428 table
429 .into_iter()
430 .map(|(key, value)| (key.as_str().to_string(), json_value(value)))
431 .collect(),
432 ),
433 AMQPValue::FieldArray(array) => {
434 serde_json::Value::Array(array.as_slice().iter().map(json_value).collect())
435 }
436 AMQPValue::Void => serde_json::Value::Null,
437 AMQPValue::Boolean(flag) => serde_json::Value::Bool(*flag),
438 other => integer(other).map_or_else(
439 || {
440 render_header(other).map_or(serde_json::Value::Null, |bytes| {
441 serde_json::Value::String(String::from_utf8_lossy(&bytes).into_owned())
442 })
443 },
444 serde_json::Value::from,
445 ),
446 }
447}
448
449#[cfg(test)]
450mod tests {
451 use super::*;
452 use lapin::types::{DecimalValue, FieldArray, FieldTable, LongString};
453
454 #[test]
455 fn parse_is_canonical_and_rejects_unknown_names() {
456 let columns = MetadataColumns::parse(Some("headers, exchange,exchange")).unwrap();
457 assert_eq!(columns.to_string(), "exchange,headers");
458 assert!(MetadataColumns::parse(Some("partition")).is_err());
459 assert!(MetadataColumns::parse(None).unwrap().is_empty());
460 }
461
462 #[test]
463 fn output_schema_round_trips_and_rejects_collisions() {
464 let columns = MetadataColumns::parse(Some("routing_key,priority")).unwrap();
465 let payload = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
466 let output = columns.output_schema(&payload);
467 assert_eq!(output.fields().len(), 3);
468 assert_eq!(columns.payload_schema(&output).unwrap(), payload);
469 assert!(columns.payload_schema(&payload).is_err());
470 let colliding = Schema::new(vec![Field::new("_ROUTING_KEY", DataType::Utf8, false)]);
471 assert!(columns.validate_payload(&colliding).is_err());
472 }
473
474 #[test]
475 fn header_rendering_keeps_bytes_and_native_scalars() {
476 assert_eq!(
477 render_header(&AMQPValue::LongString(LongString::from(vec![0xff, 0x00]))),
478 Some(vec![0xff, 0x00])
479 );
480 assert_eq!(
481 render_header(&AMQPValue::LongLongInt(-7)),
482 Some(b"-7".to_vec())
483 );
484 assert_eq!(
485 render_header(&AMQPValue::Boolean(true)),
486 Some(b"true".to_vec())
487 );
488 assert_eq!(render_header(&AMQPValue::Void), None);
489 assert_eq!(
490 render_header(&AMQPValue::DecimalValue(DecimalValue {
491 scale: 2,
492 value: 12345
493 })),
494 Some(b"12345e-2".to_vec())
495 );
496 let mut nested = FieldTable::default();
497 nested.insert("n".into(), AMQPValue::LongInt(3));
498 nested.insert(
499 "list".into(),
500 AMQPValue::FieldArray(FieldArray::from(vec![AMQPValue::Boolean(false)])),
501 );
502 assert_eq!(
503 render_header(&AMQPValue::FieldTable(nested)),
504 Some(br#"{"list":[false],"n":3}"#.to_vec())
505 );
506 }
507}