Fix marshal fragment with id and extension.

PiperOrigin-RevId: 811154837

Author: FHIR Team

go/jsonformat/marshaller.go

@@ -690,16 +690,36 @@ func (m *Marshaller) marshalNonPrimitiveFieldValue(f protoreflect.FieldDescripto
 }
 
 func (m *Marshaller) marshalReference(rpb protoreflect.Message) (jsonpbhelper.IsJSON, error) {
-	newRef, err := NewDenormalizedReference(rpb.Interface())
+	// NewDenormalizedReference creates a new Reference message for marshalling but
+	// doesn't transfer specific fields from the original fragment such as id and extension.
+	newRefMsg, err := NewDenormalizedReference(rpb.Interface())
 	if err != nil {
 		return nil, err
 	}
+	newRef := newRefMsg.ProtoReflect()
+
+	refOneofDesc := rpb.Descriptor().Oneofs().ByName("reference")
+	if refOneofDesc != nil && m.jsonFormat == formatPure {
+		if f := rpb.WhichOneof(refOneofDesc); f != nil && f.Name() == "fragment" {
+			fragmentMsg := rpb.Get(f).Message()
+
+			idField := fragmentMsg.Descriptor().Fields().ByName("id")
+			if idField != nil && fragmentMsg.Has(idField) {
+				newRef.Set(newRef.Descriptor().Fields().ByName("id"), fragmentMsg.Get(idField))
+			}
+
+			extField := fragmentMsg.Descriptor().Fields().ByName("extension")
+			if extField != nil && fragmentMsg.Has(extField) {
+				newRef.Set(newRef.Descriptor().Fields().ByName("extension"), fragmentMsg.Get(extField))
+			}
+		}
+	}
 	if m.jsonFormat != formatPure {
-		if err := normalizeRelativeReferenceAndIgnoreHistory(newRef); err != nil {
+		if err := normalizeRelativeReferenceAndIgnoreHistory(newRefMsg); err != nil {
 			return nil, err
 		}
 	}
-	return m.marshalMessageToMap(newRef.ProtoReflect())
+	return m.marshalMessageToMap(newRef)
 }
 
 func (m *Marshaller) marshalMessageToMap(pb protoreflect.Message) (jsonpbhelper.JSONObject, error) {
@@ -738,7 +758,7 @@ func (m *Marshaller) marshalMessageToMap(pb protoreflect.Message) (jsonpbhelper.
 	if err != nil {
 		return nil, err
 	}
-	if m.jsonFormat != formatPure && !jsonpbhelper.IsResourceType(pb.Descriptor()) && !jsonpbhelper.IsChoice(pb.Descriptor()){
+	if m.jsonFormat != formatPure && !jsonpbhelper.IsResourceType(pb.Descriptor()) && !jsonpbhelper.IsChoice(pb.Descriptor()) {
 		// Omit FHIR element ID fields for analytics json.
 		// See https://github.com/rbrush/sql-on-fhir/blob/master/sql-on-fhir.md#id-fields-omitted.
 		delete(decmap, "id")

go/jsonformat/marshaller_test.go

@@ -36,7 +36,9 @@ import (
 	r4codesystempb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/code_system_go_proto"
 	r4conditionpb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/condition_go_proto"
 	r4devicepb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/device_go_proto"
+	r4observationpb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/observation_go_proto"
 	r4patientpb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/patient_go_proto"
+	r4questionnairepb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/questionnaire_go_proto"
 	r4researchstudypb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/research_study_go_proto"
 	r4searchparampb "github.com/google/fhir/go/proto/google/fhir/proto/r4/core/resources/search_parameter_go_proto"
 	c5pb "github.com/google/fhir/go/proto/google/fhir/proto/r5/core/codes_go_proto"
@@ -750,6 +752,83 @@ func TestMarshalResource(t *testing.T) {
 			pretty: false,
 			want:   []byte(`{"id":"example","resourceType":"ResearchStudy","status":"draft","text":{"div":"<div xmlns=\"http://www.w3.org/1999/xhtml\">[Put rendering here]</div>","status":"generated"}}`),
 		},
+		{
+			name: "Fragment_extension",
+			inputs: []mvr{{
+				ver: fhirversion.R4,
+				r: &r4questionnairepb.Questionnaire{
+					Status: &r4questionnairepb.Questionnaire_StatusCode{
+						Value: c4pb.PublicationStatusCode_ACTIVE,
+					},
+					Contained: []*anypb.Any{
+						marshalToAny(t, &r4pb.ContainedResource{
+							OneofResource: &r4pb.ContainedResource_Observation{
+								Observation: &r4observationpb.Observation{
+									Id: &d4pb.Id{
+										Value: "Observation1",
+									},
+									Code: &d4pb.CodeableConcept{
+										Text: &d4pb.String{
+											Value: "test",
+										},
+									},
+									Status: &r4observationpb.Observation_StatusCode{
+										Value: c4pb.ObservationStatusCode_FINAL,
+									},
+									Subject: &d4pb.Reference{
+										Reference: &d4pb.Reference_Fragment{
+											Fragment: &d4pb.String{
+												Value: "Patient1",
+												Id:    &d4pb.String{Value: "123"},
+												Extension: []*d4pb.Extension{
+													&d4pb.Extension{
+														Url: &d4pb.Uri{
+															Value: "https://example.org",
+														},
+														Value: &d4pb.Extension_ValueX{
+															Choice: &d4pb.Extension_ValueX_StringValue{
+																StringValue: &d4pb.String{
+																	Value: "extension-value",
+																},
+															},
+														},
+													},
+												},
+											},
+										},
+									},
+								},
+							},
+						}),
+					},
+				},
+			}},
+			pretty: true,
+			want: []byte(`{
+        "contained": [
+          {
+            "code": {
+              "text": "test"
+            },
+            "id": "Observation1",
+            "resourceType": "Observation",
+            "status": "final",
+            "subject": {
+              "extension": [
+                {
+                  "url": "https://example.org",
+                  "valueString": "extension-value"
+                }
+              ],
+              "id": "123",
+              "reference": "#Patient1"
+            }
+          }
+        ],
+        "resourceType": "Questionnaire",
+        "status": "active"
+      }`),
+		},
 	}
 
 	for _, tt := range tests {