1use std::pin::Pin;
2use std::sync::Arc;
3
4use futures::Stream;
5
6use crate::body::Body;
7use crate::error::CamelError;
8use crate::exchange::Exchange;
9use crate::message::Message;
10
11pub type SplitExpression = Arc<dyn Fn(&Exchange) -> Vec<Exchange> + Send + Sync>;
13
14pub type StreamingSplitExpression = Arc<
21 dyn Fn(Exchange) -> Pin<Box<dyn Stream<Item = Result<Exchange, CamelError>> + Send>>
22 + Send
23 + Sync,
24>;
25
26#[derive(Clone, Default)]
28pub enum AggregationStrategy {
29 #[default]
31 LastWins,
32 CollectAll,
34 Original,
36 Custom(Arc<dyn Fn(Exchange, Exchange) -> Exchange + Send + Sync>),
38}
39
40impl std::fmt::Debug for AggregationStrategy {
41 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42 match self {
43 AggregationStrategy::LastWins => f.write_str("LastWins"),
44 AggregationStrategy::CollectAll => f.write_str("CollectAll"),
45 AggregationStrategy::Original => f.write_str("Original"),
46 AggregationStrategy::Custom(_) => f.write_str("Custom(..)"),
47 }
48 }
49}
50
51#[derive(
53 Clone,
54 Debug,
55 Default,
56 PartialEq,
57 Eq,
58 serde::Serialize,
59 serde::Deserialize,
60 schemars::JsonSchema,
61 ts_rs::TS,
62)]
63#[serde(rename_all = "snake_case")]
64#[ts(rename_all = "snake_case")]
65pub enum StreamSplitFormat {
66 #[default]
68 Auto,
69 Ndjson,
71 Lines,
73 Chunks,
75 Zip,
77}
78
79#[derive(
84 Clone,
85 Debug,
86 PartialEq,
87 Eq,
88 serde::Serialize,
89 serde::Deserialize,
90 schemars::JsonSchema,
91 ts_rs::TS,
92)]
93#[serde(rename_all = "snake_case")]
94#[ts(rename_all = "snake_case")]
95pub struct StreamSplitConfig {
96 pub format: StreamSplitFormat,
98 pub max_record_bytes: usize,
100 pub batch_size: usize,
102 pub chunk_size: Option<usize>,
104 pub include_origin: bool,
106}
107
108impl Default for StreamSplitConfig {
109 fn default() -> Self {
110 Self {
111 format: StreamSplitFormat::Auto,
112 max_record_bytes: 1024 * 1024,
113 batch_size: 1,
114 chunk_size: None,
115 include_origin: true,
116 }
117 }
118}
119
120impl StreamSplitConfig {
121 pub fn validate(&self) -> Result<(), CamelError> {
132 if self.batch_size == 0 {
133 return Err(CamelError::Config(
134 "stream split batch_size must be > 0".into(),
135 ));
136 }
137 if self.max_record_bytes == 0 {
138 return Err(CamelError::Config(
139 "stream split max_record_bytes must be > 0".into(),
140 ));
141 }
142 if self.format == StreamSplitFormat::Chunks && self.chunk_size.is_none() {
143 return Err(CamelError::Config(
144 "stream split format=Chunks requires chunk_size".into(),
145 ));
146 }
147 if self.format == StreamSplitFormat::Zip && self.chunk_size.is_some() {
150 return Err(CamelError::Config(
151 "stream split format=Zip does not support chunk_size".into(),
152 ));
153 }
154 if let Some(cs) = self.chunk_size
155 && cs == 0
156 {
157 return Err(CamelError::Config(
158 "stream split chunk_size must be > 0".into(),
159 ));
160 }
161 if self.format == StreamSplitFormat::Chunks
162 && let Some(cs) = self.chunk_size
163 && cs > self.max_record_bytes
164 {
165 return Err(CamelError::Config(
166 "stream split chunk_size must be <= max_record_bytes".into(),
167 ));
168 }
169 Ok(())
170 }
171}
172
173pub struct SplitterConfig {
175 pub expression: SplitExpression,
177 pub aggregation: AggregationStrategy,
179 pub parallel: bool,
181 pub parallel_limit: Option<usize>,
183 pub stop_on_exception: bool,
189 pub max_fragments: usize,
195}
196
197impl std::fmt::Debug for SplitterConfig {
198 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
199 f.debug_struct("SplitterConfig")
200 .field("expression", &"<split-expression>")
201 .field("aggregation", &self.aggregation)
202 .field("parallel", &self.parallel)
203 .field("parallel_limit", &self.parallel_limit)
204 .field("stop_on_exception", &self.stop_on_exception)
205 .field("max_fragments", &self.max_fragments)
206 .finish()
207 }
208}
209
210impl SplitterConfig {
211 pub fn new(expression: SplitExpression) -> Self {
213 Self {
214 expression,
215 aggregation: AggregationStrategy::default(),
216 parallel: false,
217 parallel_limit: None,
218 stop_on_exception: true,
219 max_fragments: 100_000,
220 }
221 }
222
223 pub fn aggregation(mut self, strategy: AggregationStrategy) -> Self {
225 self.aggregation = strategy;
226 self
227 }
228
229 pub fn parallel(mut self, parallel: bool) -> Self {
231 self.parallel = parallel;
232 self
233 }
234
235 pub fn parallel_limit(mut self, limit: usize) -> Self {
237 self.parallel_limit = Some(limit);
238 self
239 }
240
241 pub fn stop_on_exception(mut self, stop: bool) -> Self {
246 self.stop_on_exception = stop;
247 self
248 }
249
250 pub fn max_fragments(mut self, max: usize) -> Self {
252 self.max_fragments = max;
253 self
254 }
255
256 pub fn validate(&self) -> Result<(), CamelError> {
261 if self.parallel && self.parallel_limit == Some(0) {
262 return Err(CamelError::Config(
263 "splitter parallel_limit must be > 0".to_string(),
264 ));
265 }
266 if self.max_fragments == 0 {
267 return Err(CamelError::Config(
268 "splitter max_fragments must be > 0".to_string(),
269 ));
270 }
271 Ok(())
272 }
273}
274
275pub fn fragment_exchange(parent: &Exchange, body: Body) -> Exchange {
300 let mut msg = Message::new(body);
301 msg.headers = parent.input.headers.clone();
302 let mut ex = Exchange::new(msg);
303 ex.properties = parent.properties.clone();
304 ex.pattern = parent.pattern;
305 ex.otel_context = parent.otel_context.clone();
307 ex
308}
309
310pub fn split_body_lines() -> SplitExpression {
313 Arc::new(|exchange: &Exchange| {
314 let text = match &exchange.input.body {
315 Body::Text(s) => s.as_str(),
316 _ => return Vec::new(),
317 };
318 text.lines()
319 .map(|line| fragment_exchange(exchange, Body::Text(line.to_string())))
320 .collect()
321 })
322}
323
324pub fn split_body_json_array() -> SplitExpression {
327 Arc::new(|exchange: &Exchange| {
328 let arr = match &exchange.input.body {
329 Body::Json(serde_json::Value::Array(arr)) => arr,
330 _ => return Vec::new(),
331 };
332 arr.iter()
333 .map(|val| fragment_exchange(exchange, Body::Json(val.clone())))
334 .collect()
335 })
336}
337
338pub fn split_body<F>(f: F) -> SplitExpression
340where
341 F: Fn(&Body) -> Vec<Body> + Send + Sync + 'static,
342{
343 Arc::new(move |exchange: &Exchange| {
344 f(&exchange.input.body)
345 .into_iter()
346 .map(|body| fragment_exchange(exchange, body))
347 .collect()
348 })
349}
350
351#[cfg(test)]
352mod tests {
353 use super::*;
354 use crate::value::Value;
355
356 #[test]
357 fn test_split_body_lines() {
358 let mut ex = Exchange::new(Message::new("a\nb\nc"));
359 ex.input.set_header("source", Value::String("test".into()));
360 ex.set_property("trace", Value::Bool(true));
361
362 let fragments = split_body_lines()(&ex);
363 assert_eq!(fragments.len(), 3);
364 assert_eq!(fragments[0].input.body.as_text(), Some("a"));
365 assert_eq!(fragments[1].input.body.as_text(), Some("b"));
366 assert_eq!(fragments[2].input.body.as_text(), Some("c"));
367
368 for frag in &fragments {
370 assert_eq!(
371 frag.input.header("source"),
372 Some(&Value::String("test".into()))
373 );
374 assert_eq!(frag.property("trace"), Some(&Value::Bool(true)));
375 }
376 }
377
378 #[test]
379 fn test_split_body_lines_empty() {
380 let ex = Exchange::new(Message::default()); let fragments = split_body_lines()(&ex);
382 assert!(fragments.is_empty());
383 }
384
385 #[test]
386 fn test_split_body_json_array() {
387 let arr = serde_json::json!([1, 2, 3]);
388 let ex = Exchange::new(Message::new(arr));
389
390 let fragments = split_body_json_array()(&ex);
391 assert_eq!(fragments.len(), 3);
392 assert!(matches!(&fragments[0].input.body, Body::Json(v) if *v == serde_json::json!(1)));
393 assert!(matches!(&fragments[1].input.body, Body::Json(v) if *v == serde_json::json!(2)));
394 assert!(matches!(&fragments[2].input.body, Body::Json(v) if *v == serde_json::json!(3)));
395 }
396
397 #[test]
398 fn test_split_body_json_array_not_array() {
399 let obj = serde_json::json!({"not": "array"});
400 let ex = Exchange::new(Message::new(obj));
401
402 let fragments = split_body_json_array()(&ex);
403 assert!(fragments.is_empty());
404 }
405
406 #[test]
407 fn test_split_body_custom() {
408 let splitter = split_body(|body: &Body| match body {
409 Body::Text(s) => s
410 .split(',')
411 .map(|part| Body::Text(part.trim().to_string()))
412 .collect(),
413 _ => Vec::new(),
414 });
415
416 let mut ex = Exchange::new(Message::new("x, y, z"));
417 ex.set_property("id", Value::from(42));
418
419 let fragments = splitter(&ex);
420 assert_eq!(fragments.len(), 3);
421 assert_eq!(fragments[0].input.body.as_text(), Some("x"));
422 assert_eq!(fragments[1].input.body.as_text(), Some("y"));
423 assert_eq!(fragments[2].input.body.as_text(), Some("z"));
424
425 for frag in &fragments {
427 assert_eq!(frag.property("id"), Some(&Value::from(42)));
428 }
429 }
430
431 #[test]
432 fn test_splitter_config_defaults() {
433 let config = SplitterConfig::new(split_body_lines());
434 assert!(matches!(config.aggregation, AggregationStrategy::LastWins));
435 assert!(!config.parallel);
436 assert!(config.parallel_limit.is_none());
437 assert!(config.stop_on_exception);
438 }
439
440 #[test]
441 fn test_splitter_config_builder() {
442 let config = SplitterConfig::new(split_body_lines())
443 .aggregation(AggregationStrategy::CollectAll)
444 .parallel(true)
445 .parallel_limit(4)
446 .stop_on_exception(false);
447
448 assert!(matches!(
449 config.aggregation,
450 AggregationStrategy::CollectAll
451 ));
452 assert!(config.parallel);
453 assert_eq!(config.parallel_limit, Some(4));
454 assert!(!config.stop_on_exception);
455 }
456
457 #[test]
458 fn test_splitter_config_default_max_fragments() {
459 let cfg = SplitterConfig::new(Arc::new(|_: &Exchange| Vec::new()) as SplitExpression);
460 assert_eq!(cfg.max_fragments, 100_000);
461 }
462
463 #[test]
464 fn test_splitter_config_rejects_zero_max_fragments() {
465 let cfg = SplitterConfig::new(Arc::new(|_: &Exchange| Vec::new()) as SplitExpression)
466 .max_fragments(0);
467 assert!(cfg.validate().is_err());
468 }
469
470 #[test]
471 fn test_fragment_exchange_inherits_otel_context() {
472 use opentelemetry::Context;
473 use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId};
474
475 let mut parent = Exchange::new(Message::new("test"));
477 let trace_id = TraceId::from_bytes([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 123]);
478 let span_id = SpanId::from_bytes([0, 0, 0, 0, 0, 0, 1, 200]);
479 let span_context = SpanContext::new(
480 trace_id,
481 span_id,
482 TraceFlags::SAMPLED,
483 true,
484 Default::default(),
485 );
486 let expected_trace_id = span_context.trace_id();
487 parent.otel_context = Context::current().with_remote_span_context(span_context);
488
489 let fragments = split_body_lines()(&parent);
491 assert!(!fragments.is_empty(), "Should have at least one fragment");
492
493 for fragment in &fragments {
495 let span = fragment.otel_context.span();
496 let frag_span_ctx = span.span_context();
497 assert!(
498 frag_span_ctx.is_valid(),
499 "Fragment should have valid span context"
500 );
501 assert_eq!(
502 frag_span_ctx.trace_id(),
503 expected_trace_id,
504 "Fragment should have same trace ID as parent"
505 );
506 }
507 }
508
509 #[test]
510 fn test_stream_split_config_defaults_valid() {
511 let config = StreamSplitConfig::default();
512 assert!(config.validate().is_ok());
513 }
514
515 #[test]
516 fn test_stream_split_config_batch_size_zero_rejected() {
517 let config = StreamSplitConfig {
518 batch_size: 0,
519 ..Default::default()
520 };
521 let err = config.validate().unwrap_err();
522 assert!(err.to_string().contains("batch_size"));
523 }
524
525 #[test]
526 fn test_stream_split_config_max_record_bytes_zero_rejected() {
527 let config = StreamSplitConfig {
528 max_record_bytes: 0,
529 ..Default::default()
530 };
531 let err = config.validate().unwrap_err();
532 assert!(err.to_string().contains("max_record_bytes"));
533 }
534
535 #[test]
536 fn test_stream_split_config_chunks_requires_chunk_size() {
537 let config = StreamSplitConfig {
538 format: StreamSplitFormat::Chunks,
539 chunk_size: None,
540 ..Default::default()
541 };
542 let err = config.validate().unwrap_err();
543 assert!(err.to_string().contains("Chunks requires chunk_size"));
544 }
545
546 #[test]
547 fn test_stream_split_config_chunk_size_zero_rejected() {
548 let config = StreamSplitConfig {
549 format: StreamSplitFormat::Chunks,
550 chunk_size: Some(0),
551 ..Default::default()
552 };
553 let err = config.validate().unwrap_err();
554 assert!(err.to_string().contains("chunk_size must be > 0"));
555 }
556
557 #[test]
558 fn test_stream_split_config_chunk_size_exceeds_max_record_bytes() {
559 let config = StreamSplitConfig {
560 format: StreamSplitFormat::Chunks,
561 chunk_size: Some(2000),
562 max_record_bytes: 1000,
563 ..Default::default()
564 };
565 let err = config.validate().unwrap_err();
566 assert!(
567 err.to_string()
568 .contains("chunk_size must be <= max_record_bytes")
569 );
570 }
571
572 #[test]
573 fn test_stream_split_config_zip_rejects_chunk_size() {
574 let config = StreamSplitConfig {
575 format: StreamSplitFormat::Zip,
576 chunk_size: Some(1024),
577 ..Default::default()
578 };
579 let err = config.validate().unwrap_err();
580 assert!(err.to_string().contains("Zip does not support chunk_size"));
581 }
582
583 #[test]
584 fn test_all_fragments_share_same_trace_context() {
585 use opentelemetry::Context;
586 use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId};
587
588 let mut parent = Exchange::new(Message::new("line1\nline2\nline3"));
590 let trace_id =
591 TraceId::from_bytes([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x3B, 0x9A, 0xCA, 0x09]);
592 let span_id = SpanId::from_bytes([0, 0, 0, 0, 0, 0, 0, 111]);
593 let span_context = SpanContext::new(
594 trace_id,
595 span_id,
596 TraceFlags::SAMPLED,
597 true,
598 Default::default(),
599 );
600 parent.otel_context = Context::current().with_remote_span_context(span_context);
601
602 let fragments = split_body_lines()(&parent);
603 assert_eq!(fragments.len(), 3);
604
605 let trace_ids: Vec<_> = fragments
607 .iter()
608 .map(|f| {
609 let span = f.otel_context.span();
610 span.span_context().trace_id()
611 })
612 .collect();
613
614 assert!(
615 trace_ids.iter().all(|&id| id == trace_id),
616 "All fragments should have the same trace ID"
617 );
618 }
619}