compiler: Port bottom of parseQuery (#510)

Author: kyleconroy

internal/compiler/parse.go

@@ -2,6 +2,7 @@ package compiler
 
 import (
 	"errors"
+	"fmt"
 	"strings"
 
 	"github.com/kyleconroy/sqlc/internal/metadata"
@@ -13,6 +14,10 @@ import (
 )
 
 type Query struct {
+	SQL      string
+	Name     string
+	Cmd      string // TODO: Pick a better name. One of: one, many, exec, execrows
+	Comments []string
 }
 
 var ErrUnsupportedStatementType = errors.New("parseQuery: unsupported statement type")
@@ -59,5 +64,30 @@ func parseQuery(p Parser, c *catalog.Catalog, stmt ast.Node, src string, rewrite
 		return nil, err
 	}
 
-	return &Query{}, nil
+	// TODO: Then a miracle occurs
+
+	var edits []source.Edit
+	expanded, err := source.Mutate(rawSQL, edits)
+	if err != nil {
+		return nil, err
+	}
+
+	// If the query string was edited, make sure the syntax is valid
+	if expanded != rawSQL {
+		if _, err := p.Parse(strings.NewReader(expanded)); err != nil {
+			return nil, fmt.Errorf("edited query syntax is invalid: %w", err)
+		}
+	}
+
+	trimmed, comments, err := source.StripComments(expanded)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Query{
+		Cmd:      cmd,
+		Comments: comments,
+		Name:     name,
+		SQL:      trimmed,
+	}, nil
 }

internal/dinosql/parser.go

@@ -1,7 +1,6 @@
 package dinosql
 
 import (
-	"bufio"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -311,7 +310,7 @@ func parseQuery(c core.Catalog, stmt nodes.Node, src string, rewriteParameters b
 	}
 	edits = append(edits, expandEdits...)
 
-	expanded, err := editQuery(rawSQL, edits)
+	expanded, err := source.Mutate(rawSQL, edits)
 	if err != nil {
 		return nil, err
 	}
@@ -323,7 +322,7 @@ func parseQuery(c core.Catalog, stmt nodes.Node, src string, rewriteParameters b
 		}
 	}
 
-	trimmed, comments, err := stripComments(strings.TrimSpace(expanded))
+	trimmed, comments, err := source.StripComments(expanded)
 	if err != nil {
 		return nil, err
 	}
@@ -350,22 +349,6 @@ func rewriteNumberedParameters(refs []paramRef, raw nodes.RawStmt, sql string) (
 	return edits, nil
 }
 
-func stripComments(sql string) (string, []string, error) {
-	s := bufio.NewScanner(strings.NewReader(sql))
-	var lines, comments []string
-	for s.Scan() {
-		if strings.HasPrefix(s.Text(), "-- name:") {
-			continue
-		}
-		if strings.HasPrefix(s.Text(), "--") {
-			comments = append(comments, strings.TrimPrefix(s.Text(), "--"))
-			continue
-		}
-		lines = append(lines, s.Text())
-	}
-	return strings.Join(lines, "\n"), comments, s.Err()
-}
-
 func expand(qc *QueryCatalog, raw nodes.RawStmt) ([]source.Edit, error) {
 	list := ast.Search(raw, func(node nodes.Node) bool {
 		switch node.(type) {
@@ -486,33 +469,6 @@ func expandStmt(qc *QueryCatalog, raw nodes.RawStmt, node nodes.Node) ([]source.
 	return edits, nil
 }
 
-func editQuery(raw string, a []source.Edit) (string, error) {
-	if len(a) == 0 {
-		return raw, nil
-	}
-	sort.Slice(a, func(i, j int) bool { return a[i].Location > a[j].Location })
-	s := raw
-	for _, edit := range a {
-		start := edit.Location
-		if start > len(s) {
-			return "", fmt.Errorf("edit start location is out of bounds")
-		}
-		if len(edit.New) <= 0 {
-			return "", fmt.Errorf("empty edit contents")
-		}
-		if len(edit.Old) <= 0 {
-			return "", fmt.Errorf("empty edit contents")
-		}
-		stop := edit.Location + len(edit.Old) - 1 // Assumes edit.New is non-empty
-		if stop < len(s) {
-			s = s[:start] + edit.New + s[stop+1:]
-		} else {
-			s = s[:start] + edit.New
-		}
-	}
-	return s, nil
-}
-
 type QueryCatalog struct {
 	catalog core.Catalog
 	ctes    map[string]core.Table

internal/dinosql/parser_test.go

@@ -147,7 +147,7 @@ func TestRewriteParameters(t *testing.T) {
 			if err != nil {
 				t.Error(err)
 			}
-			rewritten, err := editQuery(q.orig, edits)
+			rewritten, err := source.Mutate(q.orig, edits)
 			if err != nil {
 				t.Error(err)
 			}
@@ -167,7 +167,7 @@ func TestExpand(t *testing.T) {
 		{10, "*", "a, b"},
 		{13, "foo.*", "foo.a, foo.b"},
 	}
-	actual, err := editQuery(raw, edits)
+	actual, err := source.Mutate(raw, edits)
 	if err != nil {
 		t.Error(err)
 	}

internal/source/code.go

@@ -1,6 +1,12 @@
 package source
 
-import "unicode"
+import (
+	"bufio"
+	"fmt"
+	"sort"
+	"strings"
+	"unicode"
+)
 
 type Edit struct {
 	Location int
@@ -43,3 +49,46 @@ func Pluck(source string, location, length int) (string, error) {
 	tail := location + length
 	return source[head:tail], nil
 }
+
+func Mutate(raw string, a []Edit) (string, error) {
+	if len(a) == 0 {
+		return raw, nil
+	}
+	sort.Slice(a, func(i, j int) bool { return a[i].Location > a[j].Location })
+	s := raw
+	for _, edit := range a {
+		start := edit.Location
+		if start > len(s) {
+			return "", fmt.Errorf("edit start location is out of bounds")
+		}
+		if len(edit.New) <= 0 {
+			return "", fmt.Errorf("empty edit contents")
+		}
+		if len(edit.Old) <= 0 {
+			return "", fmt.Errorf("empty edit contents")
+		}
+		stop := edit.Location + len(edit.Old) - 1 // Assumes edit.New is non-empty
+		if stop < len(s) {
+			s = s[:start] + edit.New + s[stop+1:]
+		} else {
+			s = s[:start] + edit.New
+		}
+	}
+	return s, nil
+}
+
+func StripComments(sql string) (string, []string, error) {
+	s := bufio.NewScanner(strings.NewReader(strings.TrimSpace(sql)))
+	var lines, comments []string
+	for s.Scan() {
+		if strings.HasPrefix(s.Text(), "-- name:") {
+			continue
+		}
+		if strings.HasPrefix(s.Text(), "--") {
+			comments = append(comments, strings.TrimPrefix(s.Text(), "--"))
+			continue
+		}
+		lines = append(lines, s.Text())
+	}
+	return strings.Join(lines, "\n"), comments, s.Err()
+}