sift_query_sqlx/export/
mod.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use futures_util::StreamExt;
5use sift_query_core::{QueryInput, SearchEntity, ServerScope};
6use tokio::io::{AsyncWrite, AsyncWriteExt};
7use tokio::sync::{OwnedSemaphorePermit, Semaphore};
8mod csv;
9use csv::{write_record, write_row_batch, write_rows};
10
11use crate::{PgSearchError, PgSiftQuery, RowMap, ScanPolicy};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct CsvExportPolicy {
15    pub max_rows: u32,
16    pub timeout: Duration,
17    pub batch_size: u32,
18}
19
20impl Default for CsvExportPolicy {
21    fn default() -> Self {
22        Self {
23            max_rows: 100_000,
24            timeout: Duration::from_secs(120),
25            batch_size: 256,
26        }
27    }
28}
29
30/// A server-owned export column and optional trusted formatter.
31#[derive(Debug, Clone)]
32pub struct CsvColumn {
33    pub field: String,
34    pub title: String,
35    /// A declared source field for a post-query label column.
36    pub label_source: Option<String>,
37    pub formatter: Option<fn(&sift_query_core::RowValue) -> Result<String, String>>,
38}
39
40impl CsvColumn {
41    #[must_use]
42    pub fn new(field: impl Into<String>, title: impl Into<String>) -> Self {
43        Self {
44            field: field.into(),
45            title: title.into(),
46            label_source: None,
47            formatter: None,
48        }
49    }
50
51    #[must_use]
52    pub fn label(source: impl Into<String>, title: impl Into<String>) -> Self {
53        let source = source.into();
54        Self {
55            field: format!("{source}Label"),
56            title: title.into(),
57            label_source: Some(source),
58            formatter: None,
59        }
60    }
61}
62
63#[derive(Debug, Clone)]
64pub struct CsvExporter {
65    policy: CsvExportPolicy,
66    columns: Vec<CsvColumn>,
67    limiter: Option<Arc<Semaphore>>,
68}
69
70impl CsvExporter {
71    #[must_use]
72    pub const fn new(policy: CsvExportPolicy) -> Self {
73        Self {
74            policy,
75            columns: Vec::new(),
76            limiter: None,
77        }
78    }
79
80    #[must_use]
81    pub const fn policy(&self) -> CsvExportPolicy {
82        self.policy
83    }
84
85    /// Fixes the export whitelist, headings and order on the server. Empty
86    /// columns use the entity's selectable projection in metadata order.
87    #[must_use]
88    pub fn with_columns(mut self, columns: impl IntoIterator<Item = CsvColumn>) -> Self {
89        self.columns = columns.into_iter().collect();
90        self
91    }
92
93    pub fn with_concurrency(
94        mut self,
95        maximum: usize,
96    ) -> Result<Self, sift_query_core::ConfigError> {
97        if maximum == 0 || maximum > 1024 {
98            return Err(sift_query_core::ConfigError::Invalid(
99                "export concurrency must be between 1 and 1024".to_owned(),
100            ));
101        }
102        self.limiter = Some(Arc::new(Semaphore::new(maximum)));
103        Ok(self)
104    }
105
106    async fn acquire(
107        &self,
108        timeout: Duration,
109    ) -> Result<Option<OwnedSemaphorePermit>, PgSearchError> {
110        let Some(limiter) = &self.limiter else {
111            return Ok(None);
112        };
113        tokio::time::timeout(timeout, limiter.clone().acquire_owned())
114            .await
115            .map_err(|_| PgSearchError::Timeout(crate::TimeoutStage::Export))?
116            .map(Some)
117            .map_err(|_| PgSearchError::ParallelGateClosed)
118    }
119
120    pub async fn write_page<W>(
121        &self,
122        writer: &mut W,
123        page: &crate::Page<RowMap>,
124    ) -> Result<u32, PgSearchError>
125    where
126        W: AsyncWrite + Unpin,
127    {
128        csv::validate_export_policy(
129            self.policy.max_rows,
130            self.policy.timeout,
131            self.policy.batch_size,
132        )?;
133        let _permit = self.acquire(self.policy.timeout).await?;
134        if page.data.len() > self.policy.max_rows as usize {
135            return Err(PgSearchError::ScanLimitExceeded {
136                max_rows: self.policy.max_rows,
137            });
138        }
139        tokio::time::timeout(
140            self.policy.timeout,
141            write_rows(writer, &page.data, &self.columns),
142        )
143        .await
144        .map_err(|_| PgSearchError::Timeout(crate::TimeoutStage::Export))?
145    }
146
147    /// Runs a bounded dynamic scan and writes escaped CSV to the caller's
148    /// response or file writer.
149    pub async fn export<E, W>(
150        &self,
151        search: &PgSiftQuery,
152        query: impl QueryInput<E>,
153        scope: &ServerScope<E>,
154        writer: &mut W,
155    ) -> Result<u32, PgSearchError>
156    where
157        E: SearchEntity + Send + Unpin + 'static,
158        W: AsyncWrite + Unpin,
159    {
160        csv::validate_export_policy(
161            self.policy.max_rows,
162            self.policy.timeout,
163            self.policy.batch_size,
164        )?;
165        let _permit = self.acquire(self.policy.timeout).await?;
166        search
167            .observe::<E, _>(
168                "csv_export",
169                async {
170                    let query = query.into_user_query(&search.policy)?;
171                    let plan = sift_query_core::QueryPlan::new(
172                        &query,
173                        scope,
174                        sift_query_core::QueryTarget::MapPage,
175                    )?;
176                    let columns = if self.columns.is_empty() {
177                        plan.projection
178                            .iter()
179                            .filter_map(|id| E::META.field(*id))
180                            .map(|field| CsvColumn::new(field.public_name, field.public_name))
181                            .collect()
182                    } else {
183                        self.columns.clone()
184                    };
185                    let mut unique = std::collections::HashSet::new();
186                    for column in &columns {
187                        let field = E::META
188                            .resolve_field(column.label_source.as_deref().unwrap_or(&column.field))
189                            .ok_or_else(|| {
190                                sift_query_core::ParamError::UnknownField(column.field.clone())
191                            })?;
192                        if !field.selectable
193                            || !plan.projection.contains(&field.id)
194                            || !unique.insert(&column.field)
195                        {
196                            return Err(sift_query_core::ParamError::NotSelectable(
197                                column.field.clone(),
198                            )
199                            .into());
200                        }
201                    }
202                    let mut rows = search.scan_map_stream(
203                        &query,
204                        scope,
205                        ScanPolicy {
206                            max_rows: self.policy.max_rows,
207                            batch_size: self.policy.batch_size,
208                            timeout: self.policy.timeout,
209                        },
210                    )?;
211                    let future = async {
212                        tracing::Span::current().record("rows_written", 0u64);
213                        write_record(writer, columns.iter().map(|column| column.title.clone()))
214                            .await?;
215                        let mut count = 0_u32;
216                        let mut batch = Vec::with_capacity(self.policy.batch_size as usize);
217                        while let Some(row) = rows.next().await {
218                            batch.push(row?);
219                            if batch.len() == self.policy.batch_size as usize {
220                                count += self
221                                    .write_batch::<E, W>(
222                                        search,
223                                        writer,
224                                        std::mem::take(&mut batch),
225                                        &columns,
226                                    )
227                                    .await?;
228                                tracing::Span::current().record("rows_written", count);
229                            }
230                        }
231                        if !batch.is_empty() {
232                            count += self
233                                .write_batch::<E, W>(search, writer, batch, &columns)
234                                .await?;
235                            tracing::Span::current().record("rows_written", count);
236                        }
237                        writer.flush().await.map_err(PgSearchError::ExportIo)?;
238                        Ok(count)
239                    };
240                    let result = tokio::time::timeout(self.policy.timeout, future)
241                        .await
242                        .map_err(|_| PgSearchError::Timeout(crate::TimeoutStage::Export))?;
243                    if result.is_err() {
244                        tracing::warn!(
245                            operation = "csv_export",
246                            outcome = "partial_failure",
247                            "CSV export failed; output may be partial"
248                        );
249                    }
250                    result
251                },
252                |count| *count as usize,
253            )
254            .await
255    }
256
257    #[cfg(feature = "rls")]
258    pub async fn export_rls<E, W>(
259        &self,
260        search: &PgSiftQuery,
261        context: &crate::RlsContext,
262        query: impl QueryInput<E>,
263        scope: &ServerScope<E>,
264        writer: &mut W,
265    ) -> Result<u32, PgSearchError>
266    where
267        E: SearchEntity + Unpin,
268        W: AsyncWrite + Unpin,
269    {
270        csv::validate_export_policy(
271            self.policy.max_rows,
272            self.policy.timeout,
273            self.policy.batch_size,
274        )?;
275        let _permit = self.acquire(self.policy.timeout).await?;
276        let query = query.into_user_query(&search.policy)?;
277        let plan =
278            sift_query_core::QueryPlan::new(&query, scope, sift_query_core::QueryTarget::MapPage)?;
279        let columns = if self.columns.is_empty() {
280            plan.projection
281                .iter()
282                .filter_map(|id| E::META.field(*id))
283                .map(|field| CsvColumn::new(field.public_name, field.public_name))
284                .collect()
285        } else {
286            self.columns.clone()
287        };
288        let mut seen = std::collections::HashSet::new();
289        for column in &columns {
290            let field = E::META
291                .resolve_field(column.label_source.as_deref().unwrap_or(&column.field))
292                .ok_or_else(|| sift_query_core::ParamError::UnknownField(column.field.clone()))?;
293            if !field.selectable
294                || !plan.projection.contains(&field.id)
295                || !seen.insert(&column.field)
296            {
297                return Err(
298                    sift_query_core::ParamError::NotSelectable(column.field.clone()).into(),
299                );
300            }
301        }
302        let mut rows = search.scan_map_stream_rls(
303            context,
304            &query,
305            scope,
306            ScanPolicy {
307                max_rows: self.policy.max_rows,
308                batch_size: self.policy.batch_size,
309                timeout: self.policy.timeout,
310            },
311        )?;
312        tokio::time::timeout(self.policy.timeout, async {
313            write_record(writer, columns.iter().map(|column| column.title.clone())).await?;
314            let mut count = 0;
315            let mut batch = Vec::with_capacity(self.policy.batch_size as usize);
316            while let Some(row) = rows.next().await {
317                batch.push(row?);
318                if batch.len() == self.policy.batch_size as usize {
319                    count += self
320                        .write_batch::<E, W>(search, writer, std::mem::take(&mut batch), &columns)
321                        .await?;
322                }
323            }
324            if !batch.is_empty() {
325                count += self
326                    .write_batch::<E, W>(search, writer, batch, &columns)
327                    .await?;
328            }
329            writer.flush().await.map_err(PgSearchError::ExportIo)?;
330            Ok(count)
331        })
332        .await
333        .map_err(|_| PgSearchError::Timeout(crate::TimeoutStage::Export))?
334    }
335
336    async fn write_batch<E: SearchEntity, W: AsyncWrite + Unpin>(
337        &self,
338        search: &PgSiftQuery,
339        writer: &mut W,
340        mut data: Vec<RowMap>,
341        columns: &[CsvColumn],
342    ) -> Result<u32, PgSearchError> {
343        search
344            .extensions
345            .filter_rows::<E>("csv_export", &mut data)?;
346        let page = search
347            .extensions
348            .enrich_page(
349                E::META.entity_name,
350                "csv_export",
351                crate::Page {
352                    data,
353                    total_count: 0,
354                    summaries: Vec::new(),
355                    next_cursor: None,
356                },
357            )
358            .await?;
359        write_row_batch(writer, &page.data, columns).await
360    }
361}
sift_query_sqlx/export/mod.rs

sift_query_sqlx/export/
mod.rs
