diff --git a/go/vt/sqlparser/ast_format.go b/go/vt/sqlparser/ast_format.go index 0e7ca0a8231..e90b8b6344b 100644 --- a/go/vt/sqlparser/ast_format.go +++ b/go/vt/sqlparser/ast_format.go @@ -74,6 +74,10 @@ func (node *CommentOnly) Format(buf *TrackedBuffer) { // Format formats the node. func (node *Union) Format(buf *TrackedBuffer) { + if node.With != nil { + buf.astPrintf(node, "%v", node.With) + } + if requiresParen(node.Left) { buf.astPrintf(node, "(%v)", node.Left) } else { diff --git a/go/vt/sqlparser/ast_format_fast.go b/go/vt/sqlparser/ast_format_fast.go index a7f1a3c1e93..4ebdbcf8475 100644 --- a/go/vt/sqlparser/ast_format_fast.go +++ b/go/vt/sqlparser/ast_format_fast.go @@ -86,6 +86,10 @@ func (node *CommentOnly) FormatFast(buf *TrackedBuffer) { // FormatFast formats the node. func (node *Union) FormatFast(buf *TrackedBuffer) { + if node.With != nil { + node.With.FormatFast(buf) + } + if requiresParen(node.Left) { buf.WriteByte('(') node.Left.FormatFast(buf) diff --git a/go/vt/sqlparser/parse_test.go b/go/vt/sqlparser/parse_test.go index f90ae16606e..0dce5eae5bc 100644 --- a/go/vt/sqlparser/parse_test.go +++ b/go/vt/sqlparser/parse_test.go @@ -505,6 +505,9 @@ var ( }, { input: "WITH topsales2003 AS (SELECT salesRepEmployeeNumber employeeNumber, SUM(quantityOrdered * priceEach) sales FROM orders INNER JOIN orderdetails USING (orderNumber) INNER JOIN customers USING (customerNumber) WHERE YEAR(shippedDate) = 2003 AND status = 'Shipped' GROUP BY salesRepEmployeeNumber ORDER BY sales DESC LIMIT 5)SELECT employeeNumber, firstName, lastName, sales FROM employees JOIN topsales2003 USING (employeeNumber)", output: "with topsales2003 as (select salesRepEmployeeNumber as employeeNumber, sum(quantityOrdered * priceEach) as sales from orders join orderdetails using (orderNumber) join customers using (customerNumber) where YEAR(shippedDate) = 2003 and `status` = 'Shipped' group by salesRepEmployeeNumber order by sales desc limit 5) select employeeNumber, firstName, lastName, sales from employees join topsales2003 using (employeeNumber)", + }, { + input: "WITH count_a AS (SELECT COUNT(`id`) AS `num` FROM `tbl_a`), count_b AS (SELECT COUNT(`id`) AS `num` FROM tbl_b) SELECT 'a', `num` FROM `count_a` UNION SELECT 'b', `num` FROM `count_b`", + output: "with count_a as (select count(id) as num from tbl_a) , count_b as (select count(id) as num from tbl_b) select 'a', num from count_a union select 'b', num from count_b", }, { input: "select 1 from t", }, { diff --git a/go/vt/vtgate/planbuilder/testdata/cte_cases.json b/go/vt/vtgate/planbuilder/testdata/cte_cases.json index e43b6320340..a7027f80348 100644 --- a/go/vt/vtgate/planbuilder/testdata/cte_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/cte_cases.json @@ -1863,5 +1863,108 @@ "main.unsharded" ] } + }, + { + "comment": "WITH two common expressions against an unsharded datbase and a SELECT UNION against those expressions", + "query": "WITH `count_a` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_a`), `count_b` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_b`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`", + "plan": { + "QueryType": "SELECT", + "Original": "WITH `count_a` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_a`), `count_b` AS (SELECT COUNT(`id`) AS `num` FROM `unsharded_b`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`", + "Instructions": { + "OperatorType": "Route", + "Variant": "Unsharded", + "Keyspace": { + "Name": "main", + "Sharded": false + }, + "FieldQuery": "select 'count_a' as tab, num from count_a where 1 != 1 union select 'count_b' as tab, num from count_b where 1 != 1", + "Query": "with count_a as (select count(id) as num from unsharded_a) , count_b as (select count(id) as num from unsharded_b) select 'count_a' as tab, num from count_a union select 'count_b' as tab, num from count_b", + "Table": "unsharded_a, unsharded_b" + }, + "TablesUsed": [ + "main.unsharded_a", + "main.unsharded_b" + ] + } + }, + { + "comment": "WITH two common expressions against a sharded datbase and a SELECT UNION against those expressions", + "query": "WITH `count_a` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_metadata`), `count_b` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_extra`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`", + "plan": { + "QueryType": "SELECT", + "Original": "WITH `count_a` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_metadata`), `count_b` AS (SELECT COUNT(`user_id`) AS `num` FROM `user_extra`) SELECT 'count_a' AS `tab`, `num` FROM `count_a` UNION SELECT 'count_b' AS `tab`, `num` FROM `count_b`", + "Instructions": { + "OperatorType": "Distinct", + "Collations": [ + "0: utf8mb4_0900_ai_ci", + "1" + ], + "Inputs": [ + { + "OperatorType": "Concatenate", + "Inputs": [ + { + "OperatorType": "Projection", + "Expressions": [ + "'count_a' as tab", + ":0 as num" + ], + "Inputs": [ + { + "OperatorType": "Aggregate", + "Variant": "Scalar", + "Aggregates": "sum_count(0) AS num", + "Inputs": [ + { + "OperatorType": "Route", + "Variant": "Scatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select count(user_id) as num from user_metadata where 1 != 1", + "Query": "select count(user_id) as num from user_metadata", + "Table": "user_metadata" + } + ] + } + ] + }, + { + "OperatorType": "Projection", + "Expressions": [ + "'count_b' as tab", + ":0 as num" + ], + "Inputs": [ + { + "OperatorType": "Aggregate", + "Variant": "Scalar", + "Aggregates": "sum_count(0) AS num", + "Inputs": [ + { + "OperatorType": "Route", + "Variant": "Scatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select count(user_id) as num from user_extra where 1 != 1", + "Query": "select count(user_id) as num from user_extra", + "Table": "user_extra" + } + ] + } + ] + } + ] + } + ] + }, + "TablesUsed": [ + "user.user_extra", + "user.user_metadata" + ] + } } ]