feat: Support "LIMIT ?" in UPDATE and DELETE for MySQL (#2365)

PostgreSQL doesn't support UPDATE-LIMIT issue #2131
sqlc-dev · Jun 27, 2023 · d2559e8 · d2559e8
1 parent 1357f50
commit d2559e8
Show file tree

Hide file tree

Showing 10 changed files with 115 additions and 2 deletions.
diff --git a/internal/compiler/find_params.go b/internal/compiler/find_params.go
@@ -69,6 +69,11 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 	case *ast.CallStmt:
 		p.parent = n.FuncCall
 
+	case *ast.DeleteStmt:
+		if n.LimitCount != nil {
+			p.limitCount = n.LimitCount
+		}
+
 	case *ast.FuncCall:
 		p.parent = node
 
@@ -129,6 +134,9 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 			}
 			p.seen[ref.Location] = struct{}{}
 		}
+		if n.LimitCount != nil {
+			p.limitCount = n.LimitCount
+		}
 
 	case *ast.RangeVar:
 		p.rangeVar = n

diff --git a/internal/endtoend/testdata/limit/mysql/go/db.go b/internal/endtoend/testdata/limit/mysql/go/db.go
diff --git a/internal/endtoend/testdata/limit/mysql/go/models.go b/internal/endtoend/testdata/limit/mysql/go/models.go
diff --git a/internal/endtoend/testdata/limit/mysql/go/query.sql.go b/internal/endtoend/testdata/limit/mysql/go/query.sql.go
diff --git a/internal/endtoend/testdata/limit/mysql/query.sql b/internal/endtoend/testdata/limit/mysql/query.sql
@@ -0,0 +1,7 @@
+CREATE TABLE foo (bar bool not null);
+
+-- name: LimitMe :exec
+UPDATE foo SET bar='baz' LIMIT ?;
+
+-- name: LimitMeToo :exec
+DELETE FROM foo LIMIT ?;
diff --git a/internal/endtoend/testdata/limit/mysql/sqlc.json b/internal/endtoend/testdata/limit/mysql/sqlc.json
@@ -0,0 +1,12 @@
+{
+  "version": "1",
+  "packages": [
+    {
+      "path": "go",
+			"engine": "mysql",
+      "name": "querytest",
+      "schema": "query.sql",
+      "queries": "query.sql"
+    }
+  ]
+}
diff --git a/internal/engine/dolphin/convert.go b/internal/engine/dolphin/convert.go
@@ -327,12 +327,16 @@ func (c *cc) convertDeleteStmt(n *pcast.DeleteStmt) *ast.DeleteStmt {
 	relations := &ast.List{}
 	convertToRangeVarList(rels, relations)
 
-	return &ast.DeleteStmt{
+	stmt := &ast.DeleteStmt{
 		Relations:     relations,
 		WhereClause:   c.convert(n.Where),
 		ReturningList: &ast.List{},
 		WithClause:    c.convertWithClause(n.With),
 	}
+	if n.Limit != nil {
+		stmt.LimitCount = c.convert(n.Limit.Count)
+	}
+	return stmt
 }
 
 func (c *cc) convertDropTableStmt(n *pcast.DropTableStmt) ast.Node {
@@ -574,14 +578,18 @@ func (c *cc) convertUpdateStmt(n *pcast.UpdateStmt) *ast.UpdateStmt {
 	for _, a := range n.List {
 		list.Items = append(list.Items, c.convertAssignment(a))
 	}
-	return &ast.UpdateStmt{
+	stmt := &ast.UpdateStmt{
 		Relations:     relations,
 		TargetList:    list,
 		WhereClause:   c.convert(n.Where),
 		FromClause:    &ast.List{},
 		ReturningList: &ast.List{},
 		WithClause:    c.convertWithClause(n.With),
 	}
+	if n.Limit != nil {
+		stmt.LimitCount = c.convert(n.Limit.Count)
+	}
+	return stmt
 }
 
 func (c *cc) convertValueExpr(n *driver.ValueExpr) *ast.A_Const {

diff --git a/internal/sql/ast/delete_stmt.go b/internal/sql/ast/delete_stmt.go
@@ -4,6 +4,7 @@ type DeleteStmt struct {
 	Relations     *List
 	UsingClause   *List
 	WhereClause   Node
+	LimitCount    Node
 	ReturningList *List
 	WithClause    *WithClause
 }

diff --git a/internal/sql/ast/update_stmt.go b/internal/sql/ast/update_stmt.go
@@ -5,6 +5,7 @@ type UpdateStmt struct {
 	TargetList    *List
 	WhereClause   Node
 	FromClause    *List
+	LimitCount    Node
 	ReturningList *List
 	WithClause    *WithClause
 }

diff --git a/internal/sql/astutils/walk.go b/internal/sql/astutils/walk.go
@@ -1068,6 +1068,9 @@ func Walk(f Visitor, node ast.Node) {
 		if n.WhereClause != nil {
 			Walk(f, n.WhereClause)
 		}
+		if n.LimitCount != nil {
+			Walk(f, n.LimitCount)
+		}
 		if n.ReturningList != nil {
 			Walk(f, n.ReturningList)
 		}
@@ -2038,6 +2041,9 @@ func Walk(f Visitor, node ast.Node) {
 		if n.FromClause != nil {
 			Walk(f, n.FromClause)
 		}
+		if n.LimitCount != nil {
+			Walk(f, n.LimitCount)
+		}
 		if n.ReturningList != nil {
 			Walk(f, n.ReturningList)
 		}