H2 reflection extensions via GraalVM Feature

extensions/jdbc/jdbc-h2/deployment/src/main/java/io/quarkus/jdbc/h2/deployment/H2JDBCReflections.java

@@ -1,8 +1,25 @@
 package io.quarkus.jdbc.h2.deployment;
 
+import java.nio.charset.StandardCharsets;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import org.h2.mvstore.type.DataType;
+import org.h2.mvstore.type.StatefulDataType;
+import org.jboss.jandex.ClassInfo;
+import org.jboss.jandex.DotName;
+import org.jboss.jandex.IndexView;
+
 import io.quarkus.deployment.annotations.BuildProducer;
 import io.quarkus.deployment.annotations.BuildStep;
+import io.quarkus.deployment.builditem.CombinedIndexBuildItem;
+import io.quarkus.deployment.builditem.GeneratedResourceBuildItem;
+import io.quarkus.deployment.builditem.IndexDependencyBuildItem;
+import io.quarkus.deployment.builditem.NativeImageFeatureBuildItem;
 import io.quarkus.deployment.builditem.nativeimage.ReflectiveClassBuildItem;
+import io.quarkus.deployment.pkg.steps.NativeOrNativeSourcesBuild;
+import io.quarkus.jdbc.h2.runtime.H2Reflections;
 
 /**
  * Registers the {@code org.h2.Driver} so that it can be loaded
@@ -21,4 +38,49 @@ void build(BuildProducer<ReflectiveClassBuildItem> reflectiveClass) {
         reflectiveClass.produce(new ReflectiveClassBuildItem(false, false, driverName));
     }
 
+    @BuildStep
+    NativeImageFeatureBuildItem enableH2Feature() {
+        return new NativeImageFeatureBuildItem("io.quarkus.jdbc.h2.runtime.H2Reflections");
+    }
+
+    /**
+     * We need to index the H2 database core jar so to include custom extension types it's
+     * loading via reflection.
+     */
+    @BuildStep(onlyIf = NativeOrNativeSourcesBuild.class)
+    IndexDependencyBuildItem indexH2Library() {
+        return new IndexDependencyBuildItem("com.h2database", "h2");
+    }
+
+    /**
+     * All implementors of {@link StatefulDataType.Factory} might get loaded reflectively.
+     * While we could hardcode the list included in H2, we prefer looking them up via Jandex
+     * so to also load custom implementations optionally provided by user code.
+     */
+    @BuildStep(onlyIf = NativeOrNativeSourcesBuild.class)
+    GeneratedResourceBuildItem listStatefulDataTypeFactories(CombinedIndexBuildItem index) {
+        return generateListBy(H2Reflections.REZ_NAME_STATEFUL_DATATYPES, index,
+                i -> i.getAllKnownImplementors(StatefulDataType.Factory.class).stream());
+    }
+
+    /**
+     * All implementors of {@link DataType} which have an INSTANCE field
+     * need this field to be accessible via reflection.
+     * While we could hardcode the list included in H2, we prefer looking them up via Jandex
+     * so to also load custom implementations optionally provided by user code.
+     */
+    @BuildStep(onlyIf = NativeOrNativeSourcesBuild.class)
+    GeneratedResourceBuildItem listBasicDataTypes(CombinedIndexBuildItem index) {
+        return generateListBy(H2Reflections.REZ_NAME_DATA_TYPE_SINGLETONS, index,
+                i -> i.getAllKnownImplementors(DataType.class)
+                        .stream().filter(classInfo -> classInfo.field("INSTANCE") != null));
+    }
+
+    private static GeneratedResourceBuildItem generateListBy(String resourceName, CombinedIndexBuildItem index,
+            Function<IndexView, Stream<ClassInfo>> selection) {
+        String classNames = selection.apply(index.getIndex()).map(ClassInfo::name).map(DotName::toString).sorted()
+                .collect(Collectors.joining("\n"));
+        return new GeneratedResourceBuildItem(resourceName, classNames.getBytes(StandardCharsets.UTF_8));
+    }
+
 }

extensions/jdbc/jdbc-h2/runtime/src/main/java/io/quarkus/jdbc/h2/runtime/H2Reflections.java

@@ -0,0 +1,73 @@
+package io.quarkus.jdbc.h2.runtime;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.util.Scanner;
+import java.util.function.BiConsumer;
+
+import org.graalvm.nativeimage.hosted.Feature;
+import org.graalvm.nativeimage.hosted.RuntimeReflection;
+
+/**
+ * Custom GraalVM feature to automatically register DataType and StatefulDataType
+ * implementors for reflective access.
+ * These are identified using Jandex, looking both into the H2 core jar and in
+ * user's indexed code.
+ */
+public final class H2Reflections implements Feature {
+
+    public static final String REZ_NAME_DATA_TYPE_SINGLETONS = "h2BasicDataTypeSingletons.classlist";
+    public static final String REZ_NAME_STATEFUL_DATATYPES = "h2StatefulDataType.classlist";
+
+    @Override
+    public void beforeAnalysis(BeforeAnalysisAccess access) {
+        Class<?> metaType = access.findClassByName("org.h2.mvstore.type.MetaType");
+        access.registerReachabilityHandler(this::metaTypeReachable, metaType);
+    }
+
+    private void metaTypeReachable(DuringAnalysisAccess access) {
+        //Register some common metatypes - these are dynamically loaded depending on the data content.
+        register(REZ_NAME_DATA_TYPE_SINGLETONS, this::registerSingletonAccess, access);
+        //Now register implementors of org.h2.mvstore.type.StatefulDataType.Factory
+        register(REZ_NAME_STATEFUL_DATATYPES, this::registerForReflection, access);
+    }
+
+    void register(final String resourceName, BiConsumer<String, DuringAnalysisAccess> action, DuringAnalysisAccess access) {
+        try (InputStream resource = access.getApplicationClassLoader().getResourceAsStream(resourceName)) {
+            Scanner s = new Scanner(resource);
+            while (s.hasNext()) {
+                final String className = s.next();
+                action.accept(className, access);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    void registerSingletonAccess(String className, DuringAnalysisAccess access) {
+        try {
+            final Field instance = access.findClassByName(className)
+                    .getDeclaredField("INSTANCE");
+            RuntimeReflection.register(instance);
+        } catch (NoSuchFieldException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    void registerForReflection(
+            String className,
+            DuringAnalysisAccess duringAnalysisAccess) {
+        final Class<?> aClass = duringAnalysisAccess.findClassByName(className);
+        final Constructor<?>[] z = aClass.getDeclaredConstructors();
+        RuntimeReflection.register(aClass);
+        RuntimeReflection.register(z);
+    }
+
+    @Override
+    public String getDescription() {
+        return "Support for H2 Database's extended data types";
+    }
+
+}