fix: improve trim unused types behavior

sudorandom · Oct 30, 2024 · e61e892 · e61e892
1 parent 8442f9e
commit e61e892
Show file tree

Hide file tree

Showing 41 changed files with 206 additions and 1,854 deletions.
diff --git a/internal/converter/components.go b/internal/converter/components.go
@@ -36,6 +36,7 @@ func fileToComponents(opts options.Options, fd protoreflect.FileDescriptor) (*hi
 	components.Schemas = stateToSchema(st)
 
 	hasGetRequests := false
+	hasMethods := false
 
 	// Add requestBodies and responses for methods
 	services := fd.Services()
@@ -48,23 +49,10 @@ func fileToComponents(opts options.Options, fd protoreflect.FileDescriptor) (*hi
 			if hasGet {
 				hasGetRequests = true
 			}
+			hasMethods = true
 		}
 	}
 
-	components.Schemas.Set("connect-protocol-version", base.CreateSchemaProxy(&base.Schema{
-		Title:       "Connect-Protocol-Version",
-		Description: "Define the version of the Connect protocol",
-		Type:        []string{"number"},
-		Enum:        []*yaml.Node{utils.CreateIntNode("1")},
-		Const:       utils.CreateIntNode("1"),
-	}))
-
-	components.Schemas.Set("connect-timeout-header", base.CreateSchemaProxy(&base.Schema{
-		Title:       "Connect-Timeout-Ms",
-		Description: "Define the timeout, in ms",
-		Type:        []string{"number"},
-	}))
-
 	if hasGetRequests {
 		components.Schemas.Set("encoding", base.CreateSchemaProxy(&base.Schema{
 			Title:       "encoding",
@@ -91,43 +79,58 @@ func fileToComponents(opts options.Options, fd protoreflect.FileDescriptor) (*hi
 			},
 		}))
 	}
-	connectErrorProps := orderedmap.New[string, *base.SchemaProxy]()
-	connectErrorProps.Set("code", base.CreateSchemaProxy(&base.Schema{
-		Description: "The status code, which should be an enum value of [google.rpc.Code][google.rpc.Code].",
-		Type:        []string{"string"},
-		Examples:    []*yaml.Node{utils.CreateStringNode("CodeNotFound")},
-		Enum: []*yaml.Node{
-			utils.CreateStringNode("CodeCanceled"),
-			utils.CreateStringNode("CodeUnknown"),
-			utils.CreateStringNode("CodeInvalidArgument"),
-			utils.CreateStringNode("CodeDeadlineExceeded"),
-			utils.CreateStringNode("CodeNotFound"),
-			utils.CreateStringNode("CodeAlreadyExists"),
-			utils.CreateStringNode("CodePermissionDenied"),
-			utils.CreateStringNode("CodeResourceExhausted"),
-			utils.CreateStringNode("CodeFailedPrecondition"),
-			utils.CreateStringNode("CodeAborted"),
-			utils.CreateStringNode("CodeOutOfRange"),
-			utils.CreateStringNode("CodeInternal"),
-			utils.CreateStringNode("CodeUnavailable"),
-			utils.CreateStringNode("CodeDataLoss"),
-			utils.CreateStringNode("CodeUnauthenticated"),
-		},
-	}))
-	connectErrorProps.Set("message", base.CreateSchemaProxy(&base.Schema{
-		Description: "A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the [google.rpc.Status.details][google.rpc.Status.details] field, or localized by the client.",
-		Type:        []string{"string"},
-	}))
-	connectErrorProps.Set("detail", base.CreateSchemaProxyRef("#/components/schemas/google.protobuf.Any"))
-	components.Schemas.Set("connect.error", base.CreateSchemaProxy(&base.Schema{
-		Title:                "Connect Error",
-		Description:          `Error type returned by Connect: https://connectrpc.com/docs/go/errors/#http-representation`,
-		Properties:           connectErrorProps,
-		Type:                 []string{"object"},
-		AdditionalProperties: &base.DynamicValue[*base.SchemaProxy, bool]{N: 1, B: true},
-	}))
-	anyPair := util.NewGoogleAny()
-	components.Schemas.Set(anyPair.ID, base.CreateSchemaProxy(anyPair.Schema))
+	if hasMethods {
+		components.Schemas.Set("connect-protocol-version", base.CreateSchemaProxy(&base.Schema{
+			Title:       "Connect-Protocol-Version",
+			Description: "Define the version of the Connect protocol",
+			Type:        []string{"number"},
+			Enum:        []*yaml.Node{utils.CreateIntNode("1")},
+			Const:       utils.CreateIntNode("1"),
+		}))
+
+		components.Schemas.Set("connect-timeout-header", base.CreateSchemaProxy(&base.Schema{
+			Title:       "Connect-Timeout-Ms",
+			Description: "Define the timeout, in ms",
+			Type:        []string{"number"},
+		}))
+		connectErrorProps := orderedmap.New[string, *base.SchemaProxy]()
+		connectErrorProps.Set("code", base.CreateSchemaProxy(&base.Schema{
+			Description: "The status code, which should be an enum value of [google.rpc.Code][google.rpc.Code].",
+			Type:        []string{"string"},
+			Examples:    []*yaml.Node{utils.CreateStringNode("CodeNotFound")},
+			Enum: []*yaml.Node{
+				utils.CreateStringNode("CodeCanceled"),
+				utils.CreateStringNode("CodeUnknown"),
+				utils.CreateStringNode("CodeInvalidArgument"),
+				utils.CreateStringNode("CodeDeadlineExceeded"),
+				utils.CreateStringNode("CodeNotFound"),
+				utils.CreateStringNode("CodeAlreadyExists"),
+				utils.CreateStringNode("CodePermissionDenied"),
+				utils.CreateStringNode("CodeResourceExhausted"),
+				utils.CreateStringNode("CodeFailedPrecondition"),
+				utils.CreateStringNode("CodeAborted"),
+				utils.CreateStringNode("CodeOutOfRange"),
+				utils.CreateStringNode("CodeInternal"),
+				utils.CreateStringNode("CodeUnavailable"),
+				utils.CreateStringNode("CodeDataLoss"),
+				utils.CreateStringNode("CodeUnauthenticated"),
+			},
+		}))
+		connectErrorProps.Set("message", base.CreateSchemaProxy(&base.Schema{
+			Description: "A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the [google.rpc.Status.details][google.rpc.Status.details] field, or localized by the client.",
+			Type:        []string{"string"},
+		}))
+		connectErrorProps.Set("detail", base.CreateSchemaProxyRef("#/components/schemas/google.protobuf.Any"))
+		components.Schemas.Set("connect.error", base.CreateSchemaProxy(&base.Schema{
+			Title:                "Connect Error",
+			Description:          `Error type returned by Connect: https://connectrpc.com/docs/go/errors/#http-representation`,
+			Properties:           connectErrorProps,
+			Type:                 []string{"object"},
+			AdditionalProperties: &base.DynamicValue[*base.SchemaProxy, bool]{N: 1, B: true},
+		}))
+		anyPair := util.NewGoogleAny()
+		components.Schemas.Set(anyPair.ID, base.CreateSchemaProxy(anyPair.Schema))
+	}
 
 	return components, nil
 }
diff --git a/internal/converter/converter.go b/internal/converter/converter.go
@@ -181,11 +181,6 @@ func ConvertWithOptions(req *pluginpb.CodeGeneratorRequest, opts options.Options
 }
 
 func specToFile(opts options.Options, spec *v3.Document) (string, error) {
-	if opts.TrimUnusedTypes {
-		if err := trimUnusedTypes(spec); err != nil {
-			return "", err
-		}
-	}
 	switch opts.Format {
 	case "yaml":
 		return string(spec.RenderWithIndention(2)), nil
@@ -200,32 +195,6 @@ func specToFile(opts options.Options, spec *v3.Document) (string, error) {
 	}
 }
 
-func trimUnusedTypes(spec *v3.Document) error {
-	slog.Debug("trimming unused types")
-	b, err := spec.Render()
-	if err != nil {
-		return err
-	}
-	doc, err := libopenapi.NewDocument(b)
-	if err != nil {
-		return err
-	}
-	model, errs := doc.BuildV3Model()
-	if errs != nil {
-		return errors.Join(errs...)
-	}
-	model.Index.BuildIndex()
-	references := model.Model.Rolodex.GetRootIndex().GetAllReferences()
-	for pair := spec.Components.Schemas.First(); pair != nil; pair = pair.Next() {
-		ref := fmt.Sprintf("#/components/schemas/%s", pair.Key())
-		if _, ok := references[ref]; !ok {
-			slog.Debug("trimming unused type", "name", pair.Key())
-			spec.Components.Schemas.Delete(pair.Key())
-		}
-	}
-	return nil
-}
-
 func appendToSpec(opts options.Options, spec *v3.Document, fd protoreflect.FileDescriptor) error {
 	gnostic.SpecWithFileAnnotations(spec, fd)
 	components, err := fileToComponents(opts, fd)

diff --git a/internal/converter/schema.go b/internal/converter/schema.go
@@ -34,16 +34,19 @@ func NewState(opts options.Options) *State {
 func (st *State) CollectFile(tt protoreflect.FileDescriptor) {
 	st.CurrentFile = tt
 
-	// Files can have enums
-	enums := tt.Enums()
-	for i := 0; i < enums.Len(); i++ {
-		st.CollectEnum(enums.Get(i))
-	}
+	// Only collect types from the root if TrimUnusedTypes is off
+	if !st.Opts.TrimUnusedTypes {
+		// Files can have enums
+		enums := tt.Enums()
+		for i := 0; i < enums.Len(); i++ {
+			st.CollectEnum(enums.Get(i))
+		}
 
-	// Files can have messages
-	messages := tt.Messages()
-	for i := 0; i < messages.Len(); i++ {
-		st.CollectMessage(messages.Get(i))
+		// Files can have messages
+		messages := tt.Messages()
+		for i := 0; i < messages.Len(); i++ {
+			st.CollectMessage(messages.Get(i))
+		}
 	}
 
 	// Also make sure to pick up messages referenced in service methods

diff --git a/internal/converter/testdata/fileset.binpb b/internal/converter/testdata/fileset.binpb
diff --git a/internal/converter/testdata/proto_names/output/proto_names.openapi.json b/internal/converter/testdata/proto_names/output/proto_names.openapi.json
@@ -32,77 +32,6 @@
         },
         "title": "Message",
         "additionalProperties": false
-      },
-      "connect-protocol-version": {
-        "type": "number",
-        "title": "Connect-Protocol-Version",
-        "enum": [
-          1
-        ],
-        "description": "Define the version of the Connect protocol",
-        "const": 1
-      },
-      "connect-timeout-header": {
-        "type": "number",
-        "title": "Connect-Timeout-Ms",
-        "description": "Define the timeout, in ms"
-      },
-      "connect.error": {
-        "type": "object",
-        "properties": {
-          "code": {
-            "type": "string",
-            "examples": [
-              "CodeNotFound"
-            ],
-            "enum": [
-              "CodeCanceled",
-              "CodeUnknown",
-              "CodeInvalidArgument",
-              "CodeDeadlineExceeded",
-              "CodeNotFound",
-              "CodeAlreadyExists",
-              "CodePermissionDenied",
-              "CodeResourceExhausted",
-              "CodeFailedPrecondition",
-              "CodeAborted",
-              "CodeOutOfRange",
-              "CodeInternal",
-              "CodeUnavailable",
-              "CodeDataLoss",
-              "CodeUnauthenticated"
-            ],
-            "description": "The status code, which should be an enum value of [google.rpc.Code][google.rpc.Code]."
-          },
-          "message": {
-            "type": "string",
-            "description": "A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the [google.rpc.Status.details][google.rpc.Status.details] field, or localized by the client."
-          },
-          "detail": {
-            "$ref": "#/components/schemas/google.protobuf.Any"
-          }
-        },
-        "title": "Connect Error",
-        "additionalProperties": true,
-        "description": "Error type returned by Connect: https://connectrpc.com/docs/go/errors/#http-representation"
-      },
-      "google.protobuf.Any": {
-        "type": "object",
-        "properties": {
-          "type": {
-            "type": "string"
-          },
-          "value": {
-            "type": "string",
-            "format": "binary"
-          },
-          "debug": {
-            "type": "object",
-            "additionalProperties": true
-          }
-        },
-        "additionalProperties": true,
-        "description": "Contains an arbitrary serialized message along with a @type that describes the type of the serialized message."
       }
     }
   },

diff --git a/internal/converter/testdata/proto_names/output/proto_names.openapi.yaml b/internal/converter/testdata/proto_names/output/proto_names.openapi.yaml
@@ -24,60 +24,4 @@ components:
           title: UPPER_SNAKE_CASE
       title: Message
       additionalProperties: false
-    connect-protocol-version:
-      type: number
-      title: Connect-Protocol-Version
-      enum:
-        - 1
-      description: Define the version of the Connect protocol
-      const: 1
-    connect-timeout-header:
-      type: number
-      title: Connect-Timeout-Ms
-      description: Define the timeout, in ms
-    connect.error:
-      type: object
-      properties:
-        code:
-          type: string
-          examples:
-            - CodeNotFound
-          enum:
-            - CodeCanceled
-            - CodeUnknown
-            - CodeInvalidArgument
-            - CodeDeadlineExceeded
-            - CodeNotFound
-            - CodeAlreadyExists
-            - CodePermissionDenied
-            - CodeResourceExhausted
-            - CodeFailedPrecondition
-            - CodeAborted
-            - CodeOutOfRange
-            - CodeInternal
-            - CodeUnavailable
-            - CodeDataLoss
-            - CodeUnauthenticated
-          description: The status code, which should be an enum value of [google.rpc.Code][google.rpc.Code].
-        message:
-          type: string
-          description: A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the [google.rpc.Status.details][google.rpc.Status.details] field, or localized by the client.
-        detail:
-          $ref: '#/components/schemas/google.protobuf.Any'
-      title: Connect Error
-      additionalProperties: true
-      description: 'Error type returned by Connect: https://connectrpc.com/docs/go/errors/#http-representation'
-    google.protobuf.Any:
-      type: object
-      properties:
-        type:
-          type: string
-        value:
-          type: string
-          format: binary
-        debug:
-          type: object
-          additionalProperties: true
-      additionalProperties: true
-      description: Contains an arbitrary serialized message along with a @type that describes the type of the serialized message.
 security: []