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

Jille · web-flow · commit 77aadba161e6 · 2023-06-27T13:16:56.000-07:00
PostgreSQL doesn't support UPDATE-LIMIT issue #2131
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)
 		}

Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@ type DeleteStmt struct {`
`4`	`4`	`Relations *List`
`5`	`5`	`UsingClause *List`
`6`	`6`	`WhereClause Node`
	`7`	`+ LimitCount Node`
`7`	`8`	`ReturningList *List`
`8`	`9`	`WithClause *WithClause`
`9`	`10`	`}`
Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@ type UpdateStmt struct {`
`5`	`5`	`TargetList *List`
`6`	`6`	`WhereClause Node`
`7`	`7`	`FromClause *List`
	`8`	`+ LimitCount Node`
`8`	`9`	`ReturningList *List`
`9`	`10`	`WithClause *WithClause`
`10`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1068,6 +1068,9 @@ func Walk(f Visitor, node ast.Node) {`
`1068`	`1068`	`if n.WhereClause != nil {`
`1069`	`1069`	`Walk(f, n.WhereClause)`
`1070`	`1070`	`}`
	`1071`	`+ if n.LimitCount != nil {`
	`1072`	`+ Walk(f, n.LimitCount)`
	`1073`	`+ }`
`1071`	`1074`	`if n.ReturningList != nil {`
`1072`	`1075`	`Walk(f, n.ReturningList)`
`1073`	`1076`	`}`
`@@ -2038,6 +2041,9 @@ func Walk(f Visitor, node ast.Node) {`
`2038`	`2041`	`if n.FromClause != nil {`
`2039`	`2042`	`Walk(f, n.FromClause)`
`2040`	`2043`	`}`
	`2044`	`+ if n.LimitCount != nil {`
	`2045`	`+ Walk(f, n.LimitCount)`
	`2046`	`+ }`
`2041`	`2047`	`if n.ReturningList != nil {`
`2042`	`2048`	`Walk(f, n.ReturningList)`
`2043`	`2049`	`}`