From 96c47af2c08555435d4db41dc8bbc8c17dba7a67 Mon Sep 17 00:00:00 2001
From: perf3ct <jonfuller2012@gmail.com>
Date: Thu, 10 Jul 2025 19:57:25 +0000
Subject: [PATCH] feat(server/client): make sure that the documents endpoint
 isn't broken

---
 frontend/src/pages/DocumentsPage.tsx |  1 +
 src/db/documents/operations.rs       | 27 +++++++++++++++++
 src/routes/documents/crud.rs         | 43 +++++++++++++++++++++++++---
 src/routes/documents/types.rs        | 14 +++++++++
 4 files changed, 81 insertions(+), 4 deletions(-)

diff --git a/frontend/src/pages/DocumentsPage.tsx b/frontend/src/pages/DocumentsPage.tsx
index a9e4603..63d594e 100644
--- a/frontend/src/pages/DocumentsPage.tsx
+++ b/frontend/src/pages/DocumentsPage.tsx
@@ -153,6 +153,7 @@ const DocumentsPage: React.FC = () => {
         pagination.offset, 
         ocrFilter || undefined
       );
+      // Backend returns wrapped object with documents and pagination
       setDocuments(response.data.documents || []);
       setPagination(response.data.pagination || { total: 0, limit: 20, offset: 0, has_more: false });
     } catch (err) {
diff --git a/src/db/documents/operations.rs b/src/db/documents/operations.rs
index e72613c..36545b4 100644
--- a/src/db/documents/operations.rs
+++ b/src/db/documents/operations.rs
@@ -271,4 +271,31 @@ impl Database {
             row.get("failed"),
         ))
     }
+
+    /// Counts total documents for a user with role-based access control
+    pub async fn count_documents_by_user_with_role(&self, user_id: Uuid, user_role: UserRole) -> Result<i64> {
+        let mut query = QueryBuilder::<Postgres>::new("SELECT COUNT(*) as total FROM documents WHERE 1=1");
+        apply_role_based_filter(&mut query, user_id, user_role);
+        let row = query.build().fetch_one(&self.pool).await?;
+        Ok(row.get("total"))
+    }
+
+    /// Counts documents for a user with role-based access control and OCR status filtering
+    pub async fn count_documents_by_user_with_role_and_filter(
+        &self, 
+        user_id: Uuid, 
+        user_role: UserRole, 
+        ocr_status: Option<&str>
+    ) -> Result<i64> {
+        let mut query = QueryBuilder::<Postgres>::new("SELECT COUNT(*) as total FROM documents WHERE 1=1");
+        apply_role_based_filter(&mut query, user_id, user_role);
+        
+        if let Some(status) = ocr_status {
+            query.push(" AND ocr_status = ");
+            query.push_bind(status);
+        }
+        
+        let row = query.build().fetch_one(&self.pool).await?;
+        Ok(row.get("total"))
+    }
 }
\ No newline at end of file
diff --git a/src/routes/documents/crud.rs b/src/routes/documents/crud.rs
index 4542516..2b97ef1 100644
--- a/src/routes/documents/crud.rs
+++ b/src/routes/documents/crud.rs
@@ -14,7 +14,7 @@ use crate::{
     models::DocumentResponse,
     AppState,
 };
-use super::types::{PaginationQuery, DocumentUploadResponse};
+use super::types::{PaginationQuery, DocumentUploadResponse, PaginatedDocumentsResponse, DocumentPaginationInfo};
 
 /// Upload a new document
 #[utoipa::path(
@@ -201,7 +201,7 @@ pub async fn get_document_by_id(
     ),
     params(PaginationQuery),
     responses(
-        (status = 200, description = "List of documents", body = Vec<DocumentResponse>),
+        (status = 200, description = "Paginated list of documents", body = PaginatedDocumentsResponse),
         (status = 401, description = "Unauthorized"),
         (status = 500, description = "Internal server error")
     )
@@ -210,10 +210,34 @@ pub async fn list_documents(
     State(state): State<Arc<AppState>>,
     auth_user: AuthUser,
     Query(query): Query<PaginationQuery>,
-) -> Result<Json<Vec<DocumentResponse>>, StatusCode> {
+) -> Result<Json<PaginatedDocumentsResponse>, StatusCode> {
     let limit = query.limit.unwrap_or(25);
     let offset = query.offset.unwrap_or(0);
 
+    // Get total count for pagination
+    let total_count = if let Some(ocr_status) = query.ocr_status.as_deref() {
+        state
+            .db
+            .count_documents_by_user_with_role_and_filter(
+                auth_user.user.id,
+                auth_user.user.role,
+                Some(ocr_status),
+            )
+            .await
+    } else {
+        state
+            .db
+            .count_documents_by_user_with_role(
+                auth_user.user.id,
+                auth_user.user.role,
+            )
+            .await
+    }
+    .map_err(|e| {
+        error!("Database error counting documents: {}", e);
+        StatusCode::INTERNAL_SERVER_ERROR
+    })?;
+
     let documents = if let Some(ocr_status) = query.ocr_status.as_deref() {
         state
             .db
@@ -272,7 +296,18 @@ pub async fn list_documents(
         })
         .collect();
 
-    Ok(Json(responses))
+    // Create pagination info
+    let pagination = DocumentPaginationInfo {
+        total: total_count,
+        limit,
+        offset,
+        has_more: offset + limit < total_count,
+    };
+
+    Ok(Json(PaginatedDocumentsResponse {
+        documents: responses,
+        pagination,
+    }))
 }
 
 /// Delete a specific document
diff --git a/src/routes/documents/types.rs b/src/routes/documents/types.rs
index da51d2d..19bdd27 100644
--- a/src/routes/documents/types.rs
+++ b/src/routes/documents/types.rs
@@ -64,6 +64,20 @@ pub struct DocumentDebugInfo {
     pub permissions: Option<String>,
 }
 
+#[derive(Serialize, ToSchema)]
+pub struct DocumentPaginationInfo {
+    pub total: i64,
+    pub limit: i64,
+    pub offset: i64,
+    pub has_more: bool,
+}
+
+#[derive(Serialize, ToSchema)]
+pub struct PaginatedDocumentsResponse {
+    pub documents: Vec<crate::models::DocumentResponse>,
+    pub pagination: DocumentPaginationInfo,
+}
+
 impl Default for PaginationQuery {
     fn default() -> Self {
         Self {