feat: add encoders to avrogen

cmd/avrogen/main.go

@@ -19,6 +19,7 @@ type config struct {
 	Out      string
 	Tags     string
 	FullName bool
+	Encoders bool
 }
 
 func main() {
@@ -33,6 +34,7 @@ func realMain(args []string, out io.Writer) int {
 	flgs.StringVar(&cfg.Out, "o", "", "The output file path.")
 	flgs.StringVar(&cfg.Tags, "tags", "", "The additional field tags <tag-name>:{snake|camel|upper-camel|kebab}>[,...]")
 	flgs.BoolVar(&cfg.FullName, "fullname", false, "Use the full name of the Record schema to create the struct name.")
+	flgs.BoolVar(&cfg.Encoders, "encoders", false, "Generate encoders for the structs.")
 	flgs.Usage = func() {
 		_, _ = fmt.Fprintln(out, "Usage: avrogen [options] schemas")
 		_, _ = fmt.Fprintln(out, "Options:")
@@ -52,7 +54,11 @@ func realMain(args []string, out io.Writer) int {
 		return 1
 	}
 
-	g := gen.NewGenerator(cfg.Pkg, tags, gen.WithFullName(cfg.FullName))
+	opts := []gen.OptsFunc{
+		gen.WithFullName(cfg.FullName),
+		gen.WithEncoders(cfg.Encoders),
+	}
+	g := gen.NewGenerator(cfg.Pkg, tags, opts...)
 	for _, file := range flgs.Args() {
 		schema, err := avro.ParseFiles(filepath.Clean(file))
 		if err != nil {

cmd/avrogen/main_test.go

@@ -117,6 +117,29 @@ func TestAvroGen_GeneratesSchemaWithFullname(t *testing.T) {
 	assert.Equal(t, want, got)
 }
 
+func TestAvroGen_GeneratesSchemaWithEncoders(t *testing.T) {
+	path, err := os.MkdirTemp("./", "avrogen")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = os.RemoveAll(path) })
+
+	file := filepath.Join(path, "test.go")
+	args := []string{"avrogen", "-pkg", "testpkg", "-o", file, "-encoders", "testdata/schema.avsc"}
+	gotCode := realMain(args, io.Discard)
+	require.Equal(t, 0, gotCode)
+
+	got, err := os.ReadFile(file)
+	require.NoError(t, err)
+
+	if *update {
+		err = os.WriteFile("testdata/golden_encoders.go", got, 0600)
+		require.NoError(t, err)
+	}
+
+	want, err := os.ReadFile("testdata/golden_encoders.go")
+	require.NoError(t, err)
+	assert.Equal(t, want, got)
+}
+
 func TestParseTags(t *testing.T) {
 	tests := []struct {
 		name string

gen/gen.go

@@ -19,6 +19,7 @@ type Config struct {
 	PackageName string
 	Tags        map[string]TagStyle
 	FullName    bool
+	Encoders    bool
 }
 
 // TagStyle defines the styling for a tag.
@@ -39,6 +40,7 @@ const outputTemplate = `package {{ .PackageName }}
 
 // Code generated by avro/gen. DO NOT EDIT.
 
+{{- $encoders := .WithEncoders }}
 {{ if len .Imports }}
 import (
 	{{- range .Imports }}
@@ -60,11 +62,36 @@ import (
 {{ end }}
 
 {{- range .Typedefs }}
+// {{ .Name }} is a generated struct.
 type {{ .Name }} struct {
 	{{- range .Fields }}
 		{{ .Name }} {{ .Type }} {{ .Tag }}
 	{{- end }}
+{{- if $encoders }}
+
+    schema avro.Schema
+{{- end }}
+}
+
+{{- if $encoders }}
+// Schema returns the schema for {{ .Name }}.
+func (o *{{ .Name }}) Schema() avro.Schema {
+  if o.schema == nil {
+    o.schema = avro.MustParse(` + "`{{ .Schema }}`" + `)
+  }
+  return o.schema
+}
+
+// Unmarshal decodes b into the receiver.
+func (o *{{ .Name }}) Unmarshal(b []byte) error {
+	return avro.Unmarshal(o.Schema(), b, o)
+}
+
+// Marshal encodes the receiver.
+func (o *{{ .Name }}) Marshal() ([]byte, error) {
+	return avro.Marshal(o.Schema(), o)
 }
+{{- end }}
 {{ end }}`
 
 var primitiveMappings = map[avro.Type]string{
@@ -95,6 +122,7 @@ func StructFromSchema(schema avro.Schema, w io.Writer, cfg Config) error {
 
 	opts := []OptsFunc{
 		WithFullName(cfg.FullName),
+		WithEncoders(cfg.Encoders),
 	}
 	g := NewGenerator(strcase.ToSnake(cfg.PackageName), cfg.Tags, opts...)
 	g.Parse(rec)
@@ -124,11 +152,23 @@ func WithFullName(b bool) OptsFunc {
 	}
 }
 
+// WithEncoders configures the generator to generate schema and encoders on
+// all objects.
+func WithEncoders(b bool) OptsFunc {
+	return func(g *Generator) {
+		g.encoders = b
+		if b {
+			g.thirdPartyImports = append(g.thirdPartyImports, "github.com/hamba/avro/v2")
+		}
+	}
+}
+
 // Generator generates Go structs from schemas.
 type Generator struct {
 	pkg      string
 	tags     map[string]TagStyle
 	fullName bool
+	encoders bool
 
 	imports           []string
 	thirdPartyImports []string
@@ -209,7 +249,7 @@ func (g *Generator) resolveRecordSchema(schema *avro.RecordSchema) string {
 
 	typeName := g.resolveTypeName(schema)
 	if !g.hasTypeDef(typeName) {
-		g.typedefs = append(g.typedefs, newType(typeName, fields))
+		g.typedefs = append(g.typedefs, newType(typeName, fields, schema.String()))
 	}
 	return typeName
 }
@@ -325,11 +365,13 @@ func (g *Generator) Write(w io.Writer) error {
 	}
 
 	data := struct {
+		WithEncoders      bool
 		PackageName       string
 		Imports           []string
 		ThirdPartyImports []string
 		Typedefs          []typedef
 	}{
+		WithEncoders:      g.encoders,
 		PackageName:       g.pkg,
 		Imports:           g.imports,
 		ThirdPartyImports: g.thirdPartyImports,
@@ -341,12 +383,14 @@ func (g *Generator) Write(w io.Writer) error {
 type typedef struct {
 	Name   string
 	Fields []field
+	Schema string
 }
 
-func newType(name string, fields []field) typedef {
+func newType(name string, fields []field, schema string) typedef {
 	return typedef{
 		Name:   name,
 		Fields: fields,
+		Schema: schema,
 	}
 }
 

gen/gen_test.go

@@ -152,6 +152,23 @@ func TestStruct_GenFromRecordSchemaWithFullName(t *testing.T) {
 	assert.Equal(t, string(want), string(file))
 }
 
+func TestStruct_GenFromRecordSchemaWithEncoders(t *testing.T) {
+	schema, err := os.ReadFile("testdata/golden.avsc")
+	require.NoError(t, err)
+
+	gc := gen.Config{PackageName: "Something", Encoders: true}
+	file, _ := generate(t, string(schema), gc)
+
+	if *update {
+		err = os.WriteFile("testdata/golden_encoders.go", file, 0600)
+		require.NoError(t, err)
+	}
+
+	want, err := os.ReadFile("testdata/golden_encoders.go")
+	require.NoError(t, err)
+	assert.Equal(t, string(want), string(file))
+}
+
 func TestGenerator(t *testing.T) {
 	unionSchema, err := avro.ParseFiles("testdata/uniontype.avsc")
 	require.NoError(t, err)