From c10d30bd35494ea336d90d0abf9977349c38d154 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Rou=C3=ABn=C3=A9?=
 <aurelien.github.arouene@rouene.fr>
Date: Mon, 24 Apr 2023 11:27:59 +0200
Subject: [PATCH] enhancement(amqp sink): Support AMQ Properties (content-type)
 in AMQP sink (#17174)

* Support AMQ Properties in AMQP sinks

* Refacto removing properties.rs

* Update src/sinks/amqp/config.rs

Co-authored-by: Stephen Wakely <stephen@lisp.space>

* Update src/sinks/amqp/config.rs

Co-authored-by: neuronull <kyle.criddle@datadoghq.com>

* Update config.rs

Refacto for documentation consistency.

* Rustfmt

* Generate component docs

---------

Co-authored-by: Stephen Wakely <stephen@lisp.space>
Co-authored-by: neuronull <kyle.criddle@datadoghq.com>
---
 src/sinks/amqp/config.rs                      | 32 +++++++++++++++++++
 src/sinks/amqp/request_builder.rs             |  4 +++
 src/sinks/amqp/service.rs                     |  5 ++-
 src/sinks/amqp/sink.rs                        | 18 +++++++++--
 .../reference/components/sinks/base/amqp.cue  | 20 ++++++++++++
 5 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/src/sinks/amqp/config.rs b/src/sinks/amqp/config.rs
index 1a22198cea06d..3af34e943d7b2 100644
--- a/src/sinks/amqp/config.rs
+++ b/src/sinks/amqp/config.rs
@@ -8,12 +8,40 @@ use crate::{
 };
 use codecs::TextSerializerConfig;
 use futures::FutureExt;
+use lapin::{types::ShortString, BasicProperties};
 use std::sync::Arc;
 use vector_config::configurable_component;
 use vector_core::config::AcknowledgementsConfig;
 
 use super::sink::AmqpSink;
 
+/// AMQP properties configuration.
+#[configurable_component]
+#[configurable(title = "Configure the AMQP message properties.")]
+#[derive(Clone, Debug, Default)]
+pub struct AmqpPropertiesConfig {
+    /// Content-Type for the AMQP messages.
+    #[configurable(derived)]
+    pub(crate) content_type: Option<String>,
+
+    /// Content-Encoding for the AMQP messages.
+    #[configurable(derived)]
+    pub(crate) content_encoding: Option<String>,
+}
+
+impl AmqpPropertiesConfig {
+    pub(super) fn build(&self) -> BasicProperties {
+        let mut prop = BasicProperties::default();
+        if let Some(content_type) = &self.content_type {
+            prop = prop.with_content_type(ShortString::from(content_type.clone()));
+        }
+        if let Some(content_encoding) = &self.content_encoding {
+            prop = prop.with_content_encoding(ShortString::from(content_encoding.clone()));
+        }
+        prop
+    }
+}
+
 /// Configuration for the `amqp` sink.
 ///
 /// Supports AMQP version 0.9.1
@@ -26,6 +54,9 @@ pub struct AmqpSinkConfig {
     /// Template used to generate a routing key which corresponds to a queue binding.
     pub(crate) routing_key: Option<Template>,
 
+    /// AMQP message properties.
+    pub(crate) properties: Option<AmqpPropertiesConfig>,
+
     #[serde(flatten)]
     pub(crate) connection: AmqpConfig,
 
@@ -46,6 +77,7 @@ impl Default for AmqpSinkConfig {
         Self {
             exchange: Template::try_from("vector").unwrap(),
             routing_key: None,
+            properties: None,
             encoding: TextSerializerConfig::default().into(),
             connection: AmqpConfig::default(),
             acknowledgements: AcknowledgementsConfig::default(),
diff --git a/src/sinks/amqp/request_builder.rs b/src/sinks/amqp/request_builder.rs
index b4ab8c9dbe133..313b37626aec8 100644
--- a/src/sinks/amqp/request_builder.rs
+++ b/src/sinks/amqp/request_builder.rs
@@ -9,6 +9,7 @@ use crate::{
     },
 };
 use bytes::Bytes;
+use lapin::BasicProperties;
 use std::io;
 use vector_common::{
     finalization::{EventFinalizers, Finalizable},
@@ -20,6 +21,7 @@ use super::{encoder::AmqpEncoder, service::AmqpRequest, sink::AmqpEvent};
 pub(super) struct AmqpMetadata {
     exchange: String,
     routing_key: String,
+    properties: BasicProperties,
     finalizers: EventFinalizers,
 }
 
@@ -55,6 +57,7 @@ impl RequestBuilder<AmqpEvent> for AmqpRequestBuilder {
         let metadata = AmqpMetadata {
             exchange: input.exchange,
             routing_key: input.routing_key,
+            properties: input.properties,
             finalizers: input.event.take_finalizers(),
         };
 
@@ -72,6 +75,7 @@ impl RequestBuilder<AmqpEvent> for AmqpRequestBuilder {
             body,
             amqp_metadata.exchange,
             amqp_metadata.routing_key,
+            amqp_metadata.properties,
             amqp_metadata.finalizers,
             metadata,
         )
diff --git a/src/sinks/amqp/service.rs b/src/sinks/amqp/service.rs
index a569035b004fa..000c753ec97e2 100644
--- a/src/sinks/amqp/service.rs
+++ b/src/sinks/amqp/service.rs
@@ -23,6 +23,7 @@ pub(super) struct AmqpRequest {
     body: Bytes,
     exchange: String,
     routing_key: String,
+    properties: BasicProperties,
     finalizers: EventFinalizers,
     metadata: RequestMetadata,
 }
@@ -32,6 +33,7 @@ impl AmqpRequest {
         body: Bytes,
         exchange: String,
         routing_key: String,
+        properties: BasicProperties,
         finalizers: EventFinalizers,
         metadata: RequestMetadata,
     ) -> Self {
@@ -39,6 +41,7 @@ impl AmqpRequest {
             body,
             exchange,
             routing_key,
+            properties,
             finalizers,
             metadata,
         }
@@ -117,7 +120,7 @@ impl Service<AmqpRequest> for AmqpService {
                     &req.routing_key,
                     BasicPublishOptions::default(),
                     req.body.as_ref(),
-                    BasicProperties::default(),
+                    req.properties,
                 )
                 .await;
 
diff --git a/src/sinks/amqp/sink.rs b/src/sinks/amqp/sink.rs
index 21d14c02cfdea..333b325dbeab4 100644
--- a/src/sinks/amqp/sink.rs
+++ b/src/sinks/amqp/sink.rs
@@ -7,7 +7,7 @@ use crate::{
 use async_trait::async_trait;
 use futures::StreamExt;
 use futures_util::stream::BoxStream;
-use lapin::options::ConfirmSelectOptions;
+use lapin::{options::ConfirmSelectOptions, BasicProperties};
 use serde::Serialize;
 use std::sync::Arc;
 use tower::ServiceBuilder;
@@ -15,8 +15,11 @@ use vector_buffers::EventCount;
 use vector_core::{sink::StreamSink, ByteSizeOf, EstimatedJsonEncodedSizeOf};
 
 use super::{
-    config::AmqpSinkConfig, encoder::AmqpEncoder, request_builder::AmqpRequestBuilder,
-    service::AmqpService, BuildError,
+    config::{AmqpPropertiesConfig, AmqpSinkConfig},
+    encoder::AmqpEncoder,
+    request_builder::AmqpRequestBuilder,
+    service::AmqpService,
+    BuildError,
 };
 
 /// Stores the event together with the rendered exchange and routing_key values.
@@ -29,6 +32,7 @@ pub(super) struct AmqpEvent {
     pub(super) event: Event,
     pub(super) exchange: String,
     pub(super) routing_key: String,
+    pub(super) properties: BasicProperties,
 }
 
 impl EventCount for AmqpEvent {
@@ -54,6 +58,7 @@ pub(super) struct AmqpSink {
     pub(super) channel: Arc<lapin::Channel>,
     exchange: Template,
     routing_key: Option<Template>,
+    properties: Option<AmqpPropertiesConfig>,
     transformer: Transformer,
     encoder: crate::codecs::Encoder<()>,
 }
@@ -81,6 +86,7 @@ impl AmqpSink {
             channel: Arc::new(channel),
             exchange: config.exchange,
             routing_key: config.routing_key,
+            properties: config.properties,
             transformer,
             encoder,
         })
@@ -115,10 +121,16 @@ impl AmqpSink {
                 .ok()?,
         };
 
+        let properties = match &self.properties {
+            None => BasicProperties::default(),
+            Some(prop) => prop.build(),
+        };
+
         Some(AmqpEvent {
             event,
             exchange,
             routing_key,
+            properties,
         })
     }
 
diff --git a/website/cue/reference/components/sinks/base/amqp.cue b/website/cue/reference/components/sinks/base/amqp.cue
index 64e1083c4a7a9..cd201bb517f08 100644
--- a/website/cue/reference/components/sinks/base/amqp.cue
+++ b/website/cue/reference/components/sinks/base/amqp.cue
@@ -187,6 +187,26 @@ base: components: sinks: amqp: configuration: {
 		required:    true
 		type: string: syntax: "template"
 	}
+	properties: {
+		description: """
+			Configure the AMQP message properties.
+
+			AMQP message properties.
+			"""
+		required: false
+		type: object: options: {
+			content_encoding: {
+				description: "Content-Encoding for the AMQP messages."
+				required:    false
+				type: string: {}
+			}
+			content_type: {
+				description: "Content-Type for the AMQP messages."
+				required:    false
+				type: string: {}
+			}
+		}
+	}
 	routing_key: {
 		description: "Template used to generate a routing key which corresponds to a queue binding."
 		required:    false