databases_list Core api route: add total for pagination (#2926)

dust-tt · Dec 19, 2023 · 96a039c · 96a039c
1 parent 90b1aa3
commit 96a039c
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 8 deletions.
diff --git a/core/bin/dust_api.rs b/core/bin/dust_api.rs
@@ -1848,12 +1848,15 @@ async fn databases_list(
             "Failed to list databases",
             Some(e),
         ),
-        Ok(dbs) => (
+        Ok((dbs, total)) => (
             StatusCode::OK,
             Json(APIResponse {
                 error: None,
                 response: Some(json!({
-                    "databases": dbs
+                    "databases": dbs,
+                    "offset": query.offset,
+                    "limit": query.limit,
+                    "total": total,
                 })),
             }),
         ),

diff --git a/core/src/data_sources/data_source.rs b/core/src/data_sources/data_source.rs
@@ -1828,15 +1828,14 @@ impl DataSource {
         ));
 
         // Delete databases (concurrently).
-        let databases = store
+        let (databases, total) = store
             .list_databases(&self.project, &self.data_source_id, None)
             .await?;
         try_join_all(databases.iter().map(|db| db.delete(store.clone()))).await?;
 
         utils::done(&format!(
             "Deleted databases: data_source_id={} database_count={}",
-            self.data_source_id,
-            databases.len(),
+            self.data_source_id, total,
         ));
 
         // Delete data source and documents (SQL).

diff --git a/core/src/stores/postgres.rs b/core/src/stores/postgres.rs
@@ -1944,7 +1944,7 @@ impl Store for PostgresStore {
         project: &Project,
         data_source_id: &str,
         limit_offset: Option<(usize, usize)>,
-    ) -> Result<Vec<Database>> {
+    ) -> Result<(Vec<Database>, usize)> {
         let project_id = project.project_id();
         let data_source_id = data_source_id.to_string();
 
@@ -1993,7 +1993,25 @@ impl Store for PostgresStore {
             })
             .collect::<Result<Vec<_>>>()?;
 
-        Ok(databases)
+        let total = match limit_offset {
+            None => databases.len(),
+            Some(_) => {
+                let stmt = c
+                    .prepare(
+                        "SELECT COUNT(*) FROM databases \
+                           INNER JOIN data_sources ON databases.data_source = data_sources.id \
+                           WHERE data_sources.project = $1 AND data_sources.data_source_id = $2",
+                    )
+                    .await?;
+                let t: i64 = c
+                    .query_one(&stmt, &[&project_id, &data_source_id])
+                    .await?
+                    .get(0);
+                t as usize
+            }
+        };
+
+        Ok((databases, total))
     }
 
     async fn assign_live_sqlite_worker_to_database(

diff --git a/core/src/stores/store.rs b/core/src/stores/store.rs
@@ -177,7 +177,7 @@ pub trait Store {
         project: &Project,
         data_source_id: &str,
         limit_offset: Option<(usize, usize)>,
-    ) -> Result<Vec<Database>>;
+    ) -> Result<(Vec<Database>, usize)>;
     async fn assign_live_sqlite_worker_to_database(
         &self,
         project: &Project,