Repository: vitessio/vitess
Pull request #8878 (2021-09-24 13:01:27)

Gen4: SQL_CALC_FOUND_ROWS support
## Description
Adds support in the gen4 planner for SQL_CALC_FOUND_ROWS.
It does this by sharing most of the logic from the V3 planner.

## Related Issue(s)
#7280 

## Checklist
- [x] Tests were added or are not required
- [x] Documentation was added or is not required



Patch:
diff --git a/go/vt/sqlparser/rewriter_api.go b/go/vt/sqlparser/rewriter_api.go
index 6f4069fbd6..71c0660deb 100644
--- a/go/vt/sqlparser/rewriter_api.go
+++ b/go/vt/sqlparser/rewriter_api.go
@@ -52,7 +52,7 @@ func Rewrite(node SQLNode, pre, post ApplyFunc) (result SQLNode) {
 	return parent.SQLNode
 }
 
-// RootNode is the root node of the AST. It is the first element of the tree.
+// RootNode is the root node of the AST when rewriting. It is the first element of the tree.
 type RootNode struct {
 	SQLNode
 }
diff --git a/go/vt/vtgate/planbuilder/gen4_planner.go b/go/vt/vtgate/planbuilder/gen4_planner.go
index 28f61d6605..0cbd64fb57 100644
--- a/go/vt/vtgate/planbuilder/gen4_planner.go
+++ b/go/vt/vtgate/planbuilder/gen4_planner.go
@@ -1,5 +1,5 @@
 /*
-Copyright 2020 The Vitess Authors.
+Copyright 2021 The Vitess Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -27,7 +27,7 @@ import (
 
 var _ selectPlanner = gen4Planner
 
-func gen4Planner(_ string) func(sqlparser.Statement, *sqlparser.ReservedVars, ContextVSchema) (engine.Primitive, error) {
+func gen4Planner(query string) func(sqlparser.Statement, *sqlparser.ReservedVars, ContextVSchema) (engine.Primitive, error) {
 	return func(stmt sqlparser.Statement, reservedVars *sqlparser.ReservedVars, vschema ContextVSchema) (engine.Primitive, error) {
 		selStatement, ok := stmt.(sqlparser.SelectStatement)
 		if !ok {
@@ -41,6 +41,12 @@ func gen4Planner(_ string) func(sqlparser.Statement, *sqlparser.ReservedVars, Co
 			if err != nil || p != nil {
 				return p, err
 			}
+
+			if sel.SQLCalcFoundRows && sel.Limit != nil {
+				return gen4planSQLCalcFoundRows(vschema, sel, query, reservedVars)
+			}
+			// if there was no limit, we can safely ignore the SQLCalcFoundRows directive
+			sel.SQLCalcFoundRows = false
 		}
 
 		getPlan := func(selStatement sqlparser.SelectStatement) (logicalPlan, error) {
@@ -63,6 +69,34 @@ func gen4Planner(_ string) func(sqlparser.Statement, *sqlparser.ReservedVars, Co
 	}
 }
 
+func gen4planSQLCalcFoundRows(vschema ContextVSchema, sel *sqlparser.Select, query string, reservedVars *sqlparser.ReservedVars) (engine.Primitive, error) {
+	ksName := ""
+	if ks, _ := vschema.DefaultKeyspace(); ks != nil {
+		ksName = ks.Name
+	}
+	semTable, err := semantics.Analyze(sel, ksName, vschema)
+	if err != nil {
+		return nil, err
+	}
+	plan, err := buildSQLCalcFoundRowsPlan(query, sel, reservedVars, vschema, planSelectGen4)
+	if err != nil {
+		return nil, err
+	}
+	err = plan.WireupGen4(semTable)
+	if err != nil {
+		return nil, err
+	}
+	return plan.Primitive(), nil
+}
+
+func planSelectGen4(reservedVars *sqlparser.ReservedVars, vschema ContextVSchema, sel *sqlparser.Select) (*jointab, logicalPlan, error) {
+	plan, err := newBuildSelectPlan(sel, reservedVars, vschema)
+	if err != nil {
+		return nil, nil, err
+	}
+	return nil, plan, nil
+}
+
 func gen4CNFRewrite(stmt sqlparser.Statement, getPlan func(selStatement sqlparser.SelectStatement) (logicalPlan, error)) engine.Primitive {
 	rewritten, isSel := sqlparser.RewriteToCNF(stmt).(sqlparser.SelectStatement)
 	if !isSel {
@@ -175,13 +209,6 @@ func planLimit(limit *sqlparser.Limit, plan logicalPlan) (logicalPlan, error) {
 	return lPlan, nil
 }
 
-func checkUnsupportedConstructs(sel *sqlparser.Select) error {
-	if sel.SQLCalcFoundRows {
-		return semantics.Gen4NotSupportedF("sql_calc_found_rows")
-	}
-	return nil
-}
-
 func planHorizon(ctx *planningContext, plan logicalPlan, in sqlparser.SelectStatement) (logicalPlan, error) {
 	switch node := in.(type) {
 	case *sqlparser.Select:
diff --git a/go/vt/vtgate/planbuilder/horizon_planning.go b/go/vt/vtgate/planbuilder/horizon_planning.go
index 4b4e2878b4..16dd77fc3f 100644
--- a/go/vt/vtgate/planbuilder/horizon_planning.go
+++ b/go/vt/vtgate/planbuilder/horizon_planning.go
@@ -39,10 +39,6 @@ type horizonPlanning struct {
 }
 
 func (hp *horizonPlanning) planHorizon(ctx *planningContext, plan logicalPlan) (logicalPlan, error) {
-	if err := checkUnsupportedConstructs(hp.sel); err != nil {
-		return nil, err
-	}
-
 	rb, isRoute := plan.(*route)
 	if !isRoute && ctx.semTable.ProjectionErr != nil {
 		return nil, ctx.semTable.ProjectionErr
diff --git a/go/vt/vtgate/planbuilder/select.go b/go/vt/vtgate/planbuilder/select.go
index 2d09980948..ec2996712a 100644
--- a/go/vt/vtgate/planbuilder/select.go
+++ b/go/vt/vtgate/planbuilder/select.go
@@ -149,7 +149,7 @@ func (pb *primitiveBuilder) processSelect(sel *sqlparser.Select, reservedVars *s
 		}
 		sel.SQLCalcFoundRows = false
 		if sel.Limit != nil {
-			plan, err := buildSQLCalcFoundRowsPlan(query, sel, reservedVars, outer, pb.vschema)
+			plan, err := buildSQLCalcFoundRowsPlan(query, sel, reservedVars, pb.vschema, planSelectV3)
 			if err != nil {
 				return err
 			}
@@ -237,15 +237,19 @@ func setMiscFunc(in logicalPlan, sel *sqlparser.Select) error {
 	return nil
 }
 
-func buildSQLCalcFoundRowsPlan(query string, sel *sqlparser.Select, reservedVars *sqlparser.ReservedVars, outer *symtab, vschema ContextVSchema) (logicalPlan, error) {
-	ljt := newJointab(reservedVars)
-	frpb := newPrimitiveBuilder(vschema, ljt)
-	err := frpb.processSelect(sel, reservedVars, outer, "")
+func buildSQLCalcFoundRowsPlan(
+	originalQuery string,
+	sel *sqlparser.Select,
+	reservedVars *sqlparser.ReservedVars,
+	vschema ContextVSchema,
+	planSelect func(reservedVars *sqlparser.ReservedVars, vschema ContextVSchema, sel *sqlparser.Select) (*jointab, logicalPlan, error),
+) (logicalPlan, error) {
+	ljt, limitPlan, err := planSelect(reservedVars, vschema, sel)
 	if err != nil {
 		return nil, err
 	}
 
-	statement2, reserved2, err := sqlparser.Parse2(query)
+	statement2, reserved2, err := sqlparser.Parse2(originalQuery)
 	if err != nil {
 		return nil, err
 	}
@@ -282,13 +286,19 @@ func buildSQLCalcFoundRowsPlan(query string, sel *sqlparser.Select, reservedVars
 	}
 
 	reservedVars2 := sqlparser.NewReservedVars("vtg", reserved2)
-	cjt := newJointab(reservedVars2)
-	countpb := newPrimitiveBuilder(vschema, cjt)
-	err = countpb.processSelect(sel2, reservedVars2, outer, "")
+
+	cjt, countPlan, err := planSelect(reservedVars2, vschema, sel2)
 	if err != nil {
 		return nil, err
 	}
-	return &sqlCalcFoundRows{LimitQuery: frpb.plan, CountQuery: countpb.plan, ljt: ljt, cjt: cjt}, nil
+	return &sqlCalcFoundRows{LimitQuery: limitPlan, CountQuery: countPlan, ljt: ljt, cjt: cjt}, nil
+}
+
+func planSelectV3(reservedVars *sqlparser.ReservedVars, vschema ContextVSchema, sel *sqlparser.Select) (*jointab, logicalPlan, error) {
+	ljt := newJointab(reservedVars)
+	frpb := newPrimitiveBuilder(vschema, ljt)
+	err := frpb.processSelect(sel, reservedVars, nil, "")
+	return ljt, frpb.plan, err
 }
 
 func handleDualSelects(sel *sqlparser.Select, vschema ContextVSchema) (engine.Primitive, error) {
diff --git a/go/vt/vtgate/planbuilder/sql_calc_found_rows.go b/go/vt/vtgate/planbuilder/sql_calc_found_rows.go
index ac55baada8..4a4c589541 100644
--- a/go/vt/vtgate/planbuilder/sql_calc_found_rows.go
+++ b/go/vt/vtgate/planbuilder/sql_calc_found_rows.go
@@ -28,10 +28,12 @@ var _ logicalPlan = (*sqlCalcFoundRows)(nil)
 
 type sqlCalcFoundRows struct {
 	LimitQuery, CountQuery logicalPlan
-	ljt, cjt               *jointab
+
+	// only used by WireUp for V3
+	ljt, cjt *jointab
 }
 
-//Wireup implements the logicalPlan interface
+// Wireup implements the logicalPlan interface
 func (s *sqlCalcFoundRows) Wireup(logicalPlan, *jointab) error {
 	err := s.LimitQuery.Wireup(s.LimitQuery, s.ljt)
 	if err != nil {
@@ -40,7 +42,7 @@ func (s *sqlCalcFoundRows) Wireup(logicalPlan, *jointab) error {
 	return s.CountQuery.Wireup(s.CountQuery, s.cjt)
 }
 
-//Wireup2 implements the logicalPlan interface
+// WireupGen4 implements the logicalPlan interface
 func (s *sqlCalcFoundRows) WireupGen4(semTable *semantics.SemTable) error {
 	err := s.LimitQuery.WireupGen4(semTable)
 	if err != nil {
@@ -49,7 +51,7 @@ func (s *sqlCalcFoundRows) WireupGen4(semTable *semantics.SemTable) error {
 	return s.CountQuery.WireupGen4(semTable)
 }
 
-// Solves implements the logicalPlan interface
+// ContainsTables implements the logicalPlan interface
 func (s *sqlCalcFoundRows) ContainsTables() semantics.TableSet {
 	return s.LimitQuery.ContainsTables()
 }
@@ -64,32 +66,32 @@ func (s *sqlCalcFoundRows) Primitive() engine.Primitive {
 
 // All the methods below are not implemented. They should not be called on a sqlCalcFoundRows plan
 
-//Order implements the logicalPlan interface
+// Order implements the logicalPlan interface
 func (s *sqlCalcFoundRows) Order() int {
 	return s.LimitQuery.Order()
 }
 
-//ResultColumns implements the logicalPlan interface
+// ResultColumns implements the logicalPlan interface
 func (s *sqlCalcFoundRows) ResultColumns() []*resultColumn {
 	return s.LimitQuery.ResultColumns()
 }
 
-//Reorder implements the logicalPlan interface
+// Reorder implements the logicalPlan interface
 func (s *sqlCalcFoundRows) Reorder(order int) {
 	s.LimitQuery.Reorder(order)
 }
 
-//SupplyVar implements the logicalPlan interface
+// SupplyVar implements the logicalPlan interface
 func (s *sqlCalcFoundRows) SupplyVar(from, to int, col *sqlparser.ColName, varname string) {
 	s.LimitQuery.SupplyVar(from, to, col, varname)
 }
 
-//SupplyCol implements the logicalPlan interface
+// SupplyCol implements the logicalPlan interface
 func (s *sqlCalcFoundRows) SupplyCol(col *sqlparser.ColName) (*resultColumn, int) {
 	return s.LimitQuery.SupplyCol(col)
 }
 
-//SupplyWeightString implements the logicalPlan interface
+// SupplyWeightString implements the logicalPlan interface
 func (s *sqlCalcFoundRows) SupplyWeightString(int, bool) (weightcolNumber int, err error) {
 	return 0, UnsupportedSupplyWeightString{Type: "sqlCalcFoundRows"}
 }
diff --git a/go/vt/vtgate/semantics/analyzer.go b/go/vt/vtgate/semantics/analyzer.go
index c035debb5b..289c4eb1db 100644
--- a/go/vt/vtgate/semantics/analyzer.go
+++ b/go/vt/vtgate/semantics/analyzer.go
@@ -268,8 +268,12 @@ func (a *analyzer) analyze(statement sqlparser.Statement) error {
 func (a *analyzer) checkForInvalidConstructs(cursor *sqlparser.Cursor) error {
 	switch node := cursor.Node().(type) {
 	case *sqlparser.Select:
-		if a.scoper.currentScope() == nil {
-			return nil
+		parent := cursor.Parent()
+		if _, isUnion := parent.(*sqlparser.Union); isUnion && node.SQLCalcFoundRows {
+			return vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "SQL_CALC_FOUND_ROWS not supported with union")
+		}
+		if _, isRoot := parent.(*sqlparser.RootNode); !isRoot && node.SQLCalcFoundRows {
+			return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "Incorrect usage/placement of 'SQL_CALC_FOUND_ROWS'")
 		}
 		errMsg := "INTO"
 		nextVal := false


Test patch:
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.txt b/go/vt/vtgate/planbuilder/testdata/select_cases.txt
index 49b8471fdd..c6b8184701 100644
--- a/go/vt/vtgate/planbuilder/testdata/select_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/select_cases.txt
@@ -1608,6 +1608,7 @@ Gen4 plan same as above
     "Vindex": "user_index"
   }
 }
+Gen4 plan same as above
 
 # sql_calc_found_rows with limit
 "select sql_calc_found_rows * from music limit 100"
@@ -1655,6 +1656,50 @@ Gen4 plan same as above
     ]
   }
 }
+{
+  "QueryType": "SELECT",
+  "Original": "select sql_calc_found_rows * from music limit 100",
+  "Instructions": {
+    "OperatorType": "SQL_CALC_FOUND_ROWS",
+    "Inputs": [
+      {
+        "OperatorType": "Limit",
+        "Count": 100,
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select * from music where 1 != 1",
+            "Query": "select * from music limit :__upper_limit",
+            "Table": "music"
+          }
+        ]
+      },
+      {
+        "OperatorType": "Aggregate",
+        "Variant": "Ordered",
+        "Aggregates": "count(0) AS count(*)",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select count(*) from music where 1 != 1",
+            "Query": "select count(*) from music",
+            "Table": "music"
+          }
+        ]
+      }
+    ]
+  }
+}
 
 # sql_calc_found_rows with SelectEqualUnique plans
 "select sql_calc_found_rows * from music where user_id = 1 limit 2"
@@ -1697,6 +1742,7 @@ Gen4 plan same as above
     ]
   }
 }
+Gen4 plan same as above
 
 # sql_calc_found_rows with group by and having
 "select sql_calc_found_rows user_id, count(id) from music group by user_id having count(user_id) = 1 order by user_id limit 2"
@@ -1746,14 +1792,62 @@ Gen4 plan same as above
     ]
   }
 }
+{
+  "QueryType": "SELECT",
+  "Original": "select sql_calc_found_rows user_id, count(id) from music group by user_id having count(user_id) = 1 order by user_id limit 2",
+  "Instructions": {
+    "OperatorType": "SQL_CALC_FOUND_ROWS",
+    "Inputs": [
+      {
+        "OperatorType": "Limit",
+        "Count": 2,
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select user_id, count(id), weight_string(user_id) from music where 1 != 1 group by user_id",
+            "OrderBy": "(0|2) ASC",
+            "Query": "select user_id, count(id), weight_string(user_id) from music group by user_id having count(user_id) = 1 order by user_id asc limit :__upper_limit",
+            "ResultColumns": 2,
+            "Table": "music"
+          }
+        ]
+      },
+      {
+        "OperatorType": "Aggregate",
+        "Variant": "Ordered",
+        "Aggregates": "count(0) AS count(*)",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select count(*) from (select user_id, count(id) from music where 1 != 1 group by user_id) as t where 1 != 1",
+            "Query": "select count(*) from (select user_id, count(id) from music group by user_id having count(user_id) = 1) as t",
+            "Table": "music"
+          }
+        ]
+      }
+    ]
+  }
+}
 
 # sql_calc_found_rows in sub queries
 "select * from music where user_id IN (select sql_calc_found_rows * from music limit 10)"
 "Incorrect usage/placement of 'SQL_CALC_FOUND_ROWS'"
+Gen4 plan same as above
 
 # sql_calc_found_rows in derived table
 "select sql_calc_found_rows * from (select sql_calc_found_rows * from music limit 10) t limit 1"
 "Incorrect usage/placement of 'SQL_CALC_FOUND_ROWS'"
+Gen4 plan same as above
 
 # select from unsharded keyspace into dumpfile
 "select * from main.unsharded into Dumpfile 'x.txt'"
diff --git a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
index b90e773c66..67440da4b6 100644
--- a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
@@ -386,6 +386,7 @@ Gen4 plan same as above
 # union with SQL_CALC_FOUND_ROWS 
 "(select sql_calc_found_rows id from user where id = 1 limit 1) union select id from user where id = 1"
 "SQL_CALC_FOUND_ROWS not supported with union"
+Gen4 plan same as above
 
 # set with DEFAULT - vitess aware
 "set workload = default"
@@ -425,10 +426,12 @@ Gen4 plan same as above
 # create view with sql_calc_found_rows with limit
 "create view user.view_a as select sql_calc_found_rows * from music limit 100"
 "Complex select queries are not supported in create or alter view statements"
+Gen4 plan same as above
 
 # create view with sql_calc_found_rows with group by and having
 "create view user.view_a as select sql_calc_found_rows user_id, count(id) from music group by user_id having count(user_id) = 1 order by user_id limit 2"
 "Complex select queries are not supported in create or alter view statements"
+Gen4 plan same as above
 
 # create view with incompatible keyspaces
 "create view main.view_a as select * from user.user_extra"
diff --git a/go/vt/vtgate/semantics/analyzer_test.go b/go/vt/vtgate/semantics/analyzer_test.go
index 0893d83518..1f5b89095b 100644
--- a/go/vt/vtgate/semantics/analyzer_test.go
+++ b/go/vt/vtgate/semantics/analyzer_test.go
@@ -763,7 +763,7 @@ func TestUnionOrderByRewrite(t *testing.T) {
 	assert.Equal(t, "select tabl1.id from tabl1 union select 1 from dual order by id asc", sqlparser.String(stmt))
 }
 
-func TestInvalidUnion(t *testing.T) {
+func TestInvalidQueries(t *testing.T) {
 	tcases := []struct {
 		sql string
 		err string
@@ -785,6 +785,15 @@ func TestInvalidUnion(t *testing.T) {
 	}, {
 		sql: "select a.id, b.id from a, b union select 1, 2 order by id",
 		err: "Column 'id' in field list is ambiguous",
+	}, {
+		sql: "select sql_calc_found_rows id from a union select 1 limit 109",
+		err: "SQL_CALC_FOUND_ROWS not supported with union",
+	}, {
+		sql: "select * from (select sql_calc_found_rows id from a) as t",
+		err: "Incorrect usage/placement of 'SQL_CALC_FOUND_ROWS'",
+	}, {
+		sql: "select (select sql_calc_found_rows id from a) as t",
+		err: "Incorrect usage/placement of 'SQL_CALC_FOUND_ROWS'",
 	}}
 	for _, tc := range tcases {
 		t.Run(tc.sql, func(t *testing.T) {