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 011ef8eb07..ee09e09532 100644 --- a/internal/endtoend/testdata/select_union/mysql/go/query.sql.go +++ b/internal/endtoend/testdata/select_union/mysql/go/query.sql.go @@ -96,6 +96,35 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { return items, nil } +const selectUnionAliased = `-- name: SelectUnionAliased :many +(SELECT a, b FROM foo) +UNION +SELECT a, b FROM bar +` + +func (q *Queries) SelectUnionAliased(ctx context.Context) ([]Foo, error) { + rows, err := q.db.QueryContext(ctx, selectUnionAliased) + 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 +} + const selectUnionOther = `-- name: SelectUnionOther :many SELECT a, b FROM foo UNION diff --git a/internal/endtoend/testdata/select_union/mysql/query.sql b/internal/endtoend/testdata/select_union/mysql/query.sql index d607aa7222..c706921bc3 100644 --- a/internal/endtoend/testdata/select_union/mysql/query.sql +++ b/internal/endtoend/testdata/select_union/mysql/query.sql @@ -22,4 +22,9 @@ SELECT * FROM foo; -- name: SelectUnionOther :many SELECT * FROM foo UNION -SELECT * FROM bar; \ No newline at end of file +SELECT * FROM bar; + +-- name: SelectUnionAliased :many +(SELECT * FROM foo) +UNION +SELECT * FROM bar; 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 9380106f0c..0dc42e4d5b 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 @@ -87,6 +87,32 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { return items, nil } +const selectUnionAliased = `-- name: SelectUnionAliased :many +(SELECT a, b FROM foo) +UNION +SELECT a, b FROM foo +` + +func (q *Queries) SelectUnionAliased(ctx context.Context) ([]Foo, error) { + rows, err := q.db.Query(ctx, selectUnionAliased) + 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 +} + const selectUnionOther = `-- name: SelectUnionOther :many SELECT a, b FROM foo UNION 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 107aee02d2..f44f1fd2af 100644 --- a/internal/endtoend/testdata/select_union/postgres/pgx/v4/query.sql +++ b/internal/endtoend/testdata/select_union/postgres/pgx/v4/query.sql @@ -22,4 +22,9 @@ SELECT * FROM foo; -- name: SelectUnionOther :many SELECT * FROM foo UNION -SELECT * FROM bar; \ No newline at end of file +SELECT * FROM bar; + +-- name: SelectUnionAliased :many +(SELECT * FROM foo) +UNION +SELECT * FROM foo; 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 9380106f0c..0dc42e4d5b 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 @@ -87,6 +87,32 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { return items, nil } +const selectUnionAliased = `-- name: SelectUnionAliased :many +(SELECT a, b FROM foo) +UNION +SELECT a, b FROM foo +` + +func (q *Queries) SelectUnionAliased(ctx context.Context) ([]Foo, error) { + rows, err := q.db.Query(ctx, selectUnionAliased) + 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 +} + const selectUnionOther = `-- name: SelectUnionOther :many SELECT a, b FROM foo UNION 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 107aee02d2..f44f1fd2af 100644 --- a/internal/endtoend/testdata/select_union/postgres/pgx/v5/query.sql +++ b/internal/endtoend/testdata/select_union/postgres/pgx/v5/query.sql @@ -22,4 +22,9 @@ SELECT * FROM foo; -- name: SelectUnionOther :many SELECT * FROM foo UNION -SELECT * FROM bar; \ No newline at end of file +SELECT * FROM bar; + +-- name: SelectUnionAliased :many +(SELECT * FROM foo) +UNION +SELECT * FROM foo; 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 ebc0fbb363..7579360f97 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 @@ -96,6 +96,35 @@ func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { return items, nil } +const selectUnionAliased = `-- name: SelectUnionAliased :many +(SELECT a, b FROM foo) +UNION +SELECT a, b FROM foo +` + +func (q *Queries) SelectUnionAliased(ctx context.Context) ([]Foo, error) { + rows, err := q.db.QueryContext(ctx, selectUnionAliased) + 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 +} + const selectUnionOther = `-- name: SelectUnionOther :many SELECT a, b FROM foo UNION diff --git a/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql b/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql index 107aee02d2..f44f1fd2af 100644 --- a/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql +++ b/internal/endtoend/testdata/select_union/postgres/stdlib/query.sql @@ -22,4 +22,9 @@ SELECT * FROM foo; -- name: SelectUnionOther :many SELECT * FROM foo UNION -SELECT * FROM bar; \ No newline at end of file +SELECT * FROM bar; + +-- name: SelectUnionAliased :many +(SELECT * FROM foo) +UNION +SELECT * FROM foo; diff --git a/internal/endtoend/testdata/select_union/sqlite/query.sql b/internal/endtoend/testdata/select_union/sqlite/query.sql index d607aa7222..67d28d1824 100644 --- a/internal/endtoend/testdata/select_union/sqlite/query.sql +++ b/internal/endtoend/testdata/select_union/sqlite/query.sql @@ -22,4 +22,4 @@ SELECT * FROM foo; -- name: SelectUnionOther :many SELECT * FROM foo UNION -SELECT * FROM bar; \ No newline at end of file +SELECT * FROM bar; diff --git a/internal/engine/dolphin/convert.go b/internal/engine/dolphin/convert.go index 6d95a75064..31d4a70a77 100644 --- a/internal/engine/dolphin/convert.go +++ b/internal/engine/dolphin/convert.go @@ -1182,7 +1182,12 @@ func (c *cc) convertSetOprType(n *pcast.SetOprType) (op ast.SetOperation, all bo func (c *cc) convertSetOprSelectList(n *pcast.SetOprSelectList) ast.Node { selectStmts := make([]*ast.SelectStmt, len(n.Selects)) for i, node := range n.Selects { - selectStmts[i] = c.convertSelectStmt(node.(*pcast.SelectStmt)) + switch node := node.(type) { + case *pcast.SelectStmt: + selectStmts[i] = c.convertSelectStmt(node) + case *pcast.SetOprSelectList: + selectStmts[i] = c.convertSetOprSelectList(node).(*ast.SelectStmt) + } } op, all := c.convertSetOprType(n.AfterSetOperator)