internal/dinosql: Implement robust expansion

Exapnds `SELECT *` into the correct columns, no matter the location

Author: kyleconroy

internal/dinosql/parser.go

@@ -289,7 +289,7 @@ func lineno(source string, head int) (int, int) {
 func pluckQuery(source string, n nodes.RawStmt) (string, error) {
 	head := n.StmtLocation
 	tail := n.StmtLocation + n.StmtLen
-	return strings.TrimSpace(source[head:tail]), nil
+	return source[head:tail], nil
 }
 
 func rangeVars(root nodes.Node) []nodes.RangeVar {
@@ -403,7 +403,7 @@ func parseQuery(c core.Catalog, stmt nodes.Node, source string) (*Query, error)
 	if err := validateFuncCall(&c, raw); err != nil {
 		return nil, err
 	}
-	name, cmd, err := parseMetadata(rawSQL)
+	name, cmd, err := parseMetadata(strings.TrimSpace(rawSQL))
 	if err != nil {
 		return nil, err
 	}
@@ -422,19 +422,28 @@ func parseQuery(c core.Catalog, stmt nodes.Node, source string) (*Query, error)
 		return nil, err
 	}
 
-	trimmed, comments, err := stripComments(rawSQL)
+	// TODO: Limit calls to sourceTables
+	tables, err := sourceTables(c, raw.Stmt)
+	if err != nil {
+		return nil, err
+	}
+	expanded, err := expand(raw, tables, rawSQL)
+	if err != nil {
+		return nil, err
+	}
+
+	trimmed, comments, err := stripComments(strings.TrimSpace(expanded))
 	if err != nil {
 		return nil, err
 	}
 
 	return &Query{
-		Cmd:       cmd,
-		Comments:  comments,
-		Name:      name,
-		Params:    params,
-		Columns:   cols,
-		SQL:       trimmed,
-		NeedsEdit: needsEdit(stmt),
+		Cmd:      cmd,
+		Comments: comments,
+		Name:     name,
+		Params:   params,
+		Columns:  cols,
+		SQL:      trimmed,
 	}, nil
 }
 
@@ -454,6 +463,86 @@ func stripComments(sql string) (string, []string, error) {
 	return strings.Join(lines, "\n"), comments, s.Err()
 }
 
+type edit struct {
+	Location int
+	Old      string
+	New      string
+}
+
+func expand(raw nodes.RawStmt, tables []core.Table, sql string) (string, error) {
+	list := search(raw, func(node nodes.Node) bool {
+		res, ok := node.(nodes.ResTarget)
+		if !ok {
+			return false
+		}
+		ref, ok := res.Val.(nodes.ColumnRef)
+		if !ok {
+			return false
+		}
+		return HasStarRef(ref)
+	})
+	if len(list.Items) == 0 {
+		return sql, nil
+	}
+	var edits []edit
+	for _, item := range list.Items {
+		res := item.(nodes.ResTarget)
+		ref := res.Val.(nodes.ColumnRef)
+		var parts, cols []string
+		for _, f := range ref.Fields.Items {
+			switch field := f.(type) {
+			case nodes.String:
+				parts = append(parts, field.Str)
+			case nodes.A_Star:
+				parts = append(parts, "*")
+			default:
+				return "", fmt.Errorf("unknown field in ColumnRef: %T", f)
+			}
+		}
+		for _, t := range tables {
+			scope := join(ref.Fields, ".")
+			if scope != "" && scope != t.Name {
+				continue
+			}
+			for _, c := range t.Columns {
+				cname := c.Name
+				if res.Name != nil {
+					cname = *res.Name
+				}
+				if scope != "" {
+					cname = scope + "." + cname
+				}
+				cols = append(cols, cname)
+			}
+		}
+		edits = append(edits, edit{
+			Location: res.Location - raw.StmtLocation,
+			Old:      strings.Join(parts, "."),
+			New:      strings.Join(cols, ", "),
+		})
+	}
+	return editQuery(sql, edits)
+}
+
+func editQuery(raw string, a []edit) (string, error) {
+	sort.Slice(a, func(i, j int) bool { return a[i].Location > a[j].Location })
+	// TODO: Check bounds
+	s := raw
+	for _, edit := range a {
+		// fmt.Printf("edit q %q\n", s)
+		// fmt.Printf("edit e %q\n", edit)
+		start := edit.Location
+		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
+		}
+	}
+	// fmt.Printf("edit fixed %q\n", s)
+	return s, nil
+}
+
 type QueryCatalog struct {
 	catalog core.Catalog
 	ctes    map[string]core.Table
@@ -653,6 +742,7 @@ func outputColumns(c core.Catalog, node nodes.Node) ([]core.Column, error) {
 
 		case nodes.ColumnRef:
 			if HasStarRef(n) {
+				// TODO: This code is copied in func expand()
 				for _, t := range tables {
 					scope := join(n.Fields, ".")
 					if scope != "" && scope != t.Name {
@@ -916,24 +1006,6 @@ func findParameters(root nodes.Node) []paramRef {
 	return refs
 }
 
-type starWalker struct {
-	found bool
-}
-
-func (s *starWalker) Visit(node nodes.Node) Visitor {
-	if _, ok := node.(nodes.A_Star); ok {
-		s.found = true
-		return nil
-	}
-	return s
-}
-
-func needsEdit(root nodes.Node) bool {
-	v := &starWalker{}
-	Walk(v, root)
-	return v.found
-}
-
 type nodeSearch struct {
 	list  nodes.List
 	check func(nodes.Node) bool

internal/dinosql/parser_test.go

@@ -26,10 +26,10 @@ func TestPluck(t *testing.T) {
 	}
 
 	expected := []string{
-		"SELECT * FROM venue WHERE slug = $1 AND city = $2",
-		"SELECT * FROM venue WHERE slug = $1",
-		"SELECT * FROM venue LIMIT $1",
-		"SELECT * FROM venue OFFSET $1",
+		"\nSELECT * FROM venue WHERE slug = $1 AND city = $2",
+		"\nSELECT * FROM venue WHERE slug = $1",
+		"\nSELECT * FROM venue LIMIT $1",
+		"\nSELECT * FROM venue OFFSET $1",
 	}
 
 	for i, stmt := range tree.Statements {
@@ -220,3 +220,21 @@ func TestParseMetadata(t *testing.T) {
 		}
 	}
 }
+
+func TestExpand(t *testing.T) {
+	// pretend that foo has two columns, a and b
+	raw := `SELECT *, *, foo.* FROM foo`
+	expected := `SELECT a, b, a, b, foo.a, foo.b FROM foo`
+	edits := []edit{
+		{7, "*", "a, b"},
+		{10, "*", "a, b"},
+		{13, "foo.*", "foo.a, foo.b"},
+	}
+	actual, err := editQuery(raw, edits)
+	if err != nil {
+		t.Error(err)
+	}
+	if expected != actual {
+		t.Errorf("mismatch:\nexpected: %s\n  acutal: %s", expected, actual)
+	}
+}

internal/dinosql/query_test.go

@@ -7,7 +7,6 @@ import (
 
 	core "github.com/kyleconroy/sqlc/internal/pg"
 
-	"github.com/davecgh/go-spew/spew"
 	"github.com/google/go-cmp/cmp"
 	pg "github.com/lfittl/pg_query_go"
 )
@@ -26,8 +25,6 @@ func parseSQL(in string) (Query, error) {
 	if q == nil {
 		return Query{}, err
 	}
-	q.SQL = ""
-	q.NeedsEdit = false
 	return *q, err
 }
 
@@ -737,13 +734,63 @@ func TestQueries(t *testing.T) {
 				},
 			},
 		},
+		{
+			"star-expansion",
+			`
+			CREATE TABLE foo (a text, b text);
+			SELECT *, *, foo.* FROM foo;
+			`,
+			Query{
+				Columns: []core.Column{
+					{Name: "a", DataType: "text", Table: public("foo")},
+					{Name: "b", DataType: "text", Table: public("foo")},
+					{Name: "a", DataType: "text", Table: public("foo")},
+					{Name: "b", DataType: "text", Table: public("foo")},
+					{Name: "a", DataType: "text", Scope: "foo", Table: public("foo")},
+					{Name: "b", DataType: "text", Scope: "foo", Table: public("foo")},
+				},
+				SQL: "SELECT a, b, a, b, foo.a, foo.b FROM foo",
+			},
+		},
+		{
+			"star-expansion-subquery",
+			`
+			CREATE TABLE foo (a text, b text);
+			SELECT * FROM foo WHERE EXISTS (SELECT * FROM foo);
+			`,
+			Query{
+				Columns: []core.Column{
+					{Name: "a", DataType: "text", Table: public("foo")},
+					{Name: "b", DataType: "text", Table: public("foo")},
+				},
+				SQL: "SELECT a, b FROM foo WHERE EXISTS (SELECT a, b FROM foo)",
+			},
+		},
+		{
+			"star-expansion-cte",
+			`
+			CREATE TABLE foo (a text, b text);
+			CREATE TABLE bar (c text, d text);
+			WITH cte AS (SELECT * FROM foo) SELECT * FROM bar;
+			`,
+			Query{
+				Columns: []core.Column{
+					{Name: "c", DataType: "text", Table: public("bar")},
+					{Name: "d", DataType: "text", Table: public("bar")},
+				},
+				SQL: "WITH cte AS (SELECT a, b FROM foo) SELECT c, d FROM bar",
+			},
+		},
 	} {
 		test := tc
 		t.Run(test.name, func(t *testing.T) {
 			q, err := parseSQL(test.stmt)
 			if err != nil {
 				t.Fatal(err)
 			}
+			if test.query.SQL == "" {
+				q.SQL = ""
+			}
 			if diff := cmp.Diff(test.query, q); diff != "" {
 				t.Errorf("query mismatch: \n%s", diff)
 			}
@@ -765,6 +812,7 @@ func TestComparisonOperators(t *testing.T) {
 				t.Fatal(err)
 			}
 			expected := Query{
+				SQL: q.SQL,
 				Columns: []core.Column{
 					{Name: "", DataType: "bool", NotNull: true},
 				},
@@ -776,59 +824,6 @@ func TestComparisonOperators(t *testing.T) {
 	}
 }
 
-func TestStarWalker(t *testing.T) {
-	for i, tc := range []struct {
-		stmt     string
-		expected bool
-	}{
-		{
-			`
-			SELECT * FROM city ORDER BY name;
-			`,
-			true,
-		},
-		{
-			`
-			INSERT INTO city (
-				name,
-				slug
-			) VALUES (
-				$1,
-				$2
-			) RETURNING *;
-			`,
-			true,
-		},
-		{
-			`
-			UPDATE city SET name = $2 WHERE slug = $1;
-			`,
-			false,
-		},
-		{
-			`
-			UPDATE venue
-			SET name = $2
-			WHERE slug = $1
-			RETURNING *;
-			`,
-			true,
-		},
-	} {
-		test := tc
-		t.Run(strconv.Itoa(i), func(t *testing.T) {
-			tree, err := pg.Parse(test.stmt)
-			if err != nil {
-				t.Fatal(err)
-			}
-			if diff := cmp.Diff(test.expected, needsEdit(tree.Statements[0])); diff != "" {
-				spew.Dump(tree.Statements[0])
-				t.Errorf("query mismatch: \n%s", diff)
-			}
-		})
-	}
-}
-
 func TestInvalidQueries(t *testing.T) {
 	for i, tc := range []struct {
 		stmt string