- final Entry<GeneratedType, WithStatus> typeWithSchema = context.getTypeWithSchema(valueType);
- final WithStatus schema = typeWithSchema.getValue();
- checkState(schema instanceof TypeDefinition, "Unexpected schema %s", schema);
- return getCodec(valueType, (TypeDefinition<?>) schema);
+ final var typeWithSchema = context.getTypeWithSchema(valueType);
+ final var schema = typeWithSchema.statement();
+ final TypeDefinition<?> def;
+ if (schema instanceof TypeDefinitionAware) {
+ def = ((TypeDefinitionAware) schema).getTypeDefinition();
+ } else if (schema instanceof TypeAware) {
+ def = ((TypeAware) schema).getType();
+ } else {
+ throw new IllegalStateException("Unexpected schema " + schema);
+ }
+ return getCodec(valueType, def);