enh: limite max # of rows returned by SQLite worker (#3674)

Co-authored-by: Henry Fontanier <[email protected]>
dust-tt · Feb 9, 2024 · e8a7d7b · e8a7d7b
1 parent 4cc469e
commit e8a7d7b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/core/src/sqlite_workers/sqlite_database.rs b/core/src/sqlite_workers/sqlite_database.rs
@@ -19,6 +19,8 @@ pub struct SqliteDatabase {
     interrupt_handle: Option<Arc<tokio::sync::Mutex<InterruptHandle>>>,
 }
 
+const MAX_ROWS: usize = 128;
+
 impl SqliteDatabase {
     pub fn new() -> Self {
         Self {
@@ -112,6 +114,8 @@ impl SqliteDatabase {
                         })
                         .collect::<Result<serde_json::Value>>()
                 })?
+                // Limit to 128 rows.
+                .take(MAX_ROWS)
                 .collect::<Result<Vec<_>>>()?
                 .into_par_iter()
                 .map(|value| QueryResult { value })