Repository: hamba/avro
Pull request #429 (2024-08-09 08:37:12)

support custom logical types
proposal to support custom logical types (implementing TextMarshaller) when generating structs

marcowork: +1
nrwiersma: Can you explain the issue you are trying to solve? Logical types are specified in the spec.
Emptyless: @nrwiersma the spec details a set of "official" logical types that most implementations support, it is however allowed to define your own custom LogicalType's that are marshallable by their base type, this has been applied in e.g. Java [[1]](https://stackoverflow.com/questions/49034266/how-to-define-a-logicaltype-in-avro-java) and Python [[2]](https://fastavro.readthedocs.io/en/latest/logical_types.html).

But this PR also allows you to register your own type for a particular LogicalType, e.g. the UUID type. The UUID type is backed by a string but a string can be anything. By registering a UUID compliant Go type, you get added type safety when going back and forth between Avro. I've included this as an example in the tests
Emptyless: resaonable suggestions, applied the fixes. `make lint` also mentioned the `_, _ = fmt.Fprintln(stdout, schema)` line so applied so that the pipeline won't break on it

Patch:
diff --git a/cmd/avrosv/main.go b/cmd/avrosv/main.go
index ec6964a..bd32c58 100644
--- a/cmd/avrosv/main.go
+++ b/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
diff --git a/gen/gen.go b/gen/gen.go
index 250bd7a..344af81 100644
--- a/gen/gen.go
+++ b/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":


Test patch:
diff --git a/gen/gen_test.go b/gen/gen_test.go
index 786beed..5d7c913 100644
--- a/gen/gen_test.go
+++ b/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))
 }
diff --git a/gen/testdata/golden_logicaltype.go b/gen/testdata/golden_logicaltype.go
new file mode 100644
index 0000000..c2f93c5
--- /dev/null
+++ b/gen/testdata/golden_logicaltype.go
@@ -0,0 +1,85 @@
+package something
+
+// Code generated by avro/gen. DO NOT EDIT.
+
+import (
+	"math/big"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/hamba/avro/v2"
+)
+
+// InnerRecord is a generated struct.
+type InnerRecord struct {
+	InnerJustBytes                   []byte    `avro:"innerJustBytes"`
+	InnerPrimitiveNullableArrayUnion *[]string `avro:"innerPrimitiveNullableArrayUnion"`
+}
+
+// RecordInMap is a generated struct.
+type RecordInMap struct {
+	Name string `avro:"name"`
+}
+
+// RecordInArray is a generated struct.
+type RecordInArray struct {
+	AString string `avro:"aString"`
+}
+
+// RecordInNullableUnion is a generated struct.
+type RecordInNullableUnion struct {
+	AString string `avro:"aString"`
+}
+
+// Record1InNonNullableUnion is a generated struct.
+type Record1InNonNullableUnion struct {
+	AString string `avro:"aString"`
+}
+
+// Record2InNonNullableUnion is a generated struct.
+type Record2InNonNullableUnion struct {
+	AString string `avro:"aString"`
+}
+
+// Record1InNullableUnion is a generated struct.
+type Record1InNullableUnion struct {
+	AString string `avro:"aString"`
+}
+
+// Record2InNullableUnion is a generated struct.
+type Record2InNullableUnion struct {
+	AString string `avro:"aString"`
+}
+
+// Test represents a golden record.
+type Test struct {
+	// aString is just a string.
+	AString string `avro:"aString"`
+	// aBoolean is just a boolean.
+	ABoolean                        bool                   `avro:"aBoolean"`
+	AnInt                           int                    `avro:"anInt"`
+	AFloat                          float32                `avro:"aFloat"`
+	ADouble                         float64                `avro:"aDouble"`
+	ALong                           int64                  `avro:"aLong"`
+	JustBytes                       []byte                 `avro:"justBytes"`
+	PrimitiveNullableArrayUnion     *[]string              `avro:"primitiveNullableArrayUnion"`
+	InnerRecord                     InnerRecord            `avro:"innerRecord"`
+	AnEnum                          string                 `avro:"anEnum"`
+	AFixed                          [7]byte                `avro:"aFixed"`
+	ALogicalFixed                   avro.LogicalDuration   `avro:"aLogicalFixed"`
+	AnotherLogicalFixed             avro.LogicalDuration   `avro:"anotherLogicalFixed"`
+	MapOfStrings                    map[string]string      `avro:"mapOfStrings"`
+	MapOfRecords                    map[string]RecordInMap `avro:"mapOfRecords"`
+	ADate                           time.Time              `avro:"aDate"`
+	ADuration                       time.Duration          `avro:"aDuration"`
+	ALongTimeMicros                 time.Duration          `avro:"aLongTimeMicros"`
+	ALongTimestampMillis            time.Time              `avro:"aLongTimestampMillis"`
+	ALongTimestampMicro             time.Time              `avro:"aLongTimestampMicro"`
+	ABytesDecimal                   *big.Rat               `avro:"aBytesDecimal"`
+	ARecordArray                    []RecordInArray        `avro:"aRecordArray"`
+	NullableRecordUnion             *RecordInNullableUnion `avro:"nullableRecordUnion"`
+	NonNullableRecordUnion          any                    `avro:"nonNullableRecordUnion"`
+	NullableRecordUnionWith3Options any                    `avro:"nullableRecordUnionWith3Options"`
+	Ref                             Record2InNullableUnion `avro:"ref"`
+	UUID                            uuid.UUID              `avro:"uuid"`
+}