Skip to main content

laminar_connectors/rabbitmq/
metadata.rs

1//! Typed AMQP delivery metadata: selectable columns appended after the decoded payload columns.
2//!
3//! Values are the broker/publisher properties as delivered, so a redelivered message reproduces
4//! them. No receive-time value is exposed because it would change on replay.
5
6use 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/// One selectable metadata column, in canonical output order.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
21pub enum MetadataColumn {
22    /// `_exchange`: the exchange the message was published to (empty for the default exchange).
23    Exchange,
24    /// `_routing_key`: the publish routing key.
25    RoutingKey,
26    /// `_redelivered`: the broker's redelivery flag.
27    Redelivered,
28    /// `_delivery_count`: the `x-delivery-count` header quorum queues maintain, else NULL.
29    DeliveryCount,
30    /// `_message_id`: the `message-id` property.
31    MessageId,
32    /// `_correlation_id`: the `correlation-id` property.
33    CorrelationId,
34    /// `_timestamp`: the `timestamp` property (whole seconds) as milliseconds since the epoch.
35    Timestamp,
36    /// `_content_type`: the `content-type` property.
37    ContentType,
38    /// `_content_encoding`: the `content-encoding` property.
39    ContentEncoding,
40    /// `_delivery_mode`: the `delivery-mode` property (1 transient, 2 persistent).
41    DeliveryMode,
42    /// `_priority`: the `priority` property.
43    Priority,
44    /// `_expiration`: the `expiration` property, the per-message TTL in milliseconds as text.
45    Expiration,
46    /// `_reply_to`: the `reply-to` property.
47    ReplyTo,
48    /// `_type`: the `type` property.
49    Type,
50    /// `_user_id`: the broker-validated `user-id` property.
51    UserId,
52    /// `_app_id`: the `app-id` property.
53    AppId,
54    /// `_headers`: every header entry in table order with its value rendered as bytes.
55    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    /// Option name used in `metadata.columns`.
80    #[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/// The metadata columns a source appends, in canonical order without duplicates.
139#[derive(Debug, Clone, Default, PartialEq, Eq)]
140pub struct MetadataColumns(Vec<MetadataColumn>);
141
142impl MetadataColumns {
143    /// Parse a comma-separated `metadata.columns` list such as `routing_key,headers`.
144    ///
145    /// # Errors
146    /// Returns a configuration error for an unknown name.
147    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    /// Whether no metadata column is selected.
176    #[must_use]
177    pub fn is_empty(&self) -> bool {
178        self.0.is_empty()
179    }
180
181    /// Reject a payload field that collides, case-insensitively, with a selected column.
182    ///
183    /// # Errors
184    /// Returns a schema mismatch naming the colliding field.
185    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    /// The payload schema followed by the selected metadata fields.
203    #[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    /// The payload prefix of an output schema produced by [`Self::output_schema`].
217    ///
218    /// # Errors
219    /// Returns a schema mismatch when `output` does not end with the selected metadata fields.
220    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    /// Build the selected metadata arrays, one row per delivery.
242    ///
243    /// # Errors
244    /// Returns a read error when a property cannot be represented without loss.
245    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
395/// Header value bytes: strings and byte arrays verbatim, scalars as their decimal or
396/// `true`/`false` text, timestamps as seconds, decimals as `value` scaled by `10^-scale`
397/// (`12345e-2`), nested tables and arrays as JSON, and `void` as NULL.
398pub(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}