add range header option (#122)

* add range header option

* remove unused variable

* lint, docs, and minor code optimization

* fixed a few typos and made docs a bit clearer

* double backtics

docs/source/crud/piccolo_crud.rst

@@ -291,6 +291,41 @@ You can also use this on GET requests when retrieving a single row, for example:
 
 -------------------------------------------------------------------------------
 
+Content-Range header
+--------------------
+
+In some applications it can be practical to get information about the
+total number of records without invoking a separate call to
+the ``count`` endpoint. Piccolo API will supply this information in the
+``Content-Range`` response header if ``add_range_headers`` is set to ``True``.
+You can use the ``range_header_plural_name`` parameter to configure the
+"plural name" used in the ``Content-Range`` response header.
+
+The contents of the ``Content-Range`` header might look something like this
+for the "Movie" table: ``movie 0-9/100``
+
+.. code-block:: python
+
+    # app.py
+    from piccolo_api.crud.endpoints import PiccoloCRUD
+    from starlette.routing import Mount, Router
+
+    from movies.tables import Movie, Director
+
+
+    app = Router([
+        Mount(
+            path='/movie',
+            app=PiccoloCRUD(
+                table=Movie,
+                add_range_headers=True,
+                range_header_plural_name="movies"
+            )
+        )
+    ])
+
+-------------------------------------------------------------------------------
+
 Source
 ------
 

piccolo_api/crud/endpoints.py

@@ -143,6 +143,8 @@ def __init__(
         schema_extra: t.Optional[t.Dict[str, t.Any]] = None,
         max_joins: int = 0,
         hooks: t.Optional[t.List[Hook]] = None,
+        add_range_headers: bool = False,
+        range_header_plural_name: t.Optional[str] = None,
     ) -> None:
         """
         :param table:
@@ -196,6 +198,11 @@ def __init__(
             To see which fields can be filtered in this way, you can check
             the ``visible_fields_options`` value returned by the ``/schema``
             endpoint.
+        :param add_range_headers:
+            if True, will add content-range headers to GET responses returning lists of data
+        :param range_header_plural_name:
+            Specify object name which is included in the Content-Range header
+            when `add_range_headers` is True. Defaults to table name if unset.
 
         """  # noqa: E501
         self.table = table
@@ -212,6 +219,8 @@ def __init__(
             }
         else:
             self._hook_map = None  # type: ignore
+        self.add_range_headers = add_range_headers
+        self.range_header_plural_name = range_header_plural_name
 
         schema_extra = schema_extra if isinstance(schema_extra, dict) else {}
         self.visible_fields_options = get_visible_fields_options(
@@ -712,11 +721,28 @@ async def get_all(
             )
         query = query.limit(page_size)
         page = split_params.page
+        offset = 0
         if page > 1:
             offset = page_size * (page - 1)
             query = query.offset(offset).limit(page_size)
 
         rows = await query.run()
+        headers = {}
+        if self.add_range_headers:
+            plural_name = (
+                self.range_header_plural_name or self.table._meta.tablename
+            )
+            row_length = len(rows)
+            if row_length == 0:
+                curr_page_len = 0
+            else:
+                curr_page_len = row_length - 1
+            curr_page_len = curr_page_len + offset
+            count = await self.table.count().run()
+            curr_page_string = f"{offset}-{curr_page_len}"
+            headers[
+                "Content-Range"
+            ] = f"{plural_name} {curr_page_string}/{count}"
 
         # We need to serialise it ourselves, in case there are datetime
         # fields.
@@ -725,7 +751,7 @@ async def get_all(
             include_columns=tuple(visible_columns),
             nested=nested,
         )(rows=rows).json()
-        return CustomJSONResponse(json)
+        return CustomJSONResponse(json, headers=headers)
 
     ###########################################################################
 

tests/crud/test_crud_endpoints.py

@@ -1233,3 +1233,96 @@ def test_parsing(self):
         self.assertEqual(
             parsed_2, {"tags": ["horror", "scifi"], "rating": "90"}
         )
+
+
+class RangeHeaders(TestCase):
+    def setUp(self):
+        Movie.create_table(if_not_exists=True).run_sync()
+
+    def tearDown(self):
+        Movie.alter().drop_table().run_sync()
+
+    def test_plural_name(self):
+        """
+        Make sure the content-range header responds correctly for empty rows
+        """
+        client = TestClient(
+            PiccoloCRUD(
+                table=Movie,
+                read_only=False,
+                add_range_headers=True,
+                range_header_plural_name="movies",
+            )
+        )
+
+        response = client.get("/")
+        self.assertTrue(response.status_code == 200)
+        # Make sure the content is correct:
+        response_json = response.json()
+        self.assertEqual(0, len(response_json["rows"]))
+        self.assertEqual(response.headers.get("Content-Range"), "movies 0-0/0")
+
+    def test_empty_list(self):
+        """
+        Make sure the content-range header responds correctly for empty rows
+        """
+        client = TestClient(
+            PiccoloCRUD(table=Movie, read_only=False, add_range_headers=True)
+        )
+
+        response = client.get("/")
+        self.assertTrue(response.status_code == 200)
+        # Make sure the content is correct:
+        response_json = response.json()
+        self.assertEqual(0, len(response_json["rows"]))
+        self.assertEqual(response.headers.get("Content-Range"), "movie 0-0/0")
+
+    def test_unpaged_ranges(self):
+        """
+        Make sure the content-range header responds
+        correctly for unpaged results
+        """
+        client = TestClient(
+            PiccoloCRUD(table=Movie, read_only=False, add_range_headers=True)
+        )
+
+        movie = Movie(name="Star Wars", rating=93)
+        movie.save().run_sync()
+        movie2 = Movie(name="Blade Runner", rating=94)
+        movie2.save().run_sync()
+
+        response = client.get("/")
+        self.assertTrue(response.status_code == 200)
+        # Make sure the content is correct:
+        response_json = response.json()
+        self.assertEqual(2, len(response_json["rows"]))
+        self.assertTrue(2, response.headers.get("Content-Range").split("/")[1])
+        self.assertEqual(response.headers.get("Content-Range"), "movie 0-1/2")
+
+    def test_page_sized_results(self):
+        """
+        Make sure the content-range header responds
+        correctly requests with page_size
+        """
+        client = TestClient(
+            PiccoloCRUD(table=Movie, read_only=False, add_range_headers=True)
+        )
+
+        movie = Movie(name="Star Wars", rating=93)
+        movie.save().run_sync()
+        movie2 = Movie(name="Blade Runner", rating=94)
+        movie2.save().run_sync()
+        movie3 = Movie(name="The Godfather", rating=95)
+        movie3.save().run_sync()
+
+        response = client.get("/?__page_size=1")
+        self.assertEqual(response.headers.get("Content-Range"), "movie 0-0/3")
+
+        response = client.get("/?__page_size=1&__page=2")
+        self.assertEqual(response.headers.get("Content-Range"), "movie 1-1/3")
+
+        response = client.get("/?__page_size=1&__page=2")
+        self.assertEqual(response.headers.get("Content-Range"), "movie 1-1/3")
+
+        response = client.get("/?__page_size=99&__page=1")
+        self.assertEqual(response.headers.get("Content-Range"), "movie 0-2/3")