From 5fd2622e4c008d628696154b58e2bbb485d67d9e Mon Sep 17 00:00:00 2001
From: Matthew Richards <matthew.richards@stfc.ac.uk>
Date: Tue, 29 Sep 2020 15:11:57 +0000
Subject: [PATCH]  #145: Basic implementation of `/count`

- Currently breaks when you add a distinct filter so that needs fixing
---
 common/icat/backend.py |  5 ++++-
 common/icat/helpers.py | 48 ++++++++++++++++++++++++++++++++++++++----
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/common/icat/backend.py b/common/icat/backend.py
index a9cacf44..8d66eb01 100644
--- a/common/icat/backend.py
+++ b/common/icat/backend.py
@@ -14,6 +14,7 @@
     update_entity_by_id,
     delete_entity_by_id,
     get_entity_with_filters,
+    get_count_with_filters,
 )
 
 from common.config import config
@@ -74,6 +75,7 @@ def logout(self, session_id):
     @requires_session_id
     @queries_records
     def get_with_filters(self, session_id, table, filters):
+        self.client.sessionId = session_id
         return get_entity_with_filters(self.client, table.__name__, filters)
 
     @requires_session_id
@@ -94,7 +96,8 @@ def get_one_with_filters(self, session_id, table, filters):
     @requires_session_id
     @queries_records
     def count_with_filters(self, session_id, table, filters):
-        pass
+        self.client.sessionId = session_id
+        return get_count_with_filters(self.client, table.__name__, filters)
 
     @requires_session_id
     @queries_records
diff --git a/common/icat/helpers.py b/common/icat/helpers.py
index d7183ac8..6d4c34da 100644
--- a/common/icat/helpers.py
+++ b/common/icat/helpers.py
@@ -180,10 +180,16 @@ def execute_query(self, client, return_json_formattable=False):
             data = []
 
             for result in query_result:
-                dict_result = self.entity_to_dict(
-                    result, flat_query_includes, mapped_distinct_fields
-                )
-                data.append(dict_result)
+                log.debug(f"Aggregate: {self.query.aggregate}")
+                # TODO - How to deal with distinct and count as aggregate
+                if self.query.aggregate != "COUNT":
+                    
+                    dict_result = self.entity_to_dict(
+                        result, flat_query_includes, mapped_distinct_fields
+                    )
+                    data.append(dict_result)
+                else:
+                    data.append(result)
             return data
         else:
             log.info("Query results will be returned as ICAT entities")
@@ -677,3 +683,37 @@ def clear_order_filters(filters):
 
     if any(isinstance(filter, PythonICATOrderFilter) for filter in filters):
         PythonICATOrderFilter.result_order = []
+
+
+def get_count_with_filters(client, table_name, filters):
+    """
+    Get the number of results of a given entity, based on the filters provided in the
+    request. This acts very much like `get_entity_with_filters()` but returns the number
+    of results, as opposed to a JSON object of data.
+
+    :param client: ICAT client containing an authenticated user
+    :type client: :class:`icat.client.Client`
+    :param table_name: Table name to extract which entity to use
+    :type table_name: :class:`str`
+    :param filters: The list of filters to be applied to the request
+    :type filters: List of specific implementations :class:`QueryFilter`
+    :return: The number of records of the given entity (of type integer), using the
+        filters to restrict the result of the query
+    """
+
+    selected_entity_name = get_python_icat_entity_name(client, table_name)
+    query = icat_query(client, selected_entity_name, aggregate="COUNT")
+
+    filter_handler = FilterOrderHandler()
+    filter_handler.add_filters(filters)
+    merge_limit_skip_filters(filter_handler)
+    clear_order_filters(filter_handler.filters)
+    filter_handler.apply_filters(query.query)
+
+    data = query.execute_query(client, True)
+
+    if not data:
+        raise MissingRecordError("No results found")
+    else:
+        # Only ever 1 element in a count query result
+        return data[0]