Merge pull request #4 from mahenzon/feature/disable-collection-count

Feature/disable collection count

Author: mahenzon

flask_rest_jsonapi/data_layers/alchemy.py

@@ -50,6 +50,37 @@ def __init__(self, kwargs):
                 f"You must provide a model in data_layer_kwargs to use sqlalchemy data layer in {self.resource.__name__}"
             )
 
+        self.disable_collection_count: bool = False
+        self.default_collection_count: int = -1
+
+    def post_init(self):
+        """
+        Checking some props here
+        :return:
+        """
+        if self.resource is None:
+            # if working outside the resource, it's not assigned here
+            return
+
+        if not hasattr(self.resource, "disable_collection_count") or self.resource.disable_collection_count is False:
+            return
+
+        params = self.resource.disable_collection_count
+
+        if isinstance(params, (bool, int)):
+            self.disable_collection_count = bool(params)
+
+        if isinstance(params, (tuple, list)):
+            try:
+                self.disable_collection_count, self.default_collection_count = params
+            except ValueError:
+                raise ValueError(
+                    "Resource's attribute `disable_collection_count` "
+                    "has to be bool or list/tuple with exactly 2 values!\n"
+                    "For example `disable_collection_count = (True, 999)`"
+                )
+        # just ignoring other types, we don't know how to process them
+
     def create_object(self, data, view_kwargs):
         """Create an object through sqlalchemy
 
@@ -125,7 +156,7 @@ def get_object(self, view_kwargs, qs=None):
 
         query = self.retrieve_object_query(view_kwargs, filter_field, filter_value)
 
-        if hasattr(self, "resource"):
+        if self.resource is not None:
             for i_plugins in self.resource.plugins:
                 try:
                     query = i_plugins.data_layer_get_object_update_query(
@@ -146,6 +177,18 @@ def get_object(self, view_kwargs, qs=None):
 
         return obj
 
+    def get_collection_count(self, query, qs, view_kwargs) -> int:
+        """
+        :param query: SQLAlchemy query
+        :param qs: QueryString
+        :param view_kwargs: view kwargs
+        :return:
+        """
+        if self.disable_collection_count is True:
+            return self.default_collection_count
+
+        return query.count()
+
     def get_collection(self, qs, view_kwargs):
         """Retrieve a collection of objects through sqlalchemy
 
@@ -174,7 +217,7 @@ def get_collection(self, qs, view_kwargs):
         if qs.sorting:
             query = self.sort_query(query, qs.sorting)
 
-        object_count = query.count()
+        objects_count = self.get_collection_count(query, qs, view_kwargs)
 
         if getattr(self, "eagerload_includes", True):
             query = self.eagerload_includes(query, qs)
@@ -185,7 +228,7 @@ def get_collection(self, qs, view_kwargs):
 
         collection = self.after_get_collection(collection, qs, view_kwargs)
 
-        return object_count, collection
+        return objects_count, collection
 
     def update_object(self, obj, data, view_kwargs):
         """Update an object through sqlalchemy

flask_rest_jsonapi/data_layers/base.py

@@ -5,6 +5,10 @@
 """
 
 import types
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from flask_rest_jsonapi.resource import Resource
 
 
 class BaseDataLayer:
@@ -38,6 +42,11 @@ def __init__(self, kwargs):
 
         :param dict kwargs: information about data layer instance
         """
+
+        # initing this attribute here in the first place
+        # because it can be easily overridden by kwargs below
+        self.resource: Optional['Resource'] = None
+
         if kwargs.get("methods") is not None:
             self.bound_rewritable_methods(kwargs["methods"])
             kwargs.pop("methods")
@@ -47,6 +56,17 @@ def __init__(self, kwargs):
         for key, value in kwargs.items():
             setattr(self, key, value)
 
+    def post_init(self):
+        """
+        Post init stage.
+
+        At this moment self.resource is already defined
+        and the layer can do any post init stuff here
+
+        NOTE that the data layer is inited for each request
+        :return:
+        """
+
     def create_object(self, data, view_kwargs):
         """Create an object
 

flask_rest_jsonapi/resource.py

@@ -55,6 +55,7 @@ def __new__(cls):
         """Constructor of a resource instance"""
         if hasattr(cls, "_data_layer"):
             cls._data_layer.resource = cls
+            cls._data_layer.post_init()
 
         return super().__new__(cls)
 

tests/test_sqlalchemy_data_layer.py

@@ -437,6 +437,23 @@ class ComputerList(ResourceList):
     yield ComputerList
 
 
+@pytest.fixture(scope="module")
+def fixed_count_for_collection_count():
+    return 42
+
+
+@pytest.fixture(scope="module")
+def computer_list_resource_with_disable_collection_count(
+    session, computer_model, computer_schema, fixed_count_for_collection_count
+):
+    class ComputerList(ResourceList):
+        disable_collection_count = True, fixed_count_for_collection_count
+        schema = computer_schema
+        data_layer = {"model": computer_model, "session": session}
+
+    yield ComputerList
+
+
 @pytest.fixture(scope="module")
 def computer_detail(session, computer_model, dummy_decorator, computer_schema):
     class ComputerDetail(ResourceDetail):
@@ -496,6 +513,7 @@ def register_routes(
     person_list_without_schema,
     computer_list,
     computer_detail,
+    computer_list_resource_with_disable_collection_count,
     computer_owner,
     string_json_attribute_person_detail,
     string_json_attribute_person_list,
@@ -510,6 +528,12 @@ def register_routes(
     api.route(person_list_raise_exception, "person_list_exception", "/persons_exception")
     api.route(person_list_response, "person_list_response", "/persons_response")
     api.route(person_list_without_schema, "person_list_without_schema", "/persons_without_schema")
+    api.route(
+        computer_list_resource_with_disable_collection_count,
+        "computer_list_with_disabled_count",
+        "/computers_with_disabled_count",
+        "/persons/<int:person_id>/computers_with_disabled_count",
+    )
     api.route(computer_list, "computer_list", "/computers", "/persons/<int:person_id>/computers")
     api.route(computer_detail, "computer_detail", "/computers/<int:id>")
     api.route(computer_owner, "computer_owner", "/computers/<int:id>/relationships/owner")
@@ -1684,3 +1708,15 @@ def test_relationship_containing_hyphens(client, register_routes, person_compute
         f"/persons/{person.person_id}/relationships/computers-owned", content_type="application/vnd.api+json"
     )
     assert response.status_code == 200
+
+
+def test__sqlalchemy_data_layer__disable_collection_count(client, fixed_count_for_collection_count):
+    """
+    :param client:
+    :param fixed_count_for_collection_count:
+    :return:
+    """
+    response = client.get(f"/computers_with_disabled_count", content_type="application/vnd.api+json")
+    assert response.status_code == 200
+    assert response.json
+    assert response.json["meta"]["count"] == fixed_count_for_collection_count