Validate metadata from comment string (#117)

Author: kyleconroy

internal/dinosql/parser.go

@@ -291,6 +291,18 @@ func parseMetadata(t string) (string, string, error) {
 			continue
 		}
 		part := strings.Split(line, " ")
+		if len(part) == 2 {
+			return "", "", fmt.Errorf("missing query type [':one', ':many', ':exec', ':execrows']: %s", line)
+		}
+		if len(part) != 4 {
+			return "", "", fmt.Errorf("invalid query comment: %s", line)
+		}
+		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
 	}
 	return "", "", nil

internal/dinosql/query_test.go

@@ -781,3 +781,39 @@ func TestStarWalker(t *testing.T) {
 		})
 	}
 }
+
+func TestInvalidQueries(t *testing.T) {
+	for i, tc := range []struct {
+		stmt string
+	}{
+		{
+			`
+			CREATE TABLE foo (id text not null);
+			-- name: ListFoos
+			SELECT id FROM foo;
+			`,
+		},
+		{
+			`
+			CREATE TABLE foo (id text not null);
+			-- name: ListFoos :one :many
+			SELECT id FROM foo;
+			`,
+		},
+		{
+			`
+			CREATE TABLE foo (id text not null);
+			-- name: ListFoos :two
+			SELECT id FROM foo;
+			`,
+		},
+	} {
+		test := tc
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			_, err := parseSQL(test.stmt)
+			if err == nil {
+				t.Errorf("expected err, got nil")
+			}
+		})
+	}
+}