fix: support Pydantic BaseModel comparison

Signed-off-by: mao3267 <[email protected]>

flytepropeller/pkg/compiler/validators/typing.go

@@ -18,33 +18,58 @@ type trivialChecker struct {
 	literalType *flyte.LiteralType
 }
 
+func removeTitleFieldFromProperties(schema map[string]interface{}) {
+	properties, ok := schema["properties"].(*structpb.Value)
+	if !ok {
+		return
+	}
+
+	for _, p := range properties.GetStructValue().Fields {
+		delete(p.GetStructValue().Fields, "title")
+	}
+}
+
 func isSuperTypeInJSON(sourceMetaData, targetMetaData *structpb.Struct) bool {
 	// Since there are lots of field differences between draft-07 and draft 2020-12,
 	// we only support json schema with 2020-12 draft, which is generated here: https://github.com/flyteorg/flytekit/blob/ff2d0da686c82266db4dbf764a009896cf062349/flytekit/core/type_engine.py#L630-L639
 	_, upstreamIsDraft7 := sourceMetaData.Fields["$schema"]
 	_, downstreamIsDraft7 := targetMetaData.Fields["$schema"]
 
+	// We only support super type check for draft 2020-12
 	if upstreamIsDraft7 || downstreamIsDraft7 {
 		return false
 	}
 
-	copySourceSchema := make(map[string]interface{})
-	copyTargetSchema := make(map[string]interface{})
+	copySrcSchema := make(map[string]interface{})
+	copyTgtSchema := make(map[string]interface{})
 
 	for k, v := range sourceMetaData.Fields {
-		copySourceSchema[k] = v
+		copySrcSchema[k] = v
 	}
 
 	for k, v := range targetMetaData.Fields {
-		copyTargetSchema[k] = v
+		copyTgtSchema[k] = v
 	}
 
-	srcSchemaBytes, _ := json.Marshal(copySourceSchema)
-	tgtSchemaBytes, _ := json.Marshal(copyTargetSchema)
+	// The JSON schema generated by Pydantic.BaseModel includes a title field in its properties, repeatedly recording the property name.
+	// Since this title field is absent in the JSON schema generated for dataclass, we need to remove the title field from the properties to ensure equivalence.
+	removeTitleFieldFromProperties(copySrcSchema)
+	removeTitleFieldFromProperties(copyTgtSchema)
+
+	srcSchemaBytes, _ := json.Marshal(copySrcSchema)
+	tgtSchemaBytes, _ := json.Marshal(copyTgtSchema)
 
 	patch, _ := jsondiff.CompareJSON(srcSchemaBytes, tgtSchemaBytes)
 	for _, p := range patch {
-		if p.Type != jsondiff.OperationAdd {
+		// If additionalProperties is false, the field is not present in the schema from Pydantic.BaseModel.
+		// We handle this case by checking the relationships by ourselves.
+		if p.Type != jsondiff.OperationAdd && p.Path == "/additionalProperties" {
+			if p.Type == jsondiff.OperationRemove || p.Type == jsondiff.OperationReplace {
+				if p.OldValue != false {
+					return false
+				}
+			}
+		} else if p.Type != jsondiff.OperationAdd {
 			return false
 		} else if strings.Contains(p.Path, "required") {
 			return false
@@ -59,15 +84,29 @@ func isSameTypeInJSON(sourceMetaData, targetMetaData *structpb.Struct) bool {
 	_, upstreamIsDraft7 := sourceMetaData.Fields["$schema"]
 	_, downstreamIsDraft7 := targetMetaData.Fields["$schema"]
 
+	// If the schema version is different, we can't compare them.
 	if upstreamIsDraft7 != downstreamIsDraft7 {
 		return false
 	}
 
-	sourceSchema := sourceMetaData.Fields
-	targetSchema := targetMetaData.Fields
+	copySrcSchema := make(map[string]interface{})
+	copyTgtSchema := make(map[string]interface{})
+
+	for k, v := range sourceMetaData.Fields {
+		copySrcSchema[k] = v
+	}
+
+	for k, v := range targetMetaData.Fields {
+		copyTgtSchema[k] = v
+	}
+
+	// The JSON schema generated by Pydantic.BaseModel includes a title field in its properties, repeatedly recording the property name.
+	// Since this title field is absent in the JSON schema generated for dataclass, we need to remove the title field from the properties to ensure equivalence.
+	removeTitleFieldFromProperties(copySrcSchema)
+	removeTitleFieldFromProperties(copyTgtSchema)
 
-	srcSchemaBytes, _ := json.Marshal(sourceSchema)
-	tgtSchemaBytes, _ := json.Marshal(targetSchema)
+	srcSchemaBytes, _ := json.Marshal(copySrcSchema)
+	tgtSchemaBytes, _ := json.Marshal(copyTgtSchema)
 
 	patch, _ := jsondiff.CompareJSON(srcSchemaBytes, tgtSchemaBytes)
 	return len(patch) == 0