feat(gen): implement struct generation from schema registry subjects and versions (#498)

Author: ekazakas

cmd/avrogen/main.go

@@ -2,31 +2,35 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"flag"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 
 	"github.com/hamba/avro/v2"
 	"github.com/hamba/avro/v2/gen"
+	"github.com/hamba/avro/v2/registry"
 	"golang.org/x/tools/imports"
 )
 
 type config struct {
 	TemplateFileName string
 
-	Pkg         string
-	PkgDoc      string
-	Out         string
-	Tags        string
-	FullName    bool
-	Encoders    bool
-	FullSchema  bool
-	StrictTypes bool
-	Initialisms string
+	Pkg            string
+	PkgDoc         string
+	Out            string
+	Tags           string
+	FullName       bool
+	Encoders       bool
+	FullSchema     bool
+	StrictTypes    bool
+	Initialisms    string
+	SchemaRegistry string
 }
 
 func main() {
@@ -47,6 +51,7 @@ func realMain(args []string, stdout, stderr io.Writer) int {
 	flgs.BoolVar(&cfg.StrictTypes, "strict-types", false, "Use strict type sizes (e.g. int32) during generation.")
 	flgs.StringVar(&cfg.Initialisms, "initialisms", "", "Custom initialisms <VAL>[,...] for struct and field names.")
 	flgs.StringVar(&cfg.TemplateFileName, "template-filename", "", "Override output template with one loaded from file.")
+	flgs.StringVar(&cfg.SchemaRegistry, "schemaregistry", "", "The URL to schema registry, e.g.: http://localhost:8081.")
 	flgs.Usage = func() {
 		_, _ = fmt.Fprintln(stderr, "Usage: avrogen [options] schemas")
 		_, _ = fmt.Fprintln(stderr, "Options:")
@@ -88,13 +93,38 @@ func realMain(args []string, stdout, stderr io.Writer) int {
 		gen.WithStrictTypes(cfg.StrictTypes),
 		gen.WithFullSchema(cfg.FullSchema),
 	}
+
 	g := gen.NewGenerator(cfg.Pkg, tags, opts...)
-	for _, file := range flgs.Args() {
-		schema, err := avro.ParseFiles(filepath.Clean(file))
-		if err != nil {
-			_, _ = fmt.Fprintf(stderr, "Error: %v\n", err)
-			return 2
+	for _, entry := range flgs.Args() {
+		var schema avro.Schema
+
+		switch cfg.SchemaRegistry {
+		case "":
+			schema, err = avro.ParseFiles(filepath.Clean(entry))
+			if err != nil {
+				_, _ = fmt.Fprintf(stderr, "Error: %v\n", err)
+				return 2
+			}
+		default:
+			client, err := registry.NewClient(cfg.SchemaRegistry)
+			if err != nil {
+				_, _ = fmt.Fprintf(stderr, "Error: %v\n", err)
+				return 2
+			}
+
+			subject, version, err := parseSubjectVersion(entry)
+			if err != nil {
+				_, _ = fmt.Fprintf(stderr, "Error: %v\n", err)
+				return 2
+			}
+
+			schema, err = client.GetSchemaByVersion(context.Background(), subject, version)
+			if err != nil {
+				_, _ = fmt.Fprintf(stderr, "Error: %v\n", err)
+				return 2
+			}
 		}
+
 		g.Parse(schema)
 	}
 
@@ -205,3 +235,17 @@ func loadTemplate(templateFileName string) ([]byte, error) {
 	}
 	return os.ReadFile(filepath.Clean(templateFileName))
 }
+
+func parseSubjectVersion(entry string) (string, int, error) {
+	parts := strings.Split(entry, ":")
+	if len(parts) != 2 {
+		return "", -1, errors.New("entry must be of format subject:version")
+	}
+
+	version, err := strconv.Atoi(parts[1])
+	if err != nil {
+		return "", -1, err
+	}
+
+	return parts[0], version, nil
+}

gen/gen.go

@@ -27,6 +27,7 @@ type Config struct {
 	StrictTypes  bool
 	Initialisms  []string
 	LogicalTypes []LogicalType
+	Metadata     any
 }
 
 // TagStyle defines the styling for a tag.
@@ -85,6 +86,7 @@ func StructFromSchema(schema avro.Schema, w io.Writer, cfg Config) error {
 		WithInitialisms(cfg.Initialisms),
 		WithStrictTypes(cfg.StrictTypes),
 		WithFullSchema(cfg.FullSchema),
+		WithMetadata(cfg.Metadata),
 	}
 	for _, opt := range cfg.LogicalTypes {
 		opts = append(opts, WithLogicalType(opt))
@@ -168,6 +170,13 @@ func WithFullSchema(b bool) OptsFunc {
 	}
 }
 
+// WithMetadata configures the generator to store the metadata within the generation context.
+func WithMetadata(m any) OptsFunc {
+	return func(g *Generator) {
+		g.metadata = m
+	}
+}
+
 // LogicalType used when the name of the "LogicalType" field in the Avro schema matches the Name attribute.
 type LogicalType struct {
 	// Name of the LogicalType
@@ -213,6 +222,7 @@ type Generator struct {
 	strictTypes  bool
 	initialisms  []string
 	logicalTypes map[avro.LogicalType]LogicalType
+	metadata     any
 
 	imports           []string
 	thirdPartyImports []string
@@ -451,12 +461,14 @@ func (g *Generator) Write(w io.Writer) error {
 		Imports           []string
 		ThirdPartyImports []string
 		Typedefs          []typedef
+		Metadata          any
 	}{
 		WithEncoders: g.encoders,
 		PackageName:  g.pkg,
 		PackageDoc:   g.pkgdoc,
 		Imports:      append(g.imports, g.thirdPartyImports...),
 		Typedefs:     g.typedefs,
+		Metadata:     g.metadata,
 	}
 	return parsed.Execute(w, data)
 }