vector/sinks/aws_s3/

config.rs

use aws_sdk_s3::Client as S3Client;
use tower::ServiceBuilder;
#[cfg(feature = "codecs-parquet")]
use vector_lib::codecs::BatchEncoder;
#[cfg(feature = "codecs-parquet")]
use vector_lib::codecs::encoding::{BatchSerializerConfig, format::ParquetSerializerConfig};
use vector_lib::{
    TimeZone,
    codecs::{
        EncoderKind, TextSerializerConfig,
        encoding::{Framer, FramingConfig},
    },
    configurable::configurable_component,
    sink::VectorSink,
};

use super::sink::S3RequestOptions;
use crate::{
    aws::{AwsAuthentication, RegionOrEndpoint},
    codecs::{Encoder, EncodingConfigWithFraming, SinkType},
    config::{AcknowledgementsConfig, GenerateConfig, Input, ProxyConfig, SinkConfig, SinkContext},
    sinks::{
        Healthcheck,
        s3_common::{
            self,
            config::{RetryStrategy, S3Options},
            partitioner::S3KeyPartitioner,
            service::S3Service,
            sink::S3Sink,
        },
        util::{
            BatchConfig, BulkSizeBasedDefaultBatchSettings, Compression, ServiceBuilderExt,
            TowerRequestConfig, timezone_to_offset,
        },
    },
    template::Template,
    tls::TlsConfig,
};

/// Batch encoding configuration for the `aws_s3` sink.
#[cfg(feature = "codecs-parquet")]
#[configurable_component]
#[derive(Clone, Debug)]
#[serde(tag = "codec", rename_all = "snake_case")]
#[configurable(metadata(
    docs::enum_tag_description = "The codec to use for batch encoding events."
))]
pub enum S3BatchEncoding {
    /// Encodes events in [Apache Parquet][apache_parquet] columnar format.
    ///
    /// [apache_parquet]: https://parquet.apache.org/
    Parquet(ParquetSerializerConfig),
}

/// Configuration for the `aws_s3` sink.
#[configurable_component(sink(
    "aws_s3",
    "Store observability events in the AWS S3 object storage system."
))]
#[derive(Clone, Debug)]
#[serde(deny_unknown_fields)]
pub struct S3SinkConfig {
    /// The S3 bucket name.
    ///
    /// This must not include a leading `s3://` or a trailing `/`.
    #[configurable(metadata(docs::examples = "my-bucket"))]
    pub bucket: String,

    /// A prefix to apply to all object keys.
    ///
    /// Prefixes are useful for partitioning objects, such as by creating an object key that
    /// stores objects under a particular directory. If using a prefix for this purpose, it must end
    /// in `/` to act as a directory path. A trailing `/` is **not** automatically added.
    #[serde(default = "default_key_prefix")]
    #[configurable(metadata(docs::templateable))]
    #[configurable(metadata(docs::examples = "date=%F/hour=%H"))]
    #[configurable(metadata(docs::examples = "year=%Y/month=%m/day=%d"))]
    #[configurable(metadata(docs::examples = "application_id={{ application_id }}/date=%F"))]
    pub key_prefix: String,

    /// The timestamp format for the time component of the object key.
    ///
    /// By default, object keys are appended with a timestamp that reflects when the objects are
    /// sent to S3, such that the resulting object key is functionally equivalent to joining the key
    /// prefix with the formatted timestamp, such as `date=2022-07-18/1658176486`.
    ///
    /// This would represent a `key_prefix` set to `date=%F/` and the timestamp of Mon Jul 18 2022
    /// 20:34:44 GMT+0000, with the `filename_time_format` being set to `%s`, which renders
    /// timestamps in seconds since the Unix epoch.
    ///
    /// Supports the common [`strftime`][chrono_strftime_specifiers] specifiers found in most
    /// languages.
    ///
    /// When set to an empty string, no timestamp is appended to the key prefix.
    ///
    /// [chrono_strftime_specifiers]: https://docs.rs/chrono/latest/chrono/format/strftime/index.html#specifiers
    #[serde(default = "default_filename_time_format")]
    pub filename_time_format: String,

    /// Whether or not to append a UUID v4 token to the end of the object key.
    ///
    /// The UUID is appended to the timestamp portion of the object key, such that if the object key
    /// generated is `date=2022-07-18/1658176486`, setting this field to `true` results
    /// in an object key that looks like `date=2022-07-18/1658176486-30f6652c-71da-4f9f-800d-a1189c47c547`.
    ///
    /// This ensures there are no name collisions, and can be useful in high-volume workloads where
    /// object keys must be unique.
    #[serde(default = "crate::serde::default_true")]
    #[configurable(metadata(docs::human_name = "Append UUID to Filename"))]
    pub filename_append_uuid: bool,

    /// The filename extension to use in the object key.
    ///
    /// This overrides setting the extension based on the configured `compression`.
    #[configurable(metadata(docs::examples = "json"))]
    pub filename_extension: Option<String>,

    #[serde(flatten)]
    pub options: S3Options,

    #[serde(flatten)]
    pub region: RegionOrEndpoint,

    #[serde(flatten)]
    pub encoding: EncodingConfigWithFraming,

    /// Batch encoding configuration for columnar formats.
    ///
    /// When set, events are encoded together as a batch in a columnar format (Parquet)
    /// instead of the standard per-event framing-based encoding. The columnar format handles
    /// its own internal compression, so the top-level `compression` setting is bypassed.
    #[cfg(feature = "codecs-parquet")]
    #[configurable(derived)]
    #[serde(default)]
    pub batch_encoding: Option<S3BatchEncoding>,

    /// Compression configuration.
    ///
    /// All compression algorithms use the default compression level unless otherwise specified.
    ///
    /// Some cloud storage API clients and browsers handle decompression transparently, so
    /// depending on how they are accessed, files may not always appear to be compressed.
    #[configurable(derived)]
    #[serde(default = "Compression::gzip_default")]
    pub compression: Compression,

    #[configurable(derived)]
    #[serde(default)]
    pub batch: BatchConfig<BulkSizeBasedDefaultBatchSettings>,

    #[configurable(derived)]
    #[serde(default)]
    pub request: TowerRequestConfig,

    #[configurable(derived)]
    pub tls: Option<TlsConfig>,

    #[configurable(derived)]
    #[serde(default)]
    pub auth: AwsAuthentication,

    #[configurable(derived)]
    #[serde(
        default,
        deserialize_with = "crate::serde::bool_or_struct",
        skip_serializing_if = "crate::serde::is_default"
    )]
    pub acknowledgements: AcknowledgementsConfig,

    #[configurable(derived)]
    #[serde(default)]
    pub timezone: Option<TimeZone>,

    /// Specifies which addressing style to use.
    ///
    /// This controls if the bucket name is in the hostname or part of the URL.
    #[serde(default = "crate::serde::default_true")]
    pub force_path_style: bool,

    /// Specifies retry strategy for failed requests.
    ///
    /// By default, the sink only retries attempts it deems possible to retry.
    /// These settings extend the default behavior.
    #[configurable(derived)]
    #[serde(default, skip_serializing_if = "vector_lib::serde::is_default")]
    pub retry_strategy: RetryStrategy,
}

pub(super) fn default_key_prefix() -> String {
    "date=%F".to_string()
}

pub(super) fn default_filename_time_format() -> String {
    "%s".to_string()
}

impl GenerateConfig for S3SinkConfig {
    fn generate_config() -> toml::Value {
        toml::Value::try_from(Self {
            bucket: "".to_owned(),
            key_prefix: default_key_prefix(),
            filename_time_format: default_filename_time_format(),
            filename_append_uuid: true,
            filename_extension: None,
            options: S3Options::default(),
            region: RegionOrEndpoint::default(),
            encoding: (None::<FramingConfig>, TextSerializerConfig::default()).into(),
            #[cfg(feature = "codecs-parquet")]
            batch_encoding: None,
            compression: Compression::gzip_default(),
            batch: BatchConfig::default(),
            request: TowerRequestConfig::default(),
            tls: Some(TlsConfig::default()),
            auth: AwsAuthentication::default(),
            acknowledgements: Default::default(),
            timezone: Default::default(),
            force_path_style: Default::default(),
            retry_strategy: Default::default(),
        })
        .unwrap()
    }
}

#[async_trait::async_trait]
#[typetag::serde(name = "aws_s3")]
impl SinkConfig for S3SinkConfig {
    async fn build(&self, cx: SinkContext) -> crate::Result<(VectorSink, Healthcheck)> {
        let service = self.create_service(&cx.proxy).await?;
        let healthcheck = self.build_healthcheck(service.client())?;
        let sink = self.build_processor(service, cx)?;
        Ok((sink, healthcheck))
    }

    fn input(&self) -> Input {
        #[cfg(feature = "codecs-parquet")]
        if let Some(batch_encoding) = &self.batch_encoding {
            let S3BatchEncoding::Parquet(parquet_config) = batch_encoding;
            let resolved = BatchSerializerConfig::Parquet(parquet_config.clone());
            return Input::new(resolved.input_type());
        }
        Input::new(self.encoding.config().1.input_type())
    }

    fn acknowledgements(&self) -> &AcknowledgementsConfig {
        &self.acknowledgements
    }
}

impl S3SinkConfig {
    pub fn build_processor(
        &self,
        service: S3Service,
        cx: SinkContext,
    ) -> crate::Result<VectorSink> {
        // Build our S3 client/service, which is what we'll ultimately feed
        // requests into in order to ship files to S3.  We build this here in
        // order to configure the client/service with retries, concurrency
        // limits, rate limits, and whatever else the client should have.
        let request_limits = self.request.into_settings();
        let retry_strategy = self.retry_strategy.clone();
        let service = ServiceBuilder::new()
            .settings(request_limits, retry_strategy)
            .service(service);

        let offset = self
            .timezone
            .or(cx.globals.timezone)
            .and_then(timezone_to_offset);

        // Configure our partitioning/batching.
        let batch_settings = self.batch.into_batcher_settings()?;

        let key_prefix = Template::try_from(self.key_prefix.clone())?.with_tz_offset(offset);

        let ssekms_key_id = self
            .options
            .ssekms_key_id
            .as_ref()
            .cloned()
            .map(|ssekms_key_id| Template::try_from(ssekms_key_id.as_str()))
            .transpose()?;

        let partitioner = S3KeyPartitioner::new(key_prefix, ssekms_key_id, None);

        let transformer = self.encoding.transformer();

        // When batch_encoding is configured (e.g., Parquet), use batch mode
        // with internal compression and appropriate file extension.
        #[cfg(feature = "codecs-parquet")]
        if let Some(batch_encoding) = &self.batch_encoding {
            let S3BatchEncoding::Parquet(parquet_config) = batch_encoding;
            let resolved_batch_config = BatchSerializerConfig::Parquet(parquet_config.clone());

            let batch_serializer = resolved_batch_config.build_batch_serializer()?;
            let batch_encoder = BatchEncoder::new(batch_serializer);

            // Auto-detect Content-Type from batch format. Users can still
            // override via `options.content_type`; we only set it when unset.
            let mut api_options = self.options.clone();
            if api_options.content_type.is_none() {
                api_options.content_type = batch_encoder.content_type().map(|s| s.to_string());
            }

            let encoder = EncoderKind::Batch(batch_encoder);

            let filename_extension = self.filename_extension.clone().or_else(|| {
                Some(
                    match batch_encoding {
                        S3BatchEncoding::Parquet(_) => "parquet",
                    }
                    .to_string(),
                )
            });

            if self.compression != Compression::None {
                warn!("Top level compression setting ignored when batch_encoding set to parquet.")
            }

            let request_options = S3RequestOptions {
                bucket: self.bucket.clone(),
                api_options,
                filename_extension,
                filename_time_format: self.filename_time_format.clone(),
                filename_append_uuid: self.filename_append_uuid,
                encoder: (transformer, encoder),
                // Batch formats handle their own compression internally
                compression: Compression::None,
                filename_tz_offset: offset,
            };

            let sink = S3Sink::new(service, request_options, partitioner, batch_settings);
            return Ok(VectorSink::from_event_streamsink(sink));
        }

        let (framer, serializer) = self.encoding.build(SinkType::MessageBased)?;
        let encoder = EncoderKind::Framed(Box::new(Encoder::<Framer>::new(framer, serializer)));

        let request_options = S3RequestOptions {
            bucket: self.bucket.clone(),
            api_options: self.options.clone(),
            filename_extension: self.filename_extension.clone(),
            filename_time_format: self.filename_time_format.clone(),
            filename_append_uuid: self.filename_append_uuid,
            encoder: (transformer, encoder),
            compression: self.compression,
            filename_tz_offset: offset,
        };

        let sink = S3Sink::new(service, request_options, partitioner, batch_settings);

        Ok(VectorSink::from_event_streamsink(sink))
    }

    pub fn build_healthcheck(&self, client: S3Client) -> crate::Result<Healthcheck> {
        s3_common::config::build_healthcheck(self.bucket.clone(), client)
    }

    pub async fn create_service(&self, proxy: &ProxyConfig) -> crate::Result<S3Service> {
        s3_common::config::create_service(
            &self.region,
            &self.auth,
            proxy,
            self.tls.as_ref(),
            self.force_path_style,
        )
        .await
    }
}

#[cfg(test)]
mod tests {
    use super::S3SinkConfig;

    #[test]
    fn generate_config() {
        crate::test_util::test_generate_config::<S3SinkConfig>();
    }

    /// Correct TOML shape: `batch_encoding.codec = "parquet"` with `schema_mode = "auto_infer"`.
    #[cfg(feature = "codecs-parquet")]
    #[test]
    fn parquet_batch_encoding_correct_toml_shape() {
        let config: S3SinkConfig = serde_yaml::from_str(indoc::indoc! {r#"
            bucket: test-bucket
            compression: none
            encoding:
              codec: text
            batch_encoding:
              schema_mode: auto_infer
              codec: parquet
              compression:
                algorithm: snappy
            "#})
        .expect("correct batch_encoding shape should parse");

        let batch_enc = config
            .batch_encoding
            .expect("batch_encoding should be Some");
        let super::S3BatchEncoding::Parquet(ref p) = batch_enc;
        use vector_lib::codecs::encoding::format::{ParquetCompression, ParquetSchemaMode};
        assert_eq!(p.schema_mode, ParquetSchemaMode::AutoInfer);
        assert_eq!(p.compression, ParquetCompression::Snappy);
    }

    /// Content-Type must be auto-detected as `application/vnd.apache.parquet`
    /// when `batch_encoding` is set and `content_type` is not explicitly provided.
    #[cfg(feature = "codecs-parquet")]
    #[test]
    fn parquet_content_type_auto_detected() {
        use vector_lib::codecs::encoding::format::{
            ParquetCompression, ParquetSchemaMode, ParquetSerializerConfig,
        };

        use crate::sinks::s3_common::config::S3Options;
        use crate::sinks::util::{BatchConfig, BulkSizeBasedDefaultBatchSettings, Compression};
        use vector_lib::codecs::TextSerializerConfig;
        use vector_lib::codecs::encoding::{BatchSerializerConfig, FramingConfig};

        let parquet_config = ParquetSerializerConfig {
            schema_mode: ParquetSchemaMode::AutoInfer,
            compression: ParquetCompression::Snappy,
            ..Default::default()
        };

        let config = S3SinkConfig {
            bucket: "test".to_string(),
            key_prefix: super::default_key_prefix(),
            filename_time_format: super::default_filename_time_format(),
            filename_append_uuid: true,
            filename_extension: None,
            options: S3Options::default(),
            region: crate::aws::RegionOrEndpoint::with_both("us-east-1", "http://localhost:4566"),
            encoding: (None::<FramingConfig>, TextSerializerConfig::default()).into(),
            batch_encoding: Some(super::S3BatchEncoding::Parquet(parquet_config)),
            compression: Compression::None,
            batch: BatchConfig::<BulkSizeBasedDefaultBatchSettings>::default(),
            request: Default::default(),
            tls: Default::default(),
            auth: Default::default(),
            acknowledgements: Default::default(),
            timezone: Default::default(),
            force_path_style: true,
            retry_strategy: Default::default(),
        };

        let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.as_ref().unwrap();
        let batch_config = BatchSerializerConfig::Parquet(p.clone());
        let batch_serializer = batch_config.build_batch_serializer().unwrap();
        let batch_encoder = vector_lib::codecs::BatchEncoder::new(batch_serializer);

        let mut api_options = config.options.clone();
        if api_options.content_type.is_none() {
            api_options.content_type = batch_encoder.content_type().map(|s| s.to_string());
        }

        assert_eq!(
            api_options.content_type.as_deref(),
            Some("application/vnd.apache.parquet"),
            "Content-Type must be auto-detected for Parquet"
        );
    }

    /// When user explicitly sets `content_type`, the auto-detection must not override it.
    #[cfg(feature = "codecs-parquet")]
    #[test]
    fn parquet_content_type_user_override_preserved() {
        let config: S3SinkConfig = serde_yaml::from_str(indoc::indoc! {r#"
            bucket: test-bucket
            compression: none
            content_type: "application/octet-stream"
            encoding:
              codec: text
            batch_encoding:
              codec: parquet
              schema_mode: auto_infer
              compression:
                algorithm: gzip
                level: 9
            "#})
        .unwrap();

        let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.as_ref().unwrap();
        let batch_config = vector_lib::codecs::encoding::BatchSerializerConfig::Parquet(p.clone());
        let batch_serializer = batch_config.build_batch_serializer().unwrap();
        let batch_encoder = vector_lib::codecs::BatchEncoder::new(batch_serializer);

        let mut api_options = config.options.clone();
        if api_options.content_type.is_none() {
            api_options.content_type = batch_encoder.content_type().map(|s| s.to_string());
        }

        assert_eq!(
            api_options.content_type.as_deref(),
            Some("application/octet-stream"),
            "User-specified Content-Type must not be overridden"
        );
    }

    /// Codecs other than `parquet` must be rejected at parse time, since
    /// `S3BatchEncoding` only exposes the `parquet` variant.
    #[cfg(feature = "codecs-parquet")]
    #[test]
    fn parquet_batch_encoding_rejects_unsupported_codec() {
        let err = serde_yaml::from_str::<S3SinkConfig>(
            r#"
            bucket: test-bucket
            compression: none
            encoding:
              codec: text
            batch_encoding:
              codec: arrow_stream
            "#,
        )
        .unwrap_err();

        assert!(
            err.to_string().contains("arrow_stream"),
            "expected error to mention the offending codec, got: {err}"
        );
    }

    /// Explicit filename_extension overrides the `.parquet` default.
    #[cfg(feature = "codecs-parquet")]
    #[test]
    fn parquet_filename_extension_user_override() {
        let config: S3SinkConfig = serde_yaml::from_str(indoc::indoc! {r#"
            bucket: test-bucket
            compression: none
            filename_extension: pq
            encoding:
              codec: text
            batch_encoding:
              codec: parquet
              schema_mode: auto_infer
            "#})
        .unwrap();

        assert_eq!(config.filename_extension.as_deref(), Some("pq"));
    }

    /// `schema_mode` defaults to `relaxed` when not specified.
    #[cfg(feature = "codecs-parquet")]
    #[test]
    fn parquet_schema_mode_defaults_to_relaxed() {
        use vector_lib::codecs::encoding::format::ParquetSchemaMode;

        let config: S3SinkConfig = serde_yaml::from_str(indoc::indoc! {r#"
            bucket: test-bucket
            compression: none
            encoding:
              codec: text
            batch_encoding:
              codec: parquet
            "#})
        .unwrap();

        let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.unwrap();
        assert_eq!(p.schema_mode, ParquetSchemaMode::Relaxed);
    }

    /// Explicit `schema_mode = "strict"` is correctly parsed.
    #[cfg(feature = "codecs-parquet")]
    #[test]
    fn parquet_schema_mode_strict_parsed() {
        use vector_lib::codecs::encoding::format::ParquetSchemaMode;

        let config: S3SinkConfig = serde_yaml::from_str(indoc::indoc! {r#"
            bucket: test-bucket
            compression: none
            encoding:
              codec: text
            batch_encoding:
              codec: parquet
              schema_mode: strict
              schema_file: tmp/something.schema
            "#})
        .unwrap();

        let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.unwrap();
        assert_eq!(p.schema_mode, ParquetSchemaMode::Strict);
    }
}