feat: add schema custom props as gen template parameters (#445)

hamba · Sep 9, 2024 · 6edc7d3 · 6edc7d3
1 parent 883fb8f
commit 6edc7d3
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/gen/gen.go b/gen/gen.go
@@ -299,12 +299,12 @@ func (g *Generator) resolveRecordSchema(schema *avro.RecordSchema) string {
 	fields := make([]field, len(schema.Fields()))
 	for i, f := range schema.Fields() {
 		typ := g.generate(f.Type())
-		fields[i] = g.newField(g.nameCaser.ToPascal(f.Name()), typ, f.Doc(), f.Name())
+		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, schema.String()))
+		g.typedefs = append(g.typedefs, newType(typeName, schema.Doc(), fields, schema.String(), schema.Props()))
 	}
 	return typeName
 }
@@ -379,13 +379,14 @@ func (g *Generator) resolveLogicalSchema(logicalType avro.LogicalType) string {
 	return typ
 }
 
-func (g *Generator) newField(name, typ, doc, avroFieldName string) field {
+func (g *Generator) newField(name, typ, doc, avroFieldName string, props map[string]any) field {
 	return field{
 		Name:          name,
 		Type:          typ,
 		AvroFieldName: avroFieldName,
 		Doc:           ensureTrailingPeriod(doc),
 		Tags:          g.tags,
+		Props:         props,
 	}
 }
 
@@ -443,14 +444,16 @@ type typedef struct {
 	Doc    string
 	Fields []field
 	Schema string
+	Props  map[string]any
 }
 
-func newType(name, doc string, fields []field, schema string) typedef {
+func newType(name, doc string, fields []field, schema string, props map[string]any) typedef {
 	return typedef{
 		Name:   name,
 		Doc:    ensureTrailingPeriod(doc),
 		Fields: fields,
 		Schema: schema,
+		Props:  props,
 	}
 }
 
@@ -460,4 +463,5 @@ type field struct {
 	Doc           string
 	AvroFieldName string
 	Tags          map[string]TagStyle
+	Props         map[string]any
 }
diff --git a/schema.go b/schema.go
@@ -366,6 +366,12 @@ func (p properties) Prop(name string) any {
 	return p.props[name]
 }
 
+// Props returns a map that contains all schema custom properties.
+// Any accidental change to the returned map will directly modify the schema custom properties.
+func (p properties) Props() map[string]any {
+	return p.props
+}
+
 func (p properties) marshalPropertiesToJSON(buf *bytes.Buffer) error {
 	sortedPropertyKeys := make([]string, 0, len(p.props))
 	for k := range p.props {

diff --git a/schema_test.go b/schema_test.go
@@ -177,6 +177,7 @@ func TestPrimitiveSchema_HandlesProps(t *testing.T) {
 	assert.Equal(t, avro.String, s.Type())
 	assert.Equal(t, "bar", s.(*avro.PrimitiveSchema).Prop("foo"))
 	assert.Equal(t, float64(1), s.(*avro.PrimitiveSchema).Prop("baz"))
+	assert.Equal(t, map[string]any{"foo": "bar", "baz": float64(1)}, s.(*avro.PrimitiveSchema).Props())
 }
 
 func TestRecordSchema(t *testing.T) {
@@ -507,6 +508,7 @@ func TestRecordSchema_HandlesProps(t *testing.T) {
    "namespace": "org.hamba.avro",
    "doc": "foo",
    "foo": "bar1",
+   "bar": "foo1",
    "fields": [
        {"name": "intField", "doc": "bar", "type": "int", "foo": "bar2"}
    ]
@@ -520,9 +522,11 @@ func TestRecordSchema_HandlesProps(t *testing.T) {
 	assert.Equal(t, avro.Record, s.Type())
 	assert.Equal(t, "foo", rs.Doc())
 	assert.Equal(t, "bar1", rs.Prop("foo"))
+	assert.Equal(t, map[string]any{"foo": "bar1", "bar": "foo1"}, rs.Props())
 	require.Len(t, rs.Fields(), 1)
 	assert.Equal(t, "bar", rs.Fields()[0].Doc())
 	assert.Equal(t, "bar2", rs.Fields()[0].Prop("foo"))
+	assert.Equal(t, map[string]any{"foo": "bar2"}, rs.Fields()[0].Props())
 }
 
 func TestRecordSchema_WithReference(t *testing.T) {
@@ -701,6 +705,7 @@ func TestEnumSchema_HandlesProps(t *testing.T) {
 	assert.Equal(t, avro.Enum, s.Type())
 	assert.Equal(t, "hello", es.Doc())
 	assert.Equal(t, "bar", es.Prop("foo"))
+	assert.Equal(t, map[string]any{"foo": "bar"}, es.Props())
 }
 
 func TestArraySchema(t *testing.T) {
@@ -749,6 +754,7 @@ func TestArraySchema_HandlesProps(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, avro.Array, s.Type())
 	assert.Equal(t, "bar", s.(*avro.ArraySchema).Prop("foo"))
+	assert.Equal(t, map[string]any{"foo": "bar"}, s.(*avro.ArraySchema).Props())
 }
 
 func TestMapSchema(t *testing.T) {
@@ -797,6 +803,7 @@ func TestMapSchema_HandlesProps(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, avro.Map, s.Type())
 	assert.Equal(t, "bar", s.(*avro.MapSchema).Prop("foo"))
+	assert.Equal(t, map[string]any{"foo": "bar"}, s.(*avro.MapSchema).Props())
 }
 
 func TestUnionSchema(t *testing.T) {
@@ -972,6 +979,7 @@ func TestFixedSchema_HandlesProps(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, avro.Fixed, s.Type())
 	assert.Equal(t, "bar", s.(*avro.FixedSchema).Prop("foo"))
+	assert.Equal(t, map[string]any{"foo": "bar"}, s.(*avro.FixedSchema).Props())
 }
 
 func TestSchema_LogicalTypes(t *testing.T) {