From 665030b018b6b313863804f8db49b33d80fd9a1d Mon Sep 17 00:00:00 2001
From: taozhi8833998 <taozhi8833998@163.com>
Date: Fri, 29 Nov 2024 09:02:40 +0800
Subject: [PATCH] feat: support subquery in limit clause

---
 pegjs/postgresql.pegjs | 21 +++++++++++----------
 test/postgres.spec.js  |  7 +++++++
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
index a4360f4f..7cb7cd8d 100644
--- a/pegjs/postgresql.pegjs
+++ b/pegjs/postgresql.pegjs
@@ -3118,15 +3118,8 @@ comment_on_stmt
       expr: is,
     }
   }
-select_stmt
-  = KW_SELECT __ ';' {
-    // => { type: 'select'; }
-    return {
-      type: 'select',
-    }
-  }
-  / select_stmt_nake
-  / s:('(' __ select_stmt __ ')') {
+select_stmt_parentheses
+  = s:('(' __ select_stmt __ ')') {
     /*
     export interface select_stmt_node extends select_stmt_nake  {
        parentheses: true;
@@ -3138,6 +3131,14 @@ select_stmt
         parentheses_symbol: true,
       }
     }
+select_stmt
+  = KW_SELECT __ ';' {
+    // => { type: 'select'; }
+    return {
+      type: 'select',
+    }
+  }
+  / select_stmt_nake / select_stmt_parentheses
 
 with_clause
   = KW_WITH __ head:cte_definition tail:(__ COMMA __ cte_definition)* {
@@ -3820,7 +3821,7 @@ number_or_param
   / param
 
 limit_clause
-  = l:(KW_LIMIT __ (number_or_param / KW_ALL))? __ tail:(KW_OFFSET __ number_or_param)? {
+  = l:(KW_LIMIT __ (number_or_param / KW_ALL / select_stmt_parentheses))? __ tail:(KW_OFFSET __ number_or_param)? {
     // => { separator: 'offset' | ''; value: [number_or_param | { type: 'origin', value: 'all' }, number_or_param?] }
       const res = []
       if (l) res.push(typeof l[2] === 'string' ? { type: 'origin', value: 'all' } : l[2])
diff --git a/test/postgres.spec.js b/test/postgres.spec.js
index 34f8b779..9cf03fe2 100644
--- a/test/postgres.spec.js
+++ b/test/postgres.spec.js
@@ -1707,6 +1707,13 @@ describe('Postgres', () => {
         'CREATE INDEX ON "tableName" (supplier, amount) INCLUDE (id)'
       ]
     },
+    {
+      title: 'limit by select stmt',
+      sql: [
+        'SELECT * FROM user LIMIT (SELECT COUNT(*) / 2 FROM user);',
+        'SELECT * FROM "user" LIMIT (SELECT COUNT(*) / 2 FROM "user")'
+      ]
+    },
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {