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

Add support for query payload limit

go/vt/vtgate/executor_test.go

@@ -1981,6 +1981,45 @@ func TestExecutorMaxPayloadSizeExceeded(t *testing.T) {
  assert.Equal(t, warningCount+4, warnings.Counts()["WarnPayloadSizeExceeded"], "warnings count")
 }
 
+func TestOlapSelectDatabase(t *testing.T) {
+ 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)
+ if err == nil {
+ assert.EqualError(t, err, "query payload size above threshold")
+ }
+ }
+ assert.Equal(t, warningCount, warnings.Counts()["WarnPayloadSizeExceeded"], "warnings count")
+
+ testMaxPayloadSizeOverride := []string{
+ "select /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ * from main1",
+ "insert /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ into main1(id) values (1), (2)",
+ "update /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ main1 set id=1",
+ "delete /*vt+ IGNORE_MAX_PAYLOAD_SIZE=1 */ from main1 where id=1",
+ }
+ for _, query := range testMaxPayloadSizeOverride {
+ _, err := executor.Execute(context.Background(), "TestExecutorMaxPayloadSizeWithOverride", session, query, nil)
+ assert.Equal(t, nil, err, "err should be nil")
+ }
+ assert.Equal(t, warningCount, warnings.Counts()["WarnPayloadSizeExceeded"], "warnings count")
+
+ *maxPayloadSize = 1000
+ for _, query := range testMaxPayloadSizeExceeded {
+ _, err := executor.Execute(context.Background(), "TestExecutorMaxPayloadSizeExceeded", session, query, nil)
+ assert.Equal(t, nil, err, "err should be nil")
+ }
+ assert.Equal(t, warningCount+4, warnings.Counts()["WarnPayloadSizeExceeded"], "warnings count")
+}
+
 func TestOlapSelectDatabase(t *testing.T) {
  executor, _, _, _ := createLegacyExecutorEnv()
  executor.normalize = true

go/vt/vtgate/vtgate.go

@@ -63,8 +63,12 @@ var (
 
  // TODO(deepthi): change these two vars to unexported and move to healthcheck.go when LegacyHealthcheck is removed
 
- maxPayloadSize = flag.Int("max_payload_size", 0, "The threshold for query payloads in bytes. A payload greater than this threshold will result in a failure to handle the query.")
- warnPayloadSize = flag.Int("warn_payload_size", 0, "The warning threshold for query payloads in bytes. A payload greater than this threshold will cause the VtGateWarnings.WarnPayloadSizeExceeded counter to be incremented.")
+ // HealthCheckRetryDelay is the time to wait before retrying healthcheck
+ HealthCheckRetryDelay = flag.Duration("healthcheck_retry_delay", 2*time.Millisecond, "health check retry delay")
+ // HealthCheckTimeout is the timeout on the RPC call to tablets
+ HealthCheckTimeout = flag.Duration("healthcheck_timeout", time.Minute, "the health check timeout period")
+ maxPayloadSize = flag.Int("max_payload_size", 0, "The threshold for query payloads in bytes. A payload greater than this threshold will result in a failure to handle the query.")
+ warnPayloadSize = flag.Int("warn_payload_size", 0, "The warning threshold for query payloads in bytes. A payload greater than this threshold will cause the VtGateWarnings.WarnPayloadSizeExceeded counter to be incremented.")
 )
 
 func getTxMode() vtgatepb.TransactionMode {