Give JpaTransactionManagerRule more parameters (#292)

* Give JpaTransactionManagerRule more parameters

Allow users of the rule to add annotated classes and properties, both useful
for testing.

* Change in response to review.

* Changes for review.

* Move test EntityManagerFactory create method

Move the test create method into the JpaTransactionManagerRuleTest.

* Remove nomulus SQL dialect from G.S.S.Command

Remove NomulusPostgreSQLDialect from GenerateSqlSchemaCommand (it has been
moved to its own top-level class).

core/src/main/java/google/registry/persistence/NomulusPostgreSQLDialect.java

@@ -0,0 +1,27 @@
+// Copyright 2019 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package google.registry.persistence;
+
+import java.sql.Types;
+import org.hibernate.dialect.PostgreSQL95Dialect;
+
+/** Nomulus mapping rules for column types in Postgresql. */
+public class NomulusPostgreSQLDialect extends PostgreSQL95Dialect {
+  public NomulusPostgreSQLDialect() {
+    super();
+    registerColumnType(Types.VARCHAR, "text");
+    registerColumnType(Types.TIMESTAMP_WITH_TIMEZONE, "timestamptz");
+    registerColumnType(Types.TIMESTAMP, "timestamptz");
+  }
+}

core/src/main/java/google/registry/persistence/PersistenceModule.java

@@ -70,6 +70,7 @@ public class PersistenceModule {
     properties.put(HIKARI_MINIMUM_IDLE, getHibernateHikariMinimumIdle());
     properties.put(HIKARI_MAXIMUM_POOL_SIZE, getHibernateHikariMaximumPoolSize());
     properties.put(HIKARI_IDLE_TIMEOUT, getHibernateHikariIdleTimeout());
+    properties.put(Environment.DIALECT, NomulusPostgreSQLDialect.class.getName());
     return properties.build();
   }
 
@@ -102,8 +103,13 @@ public class PersistenceModule {
     properties.put(Environment.URL, jdbcUrl);
     properties.put(Environment.USER, username);
     properties.put(Environment.PASS, password);
+
+    // If there are no annotated classes, we can create the EntityManagerFactory from the generic
+    // method.  Otherwise we have to use a more tailored approach.  Note that this adds to the set
+    // of annotated classes defined in the configuration, it does not override them.
     EntityManagerFactory emf =
         Persistence.createEntityManagerFactory(PERSISTENCE_UNIT_NAME, properties);
+
     checkState(
         emf != null,
         "Persistence.createEntityManagerFactory() returns a null EntityManagerFactory");

core/src/main/java/google/registry/tools/GenerateSqlSchemaCommand.java

@@ -28,20 +28,19 @@ import google.registry.model.eppcommon.Trid;
 import google.registry.model.transfer.BaseTransferObject;
 import google.registry.model.transfer.TransferData;
 import google.registry.persistence.NomulusNamingStrategy;
+import google.registry.persistence.NomulusPostgreSQLDialect;
 import google.registry.schema.domain.RegistryLock;
 import google.registry.schema.tld.PremiumList;
 import google.registry.schema.tmch.ClaimsList;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Paths;
-import java.sql.Types;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Map;
 import org.hibernate.boot.MetadataSources;
 import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
 import org.hibernate.cfg.Environment;
-import org.hibernate.dialect.PostgreSQL95Dialect;
 import org.hibernate.tool.hbm2ddl.SchemaExport;
 import org.hibernate.tool.schema.TargetType;
 import org.joda.time.Period;
@@ -198,13 +197,4 @@ public class GenerateSqlSchemaCommand implements Command {
       }
     }
   }
-
-  /** Nomulus mapping rules for column types in Postgresql. */
-  public static class NomulusPostgreSQLDialect extends PostgreSQL95Dialect {
-    public NomulusPostgreSQLDialect() {
-      super();
-      registerColumnType(Types.VARCHAR, "text");
-      registerColumnType(Types.TIMESTAMP, "timestamptz");
-    }
-  }
 }

core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java

@@ -16,9 +16,19 @@ package google.registry.model.transaction;
 
 import static org.joda.time.DateTimeZone.UTC;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import google.registry.persistence.PersistenceModule;
 import google.registry.testing.FakeClock;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import javax.persistence.EntityManagerFactory;
+import org.hibernate.boot.MetadataSources;
+import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
+import org.hibernate.cfg.Environment;
 import org.joda.time.DateTime;
 import org.junit.rules.ExternalResource;
 import org.junit.rules.RuleChain;
@@ -41,12 +51,20 @@ public class JpaTransactionManagerRule extends ExternalResource {
   private final DateTime now = DateTime.now(UTC);
   private final FakeClock clock = new FakeClock(now);
   private final String initScript;
+  private final ImmutableList<Class> extraEntityClasses;
+  private final ImmutableMap userProperties;
+
   private JdbcDatabaseContainer database;
   private EntityManagerFactory emf;
   private JpaTransactionManager cachedTm;
 
-  private JpaTransactionManagerRule(String initScript) {
+  private JpaTransactionManagerRule(
+      String initScript,
+      ImmutableList<Class> extraEntityClasses,
+      ImmutableMap<String, String> userProperties) {
     this.initScript = initScript;
+    this.extraEntityClasses = extraEntityClasses;
+    this.userProperties = userProperties;
   }
 
   /** Wraps {@link JpaTransactionManagerRule} in a {@link PostgreSQLContainer}. */
@@ -60,12 +78,21 @@ public class JpaTransactionManagerRule extends ExternalResource {
 
   @Override
   public void before() {
+    ImmutableMap properties = PersistenceModule.providesDefaultDatabaseConfigs();
+    if (!userProperties.isEmpty()) {
+      // If there are user properties, create a new properties object with these added.
+      ImmutableMap.Builder builder = properties.builder();
+      builder.putAll(userProperties);
+      properties = builder.build();
+    }
+
     emf =
-        PersistenceModule.create(
+        createEntityManagerFactory(
             database.getJdbcUrl(),
             database.getUsername(),
             database.getPassword(),
-            PersistenceModule.providesDefaultDatabaseConfigs());
+            properties,
+            extraEntityClasses);
     JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock);
     cachedTm = TransactionManagerFactory.jpaTm;
     TransactionManagerFactory.jpaTm = txnManager;
@@ -80,6 +107,24 @@ public class JpaTransactionManagerRule extends ExternalResource {
     cachedTm = null;
   }
 
+  /** Constructs the {@link EntityManagerFactory} instance. */
+  private static EntityManagerFactory createEntityManagerFactory(
+      String jdbcUrl,
+      String username,
+      String password,
+      ImmutableMap<String, String> configs,
+      ImmutableList<Class> extraEntityClasses) {
+    HashMap<String, String> properties = Maps.newHashMap(configs);
+    properties.put(Environment.URL, jdbcUrl);
+    properties.put(Environment.USER, username);
+    properties.put(Environment.PASS, password);
+
+    MetadataSources metadataSources =
+        new MetadataSources(new StandardServiceRegistryBuilder().applySettings(properties).build());
+    extraEntityClasses.forEach(metadataSources::addAnnotatedClass);
+    return metadataSources.buildMetadata().getSessionFactoryBuilder().build();
+  }
+
   /** Returns the {@link FakeClock} used by the underlying {@link JpaTransactionManagerImpl}. */
   public FakeClock getTxnClock() {
     return clock;
@@ -88,6 +133,8 @@ public class JpaTransactionManagerRule extends ExternalResource {
   /** Builder for {@link JpaTransactionManagerRule}. */
   public static class Builder {
     private String initScript;
+    private List<Class> extraEntityClasses = new ArrayList<Class>();
+    private Map<String, String> userProperties = new HashMap<String, String>();
 
     /**
      * Sets the SQL script to be used to initialize the database. If not set,
@@ -98,12 +145,27 @@ public class JpaTransactionManagerRule extends ExternalResource {
       return this;
     }
 
+    /** Adds an annotated class to the known entities for the database. */
+    public Builder withEntityClass(Class clazz) {
+      this.extraEntityClasses.add(clazz);
+      return this;
+    }
+
+    /** Adds the specified property to those used to initialize the transaction manager. */
+    public Builder withProperty(String name, String value) {
+      this.userProperties.put(name, value);
+      return this;
+    }
+
     /** Builds a {@link JpaTransactionManagerRule} instance. */
     public JpaTransactionManagerRule build() {
       if (initScript == null) {
         initScript = SCHEMA_GOLDEN_SQL;
       }
-      return new JpaTransactionManagerRule(initScript);
+      return new JpaTransactionManagerRule(
+          initScript,
+          ImmutableList.copyOf(extraEntityClasses),
+          ImmutableMap.copyOf(userProperties));
     }
   }
 }

core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRuleTest.java

@@ -18,8 +18,12 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.JUnitBackports.assertThrows;
 
+import google.registry.model.ImmutableObject;
 import java.util.List;
+import javax.persistence.Entity;
+import javax.persistence.Id;
 import javax.persistence.PersistenceException;
+import org.hibernate.cfg.Environment;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -31,7 +35,10 @@ public class JpaTransactionManagerRuleTest {
 
   @Rule
   public final JpaTransactionManagerRule jpaTmRule =
-      new JpaTransactionManagerRule.Builder().build();
+      new JpaTransactionManagerRule.Builder()
+          .withEntityClass(TestEntity.class)
+          .withProperty(Environment.HBM2DDL_AUTO, "update")
+          .build();
 
   @Test
   public void verifiesRuleWorks() {
@@ -56,4 +63,28 @@ public class JpaTransactionManagerRuleTest {
               assertThat(results).isEmpty();
             });
   }
+
+  @Test
+  public void testExtraParameters() {
+    // This test verifies that 1) withEntityClass() has registered TestEntity and 2) The table
+    // has been created, implying withProperty(HBM2DDL_AUTO, "update") worked.
+    TestEntity original = new TestEntity("key", "value");
+    jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));
+    TestEntity retrieved =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "key"));
+    assertThat(retrieved).isEqualTo(original);
+  }
+
+  @Entity(name = "TestEntity") // Specify name to avoid nested class naming issues.
+  static class TestEntity extends ImmutableObject {
+    @Id String key;
+    String value;
+
+    TestEntity(String key, String value) {
+      this.key = key;
+      this.value = value;
+    }
+
+    TestEntity() {}
+  }
 }