diff --git a/cmd/schemagen/keyspace.tmpl b/cmd/schemagen/keyspace.tmpl index 1a463e5..cddb23d 100644 --- a/cmd/schemagen/keyspace.tmpl +++ b/cmd/schemagen/keyspace.tmpl @@ -40,10 +40,10 @@ var ( {{range .}} {{- $type_name := .Name | camelize}} {{- $field_types := .FieldTypes}} -type {{$type_name}}Type struct { +type {{$type_name}}UserType struct { {{- range $index, $element := .FieldNames}} {{- $type := index $field_types $index}} - {{. | camelize}} {{getNativeTypeSting $type | mapScyllaToGoType}} + {{. | camelize}} {{typeToString $type | mapScyllaToGoType}} {{- end}} } {{- end}} @@ -54,7 +54,9 @@ type {{$type_name}}Type struct { {{- $model_name := .Name | camelize}} type {{$model_name}}Struct struct { {{- range .Columns}} + {{- if not (eq .Validator "empty") }} {{.Name | camelize}} {{.Validator | mapScyllaToGoType}} + {{- end}} {{- end}} } {{- end}} diff --git a/cmd/schemagen/map_types.go b/cmd/schemagen/map_types.go index fd34a1f..4f59554 100644 --- a/cmd/schemagen/map_types.go +++ b/cmd/schemagen/map_types.go @@ -1,6 +1,7 @@ package main import ( + "fmt" "regexp" "strconv" "strings" @@ -33,11 +34,17 @@ var types = map[string]string{ } 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) + match = mapRegex.FindAllStringSubmatch(s, -1) if match != nil { key := match[0][1] value := match[0][2] @@ -79,9 +86,20 @@ func mapScyllaToGoType(s string) string { return t } - return camelize(s) + "Type" + return camelize(s) + "UserType" } -func getNativeTypeSting(t gocql.NativeType) string { - return t.String() +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)) + } } diff --git a/cmd/schemagen/schemagen.go b/cmd/schemagen/schemagen.go index 2c5cb48..9cff543 100644 --- a/cmd/schemagen/schemagen.go +++ b/cmd/schemagen/schemagen.go @@ -75,7 +75,7 @@ func renderTemplate(md *gocql.KeyspaceMetadata) ([]byte, error) { New("keyspace.tmpl"). Funcs(template.FuncMap{"camelize": camelize}). Funcs(template.FuncMap{"mapScyllaToGoType": mapScyllaToGoType}). - Funcs(template.FuncMap{"getNativeTypeSting": getNativeTypeSting}). + Funcs(template.FuncMap{"typeToString": typeToString}). Parse(keyspaceTmpl) if err != nil { diff --git a/cmd/schemagen/schemagen_test.go b/cmd/schemagen/schemagen_test.go index 1ca13af..382d21f 100644 --- a/cmd/schemagen/schemagen_test.go +++ b/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,)`) + 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, artist text, song_id uuid, PRIMARY KEY (id, title, album, artist))`) diff --git a/cmd/schemagen/testdata/models.go.txt b/cmd/schemagen/testdata/models.go.txt index 3504e6a..6714a49 100644 --- a/cmd/schemagen/testdata/models.go.txt +++ b/cmd/schemagen/testdata/models.go.txt @@ -44,8 +44,13 @@ var ( }) ) +type AlbumUserType struct { + Name string + Songwriters []string +} + type PlaylistsStruct struct { - Album string + Album AlbumUserType Artist string Id [16]byte SongId [16]byte