Validate query name (#126)

* tests: Fix failing tests for wrapped errors

* gen: fix import error

* parser: Validate query names

Query names must be valid Go identifiers

Author: kyleconroy

internal/dinosql/gen.go

@@ -174,18 +174,22 @@ func (r Result) UsesArrays() bool {
 	return false
 }
 
-func (r Result) Imports(filename string) [][]string {
-	if filename == "db.go" {
-		return [][]string{
-			[]string{"context", "database/sql", "fmt"},
+func (r Result) Imports(settings PackageSettings) func(string) [][]string {
+	return func(filename string) [][]string {
+		if filename == "db.go" {
+			imps := []string{"context", "database/sql"}
+			if settings.EmitPreparedQueries {
+				imps = append(imps, "fmt")
+			}
+			return [][]string{imps}
 		}
-	}
 
-	if filename == "models.go" {
-		return r.ModelImports()
-	}
+		if filename == "models.go" {
+			return r.ModelImports()
+		}
 
-	return r.QueryImports(filename)
+		return r.QueryImports(filename)
+	}
 }
 
 func (r Result) ModelImports() [][]string {
@@ -950,7 +954,7 @@ func lowerTitle(s string) string {
 func Generate(r *Result, global GenerateSettings, settings PackageSettings) (map[string]string, error) {
 	funcMap := template.FuncMap{
 		"lowerTitle": lowerTitle,
-		"imports":    r.Imports,
+		"imports":    r.Imports(settings),
 	}
 
 	pkg := settings.Name

internal/dinosql/parser.go

@@ -288,7 +288,25 @@ func rangeVars(root nodes.Node) []nodes.RangeVar {
 	return vars
 }
 
-// TODO: Validate metadata
+// A query name must be a valid Go identifier
+//
+// https://golang.org/ref/spec#Identifiers
+func validateQueryName(name string) error {
+	if len(name) == 0 {
+		return fmt.Errorf("invalid query name: %q", name)
+	}
+	for i, c := range name {
+		isLetter := unicode.IsLetter(c) || c == '_'
+		isDigit := unicode.IsDigit(c)
+		if i == 0 && !isLetter {
+			return fmt.Errorf("invalid query name: %q", name)
+		} else if !(isLetter || isDigit) {
+			return fmt.Errorf("invalid query name: %q", name)
+		}
+	}
+	return nil
+}
+
 func parseMetadata(t string) (string, string, error) {
 	for _, line := range strings.Split(t, "\n") {
 		if !strings.HasPrefix(line, "-- name:") {
@@ -301,13 +319,17 @@ func parseMetadata(t string) (string, string, error) {
 		if len(part) != 4 {
 			return "", "", fmt.Errorf("invalid query comment: %s", line)
 		}
+		queryName := part[2]
 		queryType := strings.TrimSpace(part[3])
 		switch queryType {
 		case ":one", ":many", ":exec", ":execrows":
 		default:
 			return "", "", fmt.Errorf("invalid query type: %s", queryType)
 		}
-		return part[2], strings.TrimSpace(part[3]), nil
+		if err := validateQueryName(queryName); err != nil {
+			return "", "", err
+		}
+		return queryName, queryType, nil
 	}
 	return "", "", nil
 }

internal/dinosql/parser_test.go

@@ -205,3 +205,18 @@ func TestCompile(t *testing.T) {
 		t.Errorf("%s: %s", err, string(output))
 	}
 }
+
+func TestParseMetadata(t *testing.T) {
+	for _, query := range []string{
+		`-- name: CreateFoo, :one`,
+		`-- name: 9Foo_, :one`,
+		`-- name: CreateFoo :two`,
+		`-- name: CreateFoo`,
+		`-- name: CreateFoo :one something`,
+		`-- name: `,
+	} {
+		if _, _, err := parseMetadata(query); err == nil {
+			t.Errorf("expected invalid metadata: %q", query)
+		}
+	}
+}