Repository: vitessio/vitess
Pull request #12820 (2023-04-04 09:29:47)

[release-14.0] fakedbclient: Add locking to avoid races (#12814)
## Description

This is a backport of #12814.


vitess-bot[bot]: ### Review Checklist
            
Hello reviewers! :wave: Please follow this checklist when reviewing this Pull Request.

#### General
- [ ] Ensure that the Pull Request has a descriptive title.
- [ ] If this is a change that users need to know about, please apply the `release notes (needs details)` label so that merging is blocked unless the summary release notes document is included.
- [ ] If a test is added or modified, there should be a documentation on top of the test to explain what the expected behavior is what the test does.

#### If a new flag is being introduced:
- [ ] Is it really necessary to add this flag?
- [ ] Flag names should be clear and intuitive (as far as possible)
- [ ] Help text should be descriptive.
- [ ] Flag names should use dashes (`-`) as word separators rather than underscores (`_`).

#### If a workflow is added or modified:
- [ ] Each item in `Jobs` should be named in order to mark it as `required`.
- [ ] If the workflow should be required, the maintainer team should be notified.

#### Bug fixes
- [ ] There should be at least one unit or end-to-end test.
- [ ] The Pull Request description should include a link to an issue that describes the bug.

#### Non-trivial changes
- [ ] There should be some code comments as to why things are implemented the way they are.

#### New/Existing features
- [ ] Should be documented, either by modifying the existing documentation or creating new documentation.
- [ ] New features should have a link to a feature request issue or an RFC that documents the use cases, corner cases and test cases.

#### Backward compatibility            
- [ ] Protobuf changes should be wire-compatible.
- [ ] Changes to `_vt` tables and RPCs need to be backward compatible.
- [ ] `vtctl` command output order should be stable and `awk`-able.
- [ ] RPC changes should be compatible with vitess-operator
- [ ] If a flag is removed, then it should also be removed from VTop, if used there.


Patch:
diff --git a/go/vt/vtgate/planbuilder/abstract/derived.go b/go/vt/vtgate/planbuilder/abstract/derived.go
index 13a7d8a62d..cf8ca75e02 100644
--- a/go/vt/vtgate/planbuilder/abstract/derived.go
+++ b/go/vt/vtgate/planbuilder/abstract/derived.go
@@ -17,7 +17,10 @@ limitations under the License.
 package abstract
 
 import (
+	"io"
+
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/semantics"
@@ -54,11 +57,31 @@ func (d *Derived) PushPredicate(expr sqlparser.Expr, semTable *semantics.SemTabl
 	if err != nil {
 		return nil, err
 	}
+	if !canBePushedDownIntoDerived(newExpr) {
+		// if we have an aggregation, we don't want to push it inside
+		return &Filter{Source: d, Predicates: []sqlparser.Expr{expr}}, nil
+	}
 	newSrc, err := d.Inner.PushPredicate(newExpr, semTable)
 	d.Inner = newSrc
 	return d, err
 }
 
+func canBePushedDownIntoDerived(expr sqlparser.Expr) (canBePushed bool) {
+	canBePushed = true
+	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+		switch n := node.(type) {
+		case *sqlparser.FuncExpr:
+			if n.Name.EqualString("min") || n.Name.EqualString("max") {
+				break
+			}
+			canBePushed = false
+			return false, io.EOF
+		}
+		return true, nil
+	}, expr)
+	return
+}
+
 // UnsolvedPredicates implements the Operator interface
 func (d *Derived) UnsolvedPredicates(semTable *semantics.SemTable) []sqlparser.Expr {
 	return d.Inner.UnsolvedPredicates(semTable)


Test patch:
diff --git a/go/vt/vtgate/planbuilder/abstract/operator_test.go b/go/vt/vtgate/planbuilder/abstract/operator_test.go
index f518f3ab55..e1197e8d65 100644
--- a/go/vt/vtgate/planbuilder/abstract/operator_test.go
+++ b/go/vt/vtgate/planbuilder/abstract/operator_test.go
@@ -98,10 +98,9 @@ func TestOperator(t *testing.T) {
 			require.NoError(t, err)
 			tree, err := sqlparser.Parse(tc.query)
 			require.NoError(t, err)
-			stmt := tree.(sqlparser.Statement)
-			semTable, err := semantics.Analyze(stmt, "", si)
+			semTable, err := semantics.Analyze(tree, "", si)
 			require.NoError(t, err)
-			optree, err := CreateLogicalOperatorFromAST(stmt, semTable)
+			optree, err := CreateLogicalOperatorFromAST(tree, semTable)
 			require.NoError(t, err)
 			output := testString(optree)
 			if tc.expected != output {
diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
index 6016622d2e..0235bef0e7 100644
--- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
@@ -5340,9 +5340,46 @@
           "Sharded": true
         },
         "FieldQuery": "select t1.portalId, t1.flowId from (select portalId, flowId, count(*) as count from user_extra where 1 != 1 group by user_id, flowId) as t1 where 1 != 1",
-        "Query": "select t1.portalId, t1.flowId from (select portalId, flowId, count(*) as count from user_extra where localDate > :v1 group by user_id, flowId having count(*) >= :v2 order by null) as t1",
+        "Query": "select t1.portalId, t1.flowId from (select portalId, flowId, count(*) as count from user_extra where localDate > :v1 group by user_id, flowId order by null) as t1 where count >= :v2",
         "Table": "user_extra"
       }
     }
+  },
+  {
+    "comment": "aggregation, where and derived tables - we can push extremums",
+    "query": "SELECT foo FROM (SELECT foo, max(baz) as bazo FROM (SELECT foo, baz FROM user) f GROUP BY foo) tt WHERE bazo BETWEEN 100 AND 200",
+    "v3-plan": "unsupported: filtering on results of cross-shard subquery",
+    "gen4-plan": {
+      "QueryType": "SELECT",
+      "Original": "SELECT foo FROM (SELECT foo, max(baz) as bazo FROM (SELECT foo, baz FROM user) f GROUP BY foo) tt WHERE bazo BETWEEN 100 AND 200",
+      "Instructions": {
+        "OperatorType": "SimpleProjection",
+        "Columns": [
+          0
+        ],
+        "Inputs": [
+          {
+            "OperatorType": "Aggregate",
+            "Variant": "Ordered",
+            "Aggregates": "max(1) AS bazo",
+            "GroupBy": "(0|2)",
+            "Inputs": [
+              {
+                "OperatorType": "Route",
+                "Variant": "Scatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select foo, max(baz) as bazo, weight_string(foo) from (select foo, baz from `user` where 1 != 1) as f where 1 != 1 group by foo, weight_string(foo)",
+                "OrderBy": "(0|2) ASC",
+                "Query": "select foo, max(baz) as bazo, weight_string(foo) from (select foo, baz from `user` having max(baz) between 100 and 200) as f group by foo, weight_string(foo) order by foo asc",
+                "Table": "`user`"
+              }
+            ]
+          }
+        ]
+      }
+    }
   }
 ]
diff --git a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
index 39e43ff854..909b4a7d92 100644
--- a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
@@ -471,5 +471,11 @@
     "comment": "mix lock with other expr",
     "query": "select get_lock('xyz', 10), 1 from dual",
     "plan": "unsupported: lock function and other expression in same select query"
+  },
+  {
+    "comment": "aggregation, where and derived tables - we can't push aggregations that might need a second layer of aggregation",
+    "query": "SELECT foo FROM (SELECT foo, count(baz) as bazo FROM (SELECT foo, baz FROM user) f GROUP BY foo) tt WHERE bazo BETWEEN 100 AND 200",
+    "v3-plan": "unsupported: filtering on results of cross-shard subquery",
+    "gen4-plan": "expr cannot be translated, not supported: bazo between 100 and 200"
   }
 ]
diff --git a/go/vt/wrangler/fake_dbclient_test.go b/go/vt/wrangler/fake_dbclient_test.go
index f69d95f6c5..5b7fc0c187 100644
--- a/go/vt/wrangler/fake_dbclient_test.go
+++ b/go/vt/wrangler/fake_dbclient_test.go
@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"regexp"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -63,6 +64,7 @@ func (dbrs *dbResults) exhausted() bool {
 
 // fakeDBClient fakes a binlog_player.DBClient.
 type fakeDBClient struct {
+	mu         sync.Mutex
 	name       string
 	queries    map[string]*dbResults
 	queriesRE  map[string]*dbResults
@@ -86,6 +88,8 @@ func newFakeDBClient(name string) *fakeDBClient {
 }
 
 func (dc *fakeDBClient) addQuery(query string, result *sqltypes.Result, err error) {
+	dc.mu.Lock()
+	defer dc.mu.Unlock()
 	if testMode == "debug" {
 		log.Infof("%s::addQuery %s\n\n", dc.id(), query)
 	}
@@ -98,6 +102,8 @@ func (dc *fakeDBClient) addQuery(query string, result *sqltypes.Result, err erro
 }
 
 func (dc *fakeDBClient) addQueryRE(query string, result *sqltypes.Result, err error) {
+	dc.mu.Lock()
+	defer dc.mu.Unlock()
 	if testMode == "debug" {
 		log.Infof("%s::addQueryRE %s\n\n", dc.id(), query)
 	}
@@ -110,11 +116,15 @@ func (dc *fakeDBClient) addQueryRE(query string, result *sqltypes.Result, err er
 }
 
 func (dc *fakeDBClient) getInvariant(query string) *sqltypes.Result {
+	dc.mu.Lock()
+	defer dc.mu.Unlock()
 	return dc.invariants[query]
 }
 
 // note: addInvariant will replace a previous result for a query with the provided one: this is used in the tests
 func (dc *fakeDBClient) addInvariant(query string, result *sqltypes.Result) {
+	dc.mu.Lock()
+	defer dc.mu.Unlock()
 	if testMode == "debug" {
 		log.Infof("%s::addInvariant %s\n\n", dc.id(), query)
 	}
@@ -156,6 +166,8 @@ func (dc *fakeDBClient) id() string {
 
 // ExecuteFetch is part of the DBClient interface
 func (dc *fakeDBClient) ExecuteFetch(query string, maxrows int) (*sqltypes.Result, error) {
+	dc.mu.Lock()
+	defer dc.mu.Unlock()
 	qr, err := dc.executeFetch(query, maxrows)
 	if testMode == "debug" {
 		log.Infof("%s::ExecuteFetch for >>>%s<<< returns >>>%v<<< error >>>%+v<<< ", dc.id(), query, qr, err)
@@ -187,6 +199,8 @@ func (dc *fakeDBClient) executeFetch(query string, maxrows int) (*sqltypes.Resul
 }
 
 func (dc *fakeDBClient) verifyQueries(t *testing.T) {
+	dc.mu.Lock()
+	defer dc.mu.Unlock()
 	t.Helper()
 	for query, dbrs := range dc.queries {
 		if !dbrs.exhausted() {