Repository: sql-formatter-org/sql-formatter
Pull request #257 (2022-06-28 13:33:45)

8.0: Detect LIMIT clauses in Parser
Moving the logic out of formatting phase


Patch:
diff --git a/src/core/ExpressionFormatter.ts b/src/core/ExpressionFormatter.ts
index ee4a3200..17f2e414 100644
--- a/src/core/ExpressionFormatter.ts
+++ b/src/core/ExpressionFormatter.ts
@@ -4,8 +4,8 @@ import { equalizeWhitespace } from 'src/utils';
 import Indentation from './Indentation';
 import InlineBlock from './InlineBlock';
 import Params from './Params';
-import { isReserved, isToken, type Token, TokenType, EOF_TOKEN } from './token';
-import { AstNode, BetweenPredicate, isTokenNode, Parenthesis } from './ast';
+import { isReserved, type Token, TokenType, EOF_TOKEN } from './token';
+import { AstNode, BetweenPredicate, isTokenNode, LimitClause, Parenthesis } from './ast';
 import { indentString } from './config';
 import WhitespaceBuilder, { WS } from './WhitespaceBuilder';
 
@@ -18,7 +18,6 @@ export default class ExpressionFormatter {
   private query: WhitespaceBuilder;
 
   private inline = false;
-  private previousReservedToken: Token = EOF_TOKEN;
   private nodes: AstNode[] = [];
   private index = -1;
 
@@ -43,10 +42,10 @@ export default class ExpressionFormatter {
         case 'between_predicate':
           this.formatBetweenPredicate(node);
           break;
+        case 'limit_clause':
+          this.formatLimitClause(node);
+          break;
         case 'token':
-          if (isReserved(node.token)) {
-            this.previousReservedToken = node.token;
-          }
           this.formatToken(node.token);
           break;
       }
@@ -125,6 +124,21 @@ export default class ExpressionFormatter {
     this.query.add(this.show(token), WS.NEWLINE, WS.INDENT);
   }
 
+  private formatLimitClause(node: LimitClause) {
+    this.formatCommand(node.limitToken);
+    if (node.offsetToken) {
+      this.query.add(
+        this.show(node.offsetToken),
+        ',',
+        WS.SPACE,
+        this.show(node.countToken),
+        WS.SPACE
+      );
+    } else {
+      this.query.add(this.show(node.countToken), WS.SPACE);
+    }
+  }
+
   /**
    * Formats a Reserved Binary Command onto query, joining neighbouring tokens
    */
@@ -279,7 +293,7 @@ export default class ExpressionFormatter {
    * Formats a comma Operator onto query, ending line unless in an Inline Block
    */
   private formatComma(token: Token) {
-    if (!this.inline && !isToken.LIMIT(this.getPreviousReservedToken())) {
+    if (!this.inline) {
       this.query.add(WS.NO_SPACE, this.show(token), WS.NEWLINE, WS.INDENT);
     } else {
       this.query.add(WS.NO_SPACE, this.show(token), WS.SPACE);
@@ -314,11 +328,6 @@ export default class ExpressionFormatter {
     }
   }
 
-  /** Returns the latest encountered reserved keyword token */
-  private getPreviousReservedToken(): Token {
-    return this.previousReservedToken;
-  }
-
   /** Fetches nth previous token from the token stream */
   private tokenLookBehind(n = 1): Token {
     return this.tokenLookAhead(-n);
diff --git a/src/core/InlineBlock.ts b/src/core/InlineBlock.ts
index b033365b..86c9735c 100644
--- a/src/core/InlineBlock.ts
+++ b/src/core/InlineBlock.ts
@@ -30,6 +30,8 @@ export default class InlineBlock {
         case 'between_predicate':
           length += this.betweenWidth(node);
           break;
+        case 'limit_clause':
+          return Infinity;
         case 'token':
           length += node.token.value.length;
           if (this.isForbiddenToken(node.token)) {
diff --git a/src/core/Parser.ts b/src/core/Parser.ts
index c002f8ae..6ced496f 100644
--- a/src/core/Parser.ts
+++ b/src/core/Parser.ts
@@ -1,5 +1,5 @@
 /* eslint-disable no-cond-assign */
-import { AstNode, BetweenPredicate, Parenthesis, Statement, TokenNode } from './ast';
+import { AstNode, BetweenPredicate, LimitClause, Parenthesis, Statement, TokenNode } from './ast';
 import { EOF_TOKEN, type Token, TokenType, isToken } from './token';
 
 /**
@@ -22,7 +22,6 @@ export default class Parser {
 
   private statement(): Statement | undefined {
     const children: AstNode[] = [];
-    let expr: Parenthesis | BetweenPredicate | undefined;
     while (true) {
       if (this.look().value === ';') {
         children.push(this.nextTokenNode());
@@ -33,14 +32,18 @@ export default class Parser {
         } else {
           return undefined;
         }
-      } else if ((expr = this.parenthesis() || this.betweenPredicate())) {
-        children.push(expr);
       } else {
-        children.push(this.nextTokenNode());
+        children.push(this.expression());
       }
     }
   }
 
+  private expression(): AstNode {
+    return (
+      this.parenthesis() || this.betweenPredicate() || this.limitClause() || this.nextTokenNode()
+    );
+  }
+
   private parenthesis(): Parenthesis | undefined {
     if (this.look().type === TokenType.OPEN_PAREN) {
       const children: AstNode[] = [];
@@ -49,7 +52,7 @@ export default class Parser {
       const hasWhitespaceBefore = Boolean(token.whitespaceBefore);
       let closeParen = '';
       while (this.look().type !== TokenType.CLOSE_PAREN && this.look().type !== TokenType.EOF) {
-        children.push(this.parenthesis() || this.betweenPredicate() || this.nextTokenNode());
+        children.push(this.expression());
       }
       if (this.look().type === TokenType.CLOSE_PAREN) {
         closeParen = this.next().value;
@@ -72,6 +75,25 @@ export default class Parser {
     return undefined;
   }
 
+  private limitClause(): LimitClause | undefined {
+    if (isToken.LIMIT(this.look()) && this.look(2).value === ',') {
+      return {
+        type: 'limit_clause',
+        limitToken: this.next(),
+        offsetToken: this.next(),
+        countToken: this.next() && this.next(), // Discard comma token
+      };
+    }
+    if (isToken.LIMIT(this.look())) {
+      return {
+        type: 'limit_clause',
+        limitToken: this.next(),
+        countToken: this.next(),
+      };
+    }
+    return undefined;
+  }
+
   // Returns current token without advancing the pointer
   private look(ahead = 0): Token {
     return this.tokens[this.index + ahead] || EOF_TOKEN;
diff --git a/src/core/ast.ts b/src/core/ast.ts
index 9df69dca..808aec91 100644
--- a/src/core/ast.ts
+++ b/src/core/ast.ts
@@ -19,6 +19,7 @@ export type Parenthesis = {
   hasWhitespaceBefore: boolean;
 };
 
+// BETWEEN <expr1> AND <expr2>
 export type BetweenPredicate = {
   type: 'between_predicate';
   betweenToken: Token;
@@ -27,6 +28,15 @@ export type BetweenPredicate = {
   expr2: Token;
 };
 
-export type AstNode = Parenthesis | BetweenPredicate | TokenNode;
+// LIMIT <count>
+// LIMIT <offset>, <count>
+export type LimitClause = {
+  type: 'limit_clause';
+  limitToken: Token;
+  countToken: Token;
+  offsetToken?: Token;
+};
+
+export type AstNode = Parenthesis | BetweenPredicate | LimitClause | TokenNode;
 
 export const isTokenNode = (node: AstNode): node is TokenNode => node.type === 'token';


Test patch:
diff --git a/test/unit/Parser.test.ts b/test/unit/Parser.test.ts
index 5328730f..a642dfba 100644
--- a/test/unit/Parser.test.ts
+++ b/test/unit/Parser.test.ts
@@ -4,7 +4,7 @@ import Tokenizer from 'src/core/Tokenizer';
 describe('Parser', () => {
   const parse = (sql: string) => {
     const tokens = new Tokenizer({
-      reservedCommands: ['SELECT', 'FROM', 'WHERE', 'CREATE TABLE'],
+      reservedCommands: ['SELECT', 'FROM', 'WHERE', 'LIMIT', 'CREATE TABLE'],
       reservedDependentClauses: ['WHEN', 'ELSE'],
       reservedBinaryCommands: ['UNION', 'JOIN'],
       reservedJoinConditions: ['ON', 'USING'],
@@ -210,4 +210,64 @@ describe('Parser', () => {
       ]
     `);
   });
+
+  it('parses LIMIT clause', () => {
+    expect(parse('LIMIT 10')).toMatchInlineSnapshot(`
+      Array [
+        Object {
+          "children": Array [
+            Object {
+              "countToken": Object {
+                "text": "10",
+                "type": "NUMBER",
+                "value": "10",
+                "whitespaceBefore": " ",
+              },
+              "limitToken": Object {
+                "text": "LIMIT",
+                "type": "RESERVED_COMMAND",
+                "value": "LIMIT",
+                "whitespaceBefore": "",
+              },
+              "type": "limit_clause",
+            },
+          ],
+          "type": "statement",
+        },
+      ]
+    `);
+  });
+
+  it('parses LIMIT clause with offset', () => {
+    expect(parse('LIMIT 200, 10')).toMatchInlineSnapshot(`
+      Array [
+        Object {
+          "children": Array [
+            Object {
+              "countToken": Object {
+                "text": "10",
+                "type": "NUMBER",
+                "value": "10",
+                "whitespaceBefore": " ",
+              },
+              "limitToken": Object {
+                "text": "LIMIT",
+                "type": "RESERVED_COMMAND",
+                "value": "LIMIT",
+                "whitespaceBefore": "",
+              },
+              "offsetToken": Object {
+                "text": "200",
+                "type": "NUMBER",
+                "value": "200",
+                "whitespaceBefore": " ",
+              },
+              "type": "limit_clause",
+            },
+          ],
+          "type": "statement",
+        },
+      ]
+    `);
+  });
 });