Skip to content

Commit

Permalink
[Filter] Fix syntax error that occurred for some queries using subsel…
Browse files Browse the repository at this point in the history
…ects or QUALIFY
  • Loading branch information
ohaibbq committed Feb 1, 2024
1 parent 961ce06 commit fa618fe
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 2 deletions.
17 changes: 15 additions & 2 deletions internal/formatter.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package internal
import (
"context"
"fmt"
"regexp"
"strings"

"github.com/goccy/go-json"
Expand Down Expand Up @@ -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
}
Expand All @@ -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
Expand Down
29 changes: 29 additions & 0 deletions query_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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: `
Expand Down

0 comments on commit fa618fe

Please sign in to comment.