[Draft] AVRO module

avro/avro-annotation-processor/build.gradle

@@ -0,0 +1,43 @@
+import org.apache.avro.tool.SpecificCompilerTool
+
+buildscript {
+    dependencies {
+        classpath libs.avro.tools
+    }
+}
+
+apply from: "${project.rootDir}/gradle/in-test-generated.gradle"
+
+dependencies {
+    api project(":core:annotation-processor-common")
+    api project(":core:kora-app-annotation-processor")
+
+    testImplementation project(":avro:avro-common")
+    testImplementation testFixtures(project(":core:annotation-processor-common"))
+}
+
+tasks.register("generateAvroClasses") {
+    group("build")
+
+    var inputDir = "$projectDir/src/test/resources/avro"
+    var outputDir = "$buildDir/generated/sources/avro"
+    inputs.dir(inputDir)
+    outputs.dir(outputDir)
+    logging.captureStandardOutput(LogLevel.INFO);
+    logging.captureStandardError(LogLevel.ERROR)
+
+    doFirst {
+        delete outputDir
+    }
+
+    doLast {
+        var params = ["-bigDecimal", "schema", inputDir.toString(), outputDir.toString()]
+        new SpecificCompilerTool().run(System.in, System.out, System.err, params)
+    }
+}
+
+sourceSets.test.java {
+    srcDir tasks.generateAvroClasses
+}
+
+test.dependsOn(tasks.generateAvroClasses)

avro/avro-annotation-processor/src/main/java/io/koraframework/avro/annotation/processor/AvroAnnotationProcessor.java

@@ -0,0 +1,80 @@
+package io.koraframework.avro.annotation.processor;
+
+import com.palantir.javapoet.ClassName;
+import io.koraframework.annotation.processor.common.AbstractKoraProcessor;
+import io.koraframework.annotation.processor.common.LogUtils;
+import io.koraframework.annotation.processor.common.ProcessingErrorException;
+import io.koraframework.avro.annotation.processor.reader.AvroReaderGenerator;
+import io.koraframework.avro.annotation.processor.writer.AvroWriterGenerator;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.slf4j.event.Level;
+
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.annotation.processing.RoundEnvironment;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.TypeElement;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class AvroAnnotationProcessor extends AbstractKoraProcessor {
+    private static final Logger logger = LoggerFactory.getLogger(AvroAnnotationProcessor.class);
+
+    private boolean initialized = false;
+    private TypeElement avroAnnotation;
+    private TypeElement specificRecordType;
+    private AvroWriterGenerator writerGenerator;
+    private AvroReaderGenerator readerGenerator;
+
+    @Override
+    public Set<ClassName> getSupportedAnnotationClassNames() {
+        return Set.of(AvroTypes.avro);
+    }
+
+    @Override
+    public synchronized void init(ProcessingEnvironment processingEnv) {
+        super.init(processingEnv);
+        this.avroAnnotation = processingEnv.getElementUtils().getTypeElement(AvroTypes.avro.canonicalName());
+        this.specificRecordType = processingEnv.getElementUtils().getTypeElement(AvroTypes.specificRecord.canonicalName());
+        if (this.avroAnnotation == null || this.specificRecordType == null) {
+            return;
+        }
+
+        this.initialized = true;
+        this.writerGenerator = new AvroWriterGenerator(processingEnv);
+        this.readerGenerator = new AvroReaderGenerator(processingEnv);
+    }
+
+    @Override
+    protected void process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv, Map<ClassName, List<AnnotatedElement>> annotatedElements) {
+        if (!this.initialized || roundEnv.processingOver()) {
+            return;
+        }
+
+        var avroElements = roundEnv.getElementsAnnotatedWith(this.avroAnnotation).stream()
+            .filter(e -> e.getKind().isClass() || e.getKind() == ElementKind.INTERFACE)
+            .toList();
+        if (avroElements.isEmpty()) {
+            return;
+        }
+
+        LogUtils.logElementsFull(logger, Level.DEBUG, "Generating Avro Readers & Writers for", avroElements);
+        for (var e : avroElements) {
+            if (!this.types.isAssignable(e.asType(), this.specificRecordType.asType())) {
+                new ProcessingErrorException("@Avro can only be used on org.apache.avro.specific.SpecificRecord types", e).printError(this.processingEnv);
+                continue;
+            }
+            try {
+                this.readerGenerator.generate((TypeElement) e);
+            } catch (ProcessingErrorException ex) {
+                ex.printError(this.processingEnv);
+            }
+            try {
+                this.writerGenerator.generate((TypeElement) e);
+            } catch (ProcessingErrorException ex) {
+                ex.printError(this.processingEnv);
+            }
+        }
+    }
+}

avro/avro-annotation-processor/src/main/java/io/koraframework/avro/annotation/processor/AvroTypes.java

@@ -0,0 +1,21 @@
+package io.koraframework.avro.annotation.processor;
+
+import com.palantir.javapoet.ClassName;
+
+public final class AvroTypes {
+
+    private AvroTypes() {}
+
+    public static final ClassName avro = ClassName.get("io.koraframework.avro.common.annotation", "Avro");
+
+    public static final ClassName reader = ClassName.get("io.koraframework.avro.common", "AvroReader");
+    public static final ClassName writer = ClassName.get("io.koraframework.avro.common", "AvroWriter");
+
+    public static final ClassName schema = ClassName.get("org.apache.avro", "Schema");
+    public static final ClassName specificData = ClassName.get("org.apache.avro.specific", "SpecificData");
+    public static final ClassName specificRecord = ClassName.get("org.apache.avro.specific", "SpecificRecord");
+    public static final ClassName datumReader = ClassName.get("org.apache.avro.specific", "SpecificDatumReader");
+    public static final ClassName datumWriter = ClassName.get("org.apache.avro.specific", "SpecificDatumWriter");
+    public static final ClassName decoderFactory = ClassName.get("org.apache.avro.io", "DecoderFactory");
+    public static final ClassName encoderFactory = ClassName.get("org.apache.avro.io", "EncoderFactory");
+}

avro/avro-annotation-processor/src/main/java/io/koraframework/avro/annotation/processor/AvroUtils.java

@@ -0,0 +1,37 @@
+package io.koraframework.avro.annotation.processor;
+
+import io.koraframework.annotation.processor.common.NameUtils;
+
+import javax.lang.model.element.Element;
+import javax.lang.model.element.TypeElement;
+import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.Elements;
+import javax.lang.model.util.Types;
+
+
+public final class AvroUtils {
+
+    private AvroUtils() {}
+
+    public static String classPackage(Elements elements, Element typeElement) {
+        return elements.getPackageOf(typeElement).getQualifiedName().toString();
+    }
+
+    public static String writerName(Element typeElement) {
+        return NameUtils.generatedType(typeElement, "AvroWriter");
+    }
+
+    public static String writerName(Types types, TypeMirror typeMirror) {
+        var typeElement = types.asElement(typeMirror);
+        return writerName(typeElement);
+    }
+
+    public static String readerName(TypeElement typeElement) {
+        return NameUtils.generatedType(typeElement, "AvroReader");
+    }
+
+    public static String readerName(Types types, TypeMirror typeMirror) {
+        var typeElement = types.asElement(typeMirror);
+        return readerName((TypeElement) typeElement);
+    }
+}

avro/avro-annotation-processor/src/main/java/io/koraframework/avro/annotation/processor/extension/AvroExtension.java

@@ -0,0 +1,122 @@
+package io.koraframework.avro.annotation.processor.extension;
+
+import com.palantir.javapoet.ClassName;
+import com.palantir.javapoet.CodeBlock;
+import io.koraframework.annotation.processor.common.ProcessingErrorException;
+import io.koraframework.avro.annotation.processor.AvroTypes;
+import io.koraframework.avro.annotation.processor.AvroUtils;
+import io.koraframework.avro.annotation.processor.reader.AvroReaderGenerator;
+import io.koraframework.avro.annotation.processor.writer.AvroWriterGenerator;
+import io.koraframework.kora.app.annotation.processor.extension.ExtensionResult;
+import io.koraframework.kora.app.annotation.processor.extension.KoraExtension;
+import org.jspecify.annotations.Nullable;
+
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.annotation.processing.RoundEnvironment;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.TypeElement;
+import javax.lang.model.type.DeclaredType;
+import javax.lang.model.type.TypeKind;
+import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.Elements;
+import javax.lang.model.util.Types;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+public final class AvroExtension implements KoraExtension {
+
+    private final Types types;
+    private final Elements elements;
+    private final TypeMirror writerErasure;
+    private final TypeMirror readerErasure;
+    private final AvroReaderGenerator readerGenerator;
+    private final AvroWriterGenerator writerGenerator;
+    private final Set<String> generatedMappers = new HashSet<>();
+
+    public AvroExtension(ProcessingEnvironment processingEnv) {
+        this.types = processingEnv.getTypeUtils();
+        this.elements = processingEnv.getElementUtils();
+        this.writerErasure = this.types.erasure(this.elements.getTypeElement(AvroTypes.writer.canonicalName()).asType());
+        this.readerErasure = this.types.erasure(this.elements.getTypeElement(AvroTypes.reader.canonicalName()).asType());
+        this.readerGenerator = new AvroReaderGenerator(processingEnv);
+        this.writerGenerator = new AvroWriterGenerator(processingEnv);
+    }
+
+    @Nullable
+    @Override
+    public KoraExtensionDependencyGenerator getDependencyGenerator(RoundEnvironment roundEnvironment, TypeMirror typeMirror, @Nullable String tag) {
+        if (!isBinary(tag)) {
+            return null;
+        }
+        if (typeMirror.getKind() != TypeKind.DECLARED) {
+            return null;
+        }
+
+        var erasure = this.types.erasure(typeMirror);
+        if (this.types.isSameType(erasure, this.writerErasure)) {
+            var targetElement = dependencyTarget(typeMirror);
+            if (targetElement == null) {
+                throw new ProcessingErrorException("AvroWriter<T> can only be created for concrete org.apache.avro.specific.SpecificRecord types", this.types.asElement(typeMirror));
+            }
+            return generateMapper(targetElement, AvroUtils.writerName(targetElement), typeMirror, tag, () -> this.writerGenerator.generate(targetElement));
+        }
+
+        if (this.types.isSameType(erasure, this.readerErasure)) {
+            var targetElement = dependencyTarget(typeMirror);
+            if (targetElement == null) {
+                throw new ProcessingErrorException("AvroReader<T> can only be created for concrete org.apache.avro.specific.SpecificRecord types", this.types.asElement(typeMirror));
+            }
+            return generateMapper(targetElement, AvroUtils.readerName(targetElement), typeMirror, tag, () -> this.readerGenerator.generate(targetElement));
+        }
+
+        return null;
+    }
+
+    @Nullable
+    private TypeElement dependencyTarget(TypeMirror typeMirror) {
+        var declaredType = (DeclaredType) typeMirror;
+        var targetType = declaredType.getTypeArguments().get(0);
+        if (targetType.getKind() != TypeKind.DECLARED) {
+            return null;
+        }
+        var targetElement = this.types.asElement(targetType);
+        if (!(targetElement instanceof TypeElement typeElement)
+            || targetElement.getKind() == ElementKind.ENUM
+            || targetElement.getKind().isInterface()
+            || targetElement.getModifiers().contains(Modifier.ABSTRACT)) {
+            return null;
+        }
+        return typeElement;
+    }
+
+    private KoraExtensionDependencyGenerator generateMapper(
+        TypeElement targetElement,
+        String mapperName,
+        TypeMirror requestedType,
+        @Nullable String tag,
+        Runnable generate
+    ) {
+        return () -> {
+            var packageName = AvroUtils.classPackage(this.elements, targetElement);
+            var mapperCanonicalName = packageName + "." + mapperName;
+            if (this.elements.getTypeElement(mapperCanonicalName) == null && this.generatedMappers.add(mapperCanonicalName)) {
+                generate.run();
+            }
+
+            return new ExtensionResult.CodeBlockResult(
+                targetElement,
+                params -> CodeBlock.of("new $T()", ClassName.get(packageName, mapperName)),
+                requestedType,
+                tag,
+                List.of(),
+                List.of()
+            );
+        };
+    }
+
+    private boolean isBinary(@Nullable String tag) {
+        return tag == null || tag.equals(AvroTypes.avro.canonicalName());
+    }
+}

avro/avro-annotation-processor/src/main/java/io/koraframework/avro/annotation/processor/extension/AvroExtensionFactory.java

@@ -0,0 +1,21 @@
+package io.koraframework.avro.annotation.processor.extension;
+
+import io.koraframework.avro.annotation.processor.AvroTypes;
+import io.koraframework.kora.app.annotation.processor.extension.ExtensionFactory;
+import io.koraframework.kora.app.annotation.processor.extension.KoraExtension;
+
+import javax.annotation.processing.ProcessingEnvironment;
+import java.util.Optional;
+
+public class AvroExtensionFactory implements ExtensionFactory {
+
+    @Override
+    public Optional<KoraExtension> create(ProcessingEnvironment processingEnvironment) {
+        var avro = processingEnvironment.getElementUtils().getTypeElement(AvroTypes.avro.canonicalName());
+        if (avro == null) {
+            return Optional.empty();
+        } else {
+            return Optional.of(new AvroExtension(processingEnvironment));
+        }
+    }
+}

avro/avro-annotation-processor/src/main/java/io/koraframework/avro/annotation/processor/reader/AvroReaderGenerator.java

@@ -0,0 +1,69 @@
+package io.koraframework.avro.annotation.processor.reader;
+
+import com.palantir.javapoet.*;
+import io.koraframework.annotation.processor.common.CommonClassNames;
+import io.koraframework.annotation.processor.common.CommonUtils;
+import io.koraframework.avro.annotation.processor.AvroTypes;
+import io.koraframework.avro.annotation.processor.AvroUtils;
+import org.jspecify.annotations.Nullable;
+
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.TypeParameterElement;
+import java.io.IOException;
+import java.io.InputStream;
+
+public class AvroReaderGenerator {
+    private final ProcessingEnvironment env;
+
+    public AvroReaderGenerator(ProcessingEnvironment processingEnvironment) {
+        this.env = processingEnvironment;
+    }
+
+    public void generate(TypeElement element) {
+        var typeName = TypeName.get(element.asType());
+        var typeBuilder = TypeSpec.classBuilder(AvroUtils.readerName(element))
+            .addAnnotation(AnnotationSpec.builder(CommonClassNames.koraGenerated)
+                .addMember("value", CodeBlock.of("$S", AvroReaderGenerator.class.getCanonicalName()))
+                .build())
+            .addAnnotation(AvroTypes.avro)
+            .addSuperinterface(ParameterizedTypeName.get(AvroTypes.reader, typeName))
+            .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
+            .addOriginatingElement(element);
+
+        for (TypeParameterElement typeParameter : element.getTypeParameters()) {
+            typeBuilder.addTypeVariable(TypeVariableName.get(typeParameter));
+        }
+
+        typeBuilder.addField(FieldSpec.builder(AvroTypes.schema, "SCHEMA")
+            .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
+            .initializer("$T.getClassSchema()", typeName)
+            .build());
+        typeBuilder.addField(FieldSpec.builder(AvroTypes.specificData, "SPECIFIC_DATA")
+            .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
+            .initializer("new $T().getSpecificData()", typeName)
+            .build());
+        typeBuilder.addField(FieldSpec.builder(ParameterizedTypeName.get(AvroTypes.datumReader, typeName), "READER")
+            .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
+            .initializer("new $T<>(SCHEMA, SCHEMA, SPECIFIC_DATA)", AvroTypes.datumReader)
+            .build());
+
+        var method = MethodSpec.methodBuilder("read")
+            .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
+            .addException(IOException.class)
+            .addAnnotation(Nullable.class)
+            .addAnnotation(Override.class)
+            .addParameter(TypeName.get(InputStream.class), "value")
+            .returns(typeName);
+        method.beginControlFlow("if (value == null || value.available() == 0)");
+        method.addStatement("return null");
+        method.endControlFlow();
+        method.addStatement("var decoder = $T.get().directBinaryDecoder(value, null)", AvroTypes.decoderFactory);
+        method.addStatement("return READER.read(new $T(), decoder)", typeName);
+
+        typeBuilder.addMethod(method.build());
+        var javaFile = JavaFile.builder(AvroUtils.classPackage(this.env.getElementUtils(), element), typeBuilder.build()).build();
+        CommonUtils.safeWriteTo(this.env, javaFile);
+    }
+}