schemagen cli

cmd/schemagen/keyspace.tmpl

@@ -0,0 +1,30 @@
+// Code generated by "gocqlx/cmd/schemagen"; DO NOT EDIT.
+
+package {{.PackageName}}
+
+import "github.com/scylladb/gocqlx/v2/table"
+
+{{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}}
+{{end}}

cmd/schemagen/schemagen.go

@@ -0,0 +1,162 @@
+package main
+
+import (
+	"bytes"
+	_ "embed"
+	"flag"
+	"fmt"
+	"go/format"
+	"html/template"
+	"io"
+	"log"
+	"os"
+	"path"
+	"strings"
+	"unicode"
+
+	"github.com/gocql/gocql"
+	"github.com/scylladb/gocqlx/v2"
+	_ "github.com/scylladb/gocqlx/v2/table"
+)
+
+var (
+	cmd          = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
+	flagCluster  = cmd.String("cluster", "127.0.0.1", "a comma-separated list of host:port tuples")
+	flagKeyspace = cmd.String("keyspace", "", "keyspace to inspect")
+	flagPkgname  = cmd.String("pkgname", "models", "the name you wish to assign to your generated package")
+	flagOutput   = cmd.String("output", "models", "the name of the folder to output to")
+)
+
+var (
+	//go:embed keyspace.tmpl
+	keyspaceTmpl string
+)
+
+func main() {
+	err := cmd.Parse(os.Args[1:])
+	if err != nil {
+		log.Fatalln("can't parse flags")
+	}
+
+	if *flagKeyspace == "" {
+		log.Fatalln("missing required flag: keyspace")
+	}
+
+	schemagen()
+}
+
+func schemagen() {
+	err := os.MkdirAll(*flagOutput, os.ModePerm)
+	if err != nil {
+		log.Fatalln("unable to create output directory:", 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)
+}
+
+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 {
+	t, err := template.
+		New("keyspace.tmpl").
+		Funcs(template.FuncMap{"camelize": camelize}).
+		Parse(keyspaceTmpl)
+
+	if err != nil {
+		log.Fatalln("unable to parse models template:", err)
+	}
+
+	buf := &bytes.Buffer{}
+	data := map[string]interface{}{
+		"PackageName": *flagPkgname,
+		"Tables":      md.Tables,
+	}
+
+	err = t.Execute(buf, data)
+	if err != nil {
+		log.Fatalln("unable to execute models template:", 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
+}
+
+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 createCluster() *gocql.ClusterConfig {
+	clusterHosts := getClusterHosts()
+	return gocql.NewCluster(clusterHosts...)
+}
+
+func getClusterHosts() []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

@@ -0,0 +1,159 @@
+package main
+
+import (
+	"fmt"
+	"github.com/scylladb/gocqlx/v2/gocqlxtest"
+	"io"
+	"os"
+	"os/exec"
+	"path"
+	"strings"
+	"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)
+			}
+		})
+	}
+}
+
+func Test_schemagen_defaultParams(t *testing.T) {
+	cleanup(t, "models")
+	defer cleanup(t, "models")
+	createTestSchema(t)
+	runSchemagen(t, "", "")
+	assertResult(t, "models", "models")
+}
+
+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)
+	if err != nil {
+		t.Fatalf("could not delete %s directory: %v\n", output, 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))
+	}
+}
+
+func createTestSchema(t *testing.T) {
+	session := gocqlxtest.CreateSession(t)
+	defer session.Close()
+
+	err := session.ExecStmt(`CREATE KEYSPACE IF NOT EXISTS examples 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 (
+		id uuid PRIMARY KEY,
+		title text,
+		album text,
+		artist text,
+		tags set<text>,
+		data blob)`)
+	if err != nil {
+		t.Fatal("create table:", err)
+	}
+
+	err = session.ExecStmt(`CREATE TABLE IF NOT EXISTS examples.playlists (
+		id uuid,
+		title text,
+		album text, 
+		artist text,
+		song_id uuid,
+		PRIMARY KEY (id, title, album, artist))`)
+	if err != nil {
+		t.Fatal("create table:", err)
+	}
+}
+
+func runSchemagen(t *testing.T, pkgname, output string) {
+	dir, err := os.Getwd()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	args := []string{"-keyspace=examples"}
+	for _, arg := range os.Args {
+		if strings.HasPrefix(arg, "-cluster") {
+			args = append(args, arg)
+		}
+	}
+
+	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()
+	if err != nil {
+		t.Fatal(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)
+}

cmd/schemagen/testdata/models.go.txt

@@ -0,0 +1,42 @@
+// Code generated by "gocqlx/cmd/schemagen"; DO NOT EDIT.
+
+package {{pkgname}}
+
+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)
+
+var SongsMetadata = table.Metadata{
+	Name: "songs",
+	Columns: []string{
+		"album",
+		"artist",
+		"data",
+		"id",
+		"tags",
+		"title",
+	},
+	PartKey: []string{
+		"id",
+	},
+	SortKey: []string{},
+}
+var SongsTable = table.New(SongsMetadata)