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
40#[derive(
42 Clone,
43 Debug,
44 Default,
45 PartialEq,
46 Eq,
47 serde::Serialize,
48 serde::Deserialize,
49 schemars::JsonSchema,
50 ts_rs::TS,
51)]
52#[serde(rename_all = "snake_case")]
53#[ts(rename_all = "snake_case")]
54pub enum StreamSplitFormat {
55 #[default]
57 Auto,
58 Ndjson,
60 Lines,
62 Chunks,
64 Zip,
66}
67
68#[derive(
73 Clone,
74 Debug,
75 PartialEq,
76 Eq,
77 serde::Serialize,
78 serde::Deserialize,
79 schemars::JsonSchema,
80 ts_rs::TS,
81)]
82#[serde(rename_all = "snake_case")]
83#[ts(rename_all = "snake_case")]
84pub struct StreamSplitConfig {
85 pub format: StreamSplitFormat,
87 pub max_record_bytes: usize,
89 pub batch_size: usize,
91 pub chunk_size: Option<usize>,
93 pub include_origin: bool,
95}
96
97impl Default for StreamSplitConfig {
98 fn default() -> Self {
99 Self {
100 format: StreamSplitFormat::Auto,
101 max_record_bytes: 1024 * 1024,
102 batch_size: 1,
103 chunk_size: None,
104 include_origin: true,
105 }
106 }
107}
108
109impl StreamSplitConfig {
110 pub fn validate(&self) -> Result<(), CamelError> {
121 if self.batch_size == 0 {
122 return Err(CamelError::Config(
123 "stream split batch_size must be > 0".into(),
124 ));
125 }
126 if self.max_record_bytes == 0 {
127 return Err(CamelError::Config(
128 "stream split max_record_bytes must be > 0".into(),
129 ));
130 }
131 if self.format == StreamSplitFormat::Chunks && self.chunk_size.is_none() {
132 return Err(CamelError::Config(
133 "stream split format=Chunks requires chunk_size".into(),
134 ));
135 }
136 if self.format == StreamSplitFormat::Zip && self.chunk_size.is_some() {
139 return Err(CamelError::Config(
140 "stream split format=Zip does not support chunk_size".into(),
141 ));
142 }
143 if let Some(cs) = self.chunk_size
144 && cs == 0
145 {
146 return Err(CamelError::Config(
147 "stream split chunk_size must be > 0".into(),
148 ));
149 }
150 if self.format == StreamSplitFormat::Chunks
151 && let Some(cs) = self.chunk_size
152 && cs > self.max_record_bytes
153 {
154 return Err(CamelError::Config(
155 "stream split chunk_size must be <= max_record_bytes".into(),
156 ));
157 }
158 Ok(())
159 }
160}
161
162pub struct SplitterConfig {
164 pub expression: SplitExpression,
166 pub aggregation: AggregationStrategy,
168 pub parallel: bool,
170 pub parallel_limit: Option<usize>,
172 pub stop_on_exception: bool,
178 pub max_fragments: usize,
184}
185
186impl SplitterConfig {
187 pub fn new(expression: SplitExpression) -> Self {
189 Self {
190 expression,
191 aggregation: AggregationStrategy::default(),
192 parallel: false,
193 parallel_limit: None,
194 stop_on_exception: true,
195 max_fragments: 100_000,
196 }
197 }
198
199 pub fn aggregation(mut self, strategy: AggregationStrategy) -> Self {
201 self.aggregation = strategy;
202 self
203 }
204
205 pub fn parallel(mut self, parallel: bool) -> Self {
207 self.parallel = parallel;
208 self
209 }
210
211 pub fn parallel_limit(mut self, limit: usize) -> Self {
213 self.parallel_limit = Some(limit);
214 self
215 }
216
217 pub fn stop_on_exception(mut self, stop: bool) -> Self {
222 self.stop_on_exception = stop;
223 self
224 }
225
226 pub fn max_fragments(mut self, max: usize) -> Self {
228 self.max_fragments = max;
229 self
230 }
231
232 pub fn validate(&self) -> Result<(), CamelError> {
237 if self.parallel && self.parallel_limit == Some(0) {
238 return Err(CamelError::Config(
239 "splitter parallel_limit must be > 0".to_string(),
240 ));
241 }
242 if self.max_fragments == 0 {
243 return Err(CamelError::Config(
244 "splitter max_fragments must be > 0".to_string(),
245 ));
246 }
247 Ok(())
248 }
249}
250
251pub fn fragment_exchange(parent: &Exchange, body: Body) -> Exchange {
276 let mut msg = Message::new(body);
277 msg.headers = parent.input.headers.clone();
278 let mut ex = Exchange::new(msg);
279 ex.properties = parent.properties.clone();
280 ex.pattern = parent.pattern;
281 ex.otel_context = parent.otel_context.clone();
283 ex
284}
285
286pub fn split_body_lines() -> SplitExpression {
289 Arc::new(|exchange: &Exchange| {
290 let text = match &exchange.input.body {
291 Body::Text(s) => s.as_str(),
292 _ => return Vec::new(),
293 };
294 text.lines()
295 .map(|line| fragment_exchange(exchange, Body::Text(line.to_string())))
296 .collect()
297 })
298}
299
300pub fn split_body_json_array() -> SplitExpression {
303 Arc::new(|exchange: &Exchange| {
304 let arr = match &exchange.input.body {
305 Body::Json(serde_json::Value::Array(arr)) => arr,
306 _ => return Vec::new(),
307 };
308 arr.iter()
309 .map(|val| fragment_exchange(exchange, Body::Json(val.clone())))
310 .collect()
311 })
312}
313
314pub fn split_body<F>(f: F) -> SplitExpression
316where
317 F: Fn(&Body) -> Vec<Body> + Send + Sync + 'static,
318{
319 Arc::new(move |exchange: &Exchange| {
320 f(&exchange.input.body)
321 .into_iter()
322 .map(|body| fragment_exchange(exchange, body))
323 .collect()
324 })
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330 use crate::value::Value;
331
332 #[test]
333 fn test_split_body_lines() {
334 let mut ex = Exchange::new(Message::new("a\nb\nc"));
335 ex.input.set_header("source", Value::String("test".into()));
336 ex.set_property("trace", Value::Bool(true));
337
338 let fragments = split_body_lines()(&ex);
339 assert_eq!(fragments.len(), 3);
340 assert_eq!(fragments[0].input.body.as_text(), Some("a"));
341 assert_eq!(fragments[1].input.body.as_text(), Some("b"));
342 assert_eq!(fragments[2].input.body.as_text(), Some("c"));
343
344 for frag in &fragments {
346 assert_eq!(
347 frag.input.header("source"),
348 Some(&Value::String("test".into()))
349 );
350 assert_eq!(frag.property("trace"), Some(&Value::Bool(true)));
351 }
352 }
353
354 #[test]
355 fn test_split_body_lines_empty() {
356 let ex = Exchange::new(Message::default()); let fragments = split_body_lines()(&ex);
358 assert!(fragments.is_empty());
359 }
360
361 #[test]
362 fn test_split_body_json_array() {
363 let arr = serde_json::json!([1, 2, 3]);
364 let ex = Exchange::new(Message::new(arr));
365
366 let fragments = split_body_json_array()(&ex);
367 assert_eq!(fragments.len(), 3);
368 assert!(matches!(&fragments[0].input.body, Body::Json(v) if *v == serde_json::json!(1)));
369 assert!(matches!(&fragments[1].input.body, Body::Json(v) if *v == serde_json::json!(2)));
370 assert!(matches!(&fragments[2].input.body, Body::Json(v) if *v == serde_json::json!(3)));
371 }
372
373 #[test]
374 fn test_split_body_json_array_not_array() {
375 let obj = serde_json::json!({"not": "array"});
376 let ex = Exchange::new(Message::new(obj));
377
378 let fragments = split_body_json_array()(&ex);
379 assert!(fragments.is_empty());
380 }
381
382 #[test]
383 fn test_split_body_custom() {
384 let splitter = split_body(|body: &Body| match body {
385 Body::Text(s) => s
386 .split(',')
387 .map(|part| Body::Text(part.trim().to_string()))
388 .collect(),
389 _ => Vec::new(),
390 });
391
392 let mut ex = Exchange::new(Message::new("x, y, z"));
393 ex.set_property("id", Value::from(42));
394
395 let fragments = splitter(&ex);
396 assert_eq!(fragments.len(), 3);
397 assert_eq!(fragments[0].input.body.as_text(), Some("x"));
398 assert_eq!(fragments[1].input.body.as_text(), Some("y"));
399 assert_eq!(fragments[2].input.body.as_text(), Some("z"));
400
401 for frag in &fragments {
403 assert_eq!(frag.property("id"), Some(&Value::from(42)));
404 }
405 }
406
407 #[test]
408 fn test_splitter_config_defaults() {
409 let config = SplitterConfig::new(split_body_lines());
410 assert!(matches!(config.aggregation, AggregationStrategy::LastWins));
411 assert!(!config.parallel);
412 assert!(config.parallel_limit.is_none());
413 assert!(config.stop_on_exception);
414 }
415
416 #[test]
417 fn test_splitter_config_builder() {
418 let config = SplitterConfig::new(split_body_lines())
419 .aggregation(AggregationStrategy::CollectAll)
420 .parallel(true)
421 .parallel_limit(4)
422 .stop_on_exception(false);
423
424 assert!(matches!(
425 config.aggregation,
426 AggregationStrategy::CollectAll
427 ));
428 assert!(config.parallel);
429 assert_eq!(config.parallel_limit, Some(4));
430 assert!(!config.stop_on_exception);
431 }
432
433 #[test]
434 fn test_splitter_config_default_max_fragments() {
435 let cfg = SplitterConfig::new(Arc::new(|_: &Exchange| Vec::new()) as SplitExpression);
436 assert_eq!(cfg.max_fragments, 100_000);
437 }
438
439 #[test]
440 fn test_splitter_config_rejects_zero_max_fragments() {
441 let cfg = SplitterConfig::new(Arc::new(|_: &Exchange| Vec::new()) as SplitExpression)
442 .max_fragments(0);
443 assert!(cfg.validate().is_err());
444 }
445
446 #[test]
447 fn test_fragment_exchange_inherits_otel_context() {
448 use opentelemetry::Context;
449 use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId};
450
451 let mut parent = Exchange::new(Message::new("test"));
453 let trace_id = TraceId::from_bytes([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 123]);
454 let span_id = SpanId::from_bytes([0, 0, 0, 0, 0, 0, 1, 200]);
455 let span_context = SpanContext::new(
456 trace_id,
457 span_id,
458 TraceFlags::SAMPLED,
459 true,
460 Default::default(),
461 );
462 let expected_trace_id = span_context.trace_id();
463 parent.otel_context = Context::current().with_remote_span_context(span_context);
464
465 let fragments = split_body_lines()(&parent);
467 assert!(!fragments.is_empty(), "Should have at least one fragment");
468
469 for fragment in &fragments {
471 let span = fragment.otel_context.span();
472 let frag_span_ctx = span.span_context();
473 assert!(
474 frag_span_ctx.is_valid(),
475 "Fragment should have valid span context"
476 );
477 assert_eq!(
478 frag_span_ctx.trace_id(),
479 expected_trace_id,
480 "Fragment should have same trace ID as parent"
481 );
482 }
483 }
484
485 #[test]
486 fn test_stream_split_config_defaults_valid() {
487 let config = StreamSplitConfig::default();
488 assert!(config.validate().is_ok());
489 }
490
491 #[test]
492 fn test_stream_split_config_batch_size_zero_rejected() {
493 let config = StreamSplitConfig {
494 batch_size: 0,
495 ..Default::default()
496 };
497 let err = config.validate().unwrap_err();
498 assert!(err.to_string().contains("batch_size"));
499 }
500
501 #[test]
502 fn test_stream_split_config_max_record_bytes_zero_rejected() {
503 let config = StreamSplitConfig {
504 max_record_bytes: 0,
505 ..Default::default()
506 };
507 let err = config.validate().unwrap_err();
508 assert!(err.to_string().contains("max_record_bytes"));
509 }
510
511 #[test]
512 fn test_stream_split_config_chunks_requires_chunk_size() {
513 let config = StreamSplitConfig {
514 format: StreamSplitFormat::Chunks,
515 chunk_size: None,
516 ..Default::default()
517 };
518 let err = config.validate().unwrap_err();
519 assert!(err.to_string().contains("Chunks requires chunk_size"));
520 }
521
522 #[test]
523 fn test_stream_split_config_chunk_size_zero_rejected() {
524 let config = StreamSplitConfig {
525 format: StreamSplitFormat::Chunks,
526 chunk_size: Some(0),
527 ..Default::default()
528 };
529 let err = config.validate().unwrap_err();
530 assert!(err.to_string().contains("chunk_size must be > 0"));
531 }
532
533 #[test]
534 fn test_stream_split_config_chunk_size_exceeds_max_record_bytes() {
535 let config = StreamSplitConfig {
536 format: StreamSplitFormat::Chunks,
537 chunk_size: Some(2000),
538 max_record_bytes: 1000,
539 ..Default::default()
540 };
541 let err = config.validate().unwrap_err();
542 assert!(
543 err.to_string()
544 .contains("chunk_size must be <= max_record_bytes")
545 );
546 }
547
548 #[test]
549 fn test_stream_split_config_zip_rejects_chunk_size() {
550 let config = StreamSplitConfig {
551 format: StreamSplitFormat::Zip,
552 chunk_size: Some(1024),
553 ..Default::default()
554 };
555 let err = config.validate().unwrap_err();
556 assert!(err.to_string().contains("Zip does not support chunk_size"));
557 }
558
559 #[test]
560 fn test_all_fragments_share_same_trace_context() {
561 use opentelemetry::Context;
562 use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId};
563
564 let mut parent = Exchange::new(Message::new("line1\nline2\nline3"));
566 let trace_id =
567 TraceId::from_bytes([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x3B, 0x9A, 0xCA, 0x09]);
568 let span_id = SpanId::from_bytes([0, 0, 0, 0, 0, 0, 0, 111]);
569 let span_context = SpanContext::new(
570 trace_id,
571 span_id,
572 TraceFlags::SAMPLED,
573 true,
574 Default::default(),
575 );
576 parent.otel_context = Context::current().with_remote_span_context(span_context);
577
578 let fragments = split_body_lines()(&parent);
579 assert_eq!(fragments.len(), 3);
580
581 let trace_ids: Vec<_> = fragments
583 .iter()
584 .map(|f| {
585 let span = f.otel_context.span();
586 span.span_context().trace_id()
587 })
588 .collect();
589
590 assert!(
591 trace_ids.iter().all(|&id| id == trace_id),
592 "All fragments should have the same trace ID"
593 );
594 }
595}