Closes frol#55. Introduce paginate decorator that sets custom header …

…X-Total-Count

app/extensions/api/namespace.py

@@ -286,6 +286,42 @@ def _register_access_restriction_decorator(self, func, decorator_to_register):
             func._access_restriction_decorators = []  # pylint: disable=protected-access
         func._access_restriction_decorators.append(decorator_to_register)  # pylint: disable=protected-access
 
+    def paginate(self, parameters=None, locations=None):
+        """
+        Endpoint parameters registration decorator special for pagination.
+        If ``parameters`` is not provided default PaginationParameters will be
+        used.
+
+        Also, any custom Parameters can be used, but it needs to have ``limit`` and ``offset`` fields
+        """
+        if not parameters:
+            # Use default parameters if None specified
+            from app.extensions.api.parameters import PaginationParameters
+            parameters = PaginationParameters()
+
+        if not all(
+            mandatory in parameters.declared_fields
+            for mandatory in ('limit', 'offset')
+        ):
+            raise AttributeError(
+                '`limit` and `offset` fields must be in Parameter passed to `paginate()`'
+            )
+
+        def decorator(func):
+            @wraps(func)
+            def wrapper(self_, parameters_args, *args, **kwargs):
+                queryset = func(self_, parameters_args, *args, **kwargs)
+                total_count = queryset.count()
+                return (
+                    queryset
+                        .offset(parameters_args['offset'])
+                        .limit(parameters_args['limit']),
+                    HTTPStatus.OK,
+                    {'X-Total-Count': total_count}
+                )
+            return self.parameters(parameters, locations)(wrapper)
+        return decorator
+
     @contextmanager
     def commit_or_abort(self, session, default_error_message="The operation failed to complete"):
         """

app/modules/teams/resources.py

@@ -32,17 +32,16 @@ class Teams(Resource):
     """
     Manipulations with teams.
     """
-
-    @api.parameters(PaginationParameters())
     @api.response(schemas.BaseTeamSchema(many=True))
+    @api.paginate()
     def get(self, args):
         """
         List of teams.
 
         Returns a list of teams starting from ``offset`` limited by ``limit``
         parameter.
         """
-        return Team.query.offset(args['offset']).limit(args['limit'])
+        return Team.query
 
     @api.login_required(oauth_scopes=['teams:write'])
     @api.parameters(parameters.CreateTeamParameters())
@@ -144,13 +143,13 @@ class TeamMembers(Resource):
         kwargs_on_request=lambda kwargs: {'obj': kwargs['team']}
     )
     @api.permission_required(permissions.OwnerRolePermission(partial=True))
-    @api.parameters(PaginationParameters())
     @api.response(schemas.BaseTeamMemberSchema(many=True))
+    @api.paginate()
     def get(self, args, team):
         """
         Get team members by team ID.
         """
-        return team.members[args['offset']: args['offset'] + args['limit']]
+        return TeamMember.query.filter_by(team=team)
 
     @api.login_required(oauth_scopes=['teams:write'])
     @api.permission_required(

app/modules/users/resources.py

@@ -30,8 +30,8 @@ class Users(Resource):
 
     @api.login_required(oauth_scopes=['users:read'])
     @api.permission_required(permissions.AdminRolePermission())
-    @api.parameters(PaginationParameters())
     @api.response(schemas.BaseUserSchema(many=True))
+    @api.paginate()
     def get(self, args):
         """
         List of users.

flask_restplus_patched/namespace.py

@@ -1,9 +1,10 @@
+import json
 from functools import wraps
 
 import flask
 import flask_marshmallow
 from flask_restplus import Namespace as OriginalNamespace
-from flask_restplus.utils import merge
+from flask_restplus.utils import merge, unpack
 from flask_restplus._http import HTTPStatus
 from webargs.flaskparser import parser as webargs_parser
 from werkzeug import cached_property, exceptions as http_exceptions
@@ -137,6 +138,7 @@ def response_serializer_decorator(func):
             def dump_wrapper(*args, **kwargs):
                 # pylint: disable=missing-docstring
                 response = func(*args, **kwargs)
+                extra_headers = None
 
                 if response is None:
                     if model is not None:
@@ -145,13 +147,13 @@ def dump_wrapper(*args, **kwargs):
                 elif isinstance(response, flask.Response) or model is None:
                     return response
                 elif isinstance(response, tuple):
-                    response, _code = response
+                    response, _code, extra_headers = unpack(response)
                 else:
                     _code = code
 
                 if HTTPStatus(_code) is code:
                     response = model.dump(response).data
-                return response, _code
+                return response, _code, extra_headers
 
             return dump_wrapper
 

tests/modules/teams/resources/test_getting_teams_info.py

@@ -31,8 +31,9 @@ def test_getting_list_of_teams_by_authorized_user(
 ):
     with flask_app_client.login(regular_user, auth_scopes=auth_scopes):
         response = flask_app_client.get('/api/v1/teams/')
-
     assert response.status_code == 200
+    assert 'X-Total-Count' in response.headers
+    assert int(response.headers['X-Total-Count']) == 1
     assert response.content_type == 'application/json'
     assert isinstance(response.json, list)
     assert set(response.json[0].keys()) >= {'id', 'title'}