From a1bf884ff8b799246a6b770894d64b0b6fb8e743 Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Fri, 23 Feb 2024 23:28:52 +0100
Subject: [PATCH] added database convenience functions

---
 src/database/utils.py | 48 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 src/database/utils.py

diff --git a/src/database/utils.py b/src/database/utils.py
new file mode 100644
index 0000000..51d8c33
--- /dev/null
+++ b/src/database/utils.py
@@ -0,0 +1,48 @@
+from pymongo.collection import Collection
+
+from .classes import (
+    VisualCommunication,
+    NoDocumentFoundException
+)
+
+
+def total_documents(
+    collection: Collection
+) -> int:
+    """Get total number of documents in database."""
+    return collection.count_documents(filter={})
+
+
+def total_annotated(
+    collection: Collection
+) -> int:
+    """Get total number of annotated documents in database."""
+    query = {
+        "annotation": { "$ne": None }
+    }
+    return collection.count_documents(filter=query)
+
+
+def get_visual_communication(
+        collection: Collection,
+        with_annotation: bool = False,
+        with_prediction: bool = False
+    ) -> VisualCommunication:
+    """Get a random visual communication from the database."""
+    query = {}
+    if with_annotation:
+        query["annotation"] = {"$ne": None}
+    else:
+        query["annotation"] = None
+    if with_prediction:
+        query["prediction"] = {"$ne": None}
+    else:
+        query["prediction"] = None
+    data = collection.aggregate([
+        { "$match": query },   # find using filters
+        { "$sample": { "size": 1 } }    # get one random
+    ])
+    data = list(data)   # read data from cursor object
+    if len(data) == 0:
+        raise NoDocumentFoundException
+    return VisualCommunication.model_validate(data[0])