From 44d016cd5c85c668799599c3f71a460cfaac27d4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20Kr=C3=BCger?= <mkg20001@gmail.com>
Date: Fri, 24 Feb 2023 05:33:50 +0100
Subject: [PATCH] Add support for giving back total count via header

---
 lib/action-util.js  | 11 ++++++-
 lib/actions/add.js  |  2 +-
 lib/actions/find.js | 18 +++++++++--
 lib/index.js        |  2 ++
 test/index.js       | 73 +++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 101 insertions(+), 5 deletions(-)

diff --git a/lib/action-util.js b/lib/action-util.js
index 1ebacc0..abcd7d2 100644
--- a/lib/action-util.js
+++ b/lib/action-util.js
@@ -9,7 +9,7 @@ const internals = {};
  *
  * @type {Object}
  */
-module.exports = (request, options) => {
+module.exports = (request, options, h) => {
 
     return {
 
@@ -157,6 +157,15 @@ module.exports = (request, options) => {
             const userId = Hoek.reach(request.auth.credentials, options.userIdProperty);
 
             return userId;
+        },
+
+        replyWithRange: (rangeStart, rangeEnd, total, results) => {
+            if (options.totalCountHeader === 'content-range') {
+                return h.response(results).header('content-range', `${options.model} ${rangeStart}-${rangeStart + results.length}/${total[0]['count(*)']}`);
+            }
+
+            return h.response(results).header(options.totalCountHeader, String(total[0]['count(*)']));
+
         }
     };
 };
diff --git a/lib/actions/add.js b/lib/actions/add.js
index 8c3319b..be84f08 100644
--- a/lib/actions/add.js
+++ b/lib/actions/add.js
@@ -32,7 +32,7 @@ module.exports = function addToCollection(route, options) {
             return Boom.notFound('No record found with the specified `id`.');
         }
 
-        //if a key was specified for the relation, we're linking an exsiting.
+        //if a key was specified for the relation, we're linking an existing.
         if (keys.child.value) {
 
             const updatedTokenId = await foundModel.$relatedQuery(options.associationAttr)
diff --git a/lib/actions/find.js b/lib/actions/find.js
index 7d6432a..031ab5f 100644
--- a/lib/actions/find.js
+++ b/lib/actions/find.js
@@ -21,9 +21,9 @@ const Actions = require('../action-util');
 
 module.exports = function findRecords(route, options) {
 
-    return async (request) => {
+    return async (request, reply) => {
 
-        const actionUtil = Actions(request, options);
+        const actionUtil = Actions(request, options, reply);
         // Look up the model
         const Model = actionUtil.parseModel();
         // Lookup for records that match the specified criteria.  Are we just counting?
@@ -39,18 +39,30 @@ module.exports = function findRecords(route, options) {
         const rangeEnd = rangeStart ? rangeStart + limit : limit;
         const where = actionUtil.parseWhere();
 
-        const foundModelsQuery = Model.query().range(rangeStart, rangeEnd).limit(limit);
+        const foundModelsQuery = Model.query();
 
         if (where) {
             foundModelsQuery.where(where);
         }
 
+        let total;
+
+        if (options.includeTotalCount) {
+            total = await foundModelsQuery.clone().count();
+        }
+
         if (sort) {
             foundModelsQuery.orderByRaw(sort);
         }
 
+        foundModelsQuery.range(rangeStart, rangeEnd).limit(limit);
+
         const foundModels = await foundModelsQuery;
 
+        if (total) {
+            return actionUtil.replyWithRange(rangeStart, rangeEnd, total, foundModels.results);
+        }
+
         return foundModels.results;
 
     };
diff --git a/lib/index.js b/lib/index.js
index 6a5a7e1..9ea72d0 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -72,6 +72,8 @@ internals.defaults = {
     userModel: 'Users', // since it's not in the url
     userIdProperty: 'id', // on auth credentials
     prefix: '',
+    includeTotalCount: false,
+    totalCountHeader: 'content-range',
     _private: {
         actAsUserModifiedPath: false,
         count: false
diff --git a/test/index.js b/test/index.js
index 1995b5d..de5163e 100644
--- a/test/index.js
+++ b/test/index.js
@@ -955,6 +955,79 @@ describe('Tandy', () => {
         expect(result).to.be.an.array();
         expect(result.length).to.equal(4);
     });
+    it('Fetches all users and includes total count', async () => {
+
+        const server = await getServer(getOptions({
+            tandy: {
+                includeTotalCount: true
+            }
+        }));
+        server.registerModel([
+            TestModels.Users,
+            TestModels.Tokens
+        ]);
+
+        await server.initialize();
+        const knex = server.knex();
+        // const data = await knex.seed.run({ directory: 'test/seeds' });
+        await knex.seed.run({ directory: 'test/seeds' });
+
+        server.route({
+            method: 'GET',
+            path: '/users',
+            handler: { tandy: {} }
+        });
+
+        const options = {
+            method: 'GET',
+            url: '/users'
+        };
+
+        const response = await server.inject(options);
+        const result = response.result;
+
+        expect(response.statusCode).to.equal(200);
+        expect(result).to.be.an.array();
+        expect(result.length).to.equal(4);
+        expect(response.headers['content-range']).to.equal('users 0-4/4');
+    });
+    it('Fetches all users and includes total count, custom header name', async () => {
+
+        const server = await getServer(getOptions({
+            tandy: {
+                includeTotalCount: true,
+                totalCountHeader: 'x-count'
+            }
+        }));
+        server.registerModel([
+            TestModels.Users,
+            TestModels.Tokens
+        ]);
+
+        await server.initialize();
+        const knex = server.knex();
+        // const data = await knex.seed.run({ directory: 'test/seeds' });
+        await knex.seed.run({ directory: 'test/seeds' });
+
+        server.route({
+            method: 'GET',
+            path: '/users',
+            handler: { tandy: {} }
+        });
+
+        const options = {
+            method: 'GET',
+            url: '/users'
+        };
+
+        const response = await server.inject(options);
+        const result = response.result;
+
+        expect(response.statusCode).to.equal(200);
+        expect(result).to.be.an.array();
+        expect(result.length).to.equal(4);
+        expect(response.headers['x-count']).to.equal('4');
+    });
     it('Fetches all users without actAsUser', async () => {
 
         const config = getOptions({