forked from scylladb/gocqlx
/
schemagen.go
137 lines (118 loc) · 3.42 KB
/
schemagen.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
package main
import (
"bytes"
_ "embed"
"flag"
"fmt"
"go/format"
"html/template"
"io/ioutil"
"log"
"os"
"path"
"strings"
"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")
flagUser = cmd.String("user", "", "user for password authentication")
flagPassword = cmd.String("password", "", "password for password authentication")
)
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")
}
if err := schemagen(); err != nil {
log.Fatalf("failed to generate schema: %s", 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")
return ioutil.WriteFile(outputPath, b, os.ModePerm)
}
func renderTemplate(md *gocql.KeyspaceMetadata) ([]byte, error) {
t, err := template.
New("keyspace.tmpl").
Funcs(template.FuncMap{"camelize": camelize}).
Funcs(template.FuncMap{"mapScyllaToGoType": mapScyllaToGoType}).
Funcs(template.FuncMap{"typeToString": typeToString}).
Parse(keyspaceTmpl)
if err != nil {
log.Fatalln("unable to parse models template:", err)
}
imports := make([]string, 0)
for _, t := range md.Tables {
for _, c := range t.Columns {
if (c.Validator == "timestamp" || c.Validator == "date" || c.Validator == "duration" || c.Validator == "time") && !existsInSlice(imports, "time") {
imports = append(imports, "time")
}
if c.Validator == "decimal" && !existsInSlice(imports, "gopkg.in/inf.v0") {
imports = append(imports, "gopkg.in/inf.v0")
}
if c.Validator == "duration" && !existsInSlice(imports, "github.com/gocql/gocql") {
imports = append(imports, "github.com/gocql/gocql")
}
}
}
buf := &bytes.Buffer{}
data := map[string]interface{}{
"PackageName": *flagPkgname,
"Tables": md.Tables,
"UserTypes": md.UserTypes,
"Imports": imports,
}
if err = t.Execute(buf, data); err != nil {
return nil, fmt.Errorf("template: %w", err)
}
return format.Source(buf.Bytes())
}
func createSession() (gocqlx.Session, error) {
cluster := gocql.NewCluster(clusterHosts()...)
if *flagUser != "" {
cluster.Authenticator = gocql.PasswordAuthenticator{
Username: *flagUser,
Password: *flagPassword,
}
}
return gocqlx.WrapSession(cluster.CreateSession())
}
func clusterHosts() []string {
return strings.Split(*flagCluster, ",")
}
func existsInSlice(s []string, v string) bool {
for _, i := range s {
if v == i {
return true
}
}
return false
}