fix(optimize): Fix the limit/offset optimization not working with a t…

…able with multiple primary keys
teamstarter · Mar 18, 2024 · 514fffb · 514fffb
1 parent 3ce34a6
commit 514fffb
Show file tree

Hide file tree

Showing 13 changed files with 305 additions and 27 deletions.
diff --git a/lib/createListResolver.js b/lib/createListResolver.js
@@ -62,6 +62,7 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
 };
 exports.__esModule = true;
 var graphql_sequelize_1 = require("graphql-sequelize");
+var sequelize_1 = require("sequelize");
 var removeUnusedAttributes_1 = __importDefault(require("./removeUnusedAttributes"));
 function allowOrderOnAssociations(findOptions, model) {
     if (typeof findOptions.order === 'undefined') {
@@ -172,10 +173,10 @@ var argsAdvancedProcessing = function (findOptions, args, context, info, model,
 function trimAndOptimizeFindOptions(_a) {
     var findOptions = _a.findOptions, graphqlTypeDeclaration = _a.graphqlTypeDeclaration, info = _a.info, models = _a.models, args = _a.args;
     return __awaiter(this, void 0, void 0, function () {
-        var trimedFindOptions, fetchIdsFindOptions, result;
-        var _b;
-        return __generator(this, function (_c) {
-            switch (_c.label) {
+        var trimedFindOptions, fetchIdsMultiColumnsFindOptions, result_1, fetchIdsFindOptions, result;
+        var _b, _c;
+        return __generator(this, function (_d) {
+            switch (_d.label) {
                 case 0:
                     trimedFindOptions = graphqlTypeDeclaration.list &&
                         graphqlTypeDeclaration.list.removeUnusedAttributes === false
@@ -215,19 +216,42 @@ function trimAndOptimizeFindOptions(_a) {
                                 'undefined' ||
                             graphqlTypeDeclaration.list.disableOptimizationForLimitOffset !== true))) 
                     // If we have a list with a limit and an offset
-                    return [3 /*break*/, 2];
-                    fetchIdsFindOptions = __assign(__assign({}, trimedFindOptions), { 
+                    return [3 /*break*/, 4];
+                    if (!(graphqlTypeDeclaration.model.primaryKeyAttributes &&
+                        graphqlTypeDeclaration.model.primaryKeyAttributes.length > 0)) return [3 /*break*/, 2];
+                    fetchIdsMultiColumnsFindOptions = __assign(__assign({}, trimedFindOptions), { 
                         // We only fetch the primary attribute
-                        attributes: [graphqlTypeDeclaration.model.primaryKeyAttribute] });
-                    return [4 /*yield*/, graphqlTypeDeclaration.model.findAll(fetchIdsFindOptions)];
+                        attributes: graphqlTypeDeclaration.model.primaryKeyAttributes });
+                    return [4 /*yield*/, graphqlTypeDeclaration.model.findAll(fetchIdsMultiColumnsFindOptions)];
                 case 1:
-                    result = _c.sent();
+                    result_1 = _d.sent();
                     return [2 /*return*/, __assign(__assign({}, trimedFindOptions), { offset: undefined, limit: undefined, 
                             // We override the where to only fetch the rows we want.
                             where: (_b = {},
-                                _b[graphqlTypeDeclaration.model.primaryKeyAttribute] = result.map(function (r) { return r[graphqlTypeDeclaration.model.primaryKeyAttribute]; }),
+                                _b[sequelize_1.Op.or] = result_1.map(function (r) {
+                                    var where = {};
+                                    graphqlTypeDeclaration.model.primaryKeyAttributes.forEach(function (attr) {
+                                        if (!r[attr]) {
+                                            throw new Error("Got a null value for Primary key ".concat(attr, ", for model ").concat(graphqlTypeDeclaration.model.name, ". This should never be the case. Disable the optimization for this model with disableOptimizationForLimitOffset or make sure the primary keys of the table have no null values."));
+                                        }
+                                        where[attr] = r[attr];
+                                    });
+                                    return where;
+                                }),
                                 _b) })];
-                case 2: return [2 /*return*/, trimedFindOptions];
+                case 2:
+                    fetchIdsFindOptions = __assign(__assign({}, trimedFindOptions), { 
+                        // We only fetch the primary attribute
+                        attributes: [graphqlTypeDeclaration.model.primaryKeyAttribute] });
+                    return [4 /*yield*/, graphqlTypeDeclaration.model.findAll(fetchIdsFindOptions)];
+                case 3:
+                    result = _d.sent();
+                    return [2 /*return*/, __assign(__assign({}, trimedFindOptions), { offset: undefined, limit: undefined, 
+                            // We override the where to only fetch the rows we want.
+                            where: (_c = {},
+                                _c[graphqlTypeDeclaration.model.primaryKeyAttribute] = result.map(function (r) { return r[graphqlTypeDeclaration.model.primaryKeyAttribute]; }),
+                                _c) })];
+                case 4: return [2 /*return*/, trimedFindOptions];
             }
         });
     });

diff --git a/package.json b/package.json
@@ -58,7 +58,7 @@
     "test": "NODE_ENV=test DEBUG=-* PORT=8282 node ./node_modules/.bin/jest -i",
     "test-watch": "DEBUG=* PORT=8282 jest -i -u --watch",
     "lint-check": "eslint src",
-    "start": "rm -f ./src/tests/data/main.db && DEBUG=gsg node ./src/tests/testServer.js",
+    "start": "rm -f ./src/tests/data/main.db && DEBUG=gsg TS_NODE_PROJECT=./tsconfig.json node  --require ts-node/register --require tsconfig-paths/register ./src/tests/testServer.js",
     "build": "rm -rf ./lib/* && tsc --lib es2019,dom --esModuleInterop --downlevelIteration --outDir ./lib ./src/index.ts ./src/scripts/*.ts ./src/gsg.ts",
     "release": "yarn build && standard-version",
     "gsg": "node ./lib/gsg.js"

diff --git a/src/createListResolver.ts b/src/createListResolver.ts
@@ -1,5 +1,5 @@
 import { resolver } from 'graphql-sequelize'
-import { FindOptions, Model } from 'sequelize'
+import { FindOptions, Model, Op } from 'sequelize'
 import removeUnusedAttributes from './removeUnusedAttributes'
 import { GlobalBeforeHook, ModelDeclarationType } from './types/types'
 
@@ -203,6 +203,46 @@ async function trimAndOptimizeFindOptions({
       graphqlTypeDeclaration.list.disableOptimizationForLimitOffset !== true)
   ) {
     // then we pre-fetch the ids to avoid slowness problems for big offsets.
+
+    // There might be many primary keys.
+    if (
+      graphqlTypeDeclaration.model.primaryKeyAttributes &&
+      graphqlTypeDeclaration.model.primaryKeyAttributes.length > 0
+    ) {
+      const fetchIdsMultiColumnsFindOptions = {
+        ...trimedFindOptions,
+        // We only fetch the primary attribute
+        attributes: graphqlTypeDeclaration.model.primaryKeyAttributes,
+      }
+      const result = await graphqlTypeDeclaration.model.findAll(
+        fetchIdsMultiColumnsFindOptions
+      )
+
+      return {
+        ...trimedFindOptions,
+        offset: undefined,
+        limit: undefined,
+        // We override the where to only fetch the rows we want.
+        where: {
+          [Op.or]: result.map((r: any) => {
+            const where: any = {}
+            graphqlTypeDeclaration.model.primaryKeyAttributes.forEach(
+              (attr: string) => {
+                if (!r[attr]) {
+                  throw new Error(
+                    `Got a null value for Primary key ${attr}, for model ${graphqlTypeDeclaration.model.name}. This should never be the case. Disable the optimization for this model with disableOptimizationForLimitOffset or make sure the primary keys of the table have no null values.`
+                  )
+                }
+                where[attr] = r[attr]
+              }
+            )
+            return where
+          }),
+        },
+      }
+    }
+
+    // Or a single key
     const fetchIdsFindOptions = {
       ...trimedFindOptions,
       // We only fetch the primary attribute

diff --git a/src/removeUnusedAttributes.ts b/src/removeUnusedAttributes.ts
@@ -127,7 +127,6 @@ export default function removeUnusedAttributes(
   return {
     ...findOptions,
     attributes: [
-      // @ts-ignore
       ...new Set([
         ...attributes,
         ...linkFields,

diff --git a/src/tests/__snapshots__/basic.spec.js.snap b/src/tests/__snapshots__/basic.spec.js.snap
@@ -120,6 +120,53 @@ Array [
 ]
 `;
 
+exports[`Test the API queries Limit can be enforced.: Limit and offset optimization works with multi primary key tables. 1`] = `
+Object {
+  "userLocations": Array [
+    Object {
+      "locationId": 2,
+      "userId": "1",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "2",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "3",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "4",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "5",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "6",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "7",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "8",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "9",
+    },
+    Object {
+      "locationId": 2,
+      "userId": "10",
+    },
+  ],
+}
+`;
+
 exports[`Test the API queries One can exclude a basic field from a model 1`] = `
 Array [
   Object {

diff --git a/src/tests/basic.spec.js b/src/tests/basic.spec.js
@@ -300,5 +300,32 @@ describe('Test the API queries', () => {
     expect(
       responseWorksWithLimitAndOffsetOptimization.body.data.locations.length
     ).toBe(2)
+
+    const responseWorksWithLimitAndOffsetOptimizationAndMutliFieldTables =
+      await request(server)
+        .get(
+          `/graphql?query=
+          query getUserLocations {
+            userLocations: userLocation(order: "locationId,userId", limit: 10, offset:10) {
+              userId
+              locationId
+            }
+          }
+          &operationName=getUserLocations`
+        )
+        .set('userId', 1)
+    expect(
+      responseWorksWithLimitAndOffsetOptimizationAndMutliFieldTables.body.errors
+    ).toBeUndefined()
+    expect(
+      responseWorksWithLimitAndOffsetOptimizationAndMutliFieldTables.body.data
+        .userLocations.length
+    ).toBe(10)
+
+    expect(
+      responseWorksWithLimitAndOffsetOptimizationAndMutliFieldTables.body.data
+    ).toMatchSnapshot(
+      'Limit and offset optimization works with multi primary key tables.'
+    )
   })
 })
diff --git a/src/tests/migrations/20240318153721-create-user-location.js b/src/tests/migrations/20240318153721-create-user-location.js
@@ -0,0 +1,32 @@
+'use strict'
+module.exports = {
+  up: async (queryInterface, Sequelize) => {
+    await queryInterface.createTable('userLocation', {
+      userId: {
+        type: Sequelize.INTEGER,
+        allowNull: false,
+        primaryKey: true,
+      },
+      locationId: {
+        type: Sequelize.INTEGER,
+        allowNull: false,
+        primaryKey: true,
+      },
+      createdAt: {
+        allowNull: false,
+        type: Sequelize.DATE,
+      },
+      updatedAt: {
+        allowNull: false,
+        type: Sequelize.DATE,
+      },
+      deletedAt: {
+        allowNull: true,
+        type: Sequelize.DATE,
+      },
+    })
+  },
+  down: async (queryInterface, Sequelize) => {
+    await queryInterface.dropTable('userLocation')
+  },
+}
diff --git a/src/tests/models/userLocation.js b/src/tests/models/userLocation.js
@@ -0,0 +1,36 @@
+'use strict'
+module.exports = function (sequelize, DataTypes) {
+  const UserLocation = sequelize.define(
+    'userLocation',
+    {
+      userId: {
+        type: DataTypes.STRING,
+        allowNull: false,
+        primaryKey: true,
+      },
+      locationId: {
+        type: DataTypes.INTEGER,
+        allowNull: false,
+        primaryKey: true,
+      },
+      departmentId: {
+        type: DataTypes.INTEGER,
+        allowNull: false,
+      },
+      type: {
+        type: DataTypes.STRING,
+        allowNull: true,
+      },
+    },
+    {
+      freezeTableName: true,
+    }
+  )
+
+  UserLocation.associate = function (models) {
+    // associations can be defined here
+    models.userLocation.belongsTo(models.user)
+    models.userLocation.belongsTo(models.location)
+  }
+  return UserLocation
+}
diff --git a/src/tests/schema.ts b/src/tests/schema.ts
@@ -106,6 +106,7 @@ graphqlSchemaDeclaration.user = {
           findOptions.include = []
         }
 
+        // @ts-ignore
         findOptions.include.push({
           model: models.company,
           required: true,
@@ -387,6 +388,16 @@ graphqlSchemaDeclaration.companySetting = {
   actions: ['list'],
 }
 
+graphqlSchemaDeclaration.userLocation = {
+  model: models.userLocation,
+  actions: ['list'],
+  list: {
+    before: async (findOptions) => {
+      return findOptions
+    },
+  },
+} as ModelDeclarationType<typeof models.userLocation>
+
 module.exports = (globalPreCallback: any, httpServer: any) => {
   // Creating the WebSocket server
   const wsServer = new WebSocketServer({

diff --git a/src/tests/seeders/20200415181548-location.js b/src/tests/seeders/20200415181548-location.js
@@ -1,25 +1,25 @@
 'use strict'
-const timestamp = entry =>
+const timestamp = (entry) =>
   Object.assign(entry, {
     createdAt: entry.createdAt || new Date('2007-07-12 00:04:22'),
-    updatedAt: new Date('2007-07-12 00:04:22')
+    updatedAt: new Date('2007-07-12 00:04:22'),
   })
 
-const departmentPerCompany = 3
+const locationPerCompany = 3
 
 module.exports = {
-  up: function(queryInterface, Sequelize) {
+  up: function (queryInterface, Sequelize) {
     let locations = []
     for (let companyId = 0; companyId < 50; companyId++) {
       locations = [
         ...locations,
-        ...[...Array(departmentPerCompany)].map((u, index) => ({
-          id: companyId * departmentPerCompany + 1 + index,
-          name: `Location ${companyId * departmentPerCompany +
-            1 +
-            index} c ${companyId + 1}`,
-          companyId: companyId + 1
-        }))
+        ...[...Array(locationPerCompany)].map((u, index) => ({
+          id: companyId * locationPerCompany + 1 + index,
+          name: `Location ${companyId * locationPerCompany + 1 + index} c ${
+            companyId + 1
+          }`,
+          companyId: companyId + 1,
+        })),
       ]
     }
 
@@ -28,7 +28,7 @@ module.exports = {
     return queryInterface.bulkInsert('location', locations, {})
   },
 
-  down: function(queryInterface, Sequelize) {
+  down: function (queryInterface, Sequelize) {
     return queryInterface.bulkDelete('location', null, {})
-  }
+  },
 }