Feat: Add ability to generate files in folders

Adding a new feature flag `prefix_schema_files_with_package` which results in the schema files being generated in folders based on the package of the message.
flipp-oss · Apr 7, 2024 · acd8339 · acd8339
1 parent e06a9c2
commit acd8339
Show file tree

Hide file tree

Showing 11 changed files with 612 additions and 112 deletions.
diff --git a/README.md b/README.md
@@ -146,6 +146,19 @@ message MyRecord {
 }
 ```
 
+* `prefix_schema_files_with_package` - if set to true, files will be generated into folders matching the proto package. E.g. :
+if set to true, will remove the prefixes from enum values. E.g. if you have an enum like:
+```protobuf
+package my.test.data;
+message Yowza {
+  float hoo_boy = 1;
+}
+```
+
+...with this option on, it will be generated as:
+
+`./my.test.data/Yowza.avsc`
+
 ---
 
 To Do List:

diff --git a/input/params.go b/input/params.go
@@ -1,19 +1,20 @@
 package input
 
 import (
-  "google.golang.org/protobuf/proto"
-  "google.golang.org/protobuf/types/pluginpb"
-  "io"
-  "os"
-  "strings"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/pluginpb"
+	"io"
+	"os"
+	"strings"
 )
 
 type Params struct {
-	EmitOnly []string
-	NamespaceMap map[string]string
-	CollapseFields []string
-	RemoveEnumPrefixes bool
-	PreserveNonStringMaps bool
+	EmitOnly                     []string
+	NamespaceMap                 map[string]string
+	CollapseFields               []string
+	RemoveEnumPrefixes           bool
+	PreserveNonStringMaps        bool
+	PrefixSchemaFilesWithPackage bool
 }
 
 func ReadRequest() (*pluginpb.CodeGeneratorRequest, error) {
@@ -40,6 +41,8 @@ func parseRawParams(req *pluginpb.CodeGeneratorRequest) map[string]string {
 		paramStrings := strings.Split(token, "=")
 		if len(paramStrings) == 2 {
 			paramMap[paramStrings[0]] = paramStrings[1]
+		} else {
+			paramMap[paramStrings[0]] = "true"
 		}
 	}
 	return paramMap
@@ -63,6 +66,8 @@ func ParseParams(req *pluginpb.CodeGeneratorRequest) Params {
 			params.RemoveEnumPrefixes = v == "true"
 		} else if k == "preserve_non_string_maps" {
 			params.PreserveNonStringMaps = true
+		} else if k == "prefix_schema_files_with_package" {
+			params.PrefixSchemaFilesWithPackage = true
 		}
 	}
 	return params

diff --git a/main.go b/main.go
@@ -28,9 +28,16 @@ func processEnum(proto *descriptorpb.EnumDescriptorProto, protoPackage string) {
 	typeRepo.AddType(enum)
 }
 
+func generateSchemaFilename(record avro.Record) string {
+	if params.PrefixSchemaFilesWithPackage {
+		return fmt.Sprintf("%s/%s.avsc", record.Namespace, record.Name)
+	}
+	return fmt.Sprintf("%s.avsc", record.Name)
+}
+
 func generateFileResponse(record avro.Record) (*pluginpb.CodeGeneratorResponse_File, error) {
 	typeRepo.Start()
-	fileName := fmt.Sprintf("%s.avsc", record.Name)
+	fileName := generateSchemaFilename(record)
 	jsonObj, err := record.ToJSON(typeRepo)
 	if err != nil {
 		return nil, fmt.Errorf("error parsing record %s: %w", record.Name, err)

diff --git a/main_test.go b/main_test.go
@@ -1,13 +1,13 @@
 package main
 
 import (
-    "fmt"
-    "github.com/stretchr/testify/assert"
-    "os"
-  "os/exec"
-    "path/filepath"
-    "strings"
-  "testing"
+	"fmt"
+	"github.com/stretchr/testify/assert"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"testing"
 )
 
 // Overall approach taken from https://github.com/mix-php/mix/blob/master/src/grpc/protoc-gen-mix/plugin_test.go
@@ -16,124 +16,135 @@ import (
 // tests and instead act as protoc-gen-avro. This allows the test binary to
 // pass itself to protoc.
 func init() {
-    if os.Getenv("RUN_AS_PROTOC_GEN_AVRO") != "" {
-        main()
-        os.Exit(0)
-    }
+	if os.Getenv("RUN_AS_PROTOC_GEN_AVRO") != "" {
+		main()
+		os.Exit(0)
+	}
 }
 
-func fileNames(directory string, appendDirectory bool) ([]string, error) {
-    files, err := os.ReadDir(directory)
-    if err != nil {
-        return nil, fmt.Errorf("can't read %s directory: %w", directory, err)
-    }
-    var names []string
-    for _, file := range files {
-        if file.IsDir() {
-            continue
-        }
-        if appendDirectory {
-          names = append(names, filepath.Base(directory) + "/" + file.Name())
-        } else {
-            names = append(names, file.Name())
-        }
-    }
-    return names, nil
+func fileNames(directory string, appendDirectory bool, recurse bool) ([]string, error) {
+	files, err := os.ReadDir(directory)
+	if err != nil {
+		return nil, fmt.Errorf("can't read %s directory: %w", directory, err)
+	}
+	var names []string
+	for _, file := range files {
+		if file.IsDir() {
+			if recurse {
+				r_names, err := fileNames(directory+"/"+file.Name(), true, true)
+				if err != nil {
+					return nil, err
+				}
+				names = append(names, r_names...)
+			}
+			continue
+		}
+		if appendDirectory {
+			names = append(names, filepath.Base(directory)+"/"+file.Name())
+		} else {
+			names = append(names, file.Name())
+		}
+	}
+	return names, nil
 }
 
 func runTest(t *testing.T, directory string, options map[string]string) {
-    workdir, _ := os.Getwd()
-    tmpdir, err := os.MkdirTemp(workdir, "proto-test.")
-    if err != nil {
-        t.Fatal(err)
-    }
-    defer os.RemoveAll(tmpdir)
-
-    args := []string{
-        "-I.",
-        "--avro_out=" + tmpdir,
-    }
-    names, err := fileNames(workdir + "/testdata", true)
-    if err != nil {
-        t.Fatal(fmt.Errorf("testData fileNames %w", err))
-    }
-    for _, name := range names {
-        args = append(args, name)
-    }
-    for k, v := range options {
-        args = append(args, "--avro_opt=" + k + "=" + v)
-    }
-    protoc(t, args)
-
-    testDir := workdir + "/testdata/" + directory
-    if os.Getenv("UPDATE_SNAPSHOTS") != "" {
-        cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("cp %v/* %v", tmpdir, testDir))
-        cmd.Run()
-    } else {
-        assertEqualFiles(t, testDir, tmpdir)
-    }
+	workdir, _ := os.Getwd()
+	tmpdir, err := os.MkdirTemp(workdir, "proto-test.")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpdir)
+
+	args := []string{
+		"-I.",
+		"--avro_out=" + tmpdir,
+	}
+	names, err := fileNames(workdir+"/testdata", true, false)
+	if err != nil {
+		t.Fatal(fmt.Errorf("testData fileNames %w", err))
+	}
+	for _, name := range names {
+		args = append(args, name)
+	}
+	for k, v := range options {
+		args = append(args, "--avro_opt="+k+"="+v)
+	}
+	protoc(t, args)
+
+	testDir := workdir + "/testdata/" + directory
+	if os.Getenv("UPDATE_SNAPSHOTS") != "" {
+		cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("cp -r %v/* %v", tmpdir, testDir))
+		cmd.Run()
+	} else {
+		assertEqualFiles(t, testDir, tmpdir)
+	}
 }
 
 func Test_Base(t *testing.T) {
-    runTest(t, "base", map[string]string{})
+	runTest(t, "base", map[string]string{})
 }
 
 func Test_CollapseFields(t *testing.T) {
-    runTest(t, "collapse_fields", map[string]string{"collapse_fields": "StringList"})
+	runTest(t, "collapse_fields", map[string]string{"collapse_fields": "StringList"})
 }
 
 func Test_EmitOnly(t *testing.T) {
-    runTest(t, "emit_only", map[string]string{"emit_only": "Widget"})
+	runTest(t, "emit_only", map[string]string{"emit_only": "Widget"})
 }
 
 func Test_NamespaceMap(t *testing.T) {
-    runTest(t, "namespace_map", map[string]string{"namespace_map": "testdata:mynamespace"})
+	runTest(t, "namespace_map", map[string]string{"namespace_map": "testdata:mynamespace"})
 }
 
 func Test_PreserveNonStringMaps(t *testing.T) {
-    runTest(t, "preserve_non_string_maps", map[string]string{"preserve_non_string_maps": "true"})
+	runTest(t, "preserve_non_string_maps", map[string]string{"preserve_non_string_maps": "true"})
+}
+
+func Test_PrefixSchemaFilesWithPackage(t *testing.T) {
+	runTest(t, "prefix_schema_files_with_package", map[string]string{"prefix_schema_files_with_package": "true"})
 }
 
 func assertEqualFiles(t *testing.T, original, generated string) {
-    names, err := fileNames(original, false)
-    if err != nil {
-        t.Fatal(fmt.Errorf("original fileNames %w", err))
-    }
-    generatedNames, err := fileNames(generated, false)
-    if err != nil {
-        t.Fatal(fmt.Errorf("generated fileNames %w", err))
-    }
-    assert.Equal(t, names, generatedNames)
-    for i, name := range names {
-        originalData, err := os.ReadFile(original + "/" + name)
-        if err != nil {
-            t.Fatal("Can't find original file for comparison")
-        }
-
-        generatedData, err := os.ReadFile(generated + "/" + generatedNames[i])
-        if err != nil {
-            t.Fatal("Can't find generated file for comparison")
-        }
-        r := strings.NewReplacer("\r\n", "", "\n", "")
-        assert.Equal(t, r.Replace(string(originalData)), r.Replace(string(generatedData)))
-    }
+	names, err := fileNames(original, false, true)
+	if err != nil {
+		t.Fatal(fmt.Errorf("original fileNames %w", err))
+	}
+	generatedNames, err := fileNames(generated, false, true)
+	if err != nil {
+		t.Fatal(fmt.Errorf("generated fileNames %w", err))
+	}
+	assert.Equal(t, names, generatedNames)
+	for i, name := range names {
+		originalData, err := os.ReadFile(original + "/" + name)
+		if err != nil {
+			t.Fatal("Can't find original file for comparison")
+		}
+
+		generatedData, err := os.ReadFile(generated + "/" + generatedNames[i])
+		if err != nil {
+			t.Fatal("Can't find generated file for comparison")
+		}
+		r := strings.NewReplacer("\r\n", "", "\n", "")
+		assert.Equal(t, r.Replace(string(originalData)), r.Replace(string(generatedData)))
+	}
 }
 
 func protoc(t *testing.T, args []string) {
-    cmd := exec.Command("protoc", "--plugin=protoc-gen-avro=" + os.Args[0])
-    cmd.Args = append(cmd.Args, args...)
-    cmd.Env = append(os.Environ(), "RUN_AS_PROTOC_GEN_AVRO=1")
-    out, err := cmd.CombinedOutput()
-
-    if len(out) > 0 || err != nil {
-        t.Log("RUNNING: ", strings.Join(cmd.Args, " "))
-    }
-
-    if len(out) > 0 {
-        t.Log(string(out))
-    }
-
-    if err != nil {
-        t.Fatalf("protoc: %v", err)
-    }
+	cmd := exec.Command("protoc", "--plugin=protoc-gen-avro="+os.Args[0])
+	cmd.Args = append(cmd.Args, args...)
+	cmd.Env = append(os.Environ(), "RUN_AS_PROTOC_GEN_AVRO=1")
+	out, err := cmd.CombinedOutput()
+
+	if len(out) > 0 || err != nil {
+		t.Log("RUNNING: ", strings.Join(cmd.Args, " "))
+	}
+
+	if len(out) > 0 {
+		t.Log(string(out))
+	}
+
+	if err != nil {
+		t.Fatalf("protoc: %v", err)
+	}
 }
diff --git a/testdata/prefix_schema_files_with_package/testdata/AOneOf.avsc b/testdata/prefix_schema_files_with_package/testdata/AOneOf.avsc
@@ -0,0 +1,37 @@
+{
+  "type": "record",
+  "name": "AOneOf",
+  "namespace": "testdata",
+  "fields": [
+    {
+      "name": "oneof_types",
+      "type": [
+        {
+          "type": "record",
+          "name": "TypeA",
+          "namespace": "testdata",
+          "fields": [
+            {
+              "name": "foo",
+              "type": "string",
+              "default": ""
+            }
+          ]
+        },
+        {
+          "type": "record",
+          "name": "TypeB",
+          "namespace": "testdata",
+          "fields": [
+            {
+              "name": "bar",
+              "type": "string",
+              "default": ""
+            }
+          ]
+        }
+      ],
+      "default": null
+    }
+  ]
+}