[Filter] Fix syntax error that occurred for some queries using subselects or QUALIFY

internal/formatter.go

@@ -3,6 +3,7 @@ package internal
 import (
 	"context"
 	"fmt"
+	"regexp"
 	"strings"
 
 	"github.com/goccy/go-json"
@@ -686,7 +687,10 @@ func (n *ColumnHolderNode) FormatSQL(ctx context.Context) (string, error) {
 	return "", nil
 }
 
+var tokensAfterFromClause = [...]string{"WHERE", "GROUP BY", "HAVING", "QUALIFY", "WINDOW", "ORDER BY", "COLLATE"}
+
 func (n *FilterScanNode) FormatSQL(ctx context.Context) (string, error) {
+	removeExpressions := regexp.MustCompile(`\(.+?\)`)
 	if n.node == nil {
 		return "", nil
 	}
@@ -704,8 +708,17 @@ func (n *FilterScanNode) FormatSQL(ctx context.Context) (string, error) {
 			return fmt.Sprintf("%s HAVING %s", input, filter), nil
 		}
 	}
-	if strings.Contains(input, "WHERE") && input[len(input)-1] != ')' {
-		// expected to qualify clause
+	currentQuery := string(removeExpressions.ReplaceAllString(input, ""))
+
+	// Qualify the statement if the input is not wrapped in parens
+	queryWrappedInParens := currentQuery == ""
+	containsTokens := false
+	// or the input contains a token that would result in a syntax error
+	for _, token := range tokensAfterFromClause {
+		containsTokens = containsTokens || strings.Contains(currentQuery, token)
+	}
+
+	if !queryWrappedInParens && containsTokens {
 		return fmt.Sprintf("( %s ) WHERE %s", input, filter), nil
 	}
 	return fmt.Sprintf("%s WHERE %s", input, filter), nil

query_test.go

@@ -2781,6 +2781,35 @@ FROM Produce WHERE Produce.category = 'vegetable' QUALIFY rank <= 3`,
 				{"cabbage", int64(3)},
 			},
 		},
+		// Regression test goccy/go-zetasqlite#150
+		{
+			name: "qualify group",
+			query: `
+				WITH produce AS (
+					SELECT 'kale' AS item, 23 AS purchases
+				)
+				SELECT item, sum(purchases)
+				FROM produce
+				GROUP BY item
+				QUALIFY ROW_NUMBER() OVER (PARTITION BY item ORDER BY item) = 1
+			`,
+			expectedRows: [][]interface{}{{"kale", int64(23)}},
+		},
+		// Regression test goccy/go-zetasqlite#147
+		{
+			name: "subselect qualifier",
+			query: `
+				WITH produce AS (SELECT 'banana' AS item, 3 AS purchases),
+				toks AS (
+					SELECT item FROM (
+						SELECT * FROM produce
+						WHERE item = 'banana'
+					) sub
+					WHERE purchases = 3
+				)
+				SELECT * FROM toks;`,
+			expectedRows: [][]interface{}{{"banana"}},
+		},
 		{
 			name: "qualify direct",
 			query: `