Add ability to clean the DB in DevUI

Fixes #583 Fixes #19303
quarkusio · Aug 16, 2021 · 4b2ed55 · 4b2ed55
1 parent 626b2dc
commit 4b2ed55
Show file tree

Hide file tree

Showing 17 changed files with 442 additions and 1 deletion.
diff --git a/extensions/agroal/deployment/src/main/java/io/quarkus/agroal/deployment/AgroalProcessor.java b/extensions/agroal/deployment/src/main/java/io/quarkus/agroal/deployment/AgroalProcessor.java
@@ -1,12 +1,15 @@
 package io.quarkus.agroal.deployment;
 
+import static io.quarkus.deployment.annotations.ExecutionTime.STATIC_INIT;
+
 import java.sql.Driver;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 import javax.enterprise.inject.Default;
 import javax.inject.Singleton;
@@ -48,6 +51,8 @@
 import io.quarkus.deployment.builditem.nativeimage.NativeImageResourceBuildItem;
 import io.quarkus.deployment.builditem.nativeimage.ReflectiveClassBuildItem;
 import io.quarkus.deployment.pkg.builditem.CurateOutcomeBuildItem;
+import io.quarkus.devconsole.spi.DevConsoleRouteBuildItem;
+import io.quarkus.devconsole.spi.DevConsoleTemplateInfoBuildItem;
 import io.quarkus.runtime.configuration.ConfigurationException;
 import io.quarkus.smallrye.health.deployment.spi.HealthBuildItem;
 
@@ -56,6 +61,7 @@ class AgroalProcessor {
 
     private static final Logger log = Logger.getLogger(AgroalProcessor.class);
 
+    private static final String CLEAN_DATABASE = "io.quarkus.test.ResetDatabase";
     private static final DotName DATA_SOURCE = DotName.createSimple(javax.sql.DataSource.class.getName());
 
     @BuildStep
@@ -349,4 +355,18 @@ HealthBuildItem addHealthCheck(Capabilities capabilities, DataSourcesBuildTimeCo
             return null;
         }
     }
+
+    @BuildStep
+    public DevConsoleTemplateInfoBuildItem devConsoleInfo(
+            List<AggregatedDataSourceBuildTimeConfigBuildItem> dbs) {
+        return new DevConsoleTemplateInfoBuildItem("dbs",
+                dbs.stream().map(AggregatedDataSourceBuildTimeConfigBuildItem::getName)
+                        .collect(Collectors.toList()));
+    }
+
+    @BuildStep
+    @Record(value = STATIC_INIT, optional = true)
+    DevConsoleRouteBuildItem devConsoleCleanDatabaseHandler(AgroalRecorder recorder) {
+        return new DevConsoleRouteBuildItem("clean", "POST", recorder.devConsoleCleanDatabaseHandler());
+    }
 }
diff --git a/extensions/agroal/deployment/src/main/resources/dev-templates/clean.html b/extensions/agroal/deployment/src/main/resources/dev-templates/clean.html
@@ -0,0 +1,27 @@
+{#include main}
+  {#title}Clean Databases{/title}
+  {#body}
+  <table class="table table-striped">
+   <thead class="thead-dark">
+    <tr>
+      <th scope="col">Datasource</th>
+      <th scope="col">Actions</th>
+    </tr>
+   </thead>
+   <tbody>
+   {#for db in info:dbs}
+    <tr>
+      <td>
+        {db}
+      </td>
+      <td>
+       <form method="post"  enctype="application/x-www-form-urlencoded">
+        <input type="hidden" name="name" value="{db}">
+        <input id="invoke" type="submit" value="Invoke" class="btn btn-primary btn-sm">
+       </form>
+      </td>
+    {/for}
+   </tbody>
+  </table>
+  {/body}
+{/include}
diff --git a/extensions/agroal/deployment/src/main/resources/dev-templates/embedded.html b/extensions/agroal/deployment/src/main/resources/dev-templates/embedded.html
@@ -0,0 +1,3 @@
+<a href="{urlbase}/clean" class="badge badge-light">
+  <i class="fa fa-database fa-fw"></i>
+  Reset Databases</a>
diff --git a/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AgroalRecorder.java b/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AgroalRecorder.java
@@ -3,6 +3,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.ServiceLoader;
 import java.util.function.Supplier;
 
 import io.agroal.api.AgroalDataSource;
@@ -11,7 +12,12 @@
 import io.quarkus.datasource.runtime.DataSourcesBuildTimeConfig;
 import io.quarkus.datasource.runtime.DataSourcesExcludedFromHealthChecks;
 import io.quarkus.datasource.runtime.DataSourcesRuntimeConfig;
+import io.quarkus.datasource.runtime.DatabaseSchemaProvider;
+import io.quarkus.devconsole.runtime.spi.DevConsolePostHandler;
 import io.quarkus.runtime.annotations.Recorder;
+import io.vertx.core.Handler;
+import io.vertx.core.MultiMap;
+import io.vertx.ext.web.RoutingContext;
 
 @Recorder
 public class AgroalRecorder {
@@ -56,4 +62,20 @@ public DataSourcesExcludedFromHealthChecks get() {
         };
     }
 
+    public Handler<RoutingContext> devConsoleCleanDatabaseHandler() {
+        // the usual issue of Vert.x hanging on to the first TCCL and setting it on all its threads
+        final ClassLoader currentCl = Thread.currentThread().getContextClassLoader();
+        return new DevConsolePostHandler() {
+            @Override
+            protected void handlePost(RoutingContext event, MultiMap form) throws Exception {
+                String name = form.get("name");
+                ServiceLoader<DatabaseSchemaProvider> dbs = ServiceLoader.load(DatabaseSchemaProvider.class,
+                        Thread.currentThread().getContextClassLoader());
+                for (DatabaseSchemaProvider i : dbs) {
+                    i.resetDatabase(name);
+                }
+                flashMessage(event, "Action invoked");
+            }
+        };
+    }
 }
diff --git a/...atasource/runtime/src/main/java/io/quarkus/datasource/runtime/DatabaseSchemaProvider.java b/...atasource/runtime/src/main/java/io/quarkus/datasource/runtime/DatabaseSchemaProvider.java
@@ -0,0 +1,11 @@
+package io.quarkus.datasource.runtime;
+
+/**
+ * A service interface that can be used to reset the database for dev and test mode.
+ */
+public interface DatabaseSchemaProvider {
+
+    void resetDatabase(String dbName);
+
+    void resetAllDatabases();
+}
diff --git a/extensions/flyway/runtime/src/main/java/io/quarkus/flyway/runtime/FlywaySchemaProvider.java b/extensions/flyway/runtime/src/main/java/io/quarkus/flyway/runtime/FlywaySchemaProvider.java
@@ -0,0 +1,23 @@
+package io.quarkus.flyway.runtime;
+
+import io.quarkus.datasource.runtime.DatabaseSchemaProvider;
+
+public class FlywaySchemaProvider implements DatabaseSchemaProvider {
+    @Override
+    public void resetDatabase(String dbName) {
+        for (FlywayContainer i : FlywayRecorder.FLYWAY_CONTAINERS) {
+            if (i.getDataSourceName().equals(dbName)) {
+                i.getFlyway().clean();
+                i.getFlyway().migrate();
+            }
+        }
+    }
+
+    @Override
+    public void resetAllDatabases() {
+        for (FlywayContainer i : FlywayRecorder.FLYWAY_CONTAINERS) {
+            i.getFlyway().clean();
+            i.getFlyway().migrate();
+        }
+    }
+}
diff --git a/...src/main/resources/META-INF/services/io.quarkus.datasource.runtime.DatabaseSchemaProvider b/...src/main/resources/META-INF/services/io.quarkus.datasource.runtime.DatabaseSchemaProvider
@@ -0,0 +1 @@
+io.quarkus.flyway.runtime.FlywaySchemaProvider
diff --git a/...m/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateOrmProcessor.java b/...m/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateOrmProcessor.java
@@ -120,6 +120,7 @@
 import io.quarkus.hibernate.orm.runtime.cdi.QuarkusArcBeanContainer;
 import io.quarkus.hibernate.orm.runtime.integration.HibernateOrmIntegrationStaticDescriptor;
 import io.quarkus.hibernate.orm.runtime.proxies.PreGeneratedProxies;
+import io.quarkus.hibernate.orm.runtime.schema.SchemaManagementIntegrator;
 import io.quarkus.hibernate.orm.runtime.tenant.DataSourceTenantConnectionResolver;
 import io.quarkus.hibernate.orm.runtime.tenant.TenantConnectionResolver;
 import io.quarkus.panache.common.deployment.HibernateEnhancersRegisteredBuildItem;
@@ -424,6 +425,7 @@ public void build(RecorderContext recorderContext, HibernateOrmRecorder recorder
             List<HibernateOrmIntegrationStaticConfiguredBuildItem> integrationBuildItems,
             ProxyDefinitionsBuildItem proxyDefinitions,
             BuildProducer<FeatureBuildItem> feature,
+            LaunchModeBuildItem launchModeBuildItem,
             BuildProducer<BeanContainerListenerBuildItem> beanContainerListener) throws Exception {
 
         feature.produce(new FeatureBuildItem(Feature.HIBERNATE_ORM));
@@ -455,6 +457,9 @@ public void build(RecorderContext recorderContext, HibernateOrmRecorder recorder
         for (String integratorClassName : ServiceUtil.classNamesNamedIn(classLoader, INTEGRATOR_SERVICE_FILE)) {
             integratorClasses.add((Class<? extends Integrator>) recorderContext.classProxy(integratorClassName));
         }
+        if (launchModeBuildItem.getLaunchMode().isDevOrTest()) {
+            integratorClasses.add(SchemaManagementIntegrator.class);
+        }
 
         Map<String, List<HibernateOrmIntegrationStaticDescriptor>> integrationStaticDescriptors = HibernateOrmIntegrationStaticConfiguredBuildItem
                 .collectDescriptors(integrationBuildItems);

diff --git a/...ent/src/test/java/io/quarkus/hibernate/orm/HibernateSchemaRecreateDevConsoleTestCase.java b/...ent/src/test/java/io/quarkus/hibernate/orm/HibernateSchemaRecreateDevConsoleTestCase.java
@@ -0,0 +1,37 @@
+package io.quarkus.hibernate.orm;
+
+import static org.hamcrest.Matchers.is;
+
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.test.QuarkusDevModeTest;
+import io.restassured.RestAssured;
+
+public class HibernateSchemaRecreateDevConsoleTestCase {
+    @RegisterExtension
+    final static QuarkusDevModeTest TEST = new QuarkusDevModeTest()
+            .setArchiveProducer(() -> ShrinkWrap.create(JavaArchive.class)
+                    .addClasses(MyEntity.class, MyEntityTestResource.class)
+                    .addAsResource("application.properties")
+                    .addAsResource("import.sql"));
+
+    @Test
+    public void testCleanDatabase() {
+        RestAssured.when().get("/my-entity/count").then().body(is("2"));
+        RestAssured.when().get("/my-entity/add").then().body(is("MyEntity:added"));
+        RestAssured.when().get("/my-entity/count").then().body(is("3"));
+        RestAssured.with()
+                .redirects().follow(false).formParam("name", "<default>").post("q/dev/io.quarkus.quarkus-agroal/clean")
+                .then()
+                .statusCode(303);
+        RestAssured.when().get("/my-entity/count").then().body(is("2"));
+
+    }
+
+    private void assertBodyIs(String expectedBody) {
+        RestAssured.when().get("/my-entity/2").then().body(is(expectedBody));
+    }
+}
diff --git a/...hibernate-orm/deployment/src/test/java/io/quarkus/hibernate/orm/MyEntityTestResource.java b/...hibernate-orm/deployment/src/test/java/io/quarkus/hibernate/orm/MyEntityTestResource.java
@@ -2,6 +2,7 @@
 
 import javax.inject.Inject;
 import javax.persistence.EntityManager;
+import javax.transaction.Transactional;
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
@@ -25,4 +26,22 @@ public String getName(@PathParam("id") long id) {
 
         return "no entity";
     }
+
+    @GET
+    @Path("/add")
+    @Produces(MediaType.TEXT_PLAIN)
+    @Transactional
+    public String add() {
+        MyEntity entity = new MyEntity();
+        entity.setName("added");
+        em.persist(entity);
+        return entity.toString();
+    }
+
+    @GET
+    @Path("/count")
+    @Produces(MediaType.TEXT_PLAIN)
+    public int count() {
+        return em.createQuery("from MyEntity").getResultList().size();
+    }
 }
diff --git a/extensions/hibernate-orm/deployment/src/test/resources/import.sql b/extensions/hibernate-orm/deployment/src/test/resources/import.sql
@@ -1,2 +1,3 @@
 INSERT INTO MyEntity(id, name) VALUES(1, 'default sql load script entity');
 INSERT INTO MyEntity(id, name) VALUES(2, 'import.sql load script entity');
+alter sequence myEntitySeq restart with 3;
diff --git a/...nate-orm/runtime/src/main/java/io/quarkus/hibernate/orm/runtime/HibernateOrmRecorder.java b/...nate-orm/runtime/src/main/java/io/quarkus/hibernate/orm/runtime/HibernateOrmRecorder.java
@@ -21,6 +21,7 @@
 import io.quarkus.hibernate.orm.runtime.boot.QuarkusPersistenceUnitDefinition;
 import io.quarkus.hibernate.orm.runtime.integration.HibernateOrmIntegrationRuntimeDescriptor;
 import io.quarkus.hibernate.orm.runtime.proxies.PreGeneratedProxies;
+import io.quarkus.hibernate.orm.runtime.schema.SchemaManagementIntegrator;
 import io.quarkus.hibernate.orm.runtime.session.ForwardingSession;
 import io.quarkus.hibernate.orm.runtime.tenant.DataSourceTenantConnectionResolver;
 import io.quarkus.runtime.annotations.Recorder;
@@ -55,6 +56,12 @@ public void setupPersistenceProvider(HibernateOrmRuntimeConfig hibernateOrmRunti
     public BeanContainerListener initMetadata(List<QuarkusPersistenceUnitDefinition> parsedPersistenceXmlDescriptors,
             Scanner scanner, Collection<Class<? extends Integrator>> additionalIntegrators,
             PreGeneratedProxies proxyDefinitions) {
+        SchemaManagementIntegrator.clearDsMap();
+        for (QuarkusPersistenceUnitDefinition i : parsedPersistenceXmlDescriptors) {
+            if (i.getDataSource().isPresent()) {
+                SchemaManagementIntegrator.mapDatasource(i.getDataSource().get(), i.getName());
+            }
+        }
         return new BeanContainerListener() {
             @Override
             public void created(BeanContainer beanContainer) {
@@ -119,5 +126,4 @@ protected Session delegate() {
             }
         };
     }
-
 }