gocqlx/cmd/schemagen/schemagen.go

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')
}