Fix document count on collections (SciPhi-AI#1827)

* Fix document count on collections * Fix tests
FutureProofTech · Jan 15, 2025 · 5c1b4cb · 5c1b4cb
1 parent 3c36ffe
commit 5c1b4cb
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 3 deletions.
diff --git a/js/sdk/__tests__/CollectionsIntegrationSuperUser.test.ts b/js/sdk/__tests__/CollectionsIntegrationSuperUser.test.ts
@@ -64,6 +64,7 @@ describe("r2rClient V3 Collections Integration Tests", () => {
 
   test("Retrieve updated collection", async () => {
     const response = await client.collections.retrieve({ id: collectionId });
+
     expect(response.results).toBeDefined();
     expect(response.results.id).toBe(collectionId);
     expect(response.results.name).toBe("Updated Test Collection");
@@ -192,6 +193,16 @@ describe("r2rClient V3 Collections Integration Tests", () => {
     expect(response.results).toBeDefined();
   });
 
+  test("Retrieve a collection with no documents", async () => {
+    const response = await client.collections.retrieve({ id: collectionId });
+
+    expect(response.results).toBeDefined();
+    expect(response.results.id).toBe(collectionId);
+    expect(response.results.name).toBe("Updated Test Collection");
+    expect(response.results.description).toBeDefined();
+    expect(response.results.documentCount).toBe(0);
+  });
+
   test("Delete zametov.txt", async () => {
     const response = await client.documents.delete({
       id: "69100f1e-2839-5b37-916d-5c87afe14094",

diff --git a/js/sdk/__tests__/DocumentsAndCollectionsIntegrationUser.test.ts b/js/sdk/__tests__/DocumentsAndCollectionsIntegrationUser.test.ts
@@ -236,6 +236,28 @@ describe("r2rClient V3 System Integration Tests User", () => {
     expect(response2.results).toBeDefined();
   });
 
+  // test("User 1's collection should have 2 documents", async () => {
+  //   const response = await user1Client.collections.retrieve({
+  //     id: user1CollectionId,
+  //   });
+
+  //   console.log(response);
+
+  //   expect(response.results).toBeDefined();
+  //   expect(response.results.documentCount).toBe(2);
+  // });
+
+  // test("User 2's collection should have 2 documents", async () => {
+  //   const response = await user2Client.collections.retrieve({
+  //     id: user2CollectionId,
+  //   });
+
+  //   console.log(response);
+
+  //   expect(response.results).toBeDefined();
+  //   expect(response.results.documentCount).toBe(1);
+  // });
+
   test("Delete document as user 1", async () => {
     const response = await user1Client.documents.delete({
       id: user1DocumentId,
@@ -250,6 +272,28 @@ describe("r2rClient V3 System Integration Tests User", () => {
     expect(response.results).toBeDefined();
   });
 
+  // test("User 1's collection should have 0 documents after deletion", async () => {
+  //   const response = await user1Client.collections.retrieve({
+  //     id: user1CollectionId,
+  //   });
+
+  //   console.log(response);
+
+  //   expect(response.results).toBeDefined();
+  //   expect(response.results.documentCount).toBe(0);
+  // });
+
+  // test("User 2's collection should have 0 documents after deletion", async () => {
+  //   const response = await user2Client.collections.retrieve({
+  //     id: user2CollectionId,
+  //   });
+
+  //   console.log(response);
+
+  //   expect(response.results).toBeDefined();
+  //   expect(response.results.documentCount).toBe(0);
+  // });
+
   test("Delete user 1", async () => {
     const response = await client.users.delete({
       id: user1Id,

diff --git a/js/sdk/src/v3/clients/documents.ts b/js/sdk/src/v3/clients/documents.ts
@@ -197,8 +197,6 @@ export class DocumentsClient {
    */
   @feature("documents.download")
   async download(options: { id: string }): Promise<Blob> {
-    console.log("Starting download request...");
-
     const response = await this.client.makeRequest(
       "GET",
       `documents/${options.id}/download`,
@@ -220,7 +218,6 @@ export class DocumentsClient {
     }
 
     if (response.data instanceof ArrayBuffer) {
-      console.log("Creating Blob from ArrayBuffer");
       return new Blob([response.data], { type: contentType });
     }
 

diff --git a/py/core/database/collections.py b/py/core/database/collections.py
@@ -476,6 +476,29 @@ async def remove_document_from_collection_relational(
                 message="Document not found in the specified collection",
             )
 
+        await self.decrement_collection_document_count(
+            collection_id=collection_id
+        )
+
+    async def decrement_collection_document_count(
+        self, collection_id: UUID, decrement_by: int = 1
+    ) -> None:
+        """
+        Decrement the document count for a collection.
+
+        Args:
+            collection_id (UUID): The ID of the collection to update
+            decrement_by (int): Number to decrease the count by (default: 1)
+        """
+        collection_query = f"""
+            UPDATE {self._get_table_name('collections')}
+            SET document_count = document_count - $1
+            WHERE id = $2
+        """
+        await self.connection_manager.execute_query(
+            collection_query, [decrement_by, collection_id]
+        )
+
     async def export_to_csv(
         self,
         columns: Optional[list[str]] = None,

diff --git a/py/core/main/services/management_service.py b/py/core/main/services/management_service.py
@@ -191,6 +191,11 @@ def transform_chunk_id_to_id(
 
             document = documents_overview_response["results"][0]
 
+            for collection_id in document.collection_ids:
+                await self.providers.database.collections_handler.decrement_collection_document_count(
+                    collection_id=collection_id
+                )
+
             if owner_id and str(document.owner_id) != owner_id:
                 raise R2RException(
                     status_code=404,