Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix QueryStmtNode output; wrap union statements in parentheses #205

Merged
merged 1 commit into from
Apr 6, 2024
Merged

Fix QueryStmtNode output; wrap union statements in parentheses #205

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 42 additions & 8 deletions internal/formatter.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -726,12 +726,6 @@ func (n *FilterScanNode) FormatSQL(ctx context.Context) (string, error) {
if err != nil {
return "", err
}
nodeMap := nodeMapFromContext(ctx)
for _, node := range nodeMap.FindNodeFromResolvedNode(n.node) {
if _, ok := node.(*parsed_ast.HavingNode); ok {
return fmt.Sprintf("%s HAVING %s", input, filter), nil
}
}
currentQuery := removeExpressions.ReplaceAllString(input, "")

// Qualify the statement if the input is not wrapped in parens
Expand Down Expand Up @@ -910,11 +904,27 @@ func (n *SetOperationScanNode) FormatSQL(ctx context.Context) (string, error) {
}
var queries []string
for _, item := range n.node.InputItemList() {
var outputColumns []string
for _, outputColumn := range item.OutputColumnList() {
outputColumns = append(outputColumns, fmt.Sprintf("`%s`", uniqueColumnName(ctx, outputColumn)))
}
query, err := newNode(item).FormatSQL(ctx)
if err != nil {
return "", err
}
queries = append(queries, query)

formattedInput, err := formatInput(query)
if err != nil {
return "", err
}

queries = append(
queries,
fmt.Sprintf("SELECT %s %s",
strings.Join(outputColumns, ", "),
formattedInput,
),
)
}
columnMaps := []string{}
if len(n.node.InputItemList()) != 0 {
Expand Down Expand Up @@ -1283,11 +1293,35 @@ func (n *ExplainStmtNode) FormatSQL(ctx context.Context) (string, error) {
return "", nil
}

// FormatSQL Formats the outermost query statement that runs and produces rows of output, like a SELECT
// The node's `OutputColumnList()` gives user-visible column names that should be returned. There may be duplicate names,
// and multiple output columns may reference the same column from `Query()`
// https://github.com/google/zetasql/blob/master/docs/resolved_ast.md#ResolvedQueryStmt
func (n *QueryStmtNode) FormatSQL(ctx context.Context) (string, error) {
if n.node == nil {
return "", nil
}
return newNode(n.node.Query()).FormatSQL(ctx)
input, err := newNode(n.node.Query()).FormatSQL(ctx)
if err != nil {
return "", err
}

var columns []string
for _, outputColumnNode := range n.node.OutputColumnList() {
columns = append(
columns,
fmt.Sprintf("`%s` AS `%s`",
uniqueColumnName(ctx, outputColumnNode.Column()),
outputColumnNode.Name(),
),
)
}

return fmt.Sprintf(
"SELECT %s FROM (%s)",
strings.Join(columns, ", "),
input,
), nil
}

func (n *CreateDatabaseStmtNode) FormatSQL(ctx context.Context) (string, error) {
Expand Down
21 changes: 21 additions & 0 deletions query_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,27 @@ func TestQuery(t *testing.T) {
expectedRows [][]interface{}
expectedErr string
}{
// Regression test for https://github.com/goccy/go-zetasqlite/issues/191
{
name: "distinct union",
query: `WITH toks AS (SELECT true AS x, 1 AS y)
SELECT DISTINCT x, x as y FROM toks`,
expectedRows: [][]interface{}{{true, true}},
},
{
name: "with scan union all",
query: `(WITH toks AS (SELECT 1 AS x) SELECT x FROM toks)
UNION ALL
(WITH toks2 AS (SELECT 2 AS x) SELECT x FROM toks2)`,
expectedRows: [][]interface{}{{int64(1)}, {int64(2)}},
},
{
name: "having with union all",
query: `(WITH toks AS (SELECT 1 AS x) SELECT COUNT(x) AS total_rows FROM toks WHERE x > 0 HAVING total_rows >= 0)
UNION ALL
(WITH toks2 AS (SELECT 2 AS x) SELECT COUNT(x) AS total_rows FROM toks2 WHERE x > 0 HAVING total_rows >= 0)`,
expectedRows: [][]interface{}{{int64(1)}, {int64(1)}},
},
// priority 2 operator
{
name: "unary plus operator",
Expand Down
Loading