CLOUDP-325047: add support for Extentensios to the Spec Struct

tools/cli/internal/openapi/openapi.go

@@ -16,6 +16,7 @@ package openapi
 
 //go:generate mockgen -destination=../openapi/mock_openapi.go -package=openapi github.com/mongodb/openapi/tools/cli/internal/openapi Parser,Merger
 import (
+	"encoding/json"
 	"log"
 
 	"github.com/getkin/kin-openapi/openapi3"
@@ -26,6 +27,7 @@ import (
 // Spec is a struct is a 1-to-1 copy of the Spec struct in the openapi3 package.
 // We need this to override the order of the fields in the struct.
 type Spec struct {
+	Extensions   map[string]any                `json:"-" yaml:"-"`
 	OpenAPI      string                        `json:"openapi" yaml:"openapi"`
 	Security     openapi3.SecurityRequirements `json:"security,omitempty" yaml:"security,omitempty"`
 	Servers      openapi3.Servers              `json:"servers,omitempty" yaml:"servers,omitempty"`
@@ -34,7 +36,6 @@ type Spec struct {
 	Paths        *openapi3.Paths               `json:"paths" yaml:"paths"`
 	Components   *openapi3.Components          `json:"components,omitempty" yaml:"components,omitempty"`
 	ExternalDocs *openapi3.ExternalDocs        `json:"externalDocs,omitempty" yaml:"externalDocs,omitempty"`
-	Extensions   map[string]any                `json:"-" yaml:"-"`
 }
 type Parser interface {
 	CreateOpenAPISpecFromPath(string) (*load.SpecInfo, error)
@@ -98,6 +99,7 @@ func NewOasDiffWithSpecInfo(base, external *load.SpecInfo, config *diff.Config)
 
 func newSpec(spec *openapi3.T) *Spec {
 	return &Spec{
+		Extensions:   spec.Extensions,
 		OpenAPI:      spec.OpenAPI,
 		Components:   spec.Components,
 		Info:         spec.Info,
@@ -108,3 +110,45 @@ func newSpec(spec *openapi3.T) *Spec {
 		ExternalDocs: spec.ExternalDocs,
 	}
 }
+
+// MarshalJSON returns the JSON encoding of Spec.
+// We need a custom definition of MarshalJSON to include support for
+// Extensions   map[string]any                `json:"-" yaml:"-"` where
+// we only what to serialize the value of the field.
+func (doc *Spec) MarshalJSON() ([]byte, error) {
+	x, err := doc.MarshalYAML()
+	if err != nil {
+		return nil, err
+	}
+	return json.Marshal(x)
+}
+
+// MarshalYAML returns the YAML encoding of Spec.
+func (doc *Spec) MarshalYAML() (any, error) {
+	if doc == nil {
+		return nil, nil
+	}
+	m := make(map[string]any, 4+len(doc.Extensions))
+	for k, v := range doc.Extensions {
+		m[k] = v
+	}
+	m["openapi"] = doc.OpenAPI
+	if x := doc.Components; x != nil {
+		m["components"] = x
+	}
+	m["info"] = doc.Info
+	m["paths"] = doc.Paths
+	if x := doc.Security; len(x) != 0 {
+		m["security"] = x
+	}
+	if x := doc.Servers; len(x) != 0 {
+		m["servers"] = x
+	}
+	if x := doc.Tags; len(x) != 0 {
+		m["tags"] = x
+	}
+	if x := doc.ExternalDocs; x != nil {
+		m["externalDocs"] = x
+	}
+	return m, nil
+}

tools/cli/internal/openapi/openapi_test.go

@@ -0,0 +1,220 @@
+// Copyright 2025 MongoDB Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package openapi
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	"github.com/getkin/kin-openapi/openapi3"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSpec_MarshalJSON(t *testing.T) {
+	minimalInfo := &openapi3.Info{Title: "Test API", Version: "1.0.0"}
+
+	tests := []struct {
+		name       string
+		spec       *Spec
+		jsonOutput string
+		wantErr    bool
+	}{
+		{
+			name: "spec with nil extensions",
+			spec: &Spec{
+				OpenAPI:    "3.0.3",
+				Info:       minimalInfo,
+				Paths:      &openapi3.Paths{},
+				Extensions: nil,
+			},
+			jsonOutput: `{"info":{"title":"Test API","version":"1.0.0"},"openapi":"3.0.3","paths":{}}`,
+			wantErr:    false,
+		},
+		{
+			name: "spec with empty extensions",
+			spec: &Spec{
+				OpenAPI:    "3.0.3",
+				Info:       minimalInfo,
+				Paths:      &openapi3.Paths{},
+				Extensions: map[string]any{},
+			},
+			jsonOutput: `{"info":{"title":"Test API","version":"1.0.0"},"openapi":"3.0.3","paths":{}}`,
+			wantErr:    false,
+		},
+		{
+			name: "spec with single string extension",
+			spec: &Spec{
+				OpenAPI: "3.0.3",
+				Info:    minimalInfo,
+				Paths:   &openapi3.Paths{},
+				Extensions: map[string]any{
+					"x-custom-string": "hello world",
+				},
+			},
+			jsonOutput: `{
+  "info": {
+    "title": "Test API",
+    "version": "1.0.0"
+  },
+  "openapi": "3.0.3",
+  "paths": {},
+  "x-custom-string": "hello world"
+}`,
+			wantErr: false,
+		},
+		{
+			name: "spec with multiple extensions of different types",
+			spec: &Spec{
+				OpenAPI: "3.0.3",
+				Info:    minimalInfo,
+				Paths:   &openapi3.Paths{},
+				Extensions: map[string]any{
+					"x-custom-string": "hello",
+					"x-custom-number": 123.45,
+					"x-custom-bool":   true,
+					"x-custom-null":   nil,
+				},
+			},
+			jsonOutput: `{
+  "info": {
+    "title": "Test API",
+    "version": "1.0.0"
+  },
+  "openapi": "3.0.3",
+  "paths": {},
+  "x-custom-bool": true,
+  "x-custom-null": null,
+  "x-custom-number": 123.45,
+  "x-custom-string": "hello"
+}`,
+			wantErr: false,
+		},
+		{
+			name: "spec with nested object extension",
+			spec: &Spec{
+				OpenAPI: "3.0.3",
+				Info:    minimalInfo,
+				Paths:   &openapi3.Paths{},
+				Extensions: map[string]any{
+					"x-custom-object": map[string]any{
+						"key1": "value1",
+						"key2": 100,
+					},
+				},
+			},
+			jsonOutput: `{
+  "info": {
+    "title": "Test API",
+    "version": "1.0.0"
+  },
+  "openapi": "3.0.3",
+  "paths": {},
+  "x-custom-object": {
+    "key1": "value1",
+    "key2": 100
+  }
+}`,
+			wantErr: false,
+		},
+		{
+			name: "spec with array extension",
+			spec: &Spec{
+				OpenAPI: "3.0.3",
+				Info:    minimalInfo,
+				Paths:   &openapi3.Paths{},
+				Extensions: map[string]any{
+					"x-custom-array": []any{"a", 2, true, map[string]any{"nested": "item"}},
+				},
+			},
+			jsonOutput: `{
+  "info": {
+    "title": "Test API",
+    "version": "1.0.0"
+  },
+  "openapi": "3.0.3",
+  "paths": {},
+  "x-custom-array": [
+    "a",
+    2,
+    true,
+    {
+      "nested": "item"
+    }
+  ]
+}`,
+			wantErr: false,
+		},
+		{
+			name: "spec with extensions and other optional fields (e.g., components)",
+			spec: &Spec{
+				OpenAPI: "3.0.3",
+				Info:    minimalInfo,
+				Paths:   &openapi3.Paths{},
+				Components: &openapi3.Components{
+					Schemas: openapi3.Schemas{
+						"MySchema": &openapi3.SchemaRef{
+							Value: openapi3.NewObjectSchema().WithProperty("id", openapi3.NewIntegerSchema()),
+						},
+					},
+				},
+				Extensions: map[string]any{
+					"x-marker": "present",
+				},
+			},
+			jsonOutput: `{
+  "components": {
+    "schemas": {
+      "MySchema": {
+        "properties": {
+          "id": {
+            "type": "integer"
+          }
+        },
+        "type": "object"
+      }
+    }
+  },
+  "info": {
+    "title": "Test API",
+    "version": "1.0.0"
+  },
+  "openapi": "3.0.3",
+  "paths": {},
+  "x-marker": "present"
+}`,
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gotBytes, err := tt.spec.MarshalJSON()
+			require.NoError(t, err)
+
+			var gotMap map[string]any
+			err = json.Unmarshal(gotBytes, &gotMap)
+			require.NoError(t, err)
+
+			var wantMap map[string]any
+			err = json.Unmarshal([]byte(tt.jsonOutput), &wantMap)
+			require.NoError(t, err)
+			if !reflect.DeepEqual(gotMap, wantMap) {
+				gotPretty, _ := json.MarshalIndent(gotMap, "", "  ")
+				wantPretty, _ := json.MarshalIndent(wantMap, "", "  ")
+				t.Errorf("Spec.MarshalJSON() mismatch:\nGot:\n%s\nWant:\n%s", string(gotPretty), string(wantPretty))
+			}
+		})
+	}
+}