planner: support union statements with ctes (#15312)

Signed-off-by: Max Englander <max@planetscale.com>
vitessio · Feb 21, 2024 · a9b2c18 · a9b2c18
1 parent c982a04
commit a9b2c18
Show file tree

Hide file tree

Showing 4 changed files with 114 additions and 0 deletions.
diff --git 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
diff --git 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
@@ -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"
+ ]
+ }
  }
 ]