From a268445fcfe0d64de89ff2edd11c2582cb4f3ab0 Mon Sep 17 00:00:00 2001 From: aitva Date: Sun, 14 Feb 2021 19:48:17 +0100 Subject: [PATCH] dolphin: add support for union query --- .../select_union/{ => mysql}/go/db.go | 0 .../select_union/{ => mysql}/go/models.go | 0 .../select_union/{ => mysql}/go/query.sql.go | 0 .../select_union/{ => mysql}/query.sql | 0 .../testdata/select_union/mysql/sqlc.json | 12 +++ .../testdata/select_union/postgres/go/db.go | 29 ++++++ .../select_union/postgres/go/models.go | 12 +++ .../select_union/postgres/go/query.sql.go | 37 ++++++++ .../testdata/select_union/postgres/query.sql | 6 ++ .../select_union/{ => postgres}/sqlc.json | 0 internal/engine/dolphin/convert.go | 91 ++++++++++++++++++- 11 files changed, 186 insertions(+), 1 deletion(-) rename internal/endtoend/testdata/select_union/{ => mysql}/go/db.go (100%) rename internal/endtoend/testdata/select_union/{ => mysql}/go/models.go (100%) rename internal/endtoend/testdata/select_union/{ => mysql}/go/query.sql.go (100%) rename internal/endtoend/testdata/select_union/{ => mysql}/query.sql (100%) create mode 100644 internal/endtoend/testdata/select_union/mysql/sqlc.json create mode 100644 internal/endtoend/testdata/select_union/postgres/go/db.go create mode 100644 internal/endtoend/testdata/select_union/postgres/go/models.go create mode 100644 internal/endtoend/testdata/select_union/postgres/go/query.sql.go create mode 100644 internal/endtoend/testdata/select_union/postgres/query.sql rename internal/endtoend/testdata/select_union/{ => postgres}/sqlc.json (100%) diff --git a/internal/endtoend/testdata/select_union/go/db.go b/internal/endtoend/testdata/select_union/mysql/go/db.go similarity index 100% rename from internal/endtoend/testdata/select_union/go/db.go rename to internal/endtoend/testdata/select_union/mysql/go/db.go diff --git a/internal/endtoend/testdata/select_union/go/models.go b/internal/endtoend/testdata/select_union/mysql/go/models.go similarity index 100% rename from internal/endtoend/testdata/select_union/go/models.go rename to internal/endtoend/testdata/select_union/mysql/go/models.go diff --git a/internal/endtoend/testdata/select_union/go/query.sql.go b/internal/endtoend/testdata/select_union/mysql/go/query.sql.go similarity index 100% rename from internal/endtoend/testdata/select_union/go/query.sql.go rename to internal/endtoend/testdata/select_union/mysql/go/query.sql.go diff --git a/internal/endtoend/testdata/select_union/query.sql b/internal/endtoend/testdata/select_union/mysql/query.sql similarity index 100% rename from internal/endtoend/testdata/select_union/query.sql rename to internal/endtoend/testdata/select_union/mysql/query.sql diff --git a/internal/endtoend/testdata/select_union/mysql/sqlc.json b/internal/endtoend/testdata/select_union/mysql/sqlc.json new file mode 100644 index 0000000000..0657f4db83 --- /dev/null +++ b/internal/endtoend/testdata/select_union/mysql/sqlc.json @@ -0,0 +1,12 @@ +{ + "version": "1", + "packages": [ + { + "engine": "mysql", + "path": "go", + "name": "querytest", + "schema": "query.sql", + "queries": "query.sql" + } + ] +} diff --git a/internal/endtoend/testdata/select_union/postgres/go/db.go b/internal/endtoend/testdata/select_union/postgres/go/db.go new file mode 100644 index 0000000000..6a99519302 --- /dev/null +++ b/internal/endtoend/testdata/select_union/postgres/go/db.go @@ -0,0 +1,29 @@ +// Code generated by sqlc. DO NOT EDIT. + +package querytest + +import ( + "context" + "database/sql" +) + +type DBTX interface { + ExecContext(context.Context, string, ...interface{}) (sql.Result, error) + PrepareContext(context.Context, string) (*sql.Stmt, error) + QueryContext(context.Context, string, ...interface{}) (*sql.Rows, error) + QueryRowContext(context.Context, string, ...interface{}) *sql.Row +} + +func New(db DBTX) *Queries { + return &Queries{db: db} +} + +type Queries struct { + db DBTX +} + +func (q *Queries) WithTx(tx *sql.Tx) *Queries { + return &Queries{ + db: tx, + } +} diff --git a/internal/endtoend/testdata/select_union/postgres/go/models.go b/internal/endtoend/testdata/select_union/postgres/go/models.go new file mode 100644 index 0000000000..ce7117644a --- /dev/null +++ b/internal/endtoend/testdata/select_union/postgres/go/models.go @@ -0,0 +1,12 @@ +// Code generated by sqlc. DO NOT EDIT. + +package querytest + +import ( + "database/sql" +) + +type Foo struct { + A sql.NullString + B sql.NullString +} diff --git a/internal/endtoend/testdata/select_union/postgres/go/query.sql.go b/internal/endtoend/testdata/select_union/postgres/go/query.sql.go new file mode 100644 index 0000000000..ead8e54b3b --- /dev/null +++ b/internal/endtoend/testdata/select_union/postgres/go/query.sql.go @@ -0,0 +1,37 @@ +// Code generated by sqlc. DO NOT EDIT. +// source: query.sql + +package querytest + +import ( + "context" +) + +const selectUnion = `-- name: SelectUnion :many +SELECT a, b FROM foo +UNION +SELECT a, b FROM foo +` + +func (q *Queries) SelectUnion(ctx context.Context) ([]Foo, error) { + rows, err := q.db.QueryContext(ctx, selectUnion) + 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/query.sql b/internal/endtoend/testdata/select_union/postgres/query.sql new file mode 100644 index 0000000000..f2ba3ec51b --- /dev/null +++ b/internal/endtoend/testdata/select_union/postgres/query.sql @@ -0,0 +1,6 @@ +CREATE TABLE foo (a text, b text); + +-- name: SelectUnion :many +SELECT * FROM foo +UNION +SELECT * FROM foo; diff --git a/internal/endtoend/testdata/select_union/sqlc.json b/internal/endtoend/testdata/select_union/postgres/sqlc.json similarity index 100% rename from internal/endtoend/testdata/select_union/sqlc.json rename to internal/endtoend/testdata/select_union/postgres/sqlc.json diff --git a/internal/engine/dolphin/convert.go b/internal/engine/dolphin/convert.go index 782e92e7fd..2f4db8342c 100644 --- a/internal/engine/dolphin/convert.go +++ b/internal/engine/dolphin/convert.go @@ -442,10 +442,13 @@ func (c *cc) convertSelectField(n *pcast.SelectField) *ast.ResTarget { } func (c *cc) convertSelectStmt(n *pcast.SelectStmt) *ast.SelectStmt { + op, all := c.convertSetOprType(n.AfterSetOperator) stmt := &ast.SelectStmt{ TargetList: c.convertFieldList(n.Fields), FromClause: c.convertTableRefsClause(n.From), WhereClause: c.convert(n.Where), + Op: op, + All: all, } if n.Limit != nil { stmt.LimitCount = c.convert(n.Limit.Count) @@ -937,11 +940,97 @@ func (c *cc) convertSetDefaultRoleStmt(n *pcast.SetDefaultRoleStmt) ast.Node { return todo(n) } +func (c *cc) convertSetOprType(n *pcast.SetOprType) (op ast.SetOperation, all bool) { + if n == nil { + return + } + + switch *n { + case pcast.Union: + op = ast.Union + case pcast.UnionAll: + op = ast.Union + all = true + case pcast.Intersect: + op = ast.Intersect + case pcast.IntersectAll: + op = ast.Intersect + all = true + case pcast.Except: + op = ast.Except + case pcast.ExceptAll: + op = ast.Except + all = true + } + return +} + +// convertSetOprSelectList converts a list of SELECT from the Pingcap parser +// into a tree. It is called for UNION, INTERSECT or EXCLUDE operation. +// +// Given an union with the following nodes: +// [Select{1}, Select{2}, Select{3}, Select{4}] +// +// The function will return: +// Select{ +// Larg: Select{ +// Larg: Select{ +// Larg: Select{1}, +// Rarg: Select{2}, +// Op: Union +// }, +// Rarg: Select{3}, +// Op: Union, +// }, +// Rarg: Select{4}, +// Op: Union, +// } func (c *cc) convertSetOprSelectList(n *pcast.SetOprSelectList) ast.Node { - return todo(n) + selectStmts := make([]*ast.SelectStmt, len(n.Selects)) + for i, node := range n.Selects { + selectStmts[i] = c.convertSelectStmt(node.(*pcast.SelectStmt)) + } + + op, all := c.convertSetOprType(n.AfterSetOperator) + tree := &ast.SelectStmt{ + TargetList: &ast.List{}, + FromClause: &ast.List{}, + WhereClause: nil, + Op: op, + All: all, + } + for _, stmt := range selectStmts { + // We move Op and All from the child to the parent. + op, all := stmt.Op, stmt.All + stmt.Op, stmt.All = ast.None, false + + switch { + case tree.Larg == nil: + tree.Larg = stmt + case tree.Rarg == nil: + tree.Rarg = stmt + tree.Op = op + tree.All = all + default: + tree = &ast.SelectStmt{ + TargetList: &ast.List{}, + FromClause: &ast.List{}, + WhereClause: nil, + Larg: tree, + Rarg: stmt, + Op: op, + All: all, + } + } + } + + return tree } func (c *cc) convertSetOprStmt(n *pcast.SetOprStmt) ast.Node { + if n.SelectList != nil { + return c.convertSetOprSelectList(n.SelectList) + } return todo(n) }