Add mode parameter (#29)

README.md

@@ -26,7 +26,12 @@ It ignores non-SELECT queries. It understands CTE statements. It understands str
 - `sqlText` - SQL text to enforce limits on. Multiple statements allowed. Only `SELECT` statements are targeted.
 - `limitStrategies` - Keyword or array of strategies used to restrict rows. Must be either `limit`, `first`, `top`, `fetch` for `FETCH NEXT`/`FETCH FIRST`.
 - `limitNumber` - Number of rows to allow. If number in statement is lower, it is untouched. If higher it is lowered to limit. If missing it is added.
-- `offsetNumber` - Number of rows to skip before beginning to return rows from the query. If number in statement is defined, it is untouched. If missing it is added.
+- `offsetNumber` - Number of rows to skip before beginning to return rows from the query. If number in statement is defined, it is untouched. If missing it is added. (optional)
+- `mode` - Mode for enforcing `limitNumber` or `offsetNumber`. Must be either `replace`, `insert`
+  or `cap`. The default is `cap` if not defined.
+  - "replace": Replace existing value. If not existing, it will be inserted.
+  - "insert": Insert if limit or offset are not existing.
+  - "cap": Insert if not existing and if higher it is lowered to the defined value.
 
 Returns `sqlText` with limits enforced.
 

src/index.js

@@ -11,9 +11,16 @@ const getStatements = require("./get-statements");
  * @param {Array<String>|String} limitStrategies -- First strategy value takes priority if no limit exists
  * @param {number} limitNumber -- number to enforce for limit keyword
  * @param {number} [offsetNumber] -- offset number to enforce
+ * @param {boolean} [mode] -- Mode for enforcing `limitNumber` or `offsetNumber`
  * @returns {string}
  */
-function limit(sqlText, limitStrategies, limitNumber, offsetNumber) {
+function limit(
+  sqlText,
+  limitStrategies,
+  limitNumber,
+  offsetNumber,
+  mode = "cap"
+) {
   if (typeof sqlText !== "string") {
     throw new Error("sqlText must be string");
   }
@@ -36,9 +43,9 @@ function limit(sqlText, limitStrategies, limitNumber, offsetNumber) {
 
   return getStatements(sqlText)
     .map((statement) => {
-      statement.enforceLimit(strategies, limitNumber);
+      statement.enforceLimit(strategies, limitNumber, mode);
       if (typeof offsetNumber === "number") {
-        statement.injectOffset(offsetNumber);
+        statement.enforceOffset(offsetNumber, mode);
       }
       return statement.toString();
     })

src/statement.js

@@ -88,12 +88,30 @@ class Statement {
     }
   }
 
+  updateExistingNumberToken({ mode, tokens, numberToken, value }) {
+    if (
+      (mode === "cap" && parseInt(numberToken.value, 10) > value) ||
+      mode === "replace"
+    ) {
+      const firstHalf = tokens.slice(0, numberToken.index);
+      const secondhalf = tokens.slice(numberToken.index + 1);
+      this.tokens = [
+        ...firstHalf,
+        { ...numberToken, text: value, value },
+        ...secondhalf,
+      ];
+      return;
+    }
+    return;
+  }
+
   /**
    *
    * @param {Array<String>} strategiesToEnforce
    * @param {Number} limitNumber
+   *  @param {string} [mode]
    */
-  enforceLimit(strategiesToEnforce, limitNumber) {
+  enforceLimit(strategiesToEnforce, limitNumber, mode) {
     const { statementToken, tokens } = this;
 
     strategiesToEnforce.forEach((s) => {
@@ -110,19 +128,13 @@ class Statement {
           statementToken.index
         );
 
-        // If number token, check to see if over the limit and reset it if it is
-        // Otherwise return early
         if (numberToken) {
-          if (parseInt(numberToken.value, 10) > limitNumber) {
-            const firstHalf = tokens.slice(0, numberToken.index);
-            const secondhalf = tokens.slice(numberToken.index + 1);
-            this.tokens = [
-              ...firstHalf,
-              { ...numberToken, text: limitNumber, value: limitNumber },
-              ...secondhalf,
-            ];
-            return;
-          }
+          this.updateExistingNumberToken({
+            mode,
+            tokens,
+            numberToken,
+            value: limitNumber,
+          });
           return;
         }
       }
@@ -141,14 +153,20 @@ class Statement {
 
   /**
    * @param {number} offsetNumber
+   * @param {string} [mode]
    */
-  injectOffset(offsetNumber) {
+  enforceOffset(offsetNumber, mode) {
     const { statementToken, tokens } = this;
 
     if (statementToken && statementToken.value === "select") {
-      const offsetToken = offset.has(tokens, statementToken.index);
-      // If offset token exists already, return early
-      if (offsetToken) {
+      const numberToken = offset.has(tokens, statementToken.index);
+      if (numberToken) {
+        this.updateExistingNumberToken({
+          mode,
+          tokens,
+          numberToken,
+          value: offsetNumber,
+        });
         return;
       }
 

test/api-validations.js

@@ -63,6 +63,28 @@ describe("api: limit", function () {
     const res = sqlLimiter.limit(original, "limit", 100);
     assert.equal(res, expected);
   });
+
+  it("limit, offset with insert mode", function () {
+    const res = sqlLimiter.limit(
+      `SELECT * from something limit 10000`,
+      ["limit"],
+      100,
+      10,
+      "insert"
+    );
+    assert.equal(res, `SELECT * from something limit 10000 offset 10`);
+  });
+
+  it("limit with replace mode", function () {
+    const res = sqlLimiter.limit(
+      `SELECT * from something limit 10000 offset 10`,
+      ["limit"],
+      100,
+      0,
+      "replace"
+    );
+    assert.equal(res, `SELECT * from something limit 100 offset 0`);
+  });
 });
 
 describe("api: getStatementType", function () {