Add optimizer hints support and enable 2 more tests

ast/optimize_for_optimizer_hint.go

@@ -0,0 +1,11 @@
+package ast
+
+// OptimizeForOptimizerHint represents an OPTIMIZE FOR hint.
+type OptimizeForOptimizerHint struct {
+	Pairs        []*VariableValuePair `json:"Pairs,omitempty"`
+	IsForUnknown bool                 `json:"IsForUnknown,omitempty"`
+	HintKind     string               `json:"HintKind,omitempty"`
+}
+
+func (*OptimizeForOptimizerHint) node()          {}
+func (*OptimizeForOptimizerHint) optimizerHint() {}

ast/variable_value_pair.go

@@ -0,0 +1,10 @@
+package ast
+
+// VariableValuePair represents a variable-value pair in an OPTIMIZE FOR hint.
+type VariableValuePair struct {
+	Variable     *VariableReference `json:"Variable,omitempty"`
+	Value        ScalarExpression   `json:"Value,omitempty"`
+	IsForUnknown bool               `json:"IsForUnknown,omitempty"`
+}
+
+func (*VariableValuePair) node() {}

parser/marshal.go

@@ -941,11 +941,41 @@ func optimizerHintToJSON(h ast.OptimizerHintBase) jsonNode {
 			node["HintKind"] = hint.HintKind
 		}
 		return node
+	case *ast.OptimizeForOptimizerHint:
+		node := jsonNode{
+			"$type": "OptimizeForOptimizerHint",
+		}
+		if len(hint.Pairs) > 0 {
+			pairs := make([]jsonNode, len(hint.Pairs))
+			for i, pair := range hint.Pairs {
+				pairs[i] = variableValuePairToJSON(pair)
+			}
+			node["Pairs"] = pairs
+		}
+		node["IsForUnknown"] = hint.IsForUnknown
+		if hint.HintKind != "" {
+			node["HintKind"] = hint.HintKind
+		}
+		return node
 	default:
 		return jsonNode{"$type": "UnknownOptimizerHint"}
 	}
 }
 
+func variableValuePairToJSON(p *ast.VariableValuePair) jsonNode {
+	node := jsonNode{
+		"$type": "VariableValuePair",
+	}
+	if p.Variable != nil {
+		node["Variable"] = scalarExpressionToJSON(p.Variable)
+	}
+	if p.Value != nil {
+		node["Value"] = scalarExpressionToJSON(p.Value)
+	}
+	node["IsForUnknown"] = p.IsForUnknown
+	return node
+}
+
 func queryExpressionToJSON(qe ast.QueryExpression) jsonNode {
 	switch q := qe.(type) {
 	case *ast.QuerySpecification:

parser/parse_select.go

@@ -1283,28 +1283,200 @@ func (p *Parser) parseOptionClause() ([]ast.OptimizerHintBase, error) {
 
 	// Parse hints
 	for p.curTok.Type != TokenRParen && p.curTok.Type != TokenEOF {
-		if p.curTok.Type == TokenIdent || p.curTok.Type == TokenLabel {
-			hintKind := convertHintKind(p.curTok.Literal)
+		if p.curTok.Type == TokenComma {
 			p.nextToken()
+			continue
+		}
 
-			// Check if this is a literal hint (LABEL = value, etc.)
-			if p.curTok.Type == TokenEquals {
-				p.nextToken() // consume =
-				value, err := p.parseScalarExpression()
-				if err != nil {
-					return nil, err
-				}
-				hints = append(hints, &ast.LiteralOptimizerHint{
-					HintKind: hintKind,
-					Value:    value,
-				})
-			} else {
-				hints = append(hints, &ast.OptimizerHint{HintKind: hintKind})
+		hint, err := p.parseOptimizerHint()
+		if err != nil {
+			return nil, err
+		}
+		if hint != nil {
+			hints = append(hints, hint)
+		}
+	}
+
+	// Consume )
+	if p.curTok.Type == TokenRParen {
+		p.nextToken()
+	}
+
+	return hints, nil
+}
+
+func (p *Parser) parseOptimizerHint() (ast.OptimizerHintBase, error) {
+	// Handle both identifiers and keywords that can appear as optimizer hints
+	// USE is a keyword (TokenUse), so we need to handle it specially
+	if p.curTok.Type == TokenUse {
+		p.nextToken() // consume USE
+		if p.curTok.Type == TokenIdent && strings.ToUpper(p.curTok.Literal) == "PLAN" {
+			p.nextToken() // consume PLAN
+			value, err := p.parseScalarExpression()
+			if err != nil {
+				return nil, err
 			}
-		} else if p.curTok.Type == TokenComma {
+			return &ast.LiteralOptimizerHint{HintKind: "UsePlan", Value: value}, nil
+		}
+		return &ast.OptimizerHint{HintKind: "Use"}, nil
+	}
+
+	if p.curTok.Type != TokenIdent && p.curTok.Type != TokenLabel {
+		// Skip unknown tokens to avoid infinite loop
+		p.nextToken()
+		return nil, nil
+	}
+
+	upper := strings.ToUpper(p.curTok.Literal)
+
+	switch upper {
+	case "PARAMETERIZATION":
+		p.nextToken() // consume PARAMETERIZATION
+		if p.curTok.Type == TokenIdent {
+			subUpper := strings.ToUpper(p.curTok.Literal)
 			p.nextToken()
-		} else {
+			if subUpper == "SIMPLE" {
+				return &ast.OptimizerHint{HintKind: "ParameterizationSimple"}, nil
+			} else if subUpper == "FORCED" {
+				return &ast.OptimizerHint{HintKind: "ParameterizationForced"}, nil
+			}
+		}
+		return &ast.OptimizerHint{HintKind: "Parameterization"}, nil
+
+	case "MAXRECURSION":
+		p.nextToken() // consume MAXRECURSION
+		value, err := p.parseScalarExpression()
+		if err != nil {
+			return nil, err
+		}
+		return &ast.LiteralOptimizerHint{HintKind: "MaxRecursion", Value: value}, nil
+
+	case "OPTIMIZE":
+		p.nextToken() // consume OPTIMIZE
+		if p.curTok.Type == TokenIdent {
+			subUpper := strings.ToUpper(p.curTok.Literal)
+			if subUpper == "FOR" {
+				p.nextToken() // consume FOR
+				return p.parseOptimizeForHint()
+			} else if subUpper == "CORRELATED" {
+				p.nextToken() // consume CORRELATED
+				if p.curTok.Type == TokenIdent && strings.ToUpper(p.curTok.Literal) == "UNION" {
+					p.nextToken() // consume UNION
+					if p.curTok.Type == TokenIdent && strings.ToUpper(p.curTok.Literal) == "ALL" {
+						p.nextToken() // consume ALL
+					}
+				}
+				return &ast.OptimizerHint{HintKind: "OptimizeCorrelatedUnionAll"}, nil
+			}
+		}
+		return &ast.OptimizerHint{HintKind: "Optimize"}, nil
+
+	case "CHECKCONSTRAINTS":
+		p.nextToken() // consume CHECKCONSTRAINTS
+		if p.curTok.Type == TokenIdent && strings.ToUpper(p.curTok.Literal) == "PLAN" {
+			p.nextToken() // consume PLAN
+			return &ast.OptimizerHint{HintKind: "CheckConstraintsPlan"}, nil
+		}
+		return &ast.OptimizerHint{HintKind: "CheckConstraints"}, nil
+
+	case "LABEL":
+		p.nextToken() // consume LABEL
+		if p.curTok.Type == TokenEquals {
+			p.nextToken() // consume =
+			value, err := p.parseScalarExpression()
+			if err != nil {
+				return nil, err
+			}
+			return &ast.LiteralOptimizerHint{HintKind: "Label", Value: value}, nil
+		}
+		return &ast.OptimizerHint{HintKind: "Label"}, nil
+
+	case "MAX_GRANT_PERCENT":
+		p.nextToken() // consume MAX_GRANT_PERCENT
+		if p.curTok.Type == TokenEquals {
+			p.nextToken() // consume =
+			value, err := p.parseScalarExpression()
+			if err != nil {
+				return nil, err
+			}
+			return &ast.LiteralOptimizerHint{HintKind: "MaxGrantPercent", Value: value}, nil
+		}
+		return &ast.OptimizerHint{HintKind: "MaxGrantPercent"}, nil
+
+	case "MIN_GRANT_PERCENT":
+		p.nextToken() // consume MIN_GRANT_PERCENT
+		if p.curTok.Type == TokenEquals {
+			p.nextToken() // consume =
+			value, err := p.parseScalarExpression()
+			if err != nil {
+				return nil, err
+			}
+			return &ast.LiteralOptimizerHint{HintKind: "MinGrantPercent", Value: value}, nil
+		}
+		return &ast.OptimizerHint{HintKind: "MinGrantPercent"}, nil
+
+	case "FAST":
+		p.nextToken() // consume FAST
+		// FAST can take a numeric argument
+		if p.curTok.Type == TokenNumber {
+			value, err := p.parseScalarExpression()
+			if err != nil {
+				return nil, err
+			}
+			return &ast.LiteralOptimizerHint{HintKind: "Fast", Value: value}, nil
+		}
+		return &ast.OptimizerHint{HintKind: "Fast"}, nil
+
+	default:
+		// Handle generic hints
+		hintKind := convertHintKind(p.curTok.Literal)
+		p.nextToken()
+
+		// Check if this is a literal hint (LABEL = value, etc.)
+		if p.curTok.Type == TokenEquals {
+			p.nextToken() // consume =
+			value, err := p.parseScalarExpression()
+			if err != nil {
+				return nil, err
+			}
+			return &ast.LiteralOptimizerHint{HintKind: hintKind, Value: value}, nil
+		}
+		return &ast.OptimizerHint{HintKind: hintKind}, nil
+	}
+}
+
+func (p *Parser) parseOptimizeForHint() (ast.OptimizerHintBase, error) {
+	hint := &ast.OptimizeForOptimizerHint{
+		HintKind:     "OptimizeFor",
+		IsForUnknown: false,
+	}
+
+	// Check for UNKNOWN
+	if p.curTok.Type == TokenIdent && strings.ToUpper(p.curTok.Literal) == "UNKNOWN" {
+		p.nextToken()
+		hint.IsForUnknown = true
+		return hint, nil
+	}
+
+	// Expect (
+	if p.curTok.Type != TokenLParen {
+		return nil, fmt.Errorf("expected ( after OPTIMIZE FOR, got %s", p.curTok.Literal)
+	}
+	p.nextToken()
+
+	// Parse variable-value pairs
+	for p.curTok.Type != TokenRParen && p.curTok.Type != TokenEOF {
+		if p.curTok.Type == TokenComma {
 			p.nextToken()
+			continue
+		}
+
+		pair, err := p.parseVariableValuePair()
+		if err != nil {
+			return nil, err
+		}
+		if pair != nil {
+			hint.Pairs = append(hint.Pairs, pair)
 		}
 	}
 
@@ -1313,7 +1485,43 @@ func (p *Parser) parseOptionClause() ([]ast.OptimizerHintBase, error) {
 		p.nextToken()
 	}
 
-	return hints, nil
+	return hint, nil
+}
+
+func (p *Parser) parseVariableValuePair() (*ast.VariableValuePair, error) {
+	// Expect @variable (variables are TokenIdent starting with @)
+	if p.curTok.Type != TokenIdent || !strings.HasPrefix(p.curTok.Literal, "@") {
+		return nil, nil
+	}
+
+	pair := &ast.VariableValuePair{
+		Variable: &ast.VariableReference{
+			Name: p.curTok.Literal,
+		},
+		IsForUnknown: false,
+	}
+	p.nextToken()
+
+	// Expect =
+	if p.curTok.Type != TokenEquals {
+		// Could be UNKNOWN
+		if p.curTok.Type == TokenIdent && strings.ToUpper(p.curTok.Literal) == "UNKNOWN" {
+			p.nextToken()
+			pair.IsForUnknown = true
+			return pair, nil
+		}
+		return nil, fmt.Errorf("expected = after variable, got %s", p.curTok.Literal)
+	}
+	p.nextToken() // consume =
+
+	// Parse the value
+	value, err := p.parseScalarExpression()
+	if err != nil {
+		return nil, err
+	}
+	pair.Value = value
+
+	return pair, nil
 }
 
 // convertHintKind converts hint identifiers to their canonical names

parser/testdata/Baselines90_OptimizerHintsTests90/metadata.json

@@ -1 +1 @@
-{"skip": true}
+{"skip": false}

parser/testdata/OptimizerHintsTests90/metadata.json

@@ -1 +1 @@
-{"skip": true}
+{"skip": false}

skipped_tests_by_size.txt

@@ -37,7 +37,6 @@
 273 Baselines110_ServerRoleStatementTests
 275 BaselinesCommon_TSqlParserTestScript2
 277 Baselines170_VectorTypeSecondParameterTests170
-277 Baselines90_OptimizerHintsTests90
 278 Baselines160_AlterFunctionJsonObjectTests160
 278 BaselinesFabricDW_CloneTableTestsFabricDW
 278 CloneTableTestsFabricDW
@@ -46,7 +45,6 @@
 286 SecurityStatement130Tests
 288 BaselinesCommon_TSqlParserTestScript3
 290 DumpLoadStatement90Tests
-290 OptimizerHintsTests90
 291 Baselines90_DumpLoadStatementTests
 292 Baselines90_CreateTypeStatementTests
 298 ExecuteStatementTests90