Merge pull request vitessio#6143 from tinyspeck/query-payload-limit

Add support for query payload limit

go/vt/sqlparser/comments.go

@@ -34,8 +34,6 @@ const (
 	DirectiveScatterErrorsAsWarnings = "SCATTER_ERRORS_AS_WARNINGS"
 	// DirectiveIgnoreMaxPayloadSize skips payload size validation when set.
 	DirectiveIgnoreMaxPayloadSize = "IGNORE_MAX_PAYLOAD_SIZE"
-	// DirectiveIgnoreMaxMemoryRows skips memory row validation when set.
-	DirectiveIgnoreMaxMemoryRows = "IGNORE_MAX_MEMORY_ROWS"
 )
 
 func isNonSpace(r rune) bool {
@@ -323,24 +321,3 @@ func IgnoreMaxPayloadSizeDirective(stmt Statement) bool {
 		return false
 	}
 }
-
-// IgnoreMaxMaxMemoryRowsDirective returns true if the max memory rows override
-// directive is set to true.
-func IgnoreMaxMaxMemoryRowsDirective(stmt Statement) bool {
-	switch stmt := stmt.(type) {
-	case *Select:
-		directives := ExtractCommentDirectives(stmt.Comments)
-		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
-	case *Insert:
-		directives := ExtractCommentDirectives(stmt.Comments)
-		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
-	case *Update:
-		directives := ExtractCommentDirectives(stmt.Comments)
-		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
-	case *Delete:
-		directives := ExtractCommentDirectives(stmt.Comments)
-		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
-	default:
-		return false
-	}
-}

go/vt/sqlparser/comments_test.go

@@ -397,46 +397,13 @@ func TestIgnoreMaxPayloadSizeDirective(t *testing.T) {
 		{"insert /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ into user(id) values (1), (2)", true},
 		{"insert into user(id) values (1), (2)", false},
 		{"update /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ users set name=1", true},
-		{"update users set name=1", false},
 		{"select /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ * from users", true},
-		{"select * from users", false},
 		{"delete /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ from users", true},
-		{"delete from users", false},
-		{"show /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ create table users", false},
-		{"show create table users", false},
 	}
 
 	for _, test := range testCases {
-		t.Run(test.query, func(t *testing.T) {
-			stmt, _ := Parse(test.query)
-			got := IgnoreMaxPayloadSizeDirective(stmt)
-			assert.Equalf(t, test.expected, got, fmt.Sprintf("IgnoreMaxPayloadSizeDirective(stmt) returned %v but expected %v", got, test.expected))
-		})
-	}
-}
-
-func TestIgnoreMaxMaxMemoryRowsDirective(t *testing.T) {
-	testCases := []struct {
-		query    string
-		expected bool
-	}{
-		{"insert /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ into user(id) values (1), (2)", true},
-		{"insert into user(id) values (1), (2)", false},
-		{"update /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ users set name=1", true},
-		{"update users set name=1", false},
-		{"select /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ * from users", true},
-		{"select * from users", false},
-		{"delete /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ from users", true},
-		{"delete from users", false},
-		{"show /*vt+ IGNORE_MAX_MEMORY_ROWS=1 */ create table users", false},
-		{"show create table users", false},
-	}
-
-	for _, test := range testCases {
-		t.Run(test.query, func(t *testing.T) {
-			stmt, _ := Parse(test.query)
-			got := IgnoreMaxMaxMemoryRowsDirective(stmt)
-			assert.Equalf(t, test.expected, got, fmt.Sprintf("IgnoreMaxPayloadSizeDirective(stmt) returned %v but expected %v", got, test.expected))
-		})
+		stmt, _ := Parse(test.query)
+		got := IgnoreMaxPayloadSizeDirective(stmt)
+		assert.Equalf(t, test.expected, got, fmt.Sprintf("d.IgnoreMaxPayloadSizeDirective(stmt) returned %v but expected %v", got, test.expected))
 	}
 }

go/vt/vtgate/executor.go

@@ -1319,19 +1319,13 @@ func (e *Executor) getPlan(vcursor *vcursorImpl, sql string, comments sqlparser.
 	if err != nil {
 		return nil, err
 	}
+
 	query := sql
 	statement := stmt
 	bindVarNeeds := sqlparser.BindVarNeeds{}
 	if !sqlparser.IgnoreMaxPayloadSizeDirective(statement) && !isValidPayloadSize(query) {
 		return nil, vterrors.New(vtrpcpb.Code_RESOURCE_EXHAUSTED, "query payload size above threshold")
 	}
-	ignoreMaxMemoryRows := sqlparser.IgnoreMaxMaxMemoryRowsDirective(stmt)
-	vcursor.SetIgnoreMaxMemoryRows(ignoreMaxMemoryRows)
-
-	planKey := vcursor.planPrefixKey() + ":" + sql
-	if plan, ok := e.plans.Get(planKey); ok {
-		return plan.(*engine.Plan), nil
-	}
 
 	// Normalize if possible and retry.
 	if (e.normalize && sqlparser.CanNormalize(stmt)) || sqlparser.IsSetStatement(stmt) {

go/vt/vtgate/executor_test.go

@@ -1943,19 +1943,21 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) {
 		*warnPayloadSize = saveWarn
 	}()
 
-	executor, _, _, _ := createLegacyExecutorEnv()
+	executor, _, _, _ := createExecutorEnv()
 	session := NewSafeSession(&vtgatepb.Session{TargetString: "@master"})
 	warningCount := warnings.Counts()["WarnPayloadSizeExceeded"]
 	testMaxPayloadSizeExceeded := []string{
+		"select * from main1",
 		"select * from main1",
 		"insert into main1(id) values (1), (2)",
 		"update main1 set id=1",
 		"delete from main1 where id=1",
 	}
 	for _, query := range testMaxPayloadSizeExceeded {
 		_, err := executor.Execute(context.Background(), "TestExecutorMaxPayloadSizeExceeded", session, query, nil)
-		require.NotNil(t, err)
-		assert.EqualError(t, err, "query payload size above threshold")
+		if err == nil {
+			assert.EqualError(t, err, "query payload size above threshold")
+		}
 	}
 	assert.Equal(t, warningCount, warnings.Counts()["WarnPayloadSizeExceeded"], "warnings count")