tablo-core 0.10.0

//! CSV export over the rows the list shows: a count answers the cap before any byte is sent, then
//! a cursor-chunked walk streams the file.

use topcoat::{
    Result,
    context::Cx,
    router::{Body, RouteFuture, error::forbidden},
};

use super::super::gate::gate;
use crate::{
    db::db,
    policy::Ability,
    resource::{Mounted, Resource},
    table::{Table, TableState},
};

/// Fetches this many models per cursor chunk on the export walk.
const EXPORT_CHUNK_ROWS: usize = 500;

/// Rejects an export whose row count runs past the cap.
fn enforce_export_cap_count(count: u64, cap: usize) -> Result<(), topcoat::Error> {
    if count > cap as u64 {
        Err(topcoat::router::error::content_too_large().into())
    } else {
        Ok(())
    }
}

/// Bounds the slug a `Content-Disposition` filename keeps.
const MAX_EXPORT_FILENAME_LEN: usize = 100;

/// Sanitizes the export's `Content-Disposition` filename.
fn export_filename(slug: &str) -> String {
    let safe: String = slug
        .chars()
        .filter(|c| !c.is_control() && *c != '"' && *c != '\\')
        .take(MAX_EXPORT_FILENAME_LEN)
        .collect();
    if safe.is_empty() {
        "export.csv".to_string()
    } else {
        format!("{safe}.csv")
    }
}

/// `?bom=1` opts into a UTF-8 BOM prefix on the CSV body for Excel.
fn export_wants_bom(cx: &Cx) -> bool {
    let Some(parts) = topcoat::context::try_context::<http::request::Parts>(cx) else {
        return false;
    };
    let Some(query) = parts.uri.query() else {
        return false;
    };
    form_urlencoded::parse(query.as_bytes()).any(|(k, v)| k == "bom" && v == "1")
}

/// Streams the tenant-scoped query as a CSV download.
pub(crate) fn resource_export<R: Resource>(cx: &Cx, _body: Body) -> RouteFuture<'_> {
    Box::pin(async move {
        let resource = gate::<R>(cx)?;
        if !resource.can(cx, Ability::ViewAny) {
            return Err(forbidden().into());
        }
        let state = TableState::from_cx(cx);
        let table = resource.table_for(cx);
        // Fails closed on unapplied filters, a refused one included.
        let refused = table.refused_filters(cx, &state).await?;
        let unapplied = table.unapplied_filters(cx, &state, &refused);
        // A refused request's body is for whoever built the URL, so it reads in English.
        if !unapplied.is_empty() {
            return Err(topcoat::router::error::bad_request(format!(
                "invalid filters: {}",
                unapplied
                    .iter()
                    .map(|(pair, kind)| {
                        format!("{pair} ({})", crate::table::ignored_reason(kind).english())
                    })
                    .collect::<Vec<_>>()
                    .join(", ")
            ))
            .into());
        }
        // Bounds the rows an export delivers; more is a 413, so the export never returns a
        // partial file.
        let cap = crate::limits::of(cx).export_rows;
        let count = crate::db::count(
            cx,
            table.narrow(
                cx,
                &db(cx)?,
                resource.list_query(cx, &state)?,
                &state,
                &refused,
            ),
        )
        .await?;
        enforce_export_cap_count(count, cap)?;
        let scanned = count > 0;
        let want_bom = export_wants_bom(cx);
        let (tx, body) = http_body_util::Channel::<bytes::Bytes, std::io::Error>::new(8);
        let cx2 = cx.clone();
        let filename = export_filename(&resource.slug);
        // Spawned boxed, so one task type serves every resource rather than tokio's task code
        // being compiled once per model.
        let stream: crate::BoxFuture<'static, ()> = Box::pin(async move {
            let mut tx = tx;
            let mut chunker = match export_base_query(&cx2, &resource, &table, &state, &refused) {
                // A query holding no row loads without includes, which Toasty panics planning
                // over a filter it folds to `false` (`table::holds_rows`).
                Ok(query) if scanned => ExportChunker::new(table.include_relations(query), cap),
                Ok(query) => ExportChunker::new(query, cap),
                Err(error) => {
                    tracing::error!(resource = resource.slug, error = %error, "export stream failed");
                    tx.abort(std::io::Error::other("export unavailable"));
                    return;
                }
            };
            let Ok(mut db_handle) = crate::db::db(&cx2) else {
                tx.abort(std::io::Error::other("export unavailable"));
                return;
            };
            let mut head = table.csv_header(&cx2);
            if want_bom {
                head.insert(0, '\u{FEFF}');
            }
            if tx.send_data(bytes::Bytes::from(head)).await.is_err() {
                return;
            }
            let mut streamed = 0usize;
            loop {
                let rows = match chunker.next_chunk(&mut db_handle).await {
                    Ok(Some(rows)) => rows,
                    Ok(None) => break,
                    Err(error) => {
                        tracing::error!(resource = resource.slug, error = %error, "export stream failed");
                        tx.abort(std::io::Error::other("export unavailable"));
                        return;
                    }
                };
                if let Err(error) = table.restrict_relations(&cx2, &rows, &mut db_handle).await {
                    tracing::error!(resource = resource.slug, error = %error, "export stream failed");
                    tx.abort(std::io::Error::other("export unavailable"));
                    return;
                }
                let mut fragment = String::new();
                for row in &rows {
                    streamed += 1;
                    if streamed > cap {
                        // Aborts when rows were added since the count.
                        tracing::error!(
                            resource = resource.slug,
                            "export overflowed its cap mid-stream"
                        );
                        tx.abort(std::io::Error::other("export overflowed its cap"));
                        return;
                    }
                    fragment.push_str(&table.csv_row(&cx2, row));
                }
                if !fragment.is_empty() && tx.send_data(bytes::Bytes::from(fragment)).await.is_err()
                {
                    return;
                }
            }
        });
        tokio::spawn(stream);
        let res = http::Response::builder()
            .status(200)
            .header(http::header::CONTENT_TYPE, "text/csv; charset=utf-8")
            .header(
                http::header::CONTENT_DISPOSITION,
                format!("attachment; filename=\"{}\"", filename),
            )
            .body(Body::new(body))?;
        Ok(res)
    })
}

/// Builds the export's filtered and ordered base query from the tenant-scoped query.
fn export_base_query<R: Resource>(
    cx: &Cx,
    resource: &Mounted<R>,
    table: &Table<R::Model>,
    state: &TableState,
    refused: &[String],
) -> Result<toasty::stmt::Query<toasty::stmt::List<R::Model>>> {
    Ok(table.apply_declaration(
        cx,
        &crate::db::db(cx)?,
        resource.list_query(cx, state)?,
        state,
        refused,
    ))
}

/// Walks an export base query in cursor chunks.
struct ExportChunker<M> {
    query: toasty::stmt::Query<toasty::stmt::List<M>>,
    after: Option<toasty::stmt::Value>,
    /// The most rows the export delivers.
    cap: usize,
    scanned: usize,
    exhausted: bool,
}

impl<M> ExportChunker<M>
where
    M: toasty::schema::Model + Send + Sync + 'static,
{
    fn new(query: toasty::stmt::Query<toasty::stmt::List<M>>, cap: usize) -> Self {
        Self {
            query,
            after: None,
            cap,
            scanned: 0,
            exhausted: false,
        }
    }

    async fn next_chunk(&mut self, db: &mut toasty::Db) -> Result<Option<Vec<M>>, topcoat::Error> {
        if self.exhausted {
            return Ok(None);
        }
        loop {
            // One row past the cap is enough for the stream to see it overflow.
            let remaining = (self.cap + 1).saturating_sub(self.scanned);
            if remaining == 0 {
                self.exhausted = true;
                return Ok(None);
            }
            let take = remaining.min(EXPORT_CHUNK_ROWS);
            let mut page = toasty::stmt::Paginate::new(self.query.clone(), take);
            if let Some(cursor) = self.after.take() {
                page = page.after(cursor);
            }
            let loaded = page.exec(db).await.map_err(crate::error::unavailable)?;
            if loaded.items.is_empty() {
                self.after = loaded.next_cursor;
                if self.after.is_none() {
                    self.exhausted = true;
                    return Ok(None);
                }
                continue;
            }
            self.scanned += loaded.items.len();
            self.after = loaded.next_cursor;
            let items = loaded.items;
            if self.after.is_none() {
                // Stops here; re-fetching cursor-free would rescan from the start.
                self.exhausted = true;
                self.after = None;
            }
            return Ok(Some(items));
        }
    }
}

#[cfg(test)]
mod tests;