diff --git a/internal/endtoend/testdata/select_union/mysql/go/query.sql.go b/internal/endtoend/testdata/select_union/mysql/go/query.sql.go index 96147809a8..f514165eef 100644 --- a/internal/endtoend/testdata/select_union/mysql/go/query.sql.go +++ b/internal/endtoend/testdata/select_union/mysql/go/query.sql.go @@ -95,3 +95,38 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { } return items, nil } + +const selectUnionWithLimit = `-- name: SelectUnionWithLimit :many +SELECT a, b FROM foo +UNION +SELECT a, b FROM foo +LIMIT ? OFFSET ? +` + +type SelectUnionWithLimitParams struct { + Limit int32 + Offset int32 +} + +func (q *Queries) SelectUnionWithLimit(ctx context.Context, arg SelectUnionWithLimitParams) ([]Foo, error) { + rows, err := q.db.QueryContext(ctx, selectUnionWithLimit, arg.Limit, arg.Offset) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Foo + for rows.Next() { + var i Foo + if err := rows.Scan(&i.A, &i.B); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/endtoend/testdata/select_union/mysql/query.sql b/internal/endtoend/testdata/select_union/mysql/query.sql index 80dd8cc4f2..b3a6371161 100644 --- a/internal/endtoend/testdata/select_union/mysql/query.sql +++ b/internal/endtoend/testdata/select_union/mysql/query.sql @@ -5,6 +5,12 @@ SELECT * FROM foo UNION SELECT * FROM foo; +-- name: SelectUnionWithLimit :many +SELECT * FROM foo +UNION +SELECT * FROM foo +LIMIT ? OFFSET ?; + -- name: SelectExcept :many SELECT * FROM foo EXCEPT diff --git a/internal/endtoend/testdata/select_union/postgres/pgx/v4/go/query.sql.go b/internal/endtoend/testdata/select_union/postgres/pgx/v4/go/query.sql.go index d2cfd2e0df..2b21c3e7fc 100644 --- a/internal/endtoend/testdata/select_union/postgres/pgx/v4/go/query.sql.go +++ b/internal/endtoend/testdata/select_union/postgres/pgx/v4/go/query.sql.go @@ -86,3 +86,35 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { } return items, nil } + +const selectUnionWithLimit = `-- name: SelectUnionWithLimit :many +SELECT a, b FROM foo +UNION +SELECT a, b FROM foo +LIMIT $1 OFFSET $2 +` + +type SelectUnionWithLimitParams struct { + Limit int32 + Offset int32 +} + +func (q *Queries) SelectUnionWithLimit(ctx context.Context, arg SelectUnionWithLimitParams) ([]Foo, error) { + rows, err := q.db.Query(ctx, selectUnionWithLimit, arg.Limit, arg.Offset) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Foo + for rows.Next() { + var i Foo + if err := rows.Scan(&i.A, &i.B); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/endtoend/testdata/select_union/postgres/pgx/v4/query.sql b/internal/endtoend/testdata/select_union/postgres/pgx/v4/query.sql index 80dd8cc4f2..9653e0707e 100644 --- a/internal/endtoend/testdata/select_union/postgres/pgx/v4/query.sql +++ b/internal/endtoend/testdata/select_union/postgres/pgx/v4/query.sql @@ -5,6 +5,12 @@ SELECT * FROM foo UNION SELECT * FROM foo; +-- name: SelectUnionWithLimit :many +SELECT * FROM foo +UNION +SELECT * FROM foo +LIMIT $1 OFFSET $2; + -- name: SelectExcept :many SELECT * FROM foo EXCEPT diff --git a/internal/endtoend/testdata/select_union/postgres/pgx/v5/go/query.sql.go b/internal/endtoend/testdata/select_union/postgres/pgx/v5/go/query.sql.go index d2cfd2e0df..2b21c3e7fc 100644 --- a/internal/endtoend/testdata/select_union/postgres/pgx/v5/go/query.sql.go +++ b/internal/endtoend/testdata/select_union/postgres/pgx/v5/go/query.sql.go @@ -86,3 +86,35 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { } return items, nil } + +const selectUnionWithLimit = `-- name: SelectUnionWithLimit :many +SELECT a, b FROM foo +UNION +SELECT a, b FROM foo +LIMIT $1 OFFSET $2 +` + +type SelectUnionWithLimitParams struct { + Limit int32 + Offset int32 +} + +func (q *Queries) SelectUnionWithLimit(ctx context.Context, arg SelectUnionWithLimitParams) ([]Foo, error) { + rows, err := q.db.Query(ctx, selectUnionWithLimit, arg.Limit, arg.Offset) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Foo + for rows.Next() { + var i Foo + if err := rows.Scan(&i.A, &i.B); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/endtoend/testdata/select_union/postgres/pgx/v5/query.sql b/internal/endtoend/testdata/select_union/postgres/pgx/v5/query.sql index 80dd8cc4f2..9653e0707e 100644 --- a/internal/endtoend/testdata/select_union/postgres/pgx/v5/query.sql +++ b/internal/endtoend/testdata/select_union/postgres/pgx/v5/query.sql @@ -5,6 +5,12 @@ SELECT * FROM foo UNION SELECT * FROM foo; +-- name: SelectUnionWithLimit :many +SELECT * FROM foo +UNION +SELECT * FROM foo +LIMIT $1 OFFSET $2; + -- name: SelectExcept :many SELECT * FROM foo EXCEPT diff --git a/internal/endtoend/testdata/select_union/postgres/stdlib/go/query.sql.go b/internal/endtoend/testdata/select_union/postgres/stdlib/go/query.sql.go index 96147809a8..505bf72bc9 100644 --- a/internal/endtoend/testdata/select_union/postgres/stdlib/go/query.sql.go +++ b/internal/endtoend/testdata/select_union/postgres/stdlib/go/query.sql.go @@ -95,3 +95,38 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { } return items, nil } + +const selectUnionWithLimit = `-- name: SelectUnionWithLimit :many +SELECT a, b FROM foo +UNION +SELECT a, b FROM foo +LIMIT $1 OFFSET $2 +` + +type SelectUnionWithLimitParams struct { + Limit int32 + Offset int32 +} + +func (q *Queries) SelectUnionWithLimit(ctx context.Context, arg SelectUnionWithLimitParams) ([]Foo, error) { + rows, err := q.db.QueryContext(ctx, selectUnionWithLimit, arg.Limit, arg.Offset) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Foo + for rows.Next() { + var i Foo + if err := rows.Scan(&i.A, &i.B); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql b/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql index 80dd8cc4f2..9653e0707e 100644 --- a/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql +++ b/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql @@ -5,6 +5,12 @@ SELECT * FROM foo UNION SELECT * FROM foo; +-- name: SelectUnionWithLimit :many +SELECT * FROM foo +UNION +SELECT * FROM foo +LIMIT $1 OFFSET $2; + -- name: SelectExcept :many SELECT * FROM foo EXCEPT diff --git a/internal/engine/dolphin/convert.go b/internal/engine/dolphin/convert.go index 6b03774ebc..703efa9e55 100644 --- a/internal/engine/dolphin/convert.go +++ b/internal/engine/dolphin/convert.go @@ -1254,7 +1254,12 @@ func (c *cc) convertSetOprSelectList(n *pcast.SetOprSelectList) ast.Node { func (c *cc) convertSetOprStmt(n *pcast.SetOprStmt) ast.Node { if n.SelectList != nil { - return c.convertSetOprSelectList(n.SelectList) + sn := c.convertSetOprSelectList(n.SelectList) + if ss, ok := sn.(*ast.SelectStmt); ok && n.Limit != nil { + ss.LimitOffset = c.convert(n.Limit.Offset) + ss.LimitCount = c.convert(n.Limit.Count) + } + return sn } return todo(n) }