1use crate::config::{RedisSourceConfig, RedisSourceType};
4use async_trait::async_trait;
5use faucet_core::{FaucetError, Stream, StreamPage};
6use redis::AsyncCommands;
7use serde_json::{Value, json};
8use std::pin::Pin;
9
10pub struct RedisSource {
12 config: RedisSourceConfig,
13 conn: tokio::sync::OnceCell<redis::aio::MultiplexedConnection>,
18}
19
20impl RedisSource {
21 pub fn new(config: RedisSourceConfig) -> Result<Self, FaucetError> {
25 faucet_core::validate_batch_size(config.batch_size)?;
26 Ok(Self {
27 config,
28 conn: tokio::sync::OnceCell::new(),
29 })
30 }
31
32 async fn connection(&self) -> Result<redis::aio::MultiplexedConnection, FaucetError> {
35 let conn = self
36 .conn
37 .get_or_try_init(|| async {
38 let client = redis::Client::open(self.config.url.as_str())
39 .map_err(|e| FaucetError::Config(format!("invalid Redis URL: {e}")))?;
40 client
41 .get_multiplexed_async_connection()
42 .await
43 .map_err(|e| FaucetError::Source(format!("Redis connection failed: {e}")))
44 })
45 .await?;
46 Ok(conn.clone())
47 }
48
49 pub async fn fetch_all(&self) -> Result<Vec<Value>, FaucetError> {
51 let mut conn = self.connection().await?;
52
53 let mut records = match &self.config.source_type {
54 RedisSourceType::List { key } => self.fetch_list(&mut conn, key).await?,
55 RedisSourceType::Stream {
56 key,
57 group,
58 consumer,
59 count,
60 } => {
61 self.fetch_stream(&mut conn, key, group, consumer, count)
62 .await?
63 }
64 RedisSourceType::Keys { pattern } => self.fetch_keys(&mut conn, pattern).await?,
65 };
66
67 if let Some(max) = self.config.max_records {
68 records.truncate(max);
69 }
70
71 tracing::info!(records = records.len(), "Redis fetch complete");
72 Ok(records)
73 }
74
75 async fn fetch_list(
77 &self,
78 conn: &mut redis::aio::MultiplexedConnection,
79 key: &str,
80 ) -> Result<Vec<Value>, FaucetError> {
81 let values: Vec<String> = conn
82 .lrange(key, 0, -1)
83 .await
84 .map_err(|e| FaucetError::Source(format!("LRANGE failed on '{key}': {e}")))?;
85
86 let records = values
87 .into_iter()
88 .map(|v| serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone())))
89 .collect();
90
91 Ok(records)
92 }
93
94 async fn fetch_stream(
96 &self,
97 conn: &mut redis::aio::MultiplexedConnection,
98 key: &str,
99 group: &Option<String>,
100 consumer: &Option<String>,
101 count: &Option<usize>,
102 ) -> Result<Vec<Value>, FaucetError> {
103 let mut records = Vec::new();
104 match (group, consumer) {
105 (Some(group_name), Some(consumer_name)) => {
106 let per_read = count.unwrap_or(100).max(1);
112 loop {
113 let opts = redis::streams::StreamReadOptions::default().count(per_read);
114 let reply: redis::streams::StreamReadReply = conn
115 .xread_options(&[key], &[">"], &opts.group(group_name, consumer_name))
116 .await
117 .map_err(|e| {
118 FaucetError::Source(format!("XREADGROUP failed on '{key}': {e}"))
119 })?;
120 let mut got = 0usize;
121 for stream_key in &reply.keys {
122 for entry in &stream_key.ids {
123 records.push(stream_entry_to_json(&entry.id, &entry.map));
124 got += 1;
125 }
126 }
127 if got < per_read {
131 break;
132 }
133 if let Some(max) = self.config.max_records
134 && records.len() >= max
135 {
136 break;
137 }
138 }
139 }
140 _ => {
141 let mut opts = redis::streams::StreamReadOptions::default();
144 if let Some(c) = count {
145 opts = opts.count(*c);
146 }
147 let reply: redis::streams::StreamReadReply = conn
148 .xread_options(&[key], &["0"], &opts)
149 .await
150 .map_err(|e| FaucetError::Source(format!("XREAD failed on '{key}': {e}")))?;
151 for stream_key in &reply.keys {
152 for entry in &stream_key.ids {
153 records.push(stream_entry_to_json(&entry.id, &entry.map));
154 }
155 }
156 }
157 }
158
159 Ok(records)
160 }
161
162 async fn fetch_keys(
164 &self,
165 conn: &mut redis::aio::MultiplexedConnection,
166 pattern: &str,
167 ) -> Result<Vec<Value>, FaucetError> {
168 let keys: Vec<String> = {
169 let mut collected = Vec::new();
170 let mut iter: redis::AsyncIter<String> =
171 conn.scan_match(pattern).await.map_err(|e| {
172 FaucetError::Source(format!("SCAN failed with pattern '{pattern}': {e}"))
173 })?;
174
175 while let Some(key) = iter.next_item().await {
176 collected.push(key);
177 }
178 collected
179 };
180
181 if keys.is_empty() {
182 return Ok(Vec::new());
183 }
184
185 let values: Vec<Option<String>> = redis::cmd("MGET")
186 .arg(&keys)
187 .query_async(conn)
188 .await
189 .map_err(|e| FaucetError::Source(format!("MGET failed: {e}")))?;
190
191 let mut records = Vec::new();
192 for (key, value) in keys.iter().zip(values) {
193 if let Some(v) = value {
194 let parsed =
195 serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone()));
196 records.push(json!({
197 "key": key,
198 "value": parsed,
199 }));
200 }
201 }
202
203 Ok(records)
204 }
205}
206
207fn stream_entry_to_json(id: &str, map: &std::collections::HashMap<String, redis::Value>) -> Value {
210 let mut fields = serde_json::Map::new();
211 for (field_name, field_value) in map {
212 let val = match field_value {
213 redis::Value::BulkString(bytes) => {
214 let s = String::from_utf8_lossy(bytes);
215 serde_json::from_str::<Value>(&s).unwrap_or_else(|_| Value::String(s.into_owned()))
216 }
217 redis::Value::SimpleString(s) => {
218 serde_json::from_str::<Value>(s).unwrap_or_else(|_| Value::String(s.clone()))
219 }
220 redis::Value::Int(n) => json!(n),
221 redis::Value::Double(n) => json!(n),
222 redis::Value::Boolean(b) => json!(b),
223 redis::Value::Nil => Value::Null,
224 other => Value::String(format!("{other:?}")),
225 };
226 fields.insert(field_name.clone(), val);
227 }
228 json!({
229 "id": id,
230 "fields": Value::Object(fields),
231 })
232}
233
234fn next_stream_id(id: &str) -> String {
238 if let Some((ms, seq)) = id.split_once('-')
241 && let (Ok(ms), Ok(seq)) = (ms.parse::<u64>(), seq.parse::<u64>())
242 {
243 return match seq.checked_add(1) {
244 Some(next_seq) => format!("{ms}-{next_seq}"),
245 None => format!("{}-0", ms.saturating_add(1)),
246 };
247 }
248 format!("{id}\u{0}")
252}
253
254#[async_trait]
255impl faucet_core::Source for RedisSource {
256 async fn fetch_with_context(
257 &self,
258 context: &std::collections::HashMap<String, serde_json::Value>,
259 ) -> Result<Vec<Value>, FaucetError> {
260 if context.is_empty() {
261 return RedisSource::fetch_all(self).await;
262 }
263
264 let mut conn = self.connection().await?;
265
266 let mut records = match &self.config.source_type {
268 RedisSourceType::List { key } => {
269 let resolved_key = faucet_core::util::substitute_context(key, context);
270 self.fetch_list(&mut conn, &resolved_key).await?
271 }
272 RedisSourceType::Stream {
273 key,
274 group,
275 consumer,
276 count,
277 } => {
278 let resolved_key = faucet_core::util::substitute_context(key, context);
279 self.fetch_stream(&mut conn, &resolved_key, group, consumer, count)
280 .await?
281 }
282 RedisSourceType::Keys { pattern } => {
283 let resolved_pattern = faucet_core::util::substitute_context(pattern, context);
284 self.fetch_keys(&mut conn, &resolved_pattern).await?
285 }
286 };
287
288 if let Some(max) = self.config.max_records {
289 records.truncate(max);
290 }
291
292 tracing::info!(
293 records = records.len(),
294 "Redis fetch complete (with context)"
295 );
296 Ok(records)
297 }
298
299 fn stream_pages<'a>(
313 &'a self,
314 context: &'a std::collections::HashMap<String, Value>,
315 _batch_size: usize,
316 ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
317 let batch_size = self.config.batch_size;
318 let max_records = self.config.max_records;
319
320 Box::pin(async_stream::try_stream! {
321 let mut conn = self.connection().await?;
322
323 let mut emitted: usize = 0;
324
325 match &self.config.source_type {
326 RedisSourceType::List { key } => {
327 let resolved = if context.is_empty() {
328 key.clone()
329 } else {
330 faucet_core::util::substitute_context(key, context)
331 };
332 let pages = stream_list(&mut conn, &resolved, batch_size, max_records);
333 futures::pin_mut!(pages);
334 while let Some(page) = futures::StreamExt::next(&mut pages).await {
335 let page = page?;
336 emitted += page.records.len();
337 yield page;
338 }
339 }
340 RedisSourceType::Stream { key, group, consumer, .. } => {
341 if stream_ignores_consumer_group(group.as_deref(), consumer.as_deref()) {
353 tracing::warn!(
354 stream = %key,
355 "Redis Stream source has a consumer group/consumer configured, but \
356 the streaming path (stream_pages) ignores it and uses XRANGE — it \
357 re-reads the entire stream every run. Use the fetch_all batch path \
358 for XREADGROUP consume-once semantics, or drop group/consumer to \
359 silence this warning."
360 );
361 }
362 let resolved = if context.is_empty() {
363 key.clone()
364 } else {
365 faucet_core::util::substitute_context(key, context)
366 };
367 let pages = stream_xrange(&mut conn, &resolved, batch_size, max_records);
368 futures::pin_mut!(pages);
369 while let Some(page) = futures::StreamExt::next(&mut pages).await {
370 let page = page?;
371 emitted += page.records.len();
372 yield page;
373 }
374 }
375 RedisSourceType::Keys { pattern } => {
376 let resolved = if context.is_empty() {
377 pattern.clone()
378 } else {
379 faucet_core::util::substitute_context(pattern, context)
380 };
381 let pages = stream_keys(&mut conn, &resolved, batch_size, max_records);
382 futures::pin_mut!(pages);
383 while let Some(page) = futures::StreamExt::next(&mut pages).await {
384 let page = page?;
385 emitted += page.records.len();
386 yield page;
387 }
388 }
389 }
390
391 tracing::info!(
392 records = emitted,
393 batch_size,
394 "Redis source stream complete",
395 );
396 })
397 }
398
399 fn config_schema(&self) -> serde_json::Value {
400 serde_json::to_value(faucet_core::schema_for!(RedisSourceConfig))
401 .expect("schema serialization")
402 }
403
404 fn dataset_uri(&self) -> String {
405 use crate::config::RedisSourceType;
406 let base = faucet_core::redact_uri_credentials(&self.config.url);
407 match &self.config.source_type {
408 RedisSourceType::List { key } => format!("{base}?key={key}"),
409 RedisSourceType::Stream { key, .. } => format!("{base}?stream={key}"),
410 RedisSourceType::Keys { pattern } => format!("{base}?key={pattern}"),
411 }
412 }
413}
414
415fn stream_list<'a>(
426 conn: &'a mut redis::aio::MultiplexedConnection,
427 key: &'a str,
428 batch_size: usize,
429 max_records: Option<usize>,
430) -> impl Stream<Item = Result<StreamPage, FaucetError>> + 'a {
431 async_stream::try_stream! {
432 if batch_size == 0 {
433 let values: Vec<String> = conn
434 .lrange(key, 0, -1)
435 .await
436 .map_err(|e| FaucetError::Source(format!("LRANGE failed on '{key}': {e}")))?;
437 let mut records: Vec<Value> = values
438 .into_iter()
439 .map(|v| serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone())))
440 .collect();
441 if let Some(max) = max_records {
442 records.truncate(max);
443 }
444 yield StreamPage { records, bookmark: None };
445 return;
446 }
447
448 let mut start: isize = 0;
449 let mut emitted: usize = 0;
450 loop {
451 let stop: isize = start + batch_size as isize - 1;
452 let values: Vec<String> = conn
453 .lrange(key, start, stop)
454 .await
455 .map_err(|e| FaucetError::Source(format!("LRANGE failed on '{key}': {e}")))?;
456 if values.is_empty() {
457 break;
458 }
459 let mut records: Vec<Value> = values
460 .into_iter()
461 .map(|v| serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone())))
462 .collect();
463 let returned = records.len();
464 let mut stop_after_yield = false;
466 if let Some(max) = max_records
467 && emitted + records.len() >= max
468 {
469 records.truncate(max - emitted);
470 stop_after_yield = true;
471 }
472 emitted += records.len();
473 yield StreamPage { records, bookmark: None };
474 if stop_after_yield || returned < batch_size {
475 break;
476 }
477 start += batch_size as isize;
478 }
479 }
480}
481
482fn stream_xrange<'a>(
486 conn: &'a mut redis::aio::MultiplexedConnection,
487 key: &'a str,
488 batch_size: usize,
489 max_records: Option<usize>,
490) -> impl Stream<Item = Result<StreamPage, FaucetError>> + 'a {
491 async_stream::try_stream! {
492 if batch_size == 0 {
493 let reply: redis::streams::StreamRangeReply = conn
494 .xrange_all(key)
495 .await
496 .map_err(|e| FaucetError::Source(format!("XRANGE failed on '{key}': {e}")))?;
497 let mut records: Vec<Value> = reply
498 .ids
499 .iter()
500 .map(|entry| stream_entry_to_json(&entry.id, &entry.map))
501 .collect();
502 if let Some(max) = max_records {
503 records.truncate(max);
504 }
505 yield StreamPage { records, bookmark: None };
506 return;
507 }
508
509 let mut start: String = "-".to_string();
510 let mut emitted: usize = 0;
511 loop {
512 let reply: redis::streams::StreamRangeReply = conn
513 .xrange_count(key, &start, "+", batch_size)
514 .await
515 .map_err(|e| FaucetError::Source(format!("XRANGE failed on '{key}': {e}")))?;
516
517 if reply.ids.is_empty() {
518 break;
519 }
520
521 let last_id = reply
524 .ids
525 .last()
526 .expect("non-empty checked above")
527 .id
528 .clone();
529 let returned = reply.ids.len();
530 let mut records: Vec<Value> = reply
531 .ids
532 .into_iter()
533 .map(|entry| stream_entry_to_json(&entry.id, &entry.map))
534 .collect();
535
536 let mut stop_after_yield = false;
537 if let Some(max) = max_records
538 && emitted + records.len() >= max
539 {
540 records.truncate(max - emitted);
541 stop_after_yield = true;
542 }
543 emitted += records.len();
544 yield StreamPage { records, bookmark: None };
545
546 if stop_after_yield || returned < batch_size {
547 break;
548 }
549 start = next_stream_id(&last_id);
550 }
551 }
552}
553
554fn stream_keys<'a>(
560 conn: &'a mut redis::aio::MultiplexedConnection,
561 pattern: &'a str,
562 batch_size: usize,
563 max_records: Option<usize>,
564) -> impl Stream<Item = Result<StreamPage, FaucetError>> + 'a {
565 use faucet_core::DEFAULT_BATCH_SIZE;
566 async_stream::try_stream! {
567 let scan_hint = if batch_size == 0 { DEFAULT_BATCH_SIZE } else { batch_size };
574 let chunk_size = if batch_size == 0 { usize::MAX } else { batch_size };
577 let cap = max_records.unwrap_or(usize::MAX);
578
579 let mut cursor: u64 = 0;
580 let mut buffer: Vec<String> = Vec::new();
581 let mut emitted: usize = 0;
582
583 'scan: loop {
584 let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
585 .arg(cursor)
586 .arg("MATCH")
587 .arg(pattern)
588 .arg("COUNT")
589 .arg(scan_hint)
590 .query_async(conn)
591 .await
592 .map_err(|e| FaucetError::Source(format!("SCAN failed with pattern '{pattern}': {e}")))?;
593 cursor = next_cursor;
594 buffer.extend(keys);
595
596 while emitted < cap && buffer.len() >= chunk_size {
598 let take = chunk_size.min(cap - emitted);
599 let page_keys: Vec<String> = buffer.drain(..take).collect();
600 let records = mget_records(conn, &page_keys).await?;
601 emitted += records.len();
602 yield StreamPage { records, bookmark: None };
603 }
604
605 if cursor == 0 || emitted >= cap {
606 break 'scan;
607 }
608 }
609
610 if emitted < cap && !buffer.is_empty() {
612 let take = (cap - emitted).min(buffer.len());
613 let page_keys: Vec<String> = buffer.drain(..take).collect();
614 let records = mget_records(conn, &page_keys).await?;
615 yield StreamPage { records, bookmark: None };
616 }
617 }
618}
619
620async fn mget_records(
623 conn: &mut redis::aio::MultiplexedConnection,
624 keys: &[String],
625) -> Result<Vec<Value>, FaucetError> {
626 let values: Vec<Option<String>> = redis::cmd("MGET")
627 .arg(keys)
628 .query_async(conn)
629 .await
630 .map_err(|e| FaucetError::Source(format!("MGET failed: {e}")))?;
631 Ok(collect_kv_records(keys, values))
632}
633
634fn collect_kv_records(keys: &[String], values: Vec<Option<String>>) -> Vec<Value> {
638 keys.iter()
639 .zip(values)
640 .filter_map(|(key, value)| {
641 value.map(|v| {
642 let parsed =
643 serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone()));
644 json!({ "key": key, "value": parsed })
645 })
646 })
647 .collect()
648}
649
650fn stream_ignores_consumer_group(group: Option<&str>, consumer: Option<&str>) -> bool {
655 group.is_some() || consumer.is_some()
656}
657
658#[cfg(test)]
659mod tests {
660 use super::*;
661 use crate::config::RedisSourceConfig;
662
663 #[test]
664 fn stream_ignores_consumer_group_flags_configured_group() {
665 assert!(stream_ignores_consumer_group(Some("g"), Some("c")));
667 assert!(stream_ignores_consumer_group(Some("g"), None));
668 assert!(stream_ignores_consumer_group(None, Some("c")));
669 assert!(!stream_ignores_consumer_group(None, None));
671 }
672
673 #[test]
674 fn creates_source() {
675 let config = RedisSourceConfig::new(
676 "redis://localhost",
677 RedisSourceType::List { key: "test".into() },
678 );
679 let _source = RedisSource::new(config).unwrap();
680 }
681
682 #[test]
683 fn dataset_uri_list_source() {
684 use faucet_core::Source;
685 let source = RedisSource::new(RedisSourceConfig::new(
686 "redis://u:p@localhost:6379/0",
687 RedisSourceType::List {
688 key: "my-list".into(),
689 },
690 ))
691 .unwrap();
692 assert_eq!(source.dataset_uri(), "redis://localhost:6379/0?key=my-list");
693 }
694
695 #[test]
696 fn dataset_uri_stream_source() {
697 use faucet_core::Source;
698 let config = RedisSourceConfig::new(
699 "redis://localhost",
700 RedisSourceType::Stream {
701 key: "events".into(),
702 group: None,
703 consumer: None,
704 count: None,
705 },
706 );
707 let source = RedisSource::new(config).unwrap();
708 assert_eq!(source.dataset_uri(), "redis://localhost?stream=events");
709 }
710
711 #[test]
712 fn dataset_uri_keys_source() {
713 use faucet_core::Source;
714 let source = RedisSource::new(RedisSourceConfig::new(
715 "redis://u:p@localhost:6379/0",
716 RedisSourceType::Keys {
717 pattern: "user:*".into(),
718 },
719 ))
720 .unwrap();
721 assert_eq!(
722 source.dataset_uri(),
723 "redis://localhost:6379/0?key=user:*",
724 "keys variant renders the glob pattern as the key, with credentials redacted"
725 );
726 }
727
728 #[test]
729 fn config_schema_describes_redis_source_config() {
730 use faucet_core::Source;
731 let source = RedisSource::new(RedisSourceConfig::new(
732 "redis://localhost",
733 RedisSourceType::List { key: "k".into() },
734 ))
735 .unwrap();
736 let schema = source.config_schema();
737 let props = &schema["properties"];
739 assert!(props.get("url").is_some(), "schema exposes 'url'");
740 assert!(
741 props.get("source_type").is_some(),
742 "schema exposes 'source_type'"
743 );
744 assert!(
745 props.get("batch_size").is_some(),
746 "schema exposes 'batch_size'"
747 );
748 }
749
750 #[test]
751 fn new_rejects_out_of_range_batch_size() {
752 let mut config = RedisSourceConfig::new(
753 "redis://localhost",
754 RedisSourceType::List { key: "test".into() },
755 );
756 config.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
757 match RedisSource::new(config) {
758 Err(FaucetError::Config(m)) => assert!(m.contains("batch_size"), "got: {m}"),
759 other => panic!(
760 "expected a batch_size Config error, got {:?}",
761 other.is_ok()
762 ),
763 }
764 }
765
766 #[test]
767 fn next_stream_id_increments_sequence() {
768 assert_eq!(next_stream_id("1234-0"), "1234-1");
769 assert_eq!(next_stream_id("1234-99"), "1234-100");
770 }
771
772 #[test]
773 fn next_stream_id_wraps_seq_overflow() {
774 let id = format!("5-{}", u64::MAX);
775 assert_eq!(next_stream_id(&id), "6-0");
776 }
777
778 #[test]
779 fn next_stream_id_falls_back_on_malformed_id() {
780 let next = next_stream_id("not-a-real-id");
782 assert!(next.starts_with("not-a-real-id"));
783 assert!(next.ends_with('\u{0}'));
784 }
785
786 #[test]
787 fn stream_entry_to_json_extracts_id_and_fields() {
788 let mut map = std::collections::HashMap::new();
789 map.insert(
790 "field1".to_string(),
791 redis::Value::BulkString(b"value1".to_vec()),
792 );
793 map.insert("field2".to_string(), redis::Value::Int(42));
794 let json = stream_entry_to_json("100-0", &map);
795 assert_eq!(json["id"], "100-0");
796 assert_eq!(json["fields"]["field1"], "value1");
797 assert_eq!(json["fields"]["field2"], 42);
798 }
799}