sift_query_sqlx/execution/
stream.rs

1use std::pin::Pin;
2use std::sync::Arc;
3use std::time::Duration;
4
5use async_stream::try_stream;
6use futures_core::Stream;
7use futures_util::{StreamExt, TryStreamExt};
8use sift_query_core::{
9    GroupScope, PagingMode, QueryInput, QueryPlan, QueryTarget, SearchEntity, ServerScope,
10};
11use sqlx::FromRow;
12use sqlx::postgres::PgRow;
13use tokio::time::{Instant, timeout_at};
14
15use crate::cache::RenderRequest;
16use crate::execution::io::{close_hard, configure_statement_timeout, finish_statement, prepare};
17use crate::postgres::decode::decode_map;
18use crate::{PgSearchError, PgSiftQuery, RowMap, TimeoutStage};
19
20enum Outcome {
21    Done,
22    Failed(sqlx::Error),
23    TimedOut,
24}
25
26pub type SearchStream<T> = Pin<Box<dyn Stream<Item = Result<T, PgSearchError>> + Send + 'static>>;
27pub type ScanFuture<T> =
28    Pin<Box<dyn std::future::Future<Output = Result<Vec<T>, PgSearchError>> + Send>>;
29
30/// Server-owned single-statement scan budget. PostgreSQL provides one statement
31/// snapshot; `batch_size` controls cooperative handoff, not cursor pagination.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct ScanPolicy {
34    pub max_rows: u32,
35    pub batch_size: u32,
36    pub timeout: Duration,
37}
38impl Default for ScanPolicy {
39    fn default() -> Self {
40        Self {
41            max_rows: 10_000,
42            batch_size: 256,
43            timeout: Duration::from_secs(60),
44        }
45    }
46}
47impl ScanPolicy {
48    pub fn validate(self) -> Result<Self, PgSearchError> {
49        if self.max_rows == 0
50            || self.max_rows == u32::MAX
51            || self.batch_size == 0
52            || self.batch_size > 16_384
53            || self.timeout.is_zero()
54        {
55            return Err(sift_query_core::ParamError::InvalidValue {
56                field: "scan".to_owned(),
57                reason: "scan requires finite row, batch and time budgets".to_owned(),
58            }
59            .into());
60        }
61        Ok(self)
62    }
63}
64
65impl PgSiftQuery {
66    /// Streams a page without count. Early drop closes the checked-out
67    /// connection so an active statement or timeout cannot return to the pool.
68    pub fn search_stream<E>(
69        &self,
70        query: impl QueryInput<E>,
71        scope: &ServerScope<E>,
72    ) -> Result<SearchStream<E>, PgSearchError>
73    where
74        E: SearchEntity + Send + Unpin + 'static + for<'row> FromRow<'row, PgRow>,
75    {
76        self.observe_stream::<E, _>("search_stream", || {
77            let plan = self.resolve_plan(query, scope, QueryTarget::TypedList, "search_stream")?;
78            self.stream_plan(plan, "search_stream")
79        })
80    }
81
82    pub(crate) fn stream_plan<E>(
83        &self,
84        plan: QueryPlan<E>,
85        operation: &'static str,
86    ) -> Result<SearchStream<E>, PgSearchError>
87    where
88        E: SearchEntity + Send + Unpin + 'static + for<'row> FromRow<'row, PgRow>,
89    {
90        self.typed_stream(plan, operation, None)
91    }
92
93    #[allow(clippy::needless_pass_by_value)]
94    fn typed_stream<E>(
95        &self,
96        plan: QueryPlan<E>,
97        operation: &'static str,
98        scan_timeout: Option<Duration>,
99    ) -> Result<SearchStream<E>, PgSearchError>
100    where
101        E: SearchEntity + Send + Unpin + 'static + for<'row> FromRow<'row, PgRow>,
102    {
103        if plan.force_empty || plan.paging.size == 0 {
104            return Ok(Box::pin(futures_util::stream::empty()));
105        }
106        let partition = self.extensions.resolve_partition(&plan, operation)?;
107        let prepared = prepare(self.render(
108            &plan,
109            operation,
110            "list",
111            partition.as_deref(),
112            RenderRequest::List,
113        )?)?;
114        let executor = self.clone();
115        Ok(Box::pin(try_stream! {
116            let deadline = scan_timeout.map_or_else(|| executor.deadline::<E>(), |duration| Instant::now() + duration);
117            let connection = executor.acquire(deadline).await?;
118            let mut connection = configure_statement_timeout(connection, deadline, TimeoutStage::ListQuery).await?;
119            connection.close_on_drop();
120            let mut rows = sqlx::query_as_with::<_, E, _>(prepared.sql(), prepared.arguments).fetch(&mut *connection);
121            let outcome = loop {
122                match timeout_at(deadline, rows.try_next()).await {
123                    Ok(Ok(Some(row))) => yield row,
124                    Ok(Ok(None)) => break Outcome::Done,
125                    Ok(Err(error)) => break Outcome::Failed(error),
126                    Err(_) => break Outcome::TimedOut,
127                }
128            };
129            drop(rows);
130            finish_stream(connection, outcome).await?;
131        }))
132    }
133
134    pub fn search_map_stream<E>(
135        &self,
136        query: impl QueryInput<E>,
137        scope: &ServerScope<E>,
138    ) -> Result<SearchStream<RowMap>, PgSearchError>
139    where
140        E: SearchEntity + Send + Unpin + 'static,
141    {
142        self.observe_stream::<E, _>("search_map_stream", || {
143            let plan =
144                self.resolve_plan(query, scope, QueryTarget::MapPage, "search_map_stream")?;
145            self.map_stream(plan, "search_map_stream", None)
146        })
147    }
148
149    fn map_stream<E>(
150        &self,
151        plan: QueryPlan<E>,
152        operation: &'static str,
153        scan_timeout: Option<Duration>,
154    ) -> Result<SearchStream<RowMap>, PgSearchError>
155    where
156        E: SearchEntity + Send + Unpin + 'static,
157    {
158        if plan.force_empty || plan.paging.size == 0 {
159            return Ok(Box::pin(futures_util::stream::empty()));
160        }
161        let partition = self.extensions.resolve_partition(&plan, operation)?;
162        let prepared = prepare(self.render(
163            &plan,
164            operation,
165            "list",
166            partition.as_deref(),
167            RenderRequest::List,
168        )?)?;
169        let executor = self.clone();
170        Ok(Box::pin(try_stream! {
171            let deadline = scan_timeout.map_or_else(|| executor.deadline::<E>(), |duration| Instant::now() + duration);
172            let connection = executor.acquire(deadline).await?;
173            let mut connection = configure_statement_timeout(connection, deadline, TimeoutStage::ListQuery).await?;
174            connection.close_on_drop();
175            let mut rows = sqlx::query_with(prepared.sql(), prepared.arguments).fetch(&mut *connection);
176            let outcome = loop {
177                match timeout_at(deadline, rows.try_next()).await {
178                    Ok(Ok(Some(row))) => {
179                        let mut row = decode_map(&row, &plan)?;
180                        executor.format_map_rows::<E>(std::slice::from_mut(&mut row))?;
181                        yield row;
182                    }
183                    Ok(Ok(None)) => break Outcome::Done,
184                    Ok(Err(error)) => break Outcome::Failed(error),
185                    Err(_) => break Outcome::TimedOut,
186                }
187            };
188            drop(rows);
189            finish_stream(connection, outcome).await?;
190        }))
191    }
192
193    /// Explicit collecting convenience; prefer `scan_stream` for constant-memory use.
194    pub fn scan<E>(
195        &self,
196        query: impl QueryInput<E>,
197        scope: &ServerScope<E>,
198        policy: ScanPolicy,
199    ) -> Result<ScanFuture<E>, PgSearchError>
200    where
201        E: SearchEntity + Send + Unpin + 'static + for<'row> FromRow<'row, PgRow>,
202    {
203        let stream = self.scan_stream(query, scope, policy)?;
204        Ok(Box::pin(stream.try_collect()))
205    }
206
207    pub fn scan_stream<E>(
208        &self,
209        query: impl QueryInput<E>,
210        scope: &ServerScope<E>,
211        policy: ScanPolicy,
212    ) -> Result<SearchStream<E>, PgSearchError>
213    where
214        E: SearchEntity + Send + Unpin + 'static + for<'row> FromRow<'row, PgRow>,
215    {
216        self.observe_stream::<E, _>("scan", || {
217            let policy = policy.validate()?;
218            let plan = scan_plan(
219                self.resolve_plan(query, scope, QueryTarget::TypedList, "scan")?,
220                policy,
221            )?;
222            let stream = self.typed_stream(plan, "scan", Some(policy.timeout))?;
223            Ok(self.bound_scan(stream, policy))
224        })
225    }
226
227    pub fn scan_map<E>(
228        &self,
229        query: impl QueryInput<E>,
230        scope: &ServerScope<E>,
231        policy: ScanPolicy,
232    ) -> Result<ScanFuture<RowMap>, PgSearchError>
233    where
234        E: SearchEntity + Send + Unpin + 'static,
235    {
236        let stream = self.scan_map_stream(query, scope, policy)?;
237        Ok(Box::pin(stream.try_collect()))
238    }
239
240    pub fn scan_map_stream<E>(
241        &self,
242        query: impl QueryInput<E>,
243        scope: &ServerScope<E>,
244        policy: ScanPolicy,
245    ) -> Result<SearchStream<RowMap>, PgSearchError>
246    where
247        E: SearchEntity + Send + Unpin + 'static,
248    {
249        self.scan_map_stream_with_scopes(query, scope, &GroupScope::unrestricted(), policy)
250    }
251
252    pub fn scan_map_stream_with_scopes<E>(
253        &self,
254        query: impl QueryInput<E>,
255        scope: &ServerScope<E>,
256        group_scope: &GroupScope<E>,
257        policy: ScanPolicy,
258    ) -> Result<SearchStream<RowMap>, PgSearchError>
259    where
260        E: SearchEntity + Send + Unpin + 'static,
261    {
262        self.observe_stream::<E, _>("scan_map", || {
263            let policy = policy.validate()?;
264            let plan = scan_plan(
265                self.resolve_plan_with_scopes(
266                    query,
267                    scope,
268                    group_scope,
269                    QueryTarget::MapPage,
270                    "scan_map",
271                )?,
272                policy,
273            )?;
274            let stream = self.map_stream(plan, "scan_map", Some(policy.timeout))?;
275            Ok(self.bound_scan(stream, policy))
276        })
277    }
278
279    pub(crate) fn bound_scan<T: Send + 'static>(
280        &self,
281        mut stream: SearchStream<T>,
282        policy: ScanPolicy,
283    ) -> SearchStream<T> {
284        let active = Arc::clone(&self.scan_gate);
285        let waiting = Arc::clone(&self.scan_wait_gate);
286        Box::pin(try_stream! {
287            let deadline = Instant::now() + policy.timeout;
288            let _permit = if let Ok(permit) = Arc::clone(&active).try_acquire_owned() {
289                permit
290            } else {
291                    let waiter = waiting.try_acquire_owned().map_err(|_| PgSearchError::ScanQueueFull)?;
292                    let permit = timeout_at(deadline, active.acquire_owned()).await.map_err(|_| PgSearchError::Timeout(TimeoutStage::Scan))?.map_err(|_| PgSearchError::ScanQueueFull)?;
293                    drop(waiter);
294                    permit
295            };
296            let mut count = 0u32;
297            while let Some(row) = timeout_at(deadline, stream.next()).await.map_err(|_| PgSearchError::Timeout(TimeoutStage::Scan))? {
298                let row = row?;
299                if count == policy.max_rows { Err(PgSearchError::ScanLimitExceeded { max_rows: policy.max_rows })?; }
300                count += 1;
301                yield row;
302                if count.is_multiple_of(policy.batch_size) { tokio::task::yield_now().await; }
303            }
304        })
305    }
306}
307
308pub(crate) fn scan_plan<E: SearchEntity>(
309    mut plan: QueryPlan<E>,
310    policy: ScanPolicy,
311) -> Result<QueryPlan<E>, PgSearchError> {
312    if plan.paging.explicit_page
313        || plan.paging.explicit_size
314        || plan.paging.explicit_offset
315        || plan.paging.explicit_cursor
316    {
317        return Err(sift_query_core::ParamError::PagingNotAllowed.into());
318    }
319    if plan.order.is_empty() || plan.meta.stable_order.is_empty() {
320        return Err(sift_query_core::ParamError::UnstablePagination.into());
321    }
322    plan.paging.mode = PagingMode::Page { page: 1 };
323    plan.paging.offset = 0;
324    plan.paging.size = policy.max_rows + 1;
325    Ok(plan)
326}
327
328async fn finish_stream(
329    connection: sqlx::pool::PoolConnection<sqlx::Postgres>,
330    outcome: Outcome,
331) -> Result<(), PgSearchError> {
332    match outcome {
333        Outcome::Done => {
334            let (mut connection, ()) =
335                finish_statement(connection, Ok(()), TimeoutStage::ListQuery).await?;
336            connection.return_to_pool().await;
337            Ok(())
338        }
339        Outcome::Failed(error) => {
340            finish_statement(connection, Err::<(), _>(error), TimeoutStage::ListQuery)
341                .await
342                .map(|_| ())
343        }
344        Outcome::TimedOut => {
345            close_hard(connection).await;
346            Err(PgSearchError::Timeout(TimeoutStage::ListQuery))
347        }
348    }
349}
sift_query_sqlx/execution/stream.rs

sift_query_sqlx/execution/
stream.rs
