Implement parser support for all select modifiers (#16396)

Signed-off-by: Utkarsh Mahajan <[email protected]> Signed-off-by: Manan Gupta <[email protected]> Co-authored-by: Manan Gupta <[email protected]>
vitessio · Jul 19, 2024 · efd8292 · efd8292
1 parent 7621331
commit efd8292
Show file tree

Hide file tree

Showing 11 changed files with 9,439 additions and 9,279 deletions.
diff --git a/go/vt/sqlparser/ast.go b/go/vt/sqlparser/ast.go
@@ -262,7 +262,11 @@ type (
 	Select struct {
 		Cache            *bool // a reference here so it can be nil
 		Distinct         bool
+		HighPriority     bool
 		StraightJoinHint bool
+		SQLSmallResult   bool
+		SQLBigResult     bool
+		SQLBufferResult  bool
 		SQLCalcFoundRows bool
 		// The With field needs to come before the FROM clause, so any CTEs have been handled before we analyze it
 		With        *With

diff --git a/go/vt/sqlparser/ast_equals.go b/go/vt/sqlparser/ast_equals.go
diff --git a/go/vt/sqlparser/ast_format.go b/go/vt/sqlparser/ast_format.go
@@ -37,9 +37,21 @@ func (node *Select) Format(buf *TrackedBuffer) {
 			buf.literal(SQLNoCacheStr)
 		}
 	}
+	if node.HighPriority {
+		buf.literal(HighPriorityStr)
+	}
 	if node.StraightJoinHint {
 		buf.literal(StraightJoinHint)
 	}
+	if node.SQLSmallResult {
+		buf.literal(SQLSmallResultStr)
+	}
+	if node.SQLBigResult {
+		buf.literal(SQLBigResultStr)
+	}
+	if node.SQLBufferResult {
+		buf.literal(SQLBufferResultStr)
+	}
 	if node.SQLCalcFoundRows {
 		buf.literal(SQLCalcFoundRowsStr)
 	}

diff --git a/go/vt/sqlparser/ast_format_fast.go b/go/vt/sqlparser/ast_format_fast.go
diff --git a/go/vt/sqlparser/ast_funcs.go b/go/vt/sqlparser/ast_funcs.go
@@ -800,7 +800,7 @@ func NewSelect(
 	windows NamedWindows,
 ) *Select {
 	var cache *bool
-	var distinct, straightJoinHint, sqlFoundRows bool
+	var distinct, highPriority, straightJoinHint, sqlSmallResult, sqlBigResult, SQLBufferResult, sqlFoundRows bool
 	for _, option := range selectOptions {
 		switch strings.ToLower(option) {
 		case DistinctStr:
@@ -811,8 +811,16 @@ func NewSelect(
 		case SQLNoCacheStr:
 			truth := false
 			cache = &truth
+		case HighPriorityStr:
+			highPriority = true
 		case StraightJoinHint:
 			straightJoinHint = true
+		case SQLSmallResultStr:
+			sqlSmallResult = true
+		case SQLBigResultStr:
+			sqlBigResult = true
+		case SQLBufferResultStr:
+			SQLBufferResult = true
 		case SQLCalcFoundRowsStr:
 			sqlFoundRows = true
 		}
@@ -821,7 +829,11 @@ func NewSelect(
 		Cache:            cache,
 		Comments:         comments.Parsed(),
 		Distinct:         distinct,
+		HighPriority:     highPriority,
 		StraightJoinHint: straightJoinHint,
+		SQLSmallResult:   sqlSmallResult,
+		SQLBigResult:     sqlBigResult,
+		SQLBufferResult:  SQLBufferResult,
 		SQLCalcFoundRows: sqlFoundRows,
 		SelectExprs:      exprs,
 		Into:             into,

diff --git a/go/vt/sqlparser/constants.go b/go/vt/sqlparser/constants.go
@@ -23,7 +23,11 @@ const (
 	// Select.Distinct
 	AllStr              = "all "
 	DistinctStr         = "distinct "
+	HighPriorityStr     = "high_priority "
 	StraightJoinHint    = "straight_join "
+	SQLSmallResultStr   = "sql_small_result "
+	SQLBigResultStr     = "sql_big_result "
+	SQLBufferResultStr  = "sql_buffer_result "
 	SQLCalcFoundRowsStr = "sql_calc_found_rows "
 
 	// Select.Lock

diff --git a/go/vt/sqlparser/keywords.go b/go/vt/sqlparser/keywords.go
@@ -315,7 +315,7 @@ var keywords = []keyword{
 	{"hash", HASH},
 	{"having", HAVING},
 	{"header", HEADER},
-	{"high_priority", UNUSED},
+	{"high_priority", HIGH_PRIORITY},
 	{"hosts", HOSTS},
 	{"hour", HOUR},
 	{"hour_microsecond", HOUR_MICROSECOND},
@@ -594,11 +594,12 @@ var keywords = []keyword{
 	{"sqlexception", UNUSED},
 	{"sqlstate", UNUSED},
 	{"sqlwarning", UNUSED},
-	{"sql_big_result", UNUSED},
+	{"sql_big_result", SQL_BIG_RESULT},
 	{"sql_cache", SQL_CACHE},
 	{"sql_calc_found_rows", SQL_CALC_FOUND_ROWS},
 	{"sql_no_cache", SQL_NO_CACHE},
-	{"sql_small_result", UNUSED},
+	{"sql_small_result", SQL_SMALL_RESULT},
+	{"sql_buffer_result", SQL_BUFFER_RESULT},
 	{"sql_tsi_day", SQL_TSI_DAY},
 	{"sql_tsi_week", SQL_TSI_WEEK},
 	{"sql_tsi_hour", SQL_TSI_HOUR},

diff --git a/go/vt/sqlparser/parse_test.go b/go/vt/sqlparser/parse_test.go
@@ -2646,9 +2646,24 @@ var (
 	}, {
 		input:  "select sql_no_cache straight_join distinct 'foo' from t",
 		output: "select distinct sql_no_cache straight_join 'foo' from t",
+	}, {
+		input:  "select sql_buffer_result 'foo' from t",
+		output: "select sql_buffer_result 'foo' from t",
+	}, {
+		input:  "select high_priority 'foo' from t",
+		output: "select high_priority 'foo' from t",
 	}, {
 		input:  "select straight_join distinct sql_no_cache 'foo' from t",
 		output: "select distinct sql_no_cache straight_join 'foo' from t",
+	}, {
+		input:  "select sql_small_result 'foo' from t",
+		output: "select sql_small_result 'foo' from t",
+	}, {
+		input:  "select distinct sql_small_result 'foo' from t",
+		output: "select distinct sql_small_result 'foo' from t",
+	}, {
+		input:  "select sql_big_result 'foo' from t",
+		output: "select sql_big_result 'foo' from t",
 	}, {
 		input:  "select sql_calc_found_rows 'foo' from t",
 		output: "select sql_calc_found_rows 'foo' from t",
@@ -6426,8 +6441,11 @@ func testFile(t *testing.T, filename, tempDir string) {
 					errPresent := ""
 					if err != nil {
 						errPresent = err.Error()
+						expected.WriteString(fmt.Sprintf("ERROR\n%s\nEND\n", escapeNewLines(errPresent)))
+					} else {
+						out := String(tree)
+						expected.WriteString(fmt.Sprintf("OUTPUT\n%s\nEND\n", escapeNewLines(out)))
 					}
-					expected.WriteString(fmt.Sprintf("ERROR\n%s\nEND\n", escapeNewLines(errPresent)))
 					if err == nil || tcase.errStr != err.Error() {
 						fail = true
 						t.Errorf("File: %s, Line: %d\nDiff:\n%s\n[%s] \n[%s]", filename, tcase.lineno, cmp.Diff(tcase.errStr, errPresent), tcase.errStr, errPresent)