only allow alias or table name in sqlc.embed()

Author: nickjackson

internal/sql/rewrite/embeds.go

@@ -2,7 +2,6 @@ package rewrite
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
 	"github.com/kyleconroy/sqlc/internal/sql/astutils"
@@ -50,29 +49,20 @@ func Embeds(raw *ast.RawStmt) (*ast.RawStmt, EmbedSet) {
 				return false
 			}
 
-			sw := &stringWalker{}
-			astutils.Walk(sw, fun.Args)
-			str := strings.Join(sw.Parts, ".")
-
-			tableName, err := parseTable(sw)
-			if err != nil {
-				return false
-			}
+			param, _ := flatten(fun.Args)
 
 			node := &ast.ColumnRef{
 				Fields: &ast.List{
-					Items: []ast.Node{},
+					Items: []ast.Node{
+						&ast.String{Str: param},
+						&ast.A_Star{},
+					},
 				},
 			}
 
-			for _, s := range sw.Parts {
-				node.Fields.Items = append(node.Fields.Items, &ast.String{Str: s})
-			}
-			node.Fields.Items = append(node.Fields.Items, &ast.A_Star{})
-
 			embeds = append(embeds, &Embed{
-				Table: tableName,
-				param: str,
+				Table: &ast.TableName{Name: param},
+				param: param,
 				Node:  node,
 			})
 
@@ -99,27 +89,3 @@ func isEmbed(node ast.Node) bool {
 	isValid := call.Func.Schema == "sqlc" && call.Func.Name == "embed"
 	return isValid
 }
-
-func parseTable(sw *stringWalker) (*ast.TableName, error) {
-	parts := sw.Parts
-
-	switch len(parts) {
-	case 1:
-		return &ast.TableName{
-			Name: parts[0],
-		}, nil
-	case 2:
-		return &ast.TableName{
-			Schema: parts[0],
-			Name:   parts[1],
-		}, nil
-	case 3:
-		return &ast.TableName{
-			Catalog: parts[0],
-			Schema:  parts[1],
-			Name:    parts[2],
-		}, nil
-	default:
-		return nil, fmt.Errorf("invalid table name")
-	}
-}

internal/sql/rewrite/parameters.go

@@ -19,7 +19,6 @@ func flatten(root ast.Node) (string, bool) {
 
 type stringWalker struct {
 	String  string
-	Parts   []string
 	IsConst bool
 }
 
@@ -29,7 +28,6 @@ func (s *stringWalker) Visit(node ast.Node) astutils.Visitor {
 	}
 	if n, ok := node.(*ast.String); ok {
 		s.String += n.Str
-		s.Parts = append(s.Parts, n.Str)
 	}
 	return s
 }