From 1e339cecf4da9c0dbafc21b536d97cfa7b88f704 Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Mon, 11 Mar 2024 20:38:48 +0100
Subject: [PATCH] added database util functions

---
 core/database/__init__.py |  1 +
 core/database/utils.py    | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/core/database/__init__.py b/core/database/__init__.py
index 0990ebc..e964218 100644
--- a/core/database/__init__.py
+++ b/core/database/__init__.py
@@ -5,6 +5,7 @@ from .classes import NoDocumentFoundException
 from .classes import VisualCommunication
 from .database import connect
 from .utils import get_visual_communication
+from .utils import list_names
 from .utils import total_annotated
 from .utils import total_documents
 from .utils import upsert_annotations
diff --git a/core/database/utils.py b/core/database/utils.py
index d43fa77..581b807 100644
--- a/core/database/utils.py
+++ b/core/database/utils.py
@@ -9,6 +9,44 @@ from .classes import NoDocumentFoundException
 from .classes import VisualCommunication
 
 
+def count_documents(
+    collection: Collection,
+    has_annotation: bool = False,
+) -> int:
+    """
+    Get the total number of documents
+    in database that matches the filters.
+    """
+    assert isinstance(collection, Collection)
+    assert isinstance(has_annotation, bool)
+    # build query
+    query = {}
+    if has_annotation:
+        query['annotation'] = {'$ne': None}
+    return collection.count_documents(filter=query)
+
+
+def list_names(
+    collection: Collection,
+    has_annotation: bool = True,
+) -> list[str]:
+    """List the names of entries that match the filters."""
+    assert isinstance(collection, Collection)
+    assert isinstance(has_annotation, bool)
+    # build query
+    query = {}
+    if has_annotation:
+        query['annotation'] = {'$ne': None}
+    res_list = collection.find(
+        filter=query,
+        projection={
+            '_id': False,
+            'name': True,
+        },
+    )
+    return list(res_list)
+
+
 def total_documents(
     collection: Collection,
 ) -> int: