Move to cursor (keyset) based pagination

Author: piyush-jaiswal

app/routes/category.py

@@ -2,6 +2,7 @@
 from flask_jwt_extended import jwt_required
 from flask_smorest import Blueprint, abort
 from psycopg2.errors import UniqueViolation
+from sqlakeyset import get_page
 from sqlalchemy import UniqueConstraint, exists
 from sqlalchemy.exc import IntegrityError
 
@@ -154,17 +155,14 @@ class CategoryProducts(MethodView):
     @bp.doc(summary="Get Products within a Category")
     @bp.arguments(PaginationArgs, location="query", as_kwargs=True)
     @bp.response(200, ProductsOut)
-    def get(self, id, page):
+    def get(self, id, cursor):
         category_exists = db.session.query(exists().where(Category.id == id)).scalar()
         if not category_exists:
             abort(404)
 
-        products = (
-            Product.query.filter(
-                Product.subcategories.any(Subcategory.categories.any(id=id))
-            )
-            .order_by(Product.id.asc())
-            .paginate(page=page, per_page=CategoryProducts._PER_PAGE, error_out=False)
-        )
+        products = Product.query.filter(
+            Product.subcategories.any(Subcategory.categories.any(id=id))
+        ).order_by(Product.id.asc())
+        page = get_page(products, per_page=CategoryProducts._PER_PAGE, page=cursor)
 
-        return {"products": products}
+        return { "products": page, "cursor": page.paging }

app/routes/product.py

@@ -2,6 +2,7 @@
 from flask_jwt_extended import jwt_required
 from flask_smorest import Blueprint, abort
 from psycopg2.errors import UniqueViolation
+from sqlakeyset import get_page
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.exc import IntegrityError
 
@@ -51,15 +52,14 @@ def _get_by_name(self, name):
     @bp.arguments(NameArgs, location="query", as_kwargs=True)
     @bp.arguments(PaginationArgs, location="query", as_kwargs=True)
     @bp.response(200, ProductsOut)
-    def get(self, name, page):
+    def get(self, name, cursor):
         if name is not None:
             products = self._get_by_name(name)
+            return {"products": products}
         else:
-            products = Product.query.order_by(Product.id.asc()).paginate(
-                page=page, per_page=ProductCollection._PER_PAGE, error_out=False
-            )
-
-        return {"products": products}
+            products = Product.query.order_by(Product.id.asc())
+            page = get_page(products, per_page=ProductCollection._PER_PAGE, page=cursor)
+            return {"products": page, "cursor": page.paging}
 
     @jwt_required()
     @bp.doc(summary="Create Product", security=[{"access_token": []}])

app/routes/subcategory.py

@@ -2,6 +2,7 @@
 from flask_jwt_extended import jwt_required
 from flask_smorest import Blueprint, abort
 from psycopg2.errors import UniqueViolation
+from sqlakeyset import get_page
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.exc import IntegrityError
 
@@ -169,11 +170,9 @@ class SubcategoryProducts(MethodView):
     @bp.doc(summary="Get Products within a Subcategory")
     @bp.arguments(PaginationArgs, location="query", as_kwargs=True)
     @bp.response(200, ProductsOut)
-    def get(self, id, page):
+    def get(self, id, cursor):
         subcategory = Subcategory.query.get_or_404(id)
+        products = subcategory.products.order_by(Product.id.asc())
+        page = get_page(products, per_page=SubcategoryProducts._PER_PAGE, page=cursor)
 
-        products = subcategory.products.order_by(Product.id.asc()).paginate(
-            page=page, per_page=SubcategoryProducts._PER_PAGE, error_out=False
-        )
-
-        return {"products": products}
+        return {"products": page, "cursor": page.paging}

app/schemas.py

@@ -1,9 +1,38 @@
+import base64
+
 from marshmallow import Schema, ValidationError, fields, pre_load, validate, validates
 from marshmallow_sqlalchemy import SQLAlchemyAutoSchema, SQLAlchemySchema, auto_field
+from sqlakeyset import BadBookmark, unserialize_bookmark
 
 from app.models import Category, Product, Subcategory, User
 
 
+class Cursor(fields.Field[dict]):
+    def _serialize(self, paging, attr, obj, **kwargs):
+        def encode(s):
+            bytes = s.encode("utf-8")
+            return base64.urlsafe_b64encode(bytes).decode("utf-8")
+
+        if paging is None:
+            return None
+
+        return {
+            "next": encode(paging.bookmark_next) if paging.has_next else None,
+            "prev": encode(paging.bookmark_previous) if paging.has_previous else None,
+        }
+
+    def _deserialize(self, cursor, attr, data, **kwargs):
+        if cursor is None:
+            return None
+
+        try:
+            decoded_bytes = base64.urlsafe_b64decode(cursor.encode("utf-8"))
+            marker_serialized = decoded_bytes.decode("utf-8")
+            return unserialize_bookmark(marker_serialized)
+        except (TypeError, ValueError, KeyError, BadBookmark) as ex:
+            raise ValidationError("Invalid cursor") from ex
+
+
 class CategoryOut(SQLAlchemyAutoSchema):
     class Meta:
         model = Category
@@ -70,6 +99,7 @@ class Meta:
 
 class ProductsOut(Schema):
     products = fields.List(fields.Nested(ProductOut))
+    cursor = Cursor(required=False)  # require false for getting product by name
 
 
 class ProductIn(SQLAlchemySchema):
@@ -100,7 +130,7 @@ class NameArgs(Schema):
 
 
 class PaginationArgs(Schema):
-    page = fields.Int(load_default=1)
+    cursor = Cursor(load_default=None)
 
 
 class AuthIn(SQLAlchemySchema):