Repository: hamba/avro
Pull request #526 (2025-04-22 13:21:50)

feat(gen): add ability to pass type metadata
## Goal of this PR

We're looking for ways to pass down some schema metadata to the template, e.g. we have a need for the original schema subject identifier in the generated Go struct.

This change allows to pass arbitrary metadata while parsing the schema and have that metadata available in the template.

## How did I test it?

`go test`

ljank: Unfortunately current metadata is not enough (at least very inconvenient) when it comes to mapping of schema metadata -> typedef metadata. In the template Typedefs have different names compared to schemas:

https://github.com/hamba/avro/blob/5c80dfbd1161fcec89a5cfc5a213a673b08bdbf8/gen/gen.go#L311-L316

Since Typedef names are rendered using unexported `nameCaser` and unexported `resolveTypeName()`, we don't know the final Typedef name in advance and therefore can't set corresponding metadata key for access from the template.

Exporting `nameCaser` and/or `resolveTypeName()` doesn't feel right, therefore we suggest adding metadata along with the schema when calling `(Generator) Parse()` :)
ljank: @nrwiersma what do you think? can we proceed with the proposed change? 🙏🏻 
nrwiersma: @ljank I am away at the moment, will take a look when I find some time.

Patch:
diff --git a/gen/gen.go b/gen/gen.go
index c6db874..8ac94b8 100644
--- a/gen/gen.go
+++ b/gen/gen.go
@@ -269,15 +269,21 @@ func (g *Generator) Reset() {
 
 // Parse parses an avro schema into Go types.
 func (g *Generator) Parse(schema avro.Schema) {
-	_ = g.generate(schema)
+	_ = g.generate(schema, nil)
 }
 
-func (g *Generator) generate(schema avro.Schema) string {
+// ParseWithMetadata parses an avro schema into Go types with arbitrary metadata attached.
+// The metadata is then passed to the template as `Typedefs[].Metadata`.
+func (g *Generator) ParseWithMetadata(schema avro.Schema, metadata any) {
+	_ = g.generate(schema, metadata)
+}
+
+func (g *Generator) generate(schema avro.Schema, metadata any) string {
 	switch s := schema.(type) {
 	case *avro.RefSchema:
-		return g.resolveRefSchema(s)
+		return g.resolveRefSchema(s, metadata)
 	case *avro.RecordSchema:
-		return g.resolveRecordSchema(s)
+		return g.resolveRecordSchema(s, metadata)
 	case *avro.PrimitiveSchema:
 		typ := primitiveMappings[s.Type()]
 		if ls := s.Logical(); ls != nil {
@@ -290,7 +296,7 @@ func (g *Generator) generate(schema avro.Schema) string {
 		}
 		return typ
 	case *avro.ArraySchema:
-		return "[]" + g.generate(s.Items())
+		return "[]" + g.generate(s.Items(), metadata)
 	case *avro.EnumSchema:
 		return "string"
 	case *avro.FixedSchema:
@@ -300,9 +306,9 @@ func (g *Generator) generate(schema avro.Schema) string {
 		}
 		return typ
 	case *avro.MapSchema:
-		return "map[string]" + g.generate(s.Values())
+		return "map[string]" + g.generate(s.Values(), metadata)
 	case *avro.UnionSchema:
-		return g.resolveUnionTypes(s)
+		return g.resolveUnionTypes(s, metadata)
 	default:
 		return ""
 	}
@@ -315,16 +321,19 @@ func (g *Generator) resolveTypeName(s avro.NamedSchema) string {
 	return g.nameCaser.ToPascal(s.Name())
 }
 
-func (g *Generator) resolveRecordSchema(schema *avro.RecordSchema) string {
+func (g *Generator) resolveRecordSchema(schema *avro.RecordSchema, metadata any) string {
 	fields := make([]field, len(schema.Fields()))
 	for i, f := range schema.Fields() {
-		typ := g.generate(f.Type())
+		typ := g.generate(f.Type(), metadata)
 		fields[i] = g.newField(g.nameCaser.ToPascal(f.Name()), typ, f.Doc(), f.Name(), f.Props())
 	}
 
 	typeName := g.resolveTypeName(schema)
 	if !g.hasTypeDef(typeName) {
-		g.typedefs = append(g.typedefs, newType(typeName, schema.Doc(), fields, g.rawSchema(schema), schema.Props()))
+		g.typedefs = append(
+			g.typedefs,
+			newType(typeName, schema.Doc(), fields, g.rawSchema(schema), schema.Props(), metadata),
+		)
 	}
 	return typeName
 }
@@ -350,20 +359,20 @@ func (g *Generator) hasTypeDef(name string) bool {
 	return false
 }
 
-func (g *Generator) resolveRefSchema(s *avro.RefSchema) string {
+func (g *Generator) resolveRefSchema(s *avro.RefSchema, metadata any) string {
 	if sx, ok := s.Schema().(*avro.RecordSchema); ok {
 		return g.resolveTypeName(sx)
 	}
-	return g.generate(s.Schema())
+	return g.generate(s.Schema(), metadata)
 }
 
-func (g *Generator) resolveUnionTypes(s *avro.UnionSchema) string {
+func (g *Generator) resolveUnionTypes(s *avro.UnionSchema, metadata any) string {
 	types := make([]string, 0, len(s.Types()))
 	for _, elem := range s.Types() {
 		if _, ok := elem.(*avro.NullSchema); ok {
 			continue
 		}
-		types = append(types, g.generate(elem))
+		types = append(types, g.generate(elem, metadata))
 	}
 	if s.Nullable() {
 		return "*" + types[0]
@@ -474,20 +483,22 @@ func (g *Generator) Write(w io.Writer) error {
 }
 
 type typedef struct {
-	Name   string
-	Doc    string
-	Fields []field
-	Schema string
-	Props  map[string]any
+	Name     string
+	Doc      string
+	Fields   []field
+	Schema   string
+	Props    map[string]any
+	Metadata any
 }
 
-func newType(name, doc string, fields []field, schema string, props map[string]any) typedef {
+func newType(name, doc string, fields []field, schema string, props map[string]any, metadata any) typedef {
 	return typedef{
-		Name:   name,
-		Doc:    ensureTrailingPeriod(doc),
-		Fields: fields,
-		Schema: schema,
-		Props:  props,
+		Name:     name,
+		Doc:      ensureTrailingPeriod(doc),
+		Fields:   fields,
+		Schema:   schema,
+		Props:    props,
+		Metadata: metadata,
 	}
 }
 
diff --git a/gen/output_template.tmpl b/gen/output_template.tmpl
index ff36cd6..54311f2 100644
--- a/gen/output_template.tmpl
+++ b/gen/output_template.tmpl
@@ -18,7 +18,7 @@ package {{ .PackageName }}
 	// {{ replace .Doc "\n" "\n// " -1 }}
 	{{- else }}
 	// {{ .Name }} is a generated struct.
-	{{- end }}	
+	{{- end }}
 	type {{ .Name }} struct {
 	{{- range .Fields }}
 		{{- $f := . }}


Test patch:
diff --git a/gen/gen_test.go b/gen/gen_test.go
index 4cb9d9a..ce2b024 100644
--- a/gen/gen_test.go
+++ b/gen/gen_test.go
@@ -355,6 +355,46 @@ func TestGenerator(t *testing.T) {
 	assert.Equal(t, string(want), string(formatted))
 }
 
+func TestGenerator_CustomTemplateWithMetadata(t *testing.T) {
+	unionSchema, err := avro.ParseFiles("testdata/uniontype.avsc")
+	require.NoError(t, err)
+
+	mainSchema, err := avro.ParseFiles("testdata/main.avsc")
+	require.NoError(t, err)
+
+	template := `// Code generated by avro/gen. DO NOT EDIT.
+package {{ .PackageName }}
+
+{{- range .Typedefs }}
+	{{- if .Metadata }}
+	// metadata: {{ .Metadata }}
+	{{- end }}
+	type {{ .Name }} struct {
+	// fields ommitted for brevity
+	}
+{{- end }}`
+
+	g := gen.NewGenerator("something", map[string]gen.TagStyle{}, gen.WithTemplate(template))
+	g.ParseWithMetadata(unionSchema, "metadata for union schema")
+	g.ParseWithMetadata(mainSchema, map[string]any{"metadata for": "main schema"})
+
+	var buf bytes.Buffer
+	err = g.Write(&buf)
+	require.NoError(t, err)
+
+	formatted, err := format.Source(buf.Bytes())
+	require.NoError(t, err)
+
+	if *update {
+		err = os.WriteFile("testdata/golden_metadata.go", formatted, 0600)
+		require.NoError(t, err)
+	}
+
+	want, err := os.ReadFile("testdata/golden_metadata.go")
+	require.NoError(t, err)
+	assert.Equal(t, string(want), string(formatted))
+}
+
 // generate is a utility to run the generation and return the result as a tuple
 func generate(t *testing.T, schema string, gc gen.Config) ([]byte, []string) {
 	t.Helper()
diff --git a/gen/testdata/golden_metadata.go b/gen/testdata/golden_metadata.go
new file mode 100644
index 0000000..e0d09fc
--- /dev/null
+++ b/gen/testdata/golden_metadata.go
@@ -0,0 +1,12 @@
+// Code generated by avro/gen. DO NOT EDIT.
+package something
+
+// metadata: metadata for union schema
+type TestUnionType struct {
+	// fields ommitted for brevity
+}
+
+// metadata: map[metadata for:main schema]
+type TestMain struct {
+	// fields ommitted for brevity
+}