Allow customizing / disabling PolymorphicModelConverter

[MHajoha]

Is your feature request related to a problem? Please describe.
A request or response body referencing a polymorphic base class is represented in the generated OpenAPI spec using a oneOf of the inheritors of the base class, instead of the base class itself.

See https://github.com/MHajoha/springdoc-poly-demo for an example project. Springdoc currently generates both the request and response bodies of that example as:

"schema": {
  "oneOf": [
    {
      "$ref": "#/components/schemas/Impl1"
    },
    {
      "$ref": "#/components/schemas/Impl2"
    }
  ]
}

This strikes me as unnecessarily complicated and also seems to lose the discriminator information from BaseClass. Instead, I would like SpringDoc to generate:

"schema": {
  "$ref": "#/components/schemas/BaseClass"
}

This is what the model converter chain generates up until PolymorphicModelConverter, which unwraps BaseClass into the above oneOf

Describe the solution you'd like
Some simple way of either globally of on a per-model basis disabling the PolymorphicModelConverter, such as a spring boot property.

Describe alternatives you've considered
This problem can be worked around by replacing the PolymorphicModelConverter with one which simply delegates to the rest of the chain:

@Bean
public PolymorphicModelConverter noopPolymorphicModelConverter() {
	return new PolymorphicModelConverter() {

		@Nullable
		@Override
		public Schema<?> resolve(
				final AnnotatedType type,
				final ModelConverterContext context,
				final Iterator<ModelConverter> chain
		) {
			if (chain.hasNext()) {
				return chain.next().resolve(type, context, chain);
			} else {
				return null;
			}
		}
	};
}

But a better solution would be great, ideally on a per-model basis.

[bnasslahsen]

@MHajoha,

The customizing is allowed through overriding.
For the disabling, you can use the following property that will be available with v1.5.13 :

springdoc.model-converters.polymorphic-converter.enabled=false

[MHajoha]

Nice, thank you for the quick and simple solution 👍

[westse]

Posting this here in case it helps anyone else. Disabling PolymorphicModelConverter was necessary for me as well since it doesn't handle complex inheritance hierarchies where subtypes refer to parent types or other inheritance hierarchies. After failing to produce a pull request to address the issue I realized the problem is that two passes are required (and a ModelConverter is only one part of one pass). One pass is required to discover all subtypes that have a polymorphic parent, and another pass is required to replace all parent references to polymorphic references (oneOf:....).

Here is the OpenApiCustomiser I wrote to handle this. Seems to work, but I've only done limited testing so far.

/**
 * Adds polymorphic references per OAS standard, which replaces the (now-disabled) springdoc builtin
 * PolymorphicModelConverter which does not account for complex inheritance hierarchies with models that refer to
 * parent models or models in other hierarchies.
 * <p>
 * Implementation approach borrowed from {@link io.swagger.v3.core.filter.SpecFilter}
 */
@Component
@SuppressWarnings({"rawtypes", "unchecked"})
public class PolymorphicOpenApiCustomizer implements OpenApiCustomiser, GlobalOpenApiCustomizer {

    public static final String POLYMORPHIC_SUFFIX = "_Polymorphic";
    private TypeRememberingModelConverter typeRememberingModelConverter;

    public PolymorphicOpenApiCustomizer(TypeRememberingModelConverter typeRememberingModelConverter) {
        this.typeRememberingModelConverter = typeRememberingModelConverter;
    }

    @Override
    public void customise(OpenAPI openApi) {
        createPolymorphicRefs(openApi);
    }

    private void createPolymorphicRefs(OpenAPI openApi) {
        try {
            if (openApi == null || openApi.getComponents() == null || openApi.getComponents().getSchemas() == null) {
                return;
            }

            TreeMap<String, Schema> polymorphicSchemas = extractPolymorphicSchemas(openApi);

            replaceRefsWithPolymorphicRefs(openApi, polymorphicSchemas.keySet());

            // Add all the polymorphic ref schemas
            Map<String, Schema> schemas = openApi.getComponents().getSchemas();
            polymorphicSchemas.values().forEach(schema -> schemas.put(schema.getName(), schema));

        } finally {
            typeRememberingModelConverter.reset();
        }
    }

    @NotNull
    private TreeMap<String, Schema> extractPolymorphicSchemas(OpenAPI openApi) {
        TreeMap<String, Schema> polymorphicRefSchemas = new TreeMap<>();
        openApi.getComponents().getSchemas().forEach((schemaName, schema) -> {
            if (schema.getAllOf() == null || !(schema instanceof ComposedSchema)) {
                return;
            }
            for (Schema allOfElement : (List<Schema>) schema.getAllOf()) {
                if (allOfElement.get$ref() == null) {
                    continue;
                }
                String parentSchemaName = (String) RefUtils.extractSimpleName(allOfElement.get$ref()).getLeft();

                // Add subclass schema to the parent's polymorphic ref schema
                extractPolymorphicSchema(polymorphicRefSchemas, parentSchemaName, schemaName);

                // Add parent too (if concrete) since it won't be detected otherwise
                extractPolymorphicSchema(polymorphicRefSchemas, parentSchemaName, parentSchemaName);
            }
        });
        return polymorphicRefSchemas;
    }

    private void extractPolymorphicSchema(TreeMap<String, Schema> polymorphicRefSchemas, String parentSchema,
                                          String schemaName) {
        JavaType type = typeRememberingModelConverter.getSchemaNameToType().get(schemaName);
        if (isConcreteClass(type)) {
            polymorphicRefSchemas
                    .computeIfAbsent(parentSchema, k -> new ComposedSchema().name(parentSchema + POLYMORPHIC_SUFFIX))
                    .addOneOfItem(new Schema().$ref(AnnotationsUtils.COMPONENTS_REF + schemaName));
        }
    }

    private boolean isConcreteClass(JavaType javaType) {
        if (javaType == null) {
            return false;
        }
        Class<?> clazz = javaType.getRawClass();
        return !Modifier.isAbstract(clazz.getModifiers()) && !clazz.isInterface();
    }


    protected void replaceRefsWithPolymorphicRefs(OpenAPI openApi, Set<String> polymorphables) {
        if (openApi.getPaths() != null) {
            for (String resourcePath : openApi.getPaths().keySet()) {
                PathItem pathItem = openApi.getPaths().get(resourcePath);
                replacePathItemSchemaRef(polymorphables, pathItem);
            }
        }

        for (Schema schema : openApi.getComponents().getSchemas().values()) {
            replaceSchemaRef(polymorphables, schema);
        }
    }

    private void replacePathItemSchemaRef(Set<String> polymorphables, PathItem pathItem) {
        if (pathItem.getParameters() != null) {
            for (Parameter parameter : pathItem.getParameters()) {
                replaceSchemaRef(polymorphables, parameter.getSchema());
                replaceContentSchemaRef(polymorphables, parameter.getContent());
            }
        }
        Map<PathItem.HttpMethod, Operation> ops = pathItem.readOperationsMap();
        for (Operation op : ops.values()) {
            if (op.getRequestBody() != null) {
                replaceContentSchemaRef(polymorphables, op.getRequestBody().getContent());
            }
            if (op.getResponses() != null) {
                for (String keyResponses : op.getResponses().keySet()) {
                    ApiResponse response = op.getResponses().get(keyResponses);
                    if (response.getHeaders() != null) {
                        for (String keyHeaders : response.getHeaders().keySet()) {
                            Header header = response.getHeaders().get(keyHeaders);
                            replaceSchemaRef(polymorphables, header.getSchema());
                            replaceContentSchemaRef(polymorphables, header.getContent());
                        }
                    }
                    replaceContentSchemaRef(polymorphables, response.getContent());
                }
            }
            if (op.getParameters() != null) {
                for (Parameter parameter : op.getParameters()) {
                    replaceSchemaRef(polymorphables, parameter.getSchema());
                    replaceContentSchemaRef(polymorphables, parameter.getContent());
                }
            }
            if (op.getCallbacks() != null) {
                for (String keyCallback : op.getCallbacks().keySet()) {
                    Callback callback = op.getCallbacks().get(keyCallback);
                    for (PathItem callbackPathItem : callback.values()) {
                        replacePathItemSchemaRef(polymorphables, callbackPathItem);
                    }
                }
            }
        }
    }

    private void replaceContentSchemaRef(Set<String> polymorphables, Content content) {
        if (content != null) {
            for (MediaType mediaType : content.values()) {
                replaceSchemaRef(polymorphables, mediaType.getSchema());
            }
        }
    }

    private void replaceSchemaRef(Set<String> polymorphables, Schema schema) {

        if (schema == null) {
            return;
        }
        if (!StringUtils.isBlank(schema.get$ref())) {
            toPolymorphicRef(polymorphables, schema.get$ref(), schema::set$ref);
            return;
        }
        if (schema.getDiscriminator() != null && schema.getDiscriminator().getMapping() != null) {
            Map<String, String> discMap = schema.getDiscriminator().getMapping();
            for (Map.Entry<String, String> mapping : discMap.entrySet()) {
                toPolymorphicRef(polymorphables, mapping.getValue(), polyRef -> discMap.put(mapping.getKey(), polyRef));
            }
        }

        if (schema.getProperties() != null) {
            for (Object propName : schema.getProperties().keySet()) {
                Schema property = (Schema) schema.getProperties().get(propName);
                replaceSchemaRef(polymorphables, property);
            }
        }

        if (schema.getAdditionalProperties() instanceof Schema) {
            replaceSchemaRef(polymorphables, (Schema) schema.getAdditionalProperties());
        }

        if (schema instanceof ArraySchema && ((ArraySchema) schema).getItems() != null) {
            replaceSchemaRef(polymorphables, ((ArraySchema) schema).getItems());
        } else if (schema instanceof ComposedSchema) {
            ComposedSchema composedSchema = (ComposedSchema) schema;

            if (composedSchema.getAllOf() != null) {
                for (Schema ref : composedSchema.getAllOf()) {
                    // Skip Schema.allOf $ref elements because these are subtypes pointing to their polymorphic parent
                    if (ref.get$ref() == null) {
                        replaceSchemaRef(polymorphables, ref);
                    }
                }
            }
            if (composedSchema.getAnyOf() != null) {
                for (Schema ref : composedSchema.getAnyOf()) {
                    replaceSchemaRef(polymorphables, ref);
                }
            }
            if (composedSchema.getOneOf() != null) {
                for (Schema ref : composedSchema.getOneOf()) {
                    replaceSchemaRef(polymorphables, ref);
                }
            }
        }
    }

    private void toPolymorphicRef(Set<String> polymorphables, String curRef, Consumer<String> onPolymorphable) {
        String simpleName = (String) RefUtils.extractSimpleName(curRef).getLeft();
        if (polymorphables.contains(simpleName)) {
            String polymorphicRef = RefUtils.constructRef(simpleName + POLYMORPHIC_SUFFIX);
            onPolymorphable.accept(polymorphicRef);
        }
    }

}

Additionally, Java type information is needed to determine if the parent type should be included as a oneOf option. I found it convenient to use a ModelConverter for this since it already has access to the type information. Not pretty but I think it works.

/**
 * Remembers type information needed by later logic like {@link PolymorphicOpenApiCustomizer}
 */
@Component
public class TypeRememberingModelConverter implements ModelConverter {

    private final ObjectMapperProvider springDocObjectMapper;
    private final Map<String, JavaType> schemaNameToType = new HashMap<>();


    public TypeRememberingModelConverter(ObjectMapperProvider springDocObjectMapper) {
        this.springDocObjectMapper = springDocObjectMapper;
    }

    /**
     * Called to reset the state. Retaining state between calls should be ok since the results should be the same.
     * However, if this proves problematic another approach would be to put the info in the OpenApi object itself,
     * such as on a synthetic "SchemaNameToJavaType" schema that can be removed by the customizer that needs it.
     */
    public void reset() {
        schemaNameToType.clear();
    }

    public Map<String, JavaType> getSchemaNameToType() {
        return Collections.unmodifiableMap(schemaNameToType);
    }

    @SuppressWarnings("rawtypes")
    @Override
    public Schema resolve(AnnotatedType type, ModelConverterContext context, Iterator<ModelConverter> chain) {
        if (chain.hasNext()) {
            Schema<?> schema = chain.next().resolve(type, context, chain);
            if (schema == null) {
                return null;
            }
            String schemaName = schema.getName();
            if (schemaName == null && schema.get$ref() != null) {
                schemaName = (String) RefUtils.extractSimpleName(schema.get$ref()).getLeft();
            }
            if (schemaName != null) {
                JavaType javaType = springDocObjectMapper.jsonMapper().constructType(type.getType());
                if (javaType != null) {
                    schemaNameToType.put(schemaName, javaType);
                }
            }
            return schema;
        }
        return null;
    }

}