support custom logical types

cmd/avrosv/main.go

@@ -43,7 +43,7 @@ func realMain(args []string, stdout, stderr io.Writer) int {
 	}
 
 	if cfg.Verbose {
-		fmt.Fprintln(stdout, schema)
+		_, _ = fmt.Fprintln(stdout, schema)
 	}
 
 	return 0

gen/gen.go

@@ -19,12 +19,13 @@ import (
 
 // Config configures the code generation.
 type Config struct {
-	PackageName string
-	Tags        map[string]TagStyle
-	FullName    bool
-	Encoders    bool
-	StrictTypes bool
-	Initialisms []string
+	PackageName  string
+	Tags         map[string]TagStyle
+	FullName     bool
+	Encoders     bool
+	StrictTypes  bool
+	Initialisms  []string
+	LogicalTypes []LogicalType
 }
 
 // TagStyle defines the styling for a tag.
@@ -83,6 +84,9 @@ func StructFromSchema(schema avro.Schema, w io.Writer, cfg Config) error {
 		WithInitialisms(cfg.Initialisms),
 		WithStrictTypes(cfg.StrictTypes),
 	}
+	for _, opt := range cfg.LogicalTypes {
+		opts = append(opts, WithLogicalType(opt))
+	}
 	g := NewGenerator(strcase.ToSnake(cfg.PackageName), cfg.Tags, opts...)
 	g.Parse(rec)
 
@@ -155,6 +159,29 @@ func WithPackageDoc(text string) OptsFunc {
 	}
 }
 
+// LogicalType used when the name of the "LogicalType" field in the Avro schema matches the Name attribute.
+type LogicalType struct {
+	// Name of the LogicalType
+	Name string
+	// Typ returned, has to be a valid Go type
+	Typ string
+	// Import added as import (if not empty)
+	Import string
+	// ThirdPartyImport added as import (if not empty)
+	ThirdPartyImport string
+}
+
+// WithLogicalType registers a LogicalType which takes precedence over the default logical types
+// defined by this package.
+func WithLogicalType(logicalType LogicalType) OptsFunc {
+	return func(g *Generator) {
+		if g.logicalTypes == nil {
+			g.logicalTypes = map[avro.LogicalType]LogicalType{}
+		}
+		g.logicalTypes[avro.LogicalType(logicalType.Name)] = logicalType
+	}
+}
+
 func ensureTrailingPeriod(text string) string {
 	if text == "" {
 		return text
@@ -167,14 +194,15 @@ func ensureTrailingPeriod(text string) string {
 
 // Generator generates Go structs from schemas.
 type Generator struct {
-	template    string
-	pkg         string
-	pkgdoc      string
-	tags        map[string]TagStyle
-	fullName    bool
-	encoders    bool
-	strictTypes bool
-	initialisms []string
+	template     string
+	pkg          string
+	pkgdoc       string
+	tags         map[string]TagStyle
+	fullName     bool
+	encoders     bool
+	strictTypes  bool
+	initialisms  []string
+	logicalTypes map[avro.LogicalType]LogicalType
 
 	imports           []string
 	thirdPartyImports []string
@@ -313,6 +341,19 @@ func (g *Generator) resolveUnionTypes(s *avro.UnionSchema) string {
 }
 
 func (g *Generator) resolveLogicalSchema(logicalType avro.LogicalType) string {
+	if g.logicalTypes != nil {
+		if typ, ok := g.logicalTypes[logicalType]; ok {
+			if val := typ.Import; val != "" {
+				g.addImport(val)
+			}
+			if val := typ.ThirdPartyImport; val != "" {
+				g.addThirdPartyImport(val)
+			}
+
+			return typ.Typ
+		}
+	}
+
 	var typ string
 	switch logicalType {
 	case "date", "timestamp-millis", "timestamp-micros":

gen/gen_test.go

@@ -155,19 +155,74 @@ func TestStruct_ConfigurableFieldTags(t *testing.T) {
 	}
 }
 
+func TestStruct_ConfigurableLogicalTypes(t *testing.T) {
+	schema := `{
+  "type": "record",
+  "name": "test",
+  "fields": [
+    { "name": "id", "type": {"type": "string", "logicalType": "uuid"} }
+  ]
+}`
+
+	gc := gen.Config{
+		PackageName: "Something",
+		LogicalTypes: []gen.LogicalType{{
+			Name:             "uuid",
+			Typ:              "uuid.UUID",
+			ThirdPartyImport: "github.com/google/uuid",
+		}},
+	}
+	_, lines := generate(t, schema, gc)
+
+	for _, expected := range []string{
+		"package something",
+		"import (",
+		"\"github.com/google/uuid\"",
+		"type Test struct {",
+		"ID uuid.UUID `avro:\"id\"`",
+		"}",
+	} {
+		assert.Contains(t, lines, expected)
+	}
+}
+
 func TestStruct_GenFromRecordSchema(t *testing.T) {
+	fileName := "testdata/golden.go"
+	gc := gen.Config{PackageName: "Something"}
+	schema, err := os.ReadFile("testdata/golden.avsc")
+	require.NoError(t, err)
+
+	file, _ := generate(t, string(schema), gc)
+
+	if *update {
+		err = os.WriteFile(fileName, file, 0600)
+		require.NoError(t, err)
+	}
+
+	want, err := os.ReadFile(fileName)
+	require.NoError(t, err)
+	assert.Equal(t, string(want), string(file))
+}
+
+func TestStruct_GenFromRecordSchemaWithCustomLogicalTypes(t *testing.T) {
+	fileName := "testdata/golden_logicaltype.go"
+
+	gc := gen.Config{PackageName: "Something", LogicalTypes: []gen.LogicalType{{
+		Name:             "uuid",
+		Typ:              "uuid.UUID",
+		ThirdPartyImport: "github.com/google/uuid",
+	}}}
 	schema, err := os.ReadFile("testdata/golden.avsc")
 	require.NoError(t, err)
 
-	gc := gen.Config{PackageName: "Something"}
 	file, _ := generate(t, string(schema), gc)
 
 	if *update {
-		err = os.WriteFile("testdata/golden.go", file, 0600)
+		err = os.WriteFile(fileName, file, 0600)
 		require.NoError(t, err)
 	}
 
-	want, err := os.ReadFile("testdata/golden.go")
+	want, err := os.ReadFile(fileName)
 	require.NoError(t, err)
 	assert.Equal(t, string(want), string(file))
 }