fraiseql_wire/stream/
filter.rs

1//! Filtered JSON stream
2
3use crate::Result;
4use futures::stream::Stream;
5use serde_json::Value;
6use std::pin::Pin;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::task::{Context, Poll};
9
10/// Predicate function type
11pub type Predicate = Box<dyn Fn(&Value) -> bool + Send>;
12
13/// Filtered JSON stream
14pub struct FilteredStream<S> {
15    inner: S,
16    predicate: Predicate,
17    // Sampling counter for metrics recording
18    eval_count: AtomicU64,
19}
20
21impl<S> FilteredStream<S> {
22    /// Create new filtered stream
23    pub fn new(inner: S, predicate: Predicate) -> Self {
24        Self {
25            inner,
26            predicate,
27            eval_count: AtomicU64::new(0),
28        }
29    }
30}
31
32impl<S> Stream for FilteredStream<S>
33where
34    S: Stream<Item = Result<Value>> + Unpin,
35{
36    type Item = Result<Value>;
37
38    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
39        loop {
40            match Pin::new(&mut self.inner).poll_next(cx) {
41                Poll::Ready(Some(Ok(value))) => {
42                    // Apply predicate
43                    // Sample timing: only record 1 in 1000 evaluations
44                    let eval_idx = self.eval_count.fetch_add(1, Ordering::Relaxed);
45                    let passed = if eval_idx % 1000 == 0 {
46                        // Record timing for sampled evaluation
47                        let filter_start = std::time::Instant::now();
48                        let result = (self.predicate)(&value);
49                        let filter_duration = filter_start.elapsed().as_millis() as u64;
50                        crate::metrics::histograms::filter_duration("unknown", filter_duration);
51                        result
52                    } else {
53                        // No timing, just evaluate
54                        (self.predicate)(&value)
55                    };
56
57                    if passed {
58                        return Poll::Ready(Some(Ok(value)));
59                    }
60                    // Predicate failed, try next value (filter out this row)
61                    crate::metrics::counters::rows_filtered("unknown", 1);
62                    continue;
63                }
64                Poll::Ready(Some(Err(e))) => {
65                    // Propagate errors
66                    return Poll::Ready(Some(Err(e)));
67                }
68                Poll::Ready(None) => {
69                    // End of stream
70                    return Poll::Ready(None);
71                }
72                Poll::Pending => {
73                    return Poll::Pending;
74                }
75            }
76        }
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::Error;
84    use futures::{stream, StreamExt};
85
86    #[tokio::test]
87    async fn test_filter_stream() {
88        let values = vec![
89            Ok(serde_json::json!({"id": 1, "active": true})),
90            Ok(serde_json::json!({"id": 2, "active": false})),
91            Ok(serde_json::json!({"id": 3, "active": true})),
92        ];
93
94        let inner = stream::iter(values);
95
96        let predicate: Predicate = Box::new(|v| v["active"].as_bool().unwrap_or(false));
97
98        let mut filtered = FilteredStream::new(inner, predicate);
99
100        let mut results = Vec::new();
101        while let Some(item) = filtered.next().await {
102            let value = item.unwrap();
103            results.push(value["id"].as_i64().unwrap());
104        }
105
106        assert_eq!(results, vec![1, 3]);
107    }
108
109    #[tokio::test]
110    async fn test_filter_propagates_errors() {
111        let values = vec![
112            Ok(serde_json::json!({"id": 1})),
113            Err(Error::JsonDecode(serde_json::Error::io(
114                std::io::Error::other("test error"),
115            ))),
116            Ok(serde_json::json!({"id": 2})),
117        ];
118
119        let inner = stream::iter(values);
120        let predicate: Predicate = Box::new(|_| true);
121
122        let mut filtered = FilteredStream::new(inner, predicate);
123
124        // First item OK
125        assert!(filtered.next().await.unwrap().is_ok());
126
127        // Second item is error
128        assert!(filtered.next().await.unwrap().is_err());
129
130        // Third item OK
131        assert!(filtered.next().await.unwrap().is_ok());
132    }
133
134    #[tokio::test]
135    async fn test_filter_all_filtered_out() {
136        let values = vec![
137            Ok(serde_json::json!({"id": 1})),
138            Ok(serde_json::json!({"id": 2})),
139        ];
140
141        let inner = stream::iter(values);
142        let predicate: Predicate = Box::new(|_| false); // Filter everything
143
144        let mut filtered = FilteredStream::new(inner, predicate);
145
146        // Stream should be empty
147        assert!(filtered.next().await.is_none());
148    }
149}
fraiseql_wire/stream/filter.rs

fraiseql_wire/stream/
filter.rs
