Make schemagen compilable and add test for it

Makefile

@@ -38,6 +38,7 @@ GOTEST := go test -cpu $(GOTEST_CPU) -count=1 -cover -race -tags all
 
 .PHONY: test
 test:
+	@go build cmd/schemagen/testdata/models.go # Make sure the generated code compiles
 	@$(GOTEST) .
 	@$(GOTEST) ./qb
 	@$(GOTEST) ./table

cmd/schemagen/keyspace.tmpl

@@ -41,6 +41,7 @@ var (
 {{- $type_name := .Name | camelize}}
 {{- $field_types := .FieldTypes}}
 type {{$type_name}}UserType struct {
+	gocqlx.UDT
 {{- range $index, $element := .FieldNames}}
 	{{. | camelize}} {{(index $field_types $index) | mapScyllaToGoType}}
 {{- end}}

cmd/schemagen/schemagen.go

@@ -109,6 +109,10 @@ func renderTemplate(md *gocql.KeyspaceMetadata) ([]byte, error) {
 	}
 
 	imports := make([]string, 0)
+	if len(md.Types) != 0 {
+		imports = append(imports, "github.com/scylladb/gocqlx/v2")
+	}
+
 	for _, t := range md.Tables {
 		// Ensure ordered columns are sorted alphabetically
 		sort.Strings(t.OrderedColumns)
@@ -120,8 +124,8 @@ func renderTemplate(md *gocql.KeyspaceMetadata) ([]byte, error) {
 			if c.Type == "decimal" && !existsInSlice(imports, "gopkg.in/inf.v0") {
 				imports = append(imports, "gopkg.in/inf.v0")
 			}
-			if c.Type == "duration" && !existsInSlice(imports, "github.com/gocql/gocql") {
-				imports = append(imports, "github.com/gocql/gocql")
+			if c.Type == "duration" && !existsInSlice(imports, "github.com/gocql/gocql/v2") {
+				imports = append(imports, "github.com/gocql/gocql/v2")
 			}
 		}
 	}

cmd/schemagen/schemagen_test.go

@@ -30,7 +30,7 @@ func TestSchemagen(t *testing.T) {
 
 	b := runSchemagen(t, "foobar")
 
-	const goldenFile = "testdata/models.go.txt"
+	const goldenFile = "testdata/models.go"
 	if *flagUpdate {
 		if err := ioutil.WriteFile(goldenFile, b, os.ModePerm); err != nil {
 			t.Fatal(err)