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