refactor(be): filter and paginate entities list in a subquery

Joins and selecting multiple fields for both counting and paginating
caused MySQL to count rows one by one or ignore indexes. This commit
changes the paginated query to filter, order and paginate in a subquery
that selects only the entity ID; and joins the subquery with the
original query.

Author: luis-dk

common/entity_services/helpers/list_rules.py

@@ -8,9 +8,11 @@
 
 from marshmallow import EXCLUDE, Schema
 from marshmallow.fields import Enum, Int
-from peewee import Field, Ordering, Select
+from peewee import JOIN, Field, Ordering, Select
 from werkzeug.datastructures import MultiDict
 
+from common.entities import BaseEntity
+
 DEFAULT_PAGE = 1
 DEFAULT_COUNT = 10
 T = TypeVar("T")
@@ -50,9 +52,25 @@ def __len__(self) -> int:
 
     @classmethod
     def get_paginated_results(cls, query: Select, order_by: Field, list_rules: ListRules) -> Page[T]:
+        model: BaseEntity = query.model
         ordering = list_rules.order_by_field(order_by)
-        paginated_query = query.order_by(ordering).paginate(list_rules.page, list_rules.count)
-        return cls(results=list(paginated_query), total=query.count())
+
+        paginated_subquery: Select = (
+            model.select(model.id)
+            .where(query._where)
+            .order_by(ordering)
+            .paginate(list_rules.page, list_rules.count)
+            .alias("results")
+        )
+
+        results_query = query.clone()
+        results_query._where = None
+        results_query = (
+            results_query.join(paginated_subquery, join_type=JOIN.INNER, on=(model.id == paginated_subquery.c.id))
+            .order_by(ordering)
+        )
+
+        return cls(results=list(results_query), total=paginated_subquery.count(clear_limit=True))
 
 
 @dataclass

common/tests/unit/entity_services/helpers/test_list_rules.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import ANY, Mock, MagicMock
 
 import pytest
 from werkzeug.datastructures import MultiDict
@@ -20,21 +20,29 @@ def mocked_entity_class():
 
 @pytest.fixture()
 def mocked_query():
-    # .paginate() (and other query methods in general) doesn't return a list until iterated on, but this is good enough for the test
-    paginate = Mock(return_value=list())
     query = Mock()
-    query.paginate = paginate
-    query.order_by = Mock(return_value=query)
+    subquery = Mock()
+    results_query = MagicMock()
+
+    query.model.select.return_value = subquery
+    query.clone = Mock(return_value=results_query)
+
+    subquery.where.return_value = subquery
+    subquery.order_by.return_value = subquery
+    subquery.paginate.return_value = subquery
+    subquery.alias.return_value = subquery
+
+    results_query.join.return_value = results_query
+    results_query.order_by.return_value = results_query
+
     yield query
 
 
 @pytest.fixture()
-def mocked_query_with_data():
-    # .paginate() (and other query methods in general) doesn't return a list until iterated on, but this is good enough for the test
-    paginate = Mock(return_value=[Mock(), Mock()])
-    query = Mock()
-    query.paginate = paginate
-    query.order_by = Mock(return_value=query)
+def mocked_query_with_data(mocked_query):
+    query = mocked_query
+    results_query = query.clone.return_value
+    results_query.__iter__.return_value = [Mock(), Mock()]
     yield query
 
 
@@ -80,11 +88,24 @@ def test_from_params_extras():
 
 @pytest.mark.unit
 def test_page_get_paginated_results(mocked_field, mocked_entity_class, mocked_query):
+    mocked_subquery = mocked_query.model.select.return_value
+    mocked_results_query = mocked_query.clone.return_value
     rules = ListRules.from_params(MultiDict())
+
     result = Page[mocked_entity_class].get_paginated_results(mocked_query, mocked_field, rules)
+
     mocked_field.asc.assert_called_once()
-    mocked_query.paginate.assert_called_once()
-    mocked_query.count.assert_called_once()
+
+    mocked_subquery.where.assert_called_once_with(mocked_query._where)
+    mocked_subquery.order_by.assert_called_once_with(rules.order_by_field(mocked_field))
+    mocked_subquery.paginate.assert_called_once_with(rules.page, rules.count)
+    mocked_subquery.alias.assert_called_once()
+    mocked_subquery.count.assert_called_once_with(clear_limit=True)
+
+    assert mocked_results_query._where is None
+    mocked_results_query.join.assert_called_once_with(mocked_subquery, join_type=ANY, on=ANY)
+    mocked_results_query.order_by.assert_called_once_with(rules.order_by_field(mocked_field))
+
     assert type(result) == Page
 
 
@@ -102,6 +123,6 @@ def test_page_len(mocked_field, mocked_entity_class, mocked_query_with_data):
 def test_page_iter(mocked_field, mocked_entity_class, mocked_query_with_data):
     rules = ListRules.from_params(MultiDict())
     result = Page[mocked_entity_class].get_paginated_results(mocked_query_with_data, mocked_field, rules)
-    count = len(mocked_query_with_data.paginate())
+    count = len(list(mocked_query_with_data.clone.return_value))
     r = [r for r in result]
     assert len(r) == count