plan more complicated sql_calc_found_rows queries

Signed-off-by: Andres Taylor <andres@planetscale.com>
vitessio · Sep 9, 2020 · e26f979 · e26f979
1 parent cf9f6ab
commit e26f979
Show file tree

Hide file tree

Showing 5 changed files with 153 additions and 48 deletions.
diff --git a/go/test/endtoend/vtgate/found_rows_test.go b/go/test/endtoend/vtgate/found_rows_test.go
@@ -40,6 +40,8 @@ func TestFoundRows(t *testing.T) {
 	assertFoundRowsValue(t, conn, "select * from t2", 5)
 	assertFoundRowsValue(t, conn, "select * from t2 limit 2", 2)
 	assertFoundRowsValue(t, conn, "select SQL_CALC_FOUND_ROWS * from t2 limit 2", 5)
+	assertFoundRowsValue(t, conn, "select SQL_CALC_FOUND_ROWS * from t2 where id3 = 4 limit 2", 1)
+	assertFoundRowsValue(t, conn, "select SQL_CALC_FOUND_ROWS * from t2 where id4 = 4 limit 2", 1)
 }
 
 func assertFoundRowsValue(t *testing.T, conn *mysql.Conn, query string, count int) {

diff --git a/go/vt/vtgate/engine/sqlcalfoundrows.go → go/vt/vtgate/engine/sql_calc_found_rows.go b/go/vt/vtgate/engine/sqlcalfoundrows.go → go/vt/vtgate/engine/sql_calc_found_rows.go
@@ -1,3 +1,19 @@
+/*
+Copyright 2020 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.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package engine
 
 import (
@@ -8,31 +24,31 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 )
 
-var _ Primitive = (*SQLCalFoundRows)(nil)
+var _ Primitive = (*SQLCalcFoundRows)(nil)
 
-//SQLCalFoundRows is a primitive to execute limit and count query as per their individual plan.
-type SQLCalFoundRows struct {
+//SQLCalcFoundRows is a primitive to execute limit and count query as per their individual plan.
+type SQLCalcFoundRows struct {
 	LimitPrimitive Primitive
 	CountPrimitive Primitive
 }
 
 //RouteType implements the Primitive interface
-func (s SQLCalFoundRows) RouteType() string {
+func (s SQLCalcFoundRows) RouteType() string {
 	return "SQLCalcFoundRows"
 }
 
 //GetKeyspaceName implements the Primitive interface
-func (s SQLCalFoundRows) GetKeyspaceName() string {
+func (s SQLCalcFoundRows) GetKeyspaceName() string {
 	return s.LimitPrimitive.GetKeyspaceName()
 }
 
 //GetTableName implements the Primitive interface
-func (s SQLCalFoundRows) GetTableName() string {
+func (s SQLCalcFoundRows) GetTableName() string {
 	return s.LimitPrimitive.GetTableName()
 }
 
 //Execute implements the Primitive interface
-func (s SQLCalFoundRows) Execute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+func (s SQLCalcFoundRows) Execute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
 	limitQr, err := s.LimitPrimitive.Execute(vcursor, bindVars, wantfields)
 	if err != nil {
 		return nil, err
@@ -53,26 +69,26 @@ func (s SQLCalFoundRows) Execute(vcursor VCursor, bindVars map[string]*querypb.B
 }
 
 //StreamExecute implements the Primitive interface
-func (s SQLCalFoundRows) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+func (s SQLCalcFoundRows) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
 	panic("implement me")
 }
 
 //GetFields implements the Primitive interface
-func (s SQLCalFoundRows) GetFields(vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+func (s SQLCalcFoundRows) GetFields(vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
 	return s.LimitPrimitive.GetFields(vcursor, bindVars)
 }
 
 //NeedsTransaction implements the Primitive interface
-func (s SQLCalFoundRows) NeedsTransaction() bool {
+func (s SQLCalcFoundRows) NeedsTransaction() bool {
 	return s.LimitPrimitive.NeedsTransaction()
 }
 
 //Inputs implements the Primitive interface
-func (s SQLCalFoundRows) Inputs() []Primitive {
+func (s SQLCalcFoundRows) Inputs() []Primitive {
 	return []Primitive{s.LimitPrimitive, s.CountPrimitive}
 }
 
-func (s SQLCalFoundRows) description() PrimitiveDescription {
+func (s SQLCalcFoundRows) description() PrimitiveDescription {
 	return PrimitiveDescription{
 		OperatorType: "SQL_CALC_FOUND_ROWS",
 	}

diff --git a/go/vt/vtgate/planbuilder/select.go b/go/vt/vtgate/planbuilder/select.go
@@ -100,24 +100,11 @@ func (pb *primitiveBuilder) processSelect(sel *sqlparser.Select, outer *symtab)
 	if sel.SQLCalcFoundRows {
 		sel.SQLCalcFoundRows = false
 		if sel.Limit != nil {
-			frpb := newPrimitiveBuilder(pb.vschema, pb.jt)
-			err := frpb.processSelect(sel, outer)
+			builder, err := buildSQLCalcFoundRowsPlan(sel, outer, pb.vschema)
 			if err != nil {
 				return err
 			}
-			sel.SelectExprs = []sqlparser.SelectExpr{&sqlparser.AliasedExpr{
-				Expr: &sqlparser.FuncExpr{
-					Name:  sqlparser.NewColIdent("count"),
-					Exprs: []sqlparser.SelectExpr{&sqlparser.StarExpr{}},
-				},
-			}}
-			sel.OrderBy = nil
-			sel.Limit = nil
-			err = pb.processSelect(sel, outer)
-			if err != nil {
-				return err
-			}
-			pb.bldr = sqlCalcFoundRows{LimitQuery: frpb.bldr, CountQuery: pb.bldr}
+			pb.bldr = builder
 			return nil
 		}
 	}
@@ -168,6 +155,47 @@ func (pb *primitiveBuilder) processSelect(sel *sqlparser.Select, outer *symtab)
 	return nil
 }
 
+func buildSQLCalcFoundRowsPlan(sel *sqlparser.Select, outer *symtab, vschema ContextVSchema) (builder, error) {
+	ljt := newJointab(sqlparser.GetBindvars(sel))
+	frpb := newPrimitiveBuilder(vschema, ljt)
+	err := frpb.processSelect(sel, outer)
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO systay this is a hack
+	s := sqlparser.String(sel)
+	statement, err := sqlparser.Parse(s)
+	if err != nil {
+		return nil, err
+	}
+	sel2 := statement.(*sqlparser.Select)
+
+	sel2.SelectExprs = []sqlparser.SelectExpr{&sqlparser.AliasedExpr{
+		Expr: &sqlparser.FuncExpr{
+			Name:  sqlparser.NewColIdent("count"),
+			Exprs: []sqlparser.SelectExpr{&sqlparser.StarExpr{}},
+		},
+	}}
+	sel2.OrderBy = nil
+	sel2.Limit = nil
+
+	sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+		switch col := node.(type) {
+		case *sqlparser.ColName:
+			col.Metadata = nil
+		}
+		return true, nil
+	}, sel2)
+	cjt := newJointab(sqlparser.GetBindvars(sel2))
+	countpb := newPrimitiveBuilder(vschema, cjt)
+	err = countpb.processSelect(sel2, outer)
+	if err != nil {
+		return nil, err
+	}
+	return &sqlCalcFoundRows{LimitQuery: frpb.bldr, CountQuery: countpb.bldr, ljt: ljt, cjt: cjt}, nil
+}
+
 func handleDualSelects(sel *sqlparser.Select, vschema ContextVSchema) (engine.Primitive, error) {
 	if !isOnlyDual(sel) {
 		return nil, nil

diff --git a/go/vt/vtgate/planbuilder/sql_calc_found_rows.go b/go/vt/vtgate/planbuilder/sql_calc_found_rows.go
@@ -1,3 +1,19 @@
+/*
+Copyright 2020 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.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package planbuilder
 
 import (
@@ -9,78 +25,79 @@ var _ builder = (*sqlCalcFoundRows)(nil)
 
 type sqlCalcFoundRows struct {
 	LimitQuery, CountQuery builder
+	ljt, cjt               *jointab
 }
 
-func (s sqlCalcFoundRows) Order() int {
+func (s *sqlCalcFoundRows) Order() int {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) ResultColumns() []*resultColumn {
+func (s *sqlCalcFoundRows) ResultColumns() []*resultColumn {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) Reorder(i int) {
+func (s *sqlCalcFoundRows) Reorder(i int) {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) First() builder {
+func (s *sqlCalcFoundRows) First() builder {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) PushFilter(pb *primitiveBuilder, filter sqlparser.Expr, whereType string, origin builder) error {
+func (s *sqlCalcFoundRows) PushFilter(pb *primitiveBuilder, filter sqlparser.Expr, whereType string, origin builder) error {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) PushSelect(pb *primitiveBuilder, expr *sqlparser.AliasedExpr, origin builder) (rc *resultColumn, colNumber int, err error) {
+func (s *sqlCalcFoundRows) PushSelect(pb *primitiveBuilder, expr *sqlparser.AliasedExpr, origin builder) (rc *resultColumn, colNumber int, err error) {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) MakeDistinct() error {
+func (s *sqlCalcFoundRows) MakeDistinct() error {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) PushGroupBy(by sqlparser.GroupBy) error {
+func (s *sqlCalcFoundRows) PushGroupBy(by sqlparser.GroupBy) error {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) PushOrderBy(by sqlparser.OrderBy) (builder, error) {
+func (s *sqlCalcFoundRows) PushOrderBy(by sqlparser.OrderBy) (builder, error) {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) SetUpperLimit(count sqlparser.Expr) {
+func (s *sqlCalcFoundRows) SetUpperLimit(count sqlparser.Expr) {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) PushMisc(sel *sqlparser.Select) {
+func (s *sqlCalcFoundRows) PushMisc(sel *sqlparser.Select) {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) Wireup(bldr builder, jt *jointab) error {
-	err := s.LimitQuery.Wireup(bldr, jt)
+func (s *sqlCalcFoundRows) Wireup(builder, *jointab) error {
+	err := s.LimitQuery.Wireup(s.LimitQuery, s.ljt)
 	if err != nil {
 		return err
 	}
-	return s.CountQuery.Wireup(bldr, jt)
+	return s.CountQuery.Wireup(s.CountQuery, s.cjt)
 }
 
-func (s sqlCalcFoundRows) SupplyVar(from, to int, col *sqlparser.ColName, varname string) {
-	panic("implement me")
+func (s *sqlCalcFoundRows) SupplyVar(from, to int, col *sqlparser.ColName, varname string) {
+	s.LimitQuery.SupplyVar(from, to, col, varname)
 }
 
-func (s sqlCalcFoundRows) SupplyCol(col *sqlparser.ColName) (rc *resultColumn, colNumber int) {
+func (s *sqlCalcFoundRows) SupplyCol(col *sqlparser.ColName) (rc *resultColumn, colNumber int) {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) SupplyWeightString(colNumber int) (weightcolNumber int, err error) {
+func (s *sqlCalcFoundRows) SupplyWeightString(colNumber int) (weightcolNumber int, err error) {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) PushLock(lock string) error {
+func (s *sqlCalcFoundRows) PushLock(lock string) error {
 	panic("implement me")
 }
 
-func (s sqlCalcFoundRows) Primitive() engine.Primitive {
-	return engine.SQLCalFoundRows{
+func (s *sqlCalcFoundRows) Primitive() engine.Primitive {
+	return engine.SQLCalcFoundRows{
 		LimitPrimitive: s.LimitQuery.Primitive(),
 		CountPrimitive: s.CountQuery.Primitive(),
 	}

diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.txt b/go/vt/vtgate/planbuilder/testdata/select_cases.txt
@@ -1457,3 +1457,45 @@
     ]
   }
 }
+
+# sql_calc_found_rows with SelectEqualUnique plans
+"select sql_calc_found_rows * from music where user_id = 1 limit 2"
+{
+  "QueryType": "SELECT",
+  "Original": "select sql_calc_found_rows * from music where user_id = 1 limit 2",
+  "Instructions": {
+    "OperatorType": "SQL_CALC_FOUND_ROWS",
+    "Inputs": [
+      {
+        "OperatorType": "Route",
+        "Variant": "SelectEqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select * from music where 1 != 1",
+        "Query": "select * from music where user_id = 1 limit 2",
+        "Table": "music",
+        "Values": [
+          1
+        ],
+        "Vindex": "user_index"
+      },
+      {
+        "OperatorType": "Route",
+        "Variant": "SelectEqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select count(*) from music where 1 != 1",
+        "Query": "select count(*) from music where user_id = 1",
+        "Table": "music",
+        "Values": [
+          1
+        ],
+        "Vindex": "user_index"
+      }
+    ]
+  }
+}