Structs tables

cmd/schemagen/keyspace.tmpl

@@ -2,7 +2,12 @@
 
 package {{.PackageName}}
 
-import "github.com/scylladb/gocqlx/v2/table"
+import (
+	"github.com/scylladb/gocqlx/v2/table"
+	{{- range .Imports}}
+	"{{.}}"
+	{{- end}}
+)
 
 // Table models.
 var (
@@ -30,3 +35,29 @@ var (
 {{end}}
 {{end}}
 )
+
+{{with .UserTypes}}
+{{range .}}
+{{- $type_name := .Name | camelize}}
+{{- $field_types := .FieldTypes}}
+type {{$type_name}}UserType struct {
+{{- range $index, $element := .FieldNames}}
+    {{- $type := index $field_types $index}}
+    {{. | camelize}} {{typeToString $type | mapScyllaToGoType}}
+{{- end}}
+}
+{{- end}}
+{{- end}}
+
+{{with .Tables}}
+{{range .}}
+{{- $model_name := .Name | camelize}}
+type {{$model_name}}Struct struct {
+{{- range .Columns}}
+	{{- if not (eq .Validator "empty") }}
+	{{.Name | camelize}} {{.Validator | mapScyllaToGoType}}
+	{{- end}}
+{{- end}}
+}
+{{- end}}
+{{- end}}

cmd/schemagen/map_types.go

@@ -0,0 +1,105 @@
+package main
+
+import (
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/gocql/gocql"
+)
+
+var types = map[string]string{
+	"ascii":     "string",
+	"bigint":    "int64",
+	"blob":      "[]byte",
+	"boolean":   "bool",
+	"counter":   "int",
+	"date":      "time.Time",
+	"decimal":   "inf.Dec",
+	"double":    "float64",
+	"duration":  "gocql.Duration",
+	"float":     "float32",
+	"inet":      "string",
+	"int":       "int32",
+	"smallint":  "int16",
+	"text":      "string",
+	"time":      "time.Duration",
+	"timestamp": "time.Time",
+	"timeuuid":  "[16]byte",
+	"tinyint":   "int8",
+	"uuid":      "[16]byte",
+	"varchar":   "string",
+	"varint":    "int64",
+}
+
+func mapScyllaToGoType(s string) string {
+	frozenRegex := regexp.MustCompile(`frozen<([a-z]*)>`)
+	match := frozenRegex.FindAllStringSubmatch(s, -1)
+	if match != nil {
+		s = match[0][1]
+	}
+
+	mapRegex := regexp.MustCompile(`map<([a-z]*), ([a-z]*)>`)
+	setRegex := regexp.MustCompile(`set<([a-z]*)>`)
+	listRegex := regexp.MustCompile(`list<([a-z]*)>`)
+	tupleRegex := regexp.MustCompile(`tuple<(?:([a-z]*),? ?)*>`)
+	match = mapRegex.FindAllStringSubmatch(s, -1)
+	if match != nil {
+		key := match[0][1]
+		value := match[0][2]
+
+		return "map[" + types[key] + "]" + types[value]
+	}
+
+	match = setRegex.FindAllStringSubmatch(s, -1)
+	if match != nil {
+		key := match[0][1]
+
+		return "[]" + types[key]
+	}
+
+	match = listRegex.FindAllStringSubmatch(s, -1)
+	if match != nil {
+		key := match[0][1]
+
+		return "[]" + types[key]
+	}
+
+	match = tupleRegex.FindAllStringSubmatch(s, -1)
+	if match != nil {
+		tuple := match[0][0]
+		subStr := tuple[6 : len(tuple)-1]
+		types := strings.Split(subStr, ", ")
+
+		typeStr := "struct {\n"
+		for i, t := range types {
+			typeStr = typeStr + "\t\tField" + strconv.Itoa(i+1) + " " + mapScyllaToGoType(t) + "\n"
+		}
+		typeStr = typeStr + "\t}"
+
+		return typeStr
+	}
+
+	t, exists := types[s]
+	if exists {
+		return t
+	}
+
+	return camelize(s) + "UserType"
+}
+
+func typeToString(t interface{}) string {
+	tType := fmt.Sprintf("%T", t)
+	switch tType {
+	case "gocql.NativeType":
+		return t.(gocql.NativeType).String()
+	case "gocql.CollectionType":
+		collectionType := t.(gocql.CollectionType).String()
+		collectionType = strings.Replace(collectionType, "(", "<", -1)
+		collectionType = strings.Replace(collectionType, ")", ">", -1)
+		return collectionType
+	default:
+		panic(fmt.Sprintf("Did not expect %v type in user defined type", tType))
+	}
+}

cmd/schemagen/map_types_test.go

@@ -0,0 +1,49 @@
+// 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 TestMapScyllaToGoType(t *testing.T) {
+	tests := []struct {
+		input string
+		want  string
+	}{
+		{"ascii", "string"},
+		{"bigint", "int64"},
+		{"blob", "[]byte"},
+		{"boolean", "bool"},
+		{"counter", "int"},
+		{"date", "time.Time"},
+		{"decimal", "inf.Dec"},
+		{"double", "float64"},
+		{"duration", "gocql.Duration"},
+		{"float", "float32"},
+		{"inet", "string"},
+		{"int", "int32"},
+		{"smallint", "int16"},
+		{"text", "string"},
+		{"time", "time.Duration"},
+		{"timestamp", "time.Time"},
+		{"timeuuid", "[16]byte"},
+		{"tinyint", "int8"},
+		{"uuid", "[16]byte"},
+		{"varchar", "string"},
+		{"varint", "int64"},
+		{"map<int, text>", "map[int32]string"},
+		{"list<int>", "[]int32"},
+		{"set<int>", "[]int32"},
+		{"tuple<boolean, int, smallint>", "struct {\n\t\tField1 bool\n\t\tField2 int32\n\t\tField3 int16\n\t}"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			if got := mapScyllaToGoType(tt.input); got != tt.want {
+				t.Errorf("mapScyllaToGoType() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

cmd/schemagen/schemagen.go

@@ -72,16 +72,35 @@ 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 {
@@ -98,3 +117,13 @@ func createSession() (gocqlx.Session, error) {
 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
+}

cmd/schemagen/schemagen_test.go

@@ -56,10 +56,17 @@ func createTestSchema(t *testing.T) {
 		t.Fatal("create table:", err)
 	}
 
+	err = session.ExecStmt(`CREATE TYPE IF NOT EXISTS schemagen.album (
+		name text,
+		songwriters set<text>,)`)
+	if err != nil {
+		t.Fatal("create type:", err)
+	}
+
 	err = session.ExecStmt(`CREATE TABLE IF NOT EXISTS schemagen.playlists (
 		id uuid,
 		title text,
-		album text, 
+		album frozen<album>, 
 		artist text,
 		song_id uuid,
 		PRIMARY KEY (id, title, album, artist))`)

cmd/schemagen/testdata/models.go.txt

@@ -2,7 +2,9 @@
 
 package foobar
 
-import "github.com/scylladb/gocqlx/v2/table"
+import (
+	"github.com/scylladb/gocqlx/v2/table"
+)
 
 // Table models.
 var (
@@ -41,3 +43,24 @@ var (
 		SortKey: []string{},
 	})
 )
+
+type AlbumUserType struct {
+	Name        string
+	Songwriters []string
+}
+
+type PlaylistsStruct struct {
+	Album  AlbumUserType
+	Artist string
+	Id     [16]byte
+	SongId [16]byte
+	Title  string
+}
+type SongsStruct struct {
+	Album  string
+	Artist string
+	Data   []byte
+	Id     [16]byte
+	Tags   []string
+	Title  string
+}