Also generate json schema using the JSON name as the primary name

internal/cmd/jsonschema-generate-testdata/main.go

@@ -24,6 +24,7 @@ import (
 	_ "github.com/bufbuild/protoschema-plugins/internal/gen/proto/bufext/cel/expr/conformance/proto3"
 	"github.com/bufbuild/protoschema-plugins/internal/protoschema/golden"
 	"github.com/bufbuild/protoschema-plugins/internal/protoschema/jsonschema"
+	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
 func main() {
@@ -40,37 +41,45 @@ func run() error {
 		return fmt.Errorf("usage: %s [output dir]", os.Args[0])
 	}
 	outputDir := os.Args[1]
+	// Make sure the directory exists
+	if err := os.MkdirAll(outputDir, 0755); err != nil {
+		return err
+	}
 
 	testDescs, err := golden.GetTestDescriptors("./internal/testdata")
 	if err != nil {
 		return err
 	}
 	for _, testDesc := range testDescs {
-		// Generate the JSON schema
-		result := jsonschema.Generate(testDesc)
-
-		// Make sure the directory exists
-		if err := os.MkdirAll(outputDir, 0755); err != nil {
+		// Generate the JSON schema with proto names.
+		if err := writeJSONSchema(outputDir, jsonschema.Generate(testDesc)); err != nil {
 			return err
 		}
+		// Generate the JSON schema with JSON names.
+		if err := writeJSONSchema(outputDir, jsonschema.Generate(testDesc, jsonschema.WithJSONNames())); err != nil {
+			return err
+		}
+	}
+	return nil
+}
 
-		for _, jsonSchema := range result {
-			// Serialize the JSON
-			data, err := json.MarshalIndent(jsonSchema, "", "  ")
-			if err != nil {
-				return err
-			}
-			identifier, ok := jsonSchema["$id"].(string)
-			if !ok {
-				return errors.New("expected $id to be a string")
-			}
-			if identifier == "" {
-				return errors.New("expected $id to be non-empty")
-			}
-			filePath := filepath.Join(outputDir, identifier)
-			if err := golden.GenerateGolden(filePath, string(data)+"\n"); err != nil {
-				return err
-			}
+func writeJSONSchema(outputDir string, schema map[protoreflect.FullName]map[string]interface{}) error {
+	for _, jsonSchema := range schema {
+		// Serialize the JSON
+		data, err := json.MarshalIndent(jsonSchema, "", "  ")
+		if err != nil {
+			return err
+		}
+		identifier, ok := jsonSchema["$id"].(string)
+		if !ok {
+			return errors.New("expected $id to be a string")
+		}
+		if identifier == "" {
+			return errors.New("expected $id to be non-empty")
+		}
+		filePath := filepath.Join(outputDir, identifier)
+		if err := golden.GenerateGolden(filePath, string(data)+"\n"); err != nil {
+			return err
 		}
 	}
 	return nil

internal/protoschema/jsonschema/jsonschema.go

@@ -41,22 +41,38 @@ const (
 	FieldIgnore
 )
 
-// Generate generates a JSON schema for the given message descriptor.
-func Generate(input protoreflect.MessageDescriptor) map[protoreflect.FullName]map[string]interface{} {
+type GeneratorOption func(*jsonSchemaGenerator)
+
+// WithJSONNames sets the generator to use JSON field names as the primary name.
+func WithJSONNames() GeneratorOption {
+	return func(p *jsonSchemaGenerator) {
+		p.useJSONNames = true
+	}
+}
+
+// Generate generates a JSON schema for the given message descriptor, with protobuf field names.
+func Generate(input protoreflect.MessageDescriptor, opts ...GeneratorOption) map[protoreflect.FullName]map[string]interface{} {
 	generator := &jsonSchemaGenerator{
 		result: make(map[protoreflect.FullName]map[string]interface{}),
 	}
 	generator.custom = generator.makeWktGenerators()
+	for _, opt := range opts {
+		opt(generator)
+	}
 	generator.generate(input)
 	return generator.result
 }
 
 type jsonSchemaGenerator struct {
-	result map[protoreflect.FullName]map[string]interface{}
-	custom map[protoreflect.FullName]func(map[string]interface{}, protoreflect.MessageDescriptor)
+	result       map[protoreflect.FullName]map[string]interface{}
+	custom       map[protoreflect.FullName]func(map[string]interface{}, protoreflect.MessageDescriptor)
+	useJSONNames bool
 }
 
 func (p *jsonSchemaGenerator) getID(desc protoreflect.Descriptor) string {
+	if p.useJSONNames {
+		return string(desc.FullName()) + ".jsonschema.json"
+	}
 	return string(desc.FullName()) + ".schema.json"
 }
 
@@ -94,13 +110,20 @@ func (p *jsonSchemaGenerator) generateDefault(result map[string]interface{}, des
 		// TODO: Add an option to include custom alias.
 		aliases := make([]string, 0, 1)
 
-		if visibility == FieldHide {
+		switch {
+		case visibility == FieldHide:
 			aliases = append(aliases, string(field.Name()))
 			if field.JSONName() != string(field.Name()) {
 				aliases = append(aliases, field.JSONName())
 			}
-		} else {
-			// TODO: Optionally make the json name the 'primary' name.
+		case p.useJSONNames:
+			// Use the JSON name as the primary name.
+			properties[field.JSONName()] = fieldSchema
+			if field.JSONName() != string(field.Name()) {
+				aliases = append(aliases, string(field.Name()))
+			}
+		default:
+			// Use the proto name as the primary name.
 			properties[string(field.Name())] = fieldSchema
 			if field.JSONName() != string(field.Name()) {
 				aliases = append(aliases, field.JSONName())

internal/protoschema/plugin/pluginjsonschema/pluginjsonschema.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/bufbuild/protoplugin"
 	"github.com/bufbuild/protoschema-plugins/internal/protoschema/jsonschema"
+	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
 // Handle implements protoplugin.Handler and is the main entry point for the plugin.
@@ -38,31 +39,47 @@ func Handle(
 	for _, fileDescriptor := range fileDescriptors {
 		for i := range fileDescriptor.Messages().Len() {
 			messageDescriptor := fileDescriptor.Messages().Get(i)
-
-			for _, entry := range jsonschema.Generate(messageDescriptor) {
-				data, err := json.MarshalIndent(entry, "", "  ")
-				if err != nil {
-					return err
-				}
-				identifier, ok := entry["$id"].(string)
-				if !ok {
-					return fmt.Errorf("expected unique id for message %q to be a string, got type %T", messageDescriptor.FullName(), entry["$id"])
-				}
-				if identifier == "" {
-					return fmt.Errorf("expected unique id for message %q to be a non-empty string", messageDescriptor.FullName())
-				}
-				if seenIdentifiers[identifier] {
-					continue
-				}
-				responseWriter.AddFile(
-					identifier,
-					string(data)+"\n",
-				)
-				seenIdentifiers[identifier] = true
+			protoNameSchema := jsonschema.Generate(messageDescriptor)
+			if err := writeFiles(responseWriter, messageDescriptor, protoNameSchema, seenIdentifiers); err != nil {
+				return err
+			}
+			jsonNameSchema := jsonschema.Generate(messageDescriptor, jsonschema.WithJSONNames())
+			if err := writeFiles(responseWriter, messageDescriptor, jsonNameSchema, seenIdentifiers); err != nil {
+				return err
 			}
 		}
 	}
 
 	responseWriter.SetFeatureProto3Optional()
 	return nil
 }
+
+func writeFiles(
+	responseWriter protoplugin.ResponseWriter,
+	messageDescriptor protoreflect.MessageDescriptor,
+	schema map[protoreflect.FullName]map[string]interface{},
+	seenIdentifiers map[string]bool,
+) error {
+	for _, entry := range schema {
+		data, err := json.MarshalIndent(entry, "", "  ")
+		if err != nil {
+			return err
+		}
+		identifier, ok := entry["$id"].(string)
+		if !ok {
+			return fmt.Errorf("expected unique id for message %q to be a string, got type %T", messageDescriptor.FullName(), entry["$id"])
+		}
+		if identifier == "" {
+			return fmt.Errorf("expected unique id for message %q to be a non-empty string", messageDescriptor.FullName())
+		}
+		if seenIdentifiers[identifier] {
+			continue
+		}
+		responseWriter.AddFile(
+			identifier,
+			string(data)+"\n",
+		)
+		seenIdentifiers[identifier] = true
+	}
+	return nil
+}