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 connector_name(&self) -> &'static str {
400 "redis"
401 }
402
403 fn config_schema(&self) -> serde_json::Value {
404 serde_json::to_value(faucet_core::schema_for!(RedisSourceConfig))
405 .expect("schema serialization")
406 }
407
408 fn dataset_uri(&self) -> String {
409 use crate::config::RedisSourceType;
410 let base = faucet_core::redact_uri_credentials(&self.config.url);
411 match &self.config.source_type {
412 RedisSourceType::List { key } => format!("{base}?key={key}"),
413 RedisSourceType::Stream { key, .. } => format!("{base}?stream={key}"),
414 RedisSourceType::Keys { pattern } => format!("{base}?key={pattern}"),
415 }
416 }
417}
418
419fn stream_list<'a>(
430 conn: &'a mut redis::aio::MultiplexedConnection,
431 key: &'a str,
432 batch_size: usize,
433 max_records: Option<usize>,
434) -> impl Stream<Item = Result<StreamPage, FaucetError>> + 'a {
435 async_stream::try_stream! {
436 if batch_size == 0 {
437 let values: Vec<String> = conn
438 .lrange(key, 0, -1)
439 .await
440 .map_err(|e| FaucetError::Source(format!("LRANGE failed on '{key}': {e}")))?;
441 let mut records: Vec<Value> = values
442 .into_iter()
443 .map(|v| serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone())))
444 .collect();
445 if let Some(max) = max_records {
446 records.truncate(max);
447 }
448 yield StreamPage { records, bookmark: None };
449 return;
450 }
451
452 let mut start: isize = 0;
453 let mut emitted: usize = 0;
454 loop {
455 let stop: isize = start + batch_size as isize - 1;
456 let values: Vec<String> = conn
457 .lrange(key, start, stop)
458 .await
459 .map_err(|e| FaucetError::Source(format!("LRANGE failed on '{key}': {e}")))?;
460 if values.is_empty() {
461 break;
462 }
463 let mut records: Vec<Value> = values
464 .into_iter()
465 .map(|v| serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone())))
466 .collect();
467 let returned = records.len();
468 let mut stop_after_yield = false;
470 if let Some(max) = max_records
471 && emitted + records.len() >= max
472 {
473 records.truncate(max - emitted);
474 stop_after_yield = true;
475 }
476 emitted += records.len();
477 yield StreamPage { records, bookmark: None };
478 if stop_after_yield || returned < batch_size {
479 break;
480 }
481 start += batch_size as isize;
482 }
483 }
484}
485
486fn stream_xrange<'a>(
490 conn: &'a mut redis::aio::MultiplexedConnection,
491 key: &'a str,
492 batch_size: usize,
493 max_records: Option<usize>,
494) -> impl Stream<Item = Result<StreamPage, FaucetError>> + 'a {
495 async_stream::try_stream! {
496 if batch_size == 0 {
497 let reply: redis::streams::StreamRangeReply = conn
498 .xrange_all(key)
499 .await
500 .map_err(|e| FaucetError::Source(format!("XRANGE failed on '{key}': {e}")))?;
501 let mut records: Vec<Value> = reply
502 .ids
503 .iter()
504 .map(|entry| stream_entry_to_json(&entry.id, &entry.map))
505 .collect();
506 if let Some(max) = max_records {
507 records.truncate(max);
508 }
509 yield StreamPage { records, bookmark: None };
510 return;
511 }
512
513 let mut start: String = "-".to_string();
514 let mut emitted: usize = 0;
515 loop {
516 let reply: redis::streams::StreamRangeReply = conn
517 .xrange_count(key, &start, "+", batch_size)
518 .await
519 .map_err(|e| FaucetError::Source(format!("XRANGE failed on '{key}': {e}")))?;
520
521 if reply.ids.is_empty() {
522 break;
523 }
524
525 let last_id = reply
528 .ids
529 .last()
530 .expect("non-empty checked above")
531 .id
532 .clone();
533 let returned = reply.ids.len();
534 let mut records: Vec<Value> = reply
535 .ids
536 .into_iter()
537 .map(|entry| stream_entry_to_json(&entry.id, &entry.map))
538 .collect();
539
540 let mut stop_after_yield = false;
541 if let Some(max) = max_records
542 && emitted + records.len() >= max
543 {
544 records.truncate(max - emitted);
545 stop_after_yield = true;
546 }
547 emitted += records.len();
548 yield StreamPage { records, bookmark: None };
549
550 if stop_after_yield || returned < batch_size {
551 break;
552 }
553 start = next_stream_id(&last_id);
554 }
555 }
556}
557
558fn stream_keys<'a>(
564 conn: &'a mut redis::aio::MultiplexedConnection,
565 pattern: &'a str,
566 batch_size: usize,
567 max_records: Option<usize>,
568) -> impl Stream<Item = Result<StreamPage, FaucetError>> + 'a {
569 use faucet_core::DEFAULT_BATCH_SIZE;
570 async_stream::try_stream! {
571 let scan_hint = if batch_size == 0 { DEFAULT_BATCH_SIZE } else { batch_size };
578 let chunk_size = if batch_size == 0 { usize::MAX } else { batch_size };
581 let cap = max_records.unwrap_or(usize::MAX);
582
583 let mut cursor: u64 = 0;
584 let mut buffer: Vec<String> = Vec::new();
585 let mut emitted: usize = 0;
586
587 'scan: loop {
588 let (next_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
589 .arg(cursor)
590 .arg("MATCH")
591 .arg(pattern)
592 .arg("COUNT")
593 .arg(scan_hint)
594 .query_async(conn)
595 .await
596 .map_err(|e| FaucetError::Source(format!("SCAN failed with pattern '{pattern}': {e}")))?;
597 cursor = next_cursor;
598 buffer.extend(keys);
599
600 while emitted < cap && buffer.len() >= chunk_size {
602 let take = chunk_size.min(cap - emitted);
603 let page_keys: Vec<String> = buffer.drain(..take).collect();
604 let records = mget_records(conn, &page_keys).await?;
605 emitted += records.len();
606 yield StreamPage { records, bookmark: None };
607 }
608
609 if cursor == 0 || emitted >= cap {
610 break 'scan;
611 }
612 }
613
614 if emitted < cap && !buffer.is_empty() {
616 let take = (cap - emitted).min(buffer.len());
617 let page_keys: Vec<String> = buffer.drain(..take).collect();
618 let records = mget_records(conn, &page_keys).await?;
619 yield StreamPage { records, bookmark: None };
620 }
621 }
622}
623
624async fn mget_records(
627 conn: &mut redis::aio::MultiplexedConnection,
628 keys: &[String],
629) -> Result<Vec<Value>, FaucetError> {
630 let values: Vec<Option<String>> = redis::cmd("MGET")
631 .arg(keys)
632 .query_async(conn)
633 .await
634 .map_err(|e| FaucetError::Source(format!("MGET failed: {e}")))?;
635 Ok(collect_kv_records(keys, values))
636}
637
638fn collect_kv_records(keys: &[String], values: Vec<Option<String>>) -> Vec<Value> {
642 keys.iter()
643 .zip(values)
644 .filter_map(|(key, value)| {
645 value.map(|v| {
646 let parsed =
647 serde_json::from_str::<Value>(&v).unwrap_or_else(|_| Value::String(v.clone()));
648 json!({ "key": key, "value": parsed })
649 })
650 })
651 .collect()
652}
653
654fn stream_ignores_consumer_group(group: Option<&str>, consumer: Option<&str>) -> bool {
659 group.is_some() || consumer.is_some()
660}
661
662#[cfg(test)]
663mod tests {
664 use super::*;
665 use crate::config::RedisSourceConfig;
666
667 #[test]
668 fn stream_ignores_consumer_group_flags_configured_group() {
669 assert!(stream_ignores_consumer_group(Some("g"), Some("c")));
671 assert!(stream_ignores_consumer_group(Some("g"), None));
672 assert!(stream_ignores_consumer_group(None, Some("c")));
673 assert!(!stream_ignores_consumer_group(None, None));
675 }
676
677 #[test]
678 fn creates_source() {
679 let config = RedisSourceConfig::new(
680 "redis://localhost",
681 RedisSourceType::List { key: "test".into() },
682 );
683 let _source = RedisSource::new(config).unwrap();
684 }
685
686 #[test]
687 fn dataset_uri_list_source() {
688 use faucet_core::Source;
689 let source = RedisSource::new(RedisSourceConfig::new(
690 "redis://u:p@localhost:6379/0",
691 RedisSourceType::List {
692 key: "my-list".into(),
693 },
694 ))
695 .unwrap();
696 assert_eq!(source.dataset_uri(), "redis://localhost:6379/0?key=my-list");
697 }
698
699 #[test]
700 fn dataset_uri_stream_source() {
701 use faucet_core::Source;
702 let config = RedisSourceConfig::new(
703 "redis://localhost",
704 RedisSourceType::Stream {
705 key: "events".into(),
706 group: None,
707 consumer: None,
708 count: None,
709 },
710 );
711 let source = RedisSource::new(config).unwrap();
712 assert_eq!(source.dataset_uri(), "redis://localhost?stream=events");
713 }
714
715 #[test]
716 fn dataset_uri_keys_source() {
717 use faucet_core::Source;
718 let source = RedisSource::new(RedisSourceConfig::new(
719 "redis://u:p@localhost:6379/0",
720 RedisSourceType::Keys {
721 pattern: "user:*".into(),
722 },
723 ))
724 .unwrap();
725 assert_eq!(
726 source.dataset_uri(),
727 "redis://localhost:6379/0?key=user:*",
728 "keys variant renders the glob pattern as the key, with credentials redacted"
729 );
730 }
731
732 #[test]
733 fn config_schema_describes_redis_source_config() {
734 use faucet_core::Source;
735 let source = RedisSource::new(RedisSourceConfig::new(
736 "redis://localhost",
737 RedisSourceType::List { key: "k".into() },
738 ))
739 .unwrap();
740 let schema = source.config_schema();
741 let props = &schema["properties"];
743 assert!(props.get("url").is_some(), "schema exposes 'url'");
744 assert!(
745 props.get("source_type").is_some(),
746 "schema exposes 'source_type'"
747 );
748 assert!(
749 props.get("batch_size").is_some(),
750 "schema exposes 'batch_size'"
751 );
752 }
753
754 #[test]
755 fn new_rejects_out_of_range_batch_size() {
756 let mut config = RedisSourceConfig::new(
757 "redis://localhost",
758 RedisSourceType::List { key: "test".into() },
759 );
760 config.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
761 match RedisSource::new(config) {
762 Err(FaucetError::Config(m)) => assert!(m.contains("batch_size"), "got: {m}"),
763 other => panic!(
764 "expected a batch_size Config error, got {:?}",
765 other.is_ok()
766 ),
767 }
768 }
769
770 #[test]
771 fn next_stream_id_increments_sequence() {
772 assert_eq!(next_stream_id("1234-0"), "1234-1");
773 assert_eq!(next_stream_id("1234-99"), "1234-100");
774 }
775
776 #[test]
777 fn next_stream_id_wraps_seq_overflow() {
778 let id = format!("5-{}", u64::MAX);
779 assert_eq!(next_stream_id(&id), "6-0");
780 }
781
782 #[test]
783 fn next_stream_id_falls_back_on_malformed_id() {
784 let next = next_stream_id("not-a-real-id");
786 assert!(next.starts_with("not-a-real-id"));
787 assert!(next.ends_with('\u{0}'));
788 }
789
790 #[test]
791 fn stream_entry_to_json_extracts_id_and_fields() {
792 let mut map = std::collections::HashMap::new();
793 map.insert(
794 "field1".to_string(),
795 redis::Value::BulkString(b"value1".to_vec()),
796 );
797 map.insert("field2".to_string(), redis::Value::Int(42));
798 let json = stream_entry_to_json("100-0", &map);
799 assert_eq!(json["id"], "100-0");
800 assert_eq!(json["fields"]["field1"], "value1");
801 assert_eq!(json["fields"]["field2"], 42);
802 }
803}