Skip to content

Commit 7923556

Browse files
flaviofcruzprontclaude
authored
fix(clickhouse sink, aws_s3 sink): use dedicated batch_encoding types (#25340)
* Fix the clickhouse sink so that it only accepts arrow * Do not allow arrow stream or parquet on AWS S3 * reformat * Add changelog fragment * Remove stale comments * chore(aws_s3 sink): exhaustive match for batch_encoding default extension So adding a future S3BatchEncoding variant is a compile error rather than silently defaulting to "parquet". Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * chore(aws_s3 sink): rename shadowing parquet_config bind in tests Avoid rebinding parquet_config to a borrow of itself from config.batch_encoding; use a short p binding instead. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * docs(changelog): rewrite batch_encoding fragment for users Drop internal type names and dev jargon; lead with the user-visible behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * test(aws_s3 sink, clickhouse sink): reject unsupported batch_encoding codec at parse time Add per-sink deserialization-failure tests that pin the schema-tightening behavior introduced by the dedicated wrapper enums: - aws_s3 rejects codec: arrow_stream - clickhouse rejects codec: parquet Previously both codecs were accepted by serde and rejected later at sink-build time; the new wrapper enums move rejection up to parse time, and these tests prevent silent regression of that contract. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * chore(aws_s3 sink): re-export S3BatchEncoding from aws_s3 module Programmatic users constructing S3SinkConfig directly need to be able to set the batch_encoding field. With config kept private in src/sinks/aws_s3/mod.rs, S3BatchEncoding was unnameable outside the crate, regressing callers that previously used BatchSerializerConfig::Parquet(...) at the same field. Gated by codecs-parquet to mirror the type and field. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Pavlos Rontidis <pavlos.rontidis@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 9f15e23 commit 7923556

8 files changed

Lines changed: 147 additions & 212 deletions

File tree

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
The `aws_s3` and `clickhouse` sinks now correctly advertise only the `batch_encoding.codec` values they actually support: `parquet` for `aws_s3` and `arrow_stream` for `clickhouse`. Previously the documentation and configuration schema listed both codecs for both sinks, even though picking the wrong one produced a startup error.
2+
3+
authors: flaviofcruz

src/sinks/aws_s3/config.rs

Lines changed: 60 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use tower::ServiceBuilder;
33
#[cfg(feature = "codecs-parquet")]
44
use vector_lib::codecs::BatchEncoder;
55
#[cfg(feature = "codecs-parquet")]
6-
use vector_lib::codecs::encoding::BatchSerializerConfig;
6+
use vector_lib::codecs::encoding::{BatchSerializerConfig, format::ParquetSerializerConfig};
77
use vector_lib::{
88
TimeZone,
99
codecs::{
@@ -37,6 +37,21 @@ use crate::{
3737
tls::TlsConfig,
3838
};
3939

40+
/// Batch encoding configuration for the `aws_s3` sink.
41+
#[cfg(feature = "codecs-parquet")]
42+
#[configurable_component]
43+
#[derive(Clone, Debug)]
44+
#[serde(tag = "codec", rename_all = "snake_case")]
45+
#[configurable(metadata(
46+
docs::enum_tag_description = "The codec to use for batch encoding events."
47+
))]
48+
pub enum S3BatchEncoding {
49+
/// Encodes events in [Apache Parquet][apache_parquet] columnar format.
50+
///
51+
/// [apache_parquet]: https://parquet.apache.org/
52+
Parquet(ParquetSerializerConfig),
53+
}
54+
4055
/// Configuration for the `aws_s3` sink.
4156
#[configurable_component(sink(
4257
"aws_s3",
@@ -111,15 +126,13 @@ pub struct S3SinkConfig {
111126

112127
/// Batch encoding configuration for columnar formats.
113128
///
114-
/// When set, events are encoded together as a batch in a columnar format (for example, Parquet)
129+
/// When set, events are encoded together as a batch in a columnar format (Parquet)
115130
/// instead of the standard per-event framing-based encoding. The columnar format handles
116131
/// its own internal compression, so the top-level `compression` setting is bypassed.
117-
///
118-
/// Only the `parquet` codec is supported by the AWS S3 sink.
119132
#[cfg(feature = "codecs-parquet")]
120133
#[configurable(derived)]
121134
#[serde(default)]
122-
pub batch_encoding: Option<BatchSerializerConfig>,
135+
pub batch_encoding: Option<S3BatchEncoding>,
123136

124137
/// Compression configuration.
125138
///
@@ -220,8 +233,10 @@ impl SinkConfig for S3SinkConfig {
220233

221234
fn input(&self) -> Input {
222235
#[cfg(feature = "codecs-parquet")]
223-
if let Some(batch_config) = &self.batch_encoding {
224-
return Input::new(batch_config.input_type());
236+
if let Some(batch_encoding) = &self.batch_encoding {
237+
let S3BatchEncoding::Parquet(parquet_config) = batch_encoding;
238+
let resolved = BatchSerializerConfig::Parquet(parquet_config.clone());
239+
return Input::new(resolved.input_type());
225240
}
226241
Input::new(self.encoding.config().1.input_type())
227242
}
@@ -272,15 +287,11 @@ impl S3SinkConfig {
272287
// When batch_encoding is configured (e.g., Parquet), use batch mode
273288
// with internal compression and appropriate file extension.
274289
#[cfg(feature = "codecs-parquet")]
275-
if let Some(batch_config) = &self.batch_encoding {
276-
if !matches!(batch_config, BatchSerializerConfig::Parquet(_)) {
277-
return Err(
278-
"batch_encoding only supports encoding with parquet format for amazon s3 sink"
279-
.into(),
280-
);
281-
}
290+
if let Some(batch_encoding) = &self.batch_encoding {
291+
let S3BatchEncoding::Parquet(parquet_config) = batch_encoding;
292+
let resolved_batch_config = BatchSerializerConfig::Parquet(parquet_config.clone());
282293

283-
let batch_serializer = batch_config.build_batch_serializer()?;
294+
let batch_serializer = resolved_batch_config.build_batch_serializer()?;
284295
let batch_encoder = BatchEncoder::new(batch_serializer);
285296

286297
// Auto-detect Content-Type from batch format. Users can still
@@ -292,15 +303,14 @@ impl S3SinkConfig {
292303

293304
let encoder = EncoderKind::Batch(batch_encoder);
294305

295-
// Auto-detect file extension from batch format
296-
let filename_extension =
297-
self.filename_extension
298-
.clone()
299-
.or_else(|| match batch_config {
300-
BatchSerializerConfig::Parquet(_) => Some("parquet".to_string()),
301-
#[allow(unreachable_patterns)]
302-
_ => None,
303-
});
306+
let filename_extension = self.filename_extension.clone().or_else(|| {
307+
Some(
308+
match batch_encoding {
309+
S3BatchEncoding::Parquet(_) => "parquet",
310+
}
311+
.to_string(),
312+
)
313+
});
304314

305315
if self.compression != Compression::None {
306316
warn!("Top level compression setting ignored when batch_encoding set to parquet.")
@@ -392,15 +402,10 @@ mod tests {
392402
let batch_enc = config
393403
.batch_encoding
394404
.expect("batch_encoding should be Some");
395-
match batch_enc {
396-
vector_lib::codecs::encoding::BatchSerializerConfig::Parquet(ref p) => {
397-
use vector_lib::codecs::encoding::format::{ParquetCompression, ParquetSchemaMode};
398-
assert_eq!(p.schema_mode, ParquetSchemaMode::AutoInfer);
399-
assert_eq!(p.compression, ParquetCompression::Snappy);
400-
}
401-
#[allow(unreachable_patterns)]
402-
_ => panic!("expected Parquet variant"),
403-
}
405+
let super::S3BatchEncoding::Parquet(ref p) = batch_enc;
406+
use vector_lib::codecs::encoding::format::{ParquetCompression, ParquetSchemaMode};
407+
assert_eq!(p.schema_mode, ParquetSchemaMode::AutoInfer);
408+
assert_eq!(p.compression, ParquetCompression::Snappy);
404409
}
405410

406411
/// Content-Type must be auto-detected as `application/vnd.apache.parquet`
@@ -432,7 +437,7 @@ mod tests {
432437
options: S3Options::default(),
433438
region: crate::aws::RegionOrEndpoint::with_both("us-east-1", "http://localhost:4566"),
434439
encoding: (None::<FramingConfig>, TextSerializerConfig::default()).into(),
435-
batch_encoding: Some(BatchSerializerConfig::Parquet(parquet_config)),
440+
batch_encoding: Some(super::S3BatchEncoding::Parquet(parquet_config)),
436441
compression: Compression::None,
437442
batch: BatchConfig::<BulkSizeBasedDefaultBatchSettings>::default(),
438443
request: Default::default(),
@@ -444,7 +449,8 @@ mod tests {
444449
retry_strategy: Default::default(),
445450
};
446451

447-
let batch_config = config.batch_encoding.as_ref().unwrap();
452+
let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.as_ref().unwrap();
453+
let batch_config = BatchSerializerConfig::Parquet(p.clone());
448454
let batch_serializer = batch_config.build_batch_serializer().unwrap();
449455
let batch_encoder = vector_lib::codecs::BatchEncoder::new(batch_serializer);
450456

@@ -484,7 +490,8 @@ mod tests {
484490
)
485491
.unwrap();
486492

487-
let batch_config = config.batch_encoding.as_ref().unwrap();
493+
let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.as_ref().unwrap();
494+
let batch_config = vector_lib::codecs::encoding::BatchSerializerConfig::Parquet(p.clone());
488495
let batch_serializer = batch_config.build_batch_serializer().unwrap();
489496
let batch_encoder = vector_lib::codecs::BatchEncoder::new(batch_serializer);
490497

@@ -500,43 +507,27 @@ mod tests {
500507
);
501508
}
502509

503-
/// Parquet filename extension defaults to `.parquet` when not explicitly set.
510+
/// Codecs other than `parquet` must be rejected at parse time, since
511+
/// `S3BatchEncoding` only exposes the `parquet` variant.
504512
#[cfg(feature = "codecs-parquet")]
505513
#[test]
506-
fn parquet_filename_extension_defaults_to_parquet() {
507-
let config: S3SinkConfig = toml::from_str(
514+
fn parquet_batch_encoding_rejects_unsupported_codec() {
515+
let err = serde_yaml::from_str::<S3SinkConfig>(
508516
r#"
509-
bucket = "test-bucket"
510-
compression = "none"
511-
512-
[encoding]
513-
codec = "text"
514-
515-
[batch_encoding]
516-
codec = "parquet"
517-
schema_mode = "auto_infer"
517+
bucket: test-bucket
518+
compression: none
519+
encoding:
520+
codec: text
521+
batch_encoding:
522+
codec: arrow_stream
518523
"#,
519524
)
520-
.unwrap();
525+
.unwrap_err();
521526

522527
assert!(
523-
config.filename_extension.is_none(),
524-
"fixture must not set filename_extension"
528+
err.to_string().contains("arrow_stream"),
529+
"expected error to mention the offending codec, got: {err}"
525530
);
526-
527-
let batch_config = config.batch_encoding.as_ref().unwrap();
528-
let extension = config
529-
.filename_extension
530-
.clone()
531-
.or_else(|| match batch_config {
532-
vector_lib::codecs::encoding::BatchSerializerConfig::Parquet(_) => {
533-
Some("parquet".to_string())
534-
}
535-
#[allow(unreachable_patterns)]
536-
_ => None,
537-
});
538-
539-
assert_eq!(extension.as_deref(), Some("parquet"));
540531
}
541532

542533
/// Explicit filename_extension overrides the `.parquet` default.
@@ -582,13 +573,8 @@ mod tests {
582573
)
583574
.unwrap();
584575

585-
match config.batch_encoding.unwrap() {
586-
vector_lib::codecs::encoding::BatchSerializerConfig::Parquet(p) => {
587-
assert_eq!(p.schema_mode, ParquetSchemaMode::Relaxed);
588-
}
589-
#[allow(unreachable_patterns)]
590-
_ => panic!("expected Parquet variant"),
591-
}
576+
let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.unwrap();
577+
assert_eq!(p.schema_mode, ParquetSchemaMode::Relaxed);
592578
}
593579

594580
/// Explicit `schema_mode = "strict"` is correctly parsed.
@@ -613,12 +599,7 @@ mod tests {
613599
)
614600
.unwrap();
615601

616-
match config.batch_encoding.unwrap() {
617-
vector_lib::codecs::encoding::BatchSerializerConfig::Parquet(p) => {
618-
assert_eq!(p.schema_mode, ParquetSchemaMode::Strict);
619-
}
620-
#[allow(unreachable_patterns)]
621-
_ => panic!("expected Parquet variant"),
622-
}
602+
let super::S3BatchEncoding::Parquet(p) = config.batch_encoding.unwrap();
603+
assert_eq!(p.schema_mode, ParquetSchemaMode::Strict);
623604
}
624605
}

src/sinks/aws_s3/integration_tests.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ use std::{
77
time::Duration,
88
};
99

10+
#[cfg(feature = "codecs-parquet")]
11+
use super::config::S3BatchEncoding;
1012
use aws_sdk_s3::{
1113
Client as S3Client,
1214
operation::{create_bucket::CreateBucketError, get_object::GetObjectOutput},
@@ -22,8 +24,6 @@ use futures::{Stream, stream};
2224
use similar_asserts::assert_eq;
2325
use tempfile::TempDir;
2426
use tokio_stream::StreamExt;
25-
#[cfg(feature = "codecs-parquet")]
26-
use vector_lib::codecs::encoding::BatchSerializerConfig;
2727
use vector_lib::{
2828
buffers::{BufferConfig, BufferType, WhenFull},
2929
codecs::{TextSerializerConfig, encoding::FramingConfig},
@@ -502,7 +502,7 @@ async fn s3_parquet_insert_message() {
502502
};
503503

504504
let config = S3SinkConfig {
505-
batch_encoding: Some(BatchSerializerConfig::Parquet(parquet_config)),
505+
batch_encoding: Some(S3BatchEncoding::Parquet(parquet_config)),
506506
..config(&bucket, 100, 5.0)
507507
};
508508

src/sinks/aws_s3/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,6 @@ mod sink;
33

44
mod integration_tests;
55

6+
#[cfg(feature = "codecs-parquet")]
7+
pub use config::S3BatchEncoding;
68
pub use config::S3SinkConfig;

src/sinks/clickhouse/config.rs

Lines changed: 45 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ use std::fmt;
44

55
use http::{Request, StatusCode, Uri};
66
use hyper::Body;
7+
use vector_lib::codecs::encoding::ArrowStreamSerializerConfig;
78
use vector_lib::codecs::encoding::format::SchemaProvider;
8-
use vector_lib::codecs::encoding::{ArrowStreamSerializerConfig, BatchSerializerConfig};
99

1010
use super::{
1111
request_builder::ClickhouseRequestBuilder,
@@ -45,6 +45,23 @@ pub enum Format {
4545
ArrowStream,
4646
}
4747

48+
/// Batch encoding configuration for the `clickhouse` sink.
49+
#[configurable_component]
50+
#[derive(Clone, Debug)]
51+
#[serde(tag = "codec", rename_all = "snake_case")]
52+
#[configurable(metadata(
53+
docs::enum_tag_description = "The codec to use for batch encoding events."
54+
))]
55+
pub enum ClickhouseBatchEncoding {
56+
/// Encodes events in [Apache Arrow][apache_arrow] IPC streaming format.
57+
///
58+
/// This is the streaming variant of the Arrow IPC format, which writes
59+
/// a continuous stream of record batches.
60+
///
61+
/// [apache_arrow]: https://arrow.apache.org/
62+
ArrowStream(ArrowStreamSerializerConfig),
63+
}
64+
4865
impl fmt::Display for Format {
4966
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
5067
match self {
@@ -104,11 +121,9 @@ pub struct ClickhouseConfig {
104121
///
105122
/// When specified, events are encoded together as a single batch.
106123
/// This is mutually exclusive with per-event encoding based on the `format` field.
107-
///
108-
/// Only the `arrow_stream` codec is supported by the ClickHouse sink.
109124
#[configurable(derived)]
110125
#[serde(default)]
111-
pub batch_encoding: Option<BatchSerializerConfig>,
126+
pub batch_encoding: Option<ClickhouseBatchEncoding>,
112127

113128
#[configurable(derived)]
114129
#[serde(default)]
@@ -282,6 +297,7 @@ impl ClickhouseConfig {
282297

283298
if let Some(batch_encoding) = &self.batch_encoding {
284299
use vector_lib::codecs::BatchEncoder;
300+
use vector_lib::codecs::encoding::BatchSerializerConfig;
285301

286302
// Validate that batch_encoding is only compatible with ArrowStream format
287303
if self.format != Format::ArrowStream {
@@ -292,16 +308,8 @@ impl ClickhouseConfig {
292308
.into());
293309
}
294310

295-
let mut arrow_config = match batch_encoding {
296-
BatchSerializerConfig::ArrowStream(config) => config.clone(),
297-
#[cfg(feature = "codecs-parquet")]
298-
BatchSerializerConfig::Parquet(_) => {
299-
return Err(
300-
"ClickHouse sink does not support Parquet batch encoding. Use 'arrow_stream' instead."
301-
.into(),
302-
);
303-
}
304-
};
311+
let ClickhouseBatchEncoding::ArrowStream(arrow_config) = batch_encoding;
312+
let mut arrow_config = arrow_config.clone();
305313

306314
self.resolve_arrow_schema(
307315
client,
@@ -433,10 +441,31 @@ mod tests {
433441
);
434442
}
435443

444+
/// Codecs other than `arrow_stream` must be rejected at parse time, since
445+
/// `ClickhouseBatchEncoding` only exposes the `arrow_stream` variant.
446+
#[cfg(feature = "codecs-parquet")]
447+
#[test]
448+
fn batch_encoding_rejects_unsupported_codec() {
449+
let err = serde_yaml::from_str::<ClickhouseConfig>(
450+
r#"
451+
endpoint: http://localhost:8123
452+
table: test_table
453+
batch_encoding:
454+
codec: parquet
455+
"#,
456+
)
457+
.unwrap_err();
458+
459+
assert!(
460+
err.to_string().contains("parquet"),
461+
"expected error to mention the offending codec, got: {err}"
462+
);
463+
}
464+
436465
/// Helper to create a minimal ClickhouseConfig for testing
437466
fn create_test_config(
438467
format: Format,
439-
batch_encoding: Option<BatchSerializerConfig>,
468+
batch_encoding: Option<ClickhouseBatchEncoding>,
440469
) -> ClickhouseConfig {
441470
ClickhouseConfig {
442471
endpoint: "http://localhost:8123".parse::<http::Uri>().unwrap().into(),
@@ -469,7 +498,7 @@ mod tests {
469498
for (format, format_name) in incompatible_formats {
470499
let config = create_test_config(
471500
format,
472-
Some(BatchSerializerConfig::ArrowStream(
501+
Some(ClickhouseBatchEncoding::ArrowStream(
473502
ArrowStreamSerializerConfig::default(),
474503
)),
475504
);

0 commit comments

Comments
 (0)