cmd/schemagen: refactoring

- Replace log.Faltal with error wrapping in schemagen func
- Simplify tests, use temp dir and ioutil functions, remove boilerplate code
- In test use schemagen keyspace to avoid name conflict with examples
- Change template

README.md

@@ -141,42 +141,43 @@ package models
 
 import "github.com/scylladb/gocqlx/v2/table"
 
-var PlaylistsMetadata = table.Metadata{
-	Name: "playlists",
-	Columns: []string{
-		"album",
-		"artist",
-		"id",
-		"song_id",
-		"title",
-	},
-	PartKey: []string{
-		"id",
-	},
-	SortKey: []string{
-		"title",
-		"album",
-		"artist",
-	},
-}
-var PlaylistsTable = table.New(PlaylistsMetadata)
+// Table models.
+var (
+	Playlists = table.New(table.Metadata{
+		Name: "playlists",
+		Columns: []string{
+			"album",
+			"artist",
+			"id",
+			"song_id",
+			"title",
+		},
+		PartKey: []string{
+			"id",
+		},
+		SortKey: []string{
+			"title",
+			"album",
+			"artist",
+		},
+	})
 
-var SongsMetadata = table.Metadata{
-	Name: "songs",
-	Columns: []string{
-		"album",
-		"artist",
-		"data",
-		"id",
-		"tags",
-		"title",
-	},
-	PartKey: []string{
-		"id",
-	},
-	SortKey: []string{},
-}
-var SongsTable = table.New(SongsMetadata)
+	Songs = table.New(table.Metadata{
+		Name: "songs",
+		Columns: []string{
+			"album",
+			"artist",
+			"data",
+			"id",
+			"tags",
+			"title",
+		},
+		PartKey: []string{
+			"id",
+		},
+		SortKey: []string{},
+	})
+)
 ```
 
 ## Examples

cmd/schemagen/camelize.go

@@ -0,0 +1,43 @@
+// Copyright (C) 2017 ScyllaDB
+// Use of this source code is governed by a ALv2-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"fmt"
+	"unicode"
+)
+
+func camelize(s string) string {
+	buf := []byte(s)
+	out := make([]byte, 0, len(buf))
+	underscoreSeen := false
+
+	l := len(buf)
+	for i := 0; i < l; i++ {
+		if !(allowedBindRune(buf[i]) || buf[i] == '_') {
+			panic(fmt.Sprint("not allowed name ", s))
+		}
+
+		b := rune(buf[i])
+
+		if b == '_' {
+			underscoreSeen = true
+			continue
+		}
+
+		if (i == 0 || underscoreSeen) && unicode.IsLower(b) {
+			b = unicode.ToUpper(b)
+			underscoreSeen = false
+		}
+
+		out = append(out, byte(b))
+	}
+
+	return string(out)
+}
+
+func allowedBindRune(b byte) bool {
+	return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9')
+}

cmd/schemagen/camelize_test.go

@@ -0,0 +1,31 @@
+// Copyright (C) 2017 ScyllaDB
+// Use of this source code is governed by a ALv2-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import "testing"
+
+func TestCamelize(t *testing.T) {
+	tests := []struct {
+		input string
+		want  string
+	}{
+		{"hello", "Hello"},
+		{"_hello", "Hello"},
+		{"__hello", "Hello"},
+		{"hello_", "Hello"},
+		{"hello_world", "HelloWorld"},
+		{"hello__world", "HelloWorld"},
+		{"_hello_world", "HelloWorld"},
+		{"helloWorld", "HelloWorld"},
+		{"HelloWorld", "HelloWorld"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			if got := camelize(tt.input); got != tt.want {
+				t.Errorf("camelize() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

cmd/schemagen/keyspace.tmpl

@@ -4,27 +4,29 @@ package {{.PackageName}}
 
 import "github.com/scylladb/gocqlx/v2/table"
 
+// Table models.
+var (
 {{with .Tables}}
-    {{range .}}
-        {{$model_name := .Name | camelize}}
-        var {{$model_name}}Metadata = table.Metadata {
-            Name:    "{{.Name}}",
-            Columns: []string{
-                {{- range .OrderedColumns}}
-                    "{{.}}",
-                {{- end}}
-            },
-            PartKey: []string {
-                {{- range .PartitionKey}}
-                    "{{.Name}}",
-                {{- end}}
-            },
-            SortKey: []string{
-                {{- range .ClusteringColumns}}
-                    "{{.Name}}",
-                {{- end}}
-            },
-        }
-        var {{$model_name}}Table = table.New({{$model_name}}Metadata)
-    {{end}}
+{{range .}}
+	{{$model_name := .Name | camelize}}
+		{{$model_name}} = table.New(table.Metadata {
+			Name:	"{{.Name}}",
+			Columns: []string{
+				{{- range .OrderedColumns}}
+					"{{.}}",
+				{{- end}}
+			},
+			PartKey: []string {
+				{{- range .PartitionKey}}
+					"{{.Name}}",
+				{{- end}}
+			},
+			SortKey: []string{
+				{{- range .ClusteringColumns}}
+					"{{.Name}}",
+				{{- end}}
+			},
+	})
 {{end}}
+{{end}}
+)

cmd/schemagen/schemagen.go

@@ -7,12 +7,11 @@ import (
 	"fmt"
 	"go/format"
 	"html/template"
-	"io"
+	"io/ioutil"
 	"log"
 	"os"
 	"path"
 	"strings"
-	"unicode"
 
 	"github.com/gocql/gocql"
 	"github.com/scylladb/gocqlx/v2"
@@ -42,44 +41,34 @@ func main() {
 		log.Fatalln("missing required flag: keyspace")
 	}
 
-	schemagen()
+	if err := schemagen(); err != nil {
+		log.Fatalf("failed to generate schema: %s", err)
+	}
 }
 
-func schemagen() {
-	err := os.MkdirAll(*flagOutput, os.ModePerm)
-	if err != nil {
-		log.Fatalln("unable to create output directory:", err)
+func schemagen() error {
+	if err := os.MkdirAll(*flagOutput, os.ModePerm); err != nil {
+		return fmt.Errorf("create output directory: %w", err)
 	}
 
+	session, err := createSession()
+	if err != nil {
+		return fmt.Errorf("open output file: %w", err)
+	}
+	metadata, err := session.KeyspaceMetadata(*flagKeyspace)
+	if err != nil {
+		return fmt.Errorf("fetch keyspace metadata: %w", err)
+	}
+	b, err := renderTemplate(metadata)
+	if err != nil {
+		return fmt.Errorf("render template: %w", err)
+	}
 	outputPath := path.Join(*flagOutput, *flagPkgname+".go")
-	f, err := os.OpenFile(outputPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
-	if err != nil {
-		log.Fatalln("unable to open output file:", err)
-	}
 
-	metadata := fetchMetadata(createSession())
-
-	if err = renderTemplate(f, metadata); err != nil {
-		log.Fatalln("unable to output template:", err)
-	}
-
-	if err = f.Close(); err != nil {
-		log.Fatalln("unable to close output file:", err)
-	}
-
-	log.Println("File written to", outputPath)
+	return ioutil.WriteFile(outputPath, b, os.ModePerm)
 }
 
-func fetchMetadata(s *gocqlx.Session) *gocql.KeyspaceMetadata {
-	md, err := s.KeyspaceMetadata(*flagKeyspace)
-	if err != nil {
-		log.Fatalln("unable to fetch keyspace metadata:", err)
-	}
-
-	return md
-}
-
-func renderTemplate(w io.Writer, md *gocql.KeyspaceMetadata) error {
+func renderTemplate(md *gocql.KeyspaceMetadata) ([]byte, error) {
 	t, err := template.
 		New("keyspace.tmpl").
 		Funcs(template.FuncMap{"camelize": camelize}).
@@ -95,68 +84,17 @@ func renderTemplate(w io.Writer, md *gocql.KeyspaceMetadata) error {
 		"Tables":      md.Tables,
 	}
 
-	err = t.Execute(buf, data)
-	if err != nil {
-		log.Fatalln("unable to execute models template:", err)
+	if err = t.Execute(buf, data); err != nil {
+		return nil, fmt.Errorf("template: %w", err)
 	}
-
-	res, err := format.Source(buf.Bytes())
-	if err != nil {
-		log.Fatalln("template output is not a valid go code:", err)
-	}
-
-	_, err = w.Write(res)
-
-	return err
+	return format.Source(buf.Bytes())
 }
 
-func createSession() *gocqlx.Session {
-	cluster := createCluster()
-	s, err := gocqlx.WrapSession(cluster.CreateSession())
-	if err != nil {
-		log.Fatalln("unable to create scylla session:", err)
-	}
-	return &s
+func createSession() (gocqlx.Session, error) {
+	cluster := gocql.NewCluster(clusterHosts()...)
+	return gocqlx.WrapSession(cluster.CreateSession())
 }
 
-func createCluster() *gocql.ClusterConfig {
-	clusterHosts := getClusterHosts()
-	return gocql.NewCluster(clusterHosts...)
-}
-
-func getClusterHosts() []string {
+func clusterHosts() []string {
 	return strings.Split(*flagCluster, ",")
 }
-
-func camelize(s string) string {
-	buf := []byte(s)
-	out := make([]byte, 0, len(buf))
-	underscoreSeen := false
-
-	l := len(buf)
-	for i := 0; i < l; i++ {
-		if !(allowedBindRune(buf[i]) || buf[i] == '_') {
-			panic(fmt.Sprint("not allowed name ", s))
-		}
-
-		b := rune(buf[i])
-
-		if b == '_' {
-			underscoreSeen = true
-			continue
-		}
-
-		if (i == 0 || underscoreSeen) && unicode.IsLower(b) {
-			b = unicode.ToUpper(b)
-			underscoreSeen = false
-		}
-
-		out = append(out, byte(b))
-	}
-
-	return string(out)
-}
-
-func allowedBindRune(b byte) bool {
-	return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9')
-}

cmd/schemagen/schemagen_test.go

@@ -1,84 +1,51 @@
 package main
 
 import (
+	"flag"
 	"fmt"
-	"github.com/scylladb/gocqlx/v2/gocqlxtest"
-	"io"
+	"io/ioutil"
 	"os"
-	"os/exec"
-	"path"
-	"strings"
 	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/scylladb/gocqlx/v2/gocqlxtest"
 )
 
-func TestCamelize(t *testing.T) {
-	tests := []struct {
-		input string
-		want  string
-	}{
-		{"hello", "Hello"},
-		{"_hello", "Hello"},
-		{"__hello", "Hello"},
-		{"hello_", "Hello"},
-		{"hello_world", "HelloWorld"},
-		{"hello__world", "HelloWorld"},
-		{"_hello_world", "HelloWorld"},
-		{"helloWorld", "HelloWorld"},
-		{"HelloWorld", "HelloWorld"},
-	}
-	for _, tt := range tests {
-		t.Run(tt.input, func(t *testing.T) {
-			if got := camelize(tt.input); got != tt.want {
-				t.Errorf("camelize() = %v, want %v", got, tt.want)
-			}
-		})
-	}
-}
+var flagUpdate = flag.Bool("update", false, "update golden file")
 
-func Test_schemagen_defaultParams(t *testing.T) {
-	cleanup(t, "models")
-	defer cleanup(t, "models")
+func TestSchemagen(t *testing.T) {
+	flag.Parse()
 	createTestSchema(t)
-	runSchemagen(t, "", "")
-	assertResult(t, "models", "models")
-}
+	b := runSchemagen(t, "foobar")
 
-func Test_schemagen_customParams(t *testing.T) {
-	cleanup(t, "asdf")
-	defer cleanup(t, "asdf")
-	createTestSchema(t)
-	runSchemagen(t, "qwer", "asdf")
-	assertResult(t, "qwer", "asdf")
-}
-
-func cleanup(t *testing.T, output string) {
-	err := os.RemoveAll(output)
+	const goldenFile = "testdata/models.go.txt"
+	if *flagUpdate {
+		if err := ioutil.WriteFile(goldenFile, b, os.ModePerm); err != nil {
+			t.Fatal(err)
+		}
+	}
+	golden, err := ioutil.ReadFile(goldenFile)
 	if err != nil {
-		t.Fatalf("could not delete %s directory: %v\n", output, err)
+		t.Fatal(err)
 	}
 
-	err = os.Remove("./schemagen")
-	if err != nil {
-		t.Fatalf("could not delete binary: %v\n", err)
-	}
-
-	cmd := exec.Command("go", "build")
-	out, err := cmd.CombinedOutput()
-	if err != nil {
-		t.Fatalf("could not build binary for schemagen: %v\nOutput:\n%v\n", err, string(out))
+	if diff := cmp.Diff(string(golden), string(b)); diff != "" {
+		t.Fatalf(diff)
 	}
 }
 
 func createTestSchema(t *testing.T) {
+	t.Helper()
+
 	session := gocqlxtest.CreateSession(t)
 	defer session.Close()
 
-	err := session.ExecStmt(`CREATE KEYSPACE IF NOT EXISTS examples WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}`)
+	err := session.ExecStmt(`CREATE KEYSPACE IF NOT EXISTS schemagen WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}`)
 	if err != nil {
 		t.Fatal("create keyspace:", err)
 	}
 
-	err = session.ExecStmt(`CREATE TABLE IF NOT EXISTS examples.songs (
+	err = session.ExecStmt(`CREATE TABLE IF NOT EXISTS schemagen.songs (
 		id uuid PRIMARY KEY,
 		title text,
 		album text,
@@ -89,7 +56,7 @@ func createTestSchema(t *testing.T) {
 		t.Fatal("create table:", err)
 	}
 
-	err = session.ExecStmt(`CREATE TABLE IF NOT EXISTS examples.playlists (
+	err = session.ExecStmt(`CREATE TABLE IF NOT EXISTS schemagen.playlists (
 		id uuid,
 		title text,
 		album text, 
@@ -101,59 +68,27 @@ func createTestSchema(t *testing.T) {
 	}
 }
 
-func runSchemagen(t *testing.T, pkgname, output string) {
-	dir, err := os.Getwd()
+func runSchemagen(t *testing.T, pkgname string) []byte {
+	t.Helper()
+
+	dir, err := os.MkdirTemp("", "gocqlx")
 	if err != nil {
 		t.Fatal(err)
 	}
+	keyspace := "schemagen"
 
-	args := []string{"-keyspace=examples"}
-	for _, arg := range os.Args {
-		if strings.HasPrefix(arg, "-cluster") {
-			args = append(args, arg)
-		}
+	flagKeyspace = &keyspace
+	flagPkgname = &pkgname
+	flagOutput = &dir
+
+	if err := schemagen(); err != nil {
+		t.Fatalf("schemagen() error %s", err)
 	}
 
-	if pkgname != "" {
-		args = append(args, fmt.Sprintf("-pkgname=%s", pkgname))
-	}
-
-	if output != "" {
-		args = append(args, fmt.Sprintf("-output=%s", output))
-	}
-
-	cmd := exec.Command(path.Join(dir, "schemagen"), args...)
-	err = cmd.Run()
+	f := fmt.Sprintf("%s/%s.go", dir, pkgname)
+	b, err := os.ReadFile(f)
 	if err != nil {
-		t.Fatal(err)
+		t.Fatalf("%s: %s", f, err)
 	}
-}
-
-func assertResult(t *testing.T, pkgname, output string) {
-	path := fmt.Sprintf("%s/%s.go", output, pkgname)
-	res, err := os.ReadFile(path)
-	if err != nil {
-		t.Fatalf("can't read output file (%s): %s\n", path, err)
-	}
-
-	want := resultWant(t, pkgname)
-
-	if string(res) != want {
-		t.Fatalf("unexpected result: %s\nWanted:\n%s\n", string(res), want)
-	}
-}
-
-func resultWant(t *testing.T, pkgname string) string {
-	f, err := os.Open("testdata/models.go.txt")
-	if err != nil {
-		t.Fatalf("can't open testdata/models.go.txt")
-	}
-	defer f.Close()
-
-	b, err := io.ReadAll(f)
-	if err != nil {
-		t.Fatalf("can't read testdata/models.go.txt")
-	}
-
-	return strings.Replace(string(b), "{{pkgname}}", pkgname, 1)
+	return b
 }

cmd/schemagen/testdata/models.go.txt

@@ -1,42 +1,43 @@
 // Code generated by "gocqlx/cmd/schemagen"; DO NOT EDIT.
 
-package {{pkgname}}
+package foobar
 
 import "github.com/scylladb/gocqlx/v2/table"
 
-var PlaylistsMetadata = table.Metadata{
-	Name: "playlists",
-	Columns: []string{
-		"album",
-		"artist",
-		"id",
-		"song_id",
-		"title",
-	},
-	PartKey: []string{
-		"id",
-	},
-	SortKey: []string{
-		"title",
-		"album",
-		"artist",
-	},
-}
-var PlaylistsTable = table.New(PlaylistsMetadata)
+// Table models.
+var (
+	Playlists = table.New(table.Metadata{
+		Name: "playlists",
+		Columns: []string{
+			"album",
+			"artist",
+			"id",
+			"song_id",
+			"title",
+		},
+		PartKey: []string{
+			"id",
+		},
+		SortKey: []string{
+			"title",
+			"album",
+			"artist",
+		},
+	})
 
-var SongsMetadata = table.Metadata{
-	Name: "songs",
-	Columns: []string{
-		"album",
-		"artist",
-		"data",
-		"id",
-		"tags",
-		"title",
-	},
-	PartKey: []string{
-		"id",
-	},
-	SortKey: []string{},
-}
-var SongsTable = table.New(SongsMetadata)
+	Songs = table.New(table.Metadata{
+		Name: "songs",
+		Columns: []string{
+			"album",
+			"artist",
+			"data",
+			"id",
+			"tags",
+			"title",
+		},
+		PartKey: []string{
+			"id",
+		},
+		SortKey: []string{},
+	})
+)