Merge pull request #61 from tjmlabs/collection-return

Adjust CollectionOut return value

web/api/models.py

@@ -78,6 +78,9 @@ class Collection(models.Model):
     def __str__(self) -> str:
         return self.name
 
+    async def document_count(self) -> int:
+        return await self.documents.acount()
+
     class Meta(TypedModelMeta):
         constraints = [
             models.UniqueConstraint(

web/api/tests/tests.py

@@ -144,6 +144,7 @@ async def test_create_collection(async_client, user):
         "id": 1,
         "name": "Test Collection Fixture",
         "metadata": {"key": "value"},
+        "num_documents": 0,
     }
 
 
@@ -186,6 +187,7 @@ async def test_get_collection_by_name(async_client, user, collection):
         "id": 1,
         "name": "Test Collection Fixture",
         "metadata": {"key": "value"},
+        "num_documents": 0,
     }
 
 
@@ -196,7 +198,7 @@ async def test_list_collection(async_client, user, collection):
     )
     assert response.status_code == 200
     assert response.json() == [
-        {"id": 1, "name": "Test Collection Fixture", "metadata": {"key": "value"}}
+        {"id": 1, "name": "Test Collection Fixture", "metadata": {"key": "value"}, "num_documents": 0}
     ]
 
 
@@ -212,6 +214,7 @@ async def test_patch_collection(async_client, user, collection):
         "id": 1,
         "name": "Test Collection Update",
         "metadata": {"key": "value"},
+        "num_documents": 0,
     }
 
     # now check if the collection was actually updated
@@ -225,6 +228,7 @@ async def test_patch_collection(async_client, user, collection):
         "id": 1,
         "name": "Test Collection Update",
         "metadata": {"key": "value"},
+        "num_documents": 0,
     }
 
 
@@ -258,6 +262,7 @@ async def test_patch_collection_no_metadata(async_client, user, collection):
         "id": 1,
         "name": "Test Collection Update",
         "metadata": {"key": "value"},
+        "num_documents": 0,
     }
 
 

web/api/views.py

@@ -98,6 +98,7 @@ class CollectionOut(Schema):
     id: int
     name: str
     metadata: dict
+    num_documents: int
 
 
 class GenericError(Schema):
@@ -137,7 +138,7 @@ async def create_collection(
             name=payload.name, owner=request.auth, metadata=payload.metadata
         )
         return 201, CollectionOut(
-            id=collection.id, name=collection.name, metadata=collection.metadata
+            id=collection.id, name=collection.name, metadata=collection.metadata, num_documents=0
         )
     except IntegrityError:
         return 409, GenericError(
@@ -164,7 +165,7 @@ async def list_collections(request: Request) -> List[CollectionOut]:
         HTTPException: If there is an issue with the request or authentication.
     """
     collections = [
-        CollectionOut(id=c.id, name=c.name, metadata=c.metadata)
+        CollectionOut(id=c.id, name=c.name, metadata=c.metadata, num_documents=await c.document_count())
         async for c in Collection.objects.filter(owner=request.auth)
     ]
     return collections
@@ -206,7 +207,7 @@ async def get_collection(
             name=collection_name, owner=request.auth
         )
         return 200, CollectionOut(
-            id=collection.id, name=collection.name, metadata=collection.metadata
+            id=collection.id, name=collection.name, metadata=collection.metadata, num_documents=await collection.document_count()
         )
     except Collection.DoesNotExist:
         return 404, GenericError(detail=f"Collection: {collection_name} doesn't exist")
@@ -251,7 +252,7 @@ async def partial_update_collection(
 
     await collection.asave()
     return 200, CollectionOut(
-        id=collection.id, name=collection.name, metadata=collection.metadata
+        id=collection.id, name=collection.name, metadata=collection.metadata, num_documents=await collection.document_count()
     )