Speed up test database resets and Postgres config (#3238)

- Mount the test Postgres data directory on an in-memory tmpfs and disable fsync, synchronous_commit, and full_page_writes for ephemeral test containers.
- Override the Nomulus SQL dialect to emit 'deferrable initially deferred' on foreign keys exported by HibernateSchemaExporter. This matches prod.
- Replace per-test TRUNCATES (etc) with DELETEs -- the vast majority of
  tests only change a couple tables and DELETEing a few rows is much
faster than having to reallocate all the disk space for each and every
table every time
- Cache the connection to Postgres that we use for resetting data,
  instead of having to re-create it each time
- Add a test for child entity foreign key generation

This drastically speeds up tests, cutting out like 30-40% of total test
time (not including compilation)

(we also update a couple queries to be deterministically ordered to
avoid test flakes)

core/src/main/java/google/registry/model/registrar/Registrar.java

@@ -211,7 +211,7 @@ public class Registrar extends UpdateAutoTimestampEntity implements Buildable, J
       memoizeWithShortExpiration(
           () ->
               Maps.uniqueIndex(
-                  tm().reTransact(() -> tm().loadAllOf(Registrar.class)),
+                  tm().reTransact(() -> tm().loadAllOfSorted(Registrar.class, "registrarId")),
                   Registrar::getRegistrarId));
 
   /**
@@ -1050,9 +1050,9 @@ public class Registrar extends UpdateAutoTimestampEntity implements Buildable, J
     return email;
   }
 
-  /** Loads all registrar entities directly from the database. */
+  /** Loads all registrar entities directly from the database, sorted by {@code registrarId}. */
   public static Iterable<Registrar> loadAll() {
-    return tm().transact(() -> tm().loadAllOf(Registrar.class));
+    return loadAllSorted("registrarId");
   }
 
   /** Loads all registrar entities directly from the database, sorted by the given field names. */

core/src/main/java/google/registry/model/tld/Tlds.java

@@ -63,7 +63,8 @@ public final class Tlds {
                       EntityManager entityManager = tm().getEntityManager();
                       Stream<Object[]> resultStream =
                           entityManager
-                              .createQuery("SELECT tldStr, tldType FROM Tld", Object[].class)
+                              .createQuery(
+                                  "SELECT tldStr, tldType FROM Tld ORDER BY tldStr", Object[].class)
                               .getResultStream();
                       return resultStream
                           .map(e -> Maps.immutableEntry((String) e[0], ((TldType) e[1])))

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

@@ -41,6 +41,18 @@ public class NomulusPostgreSQLDialect extends PostgreSQLDialect {
     return elementTypeName + "[]";
   }
 
+  @Override
+  public String getAddForeignKeyConstraintString(
+      String constraintName,
+      String[] foreignKey,
+      String referencedTable,
+      String[] primaryKey,
+      boolean referencesPrimaryKey) {
+    return super.getAddForeignKeyConstraintString(
+            constraintName, foreignKey, referencedTable, primaryKey, referencesPrimaryKey)
+        + " deferrable initially deferred";
+  }
+
   @Override
   public void contributeTypes(
       TypeContributions typeContributions, ServiceRegistry serviceRegistry) {

core/src/test/java/google/registry/persistence/HibernateSchemaExporterTest.java

@@ -18,7 +18,10 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.ImmutableList;
 import jakarta.persistence.Entity;
+import jakarta.persistence.ForeignKey;
 import jakarta.persistence.Id;
+import jakarta.persistence.JoinColumn;
+import jakarta.persistence.ManyToOne;
 import java.io.File;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -56,16 +59,29 @@ class HibernateSchemaExporterTest {
   @Test
   void export_succeeds() throws IOException {
     File sqlFile = Files.createFile(tmpDir.resolve("tempfile.dat")).toFile();
-    exporter.export(ImmutableList.of(HibernateSchemaTestEntity.class), sqlFile);
+    exporter.export(
+        ImmutableList.of(HibernateSchemaTestEntity.class, HibernateSchemaChildEntity.class),
+        sqlFile);
     assertThat(Files.readAllBytes(sqlFile.toPath()))
         .isEqualTo(
-            """
+"""
+
+    create table "ChildEntity" (
+        id text not null,
+        parent_name text,
+        primary key (id)
+    );
 
     create table "TestEntity" (
         name text not null,
         cu text,
         primary key (name)
     );
+
+    alter table if exists "ChildEntity"\s
+       add constraint FK_child_parent\s
+       foreign key (parent_name)\s
+       references "TestEntity" deferrable initially deferred;
 """
                 .getBytes(StandardCharsets.UTF_8));
   }
@@ -76,4 +92,13 @@ class HibernateSchemaExporterTest {
 
     CurrencyUnit cu;
   }
+
+  @Entity(name = "ChildEntity")
+  private static class HibernateSchemaChildEntity {
+    @Id String id;
+
+    @ManyToOne
+    @JoinColumn(name = "parent_name", foreignKey = @ForeignKey(name = "FK_child_parent"))
+    HibernateSchemaTestEntity parent;
+  }
 }

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java

@@ -21,7 +21,6 @@ import static google.registry.testing.DatabaseHelper.persistResources;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.testcontainers.containers.PostgreSQLContainer.POSTGRESQL_PORT;
 
-import com.google.common.base.Joiner;
 import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -104,17 +103,20 @@ public abstract class JpaTransactionManagerExtension
   // reused between test methods if the requested schema remains the same.
   private static EntityManagerFactory emf;
   // Hash of the ORM entity names in the current schema in the test db.
-
   private static int emfEntityHash;
-
-  private JpaTransactionManager cachedTm;
   // Hash of the ORM entity names requested by this extension instance.
   private final int entityHash;
+  private JpaTransactionManager cachedTm;
 
   // Whether to create nomulus tables in the test db. Right now, only the JpaTestExtensions set this
   // to false.
   private boolean includeNomulusSchema = true;
 
+  // Cached re-usable connection used for resetting tables / schema
+  private static Connection reusableConnection;
+  // Cached "reset" statement so we don't need to query SQL each time to find out what tables exist
+  private static String cachedResetSql;
+
   // Whether to pre-populate some registrars for ease of testing.
   private final boolean withCannedData;
 
@@ -144,7 +146,19 @@ public abstract class JpaTransactionManagerExtension
                 "POSTGRES_INITDB_ARGS",
                 "--encoding=UTF8 --lc-collate=en_US.UTF8 --lc-ctype=en_US.UTF8"
                     + " --locale-provider=libc --no-locale")
-            .withDatabaseName(POSTGRES_DB_NAME);
+            .withDatabaseName(POSTGRES_DB_NAME)
+            // Use tmpFs so that we keep everything in memory. No point in writing to disk
+            .withTmpFs(ImmutableMap.of("/var/lib/postgresql/data", "rw"))
+            // Testcontainers defaults to fsync=off but we want the additional options turned off as
+            // well, for speed purposes
+            .withCommand(
+                "postgres",
+                "-c",
+                "fsync=off",
+                "-c",
+                "synchronous_commit=off",
+                "-c",
+                "full_page_writes=off");
     container.start();
     return container;
   }
@@ -159,11 +173,31 @@ public abstract class JpaTransactionManagerExtension
         .hashCode();
   }
 
+  private static void closeReusableConnection() {
+    if (reusableConnection != null) {
+      try {
+        reusableConnection.close();
+      } catch (SQLException e) {
+        // Ignore close errors during cleanup
+      }
+      reusableConnection = null;
+    }
+  }
+
+  private static Connection getReusableConnection() throws SQLException {
+    if (reusableConnection == null || reusableConnection.isClosed()) {
+      reusableConnection = createConnection();
+    }
+    return reusableConnection;
+  }
+
   /**
    * Drops and recreates the 'public' schema and all tables, then creates a new {@link
    * EntityManagerFactory} and save it in {@link #emf}.
    */
   private void recreateSchema() throws Exception {
+    closeReusableConnection();
+    cachedResetSql = null;
     if (emf != null) {
       emf.close();
       emf = null;
@@ -178,6 +212,10 @@ public abstract class JpaTransactionManagerExtension
       exporter.export(extraEntityClasses, tempSqlFile);
       executeSql(Files.readString(tempSqlFile.toPath(), UTF_8));
     }
+    // Lower project_wide_unique_id_seq's non-1 MINVALUE from the golden schema so deterministic
+    // sequence resets to 1 succeed.
+    executeSql(
+        "ALTER SEQUENCE IF EXISTS project_wide_unique_id_seq START 1 MINVALUE 1 RESTART WITH 1;");
     assertReasonableNumDbConnections();
     emf = createEntityManagerFactory(getJpaProperties());
     emfEntityHash = entityHash;
@@ -228,16 +266,6 @@ public abstract class JpaTransactionManagerExtension
     cachedTm = TransactionManagerFactory.tm();
     TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(txnManager));
     TransactionManagerFactory.setReplicaJpaTm(Suppliers.ofInstance(readOnlyTxnManager));
-    // Reset SQL Sequence based id allocation so that ids are deterministic in tests.
-    TransactionManagerFactory.tm()
-        .transact(
-            () ->
-                TransactionManagerFactory.tm()
-                    .getEntityManager()
-                    .createNativeQuery(
-                        "alter sequence if exists project_wide_unique_id_seq start 1 minvalue 1"
-                            + " restart with 1")
-                    .executeUpdate());
     if (withCannedData) {
       loadInitialData();
     }
@@ -256,22 +284,34 @@ public abstract class JpaTransactionManagerExtension
   }
 
   private void resetTablesAndSequences() {
-    try (Connection conn = createConnection();
-        Statement statement = conn.createStatement()) {
-      ResultSet rs =
-          statement.executeQuery(
-              "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';");
-      ImmutableList.Builder<String> tableNamesBuilder = new ImmutableList.Builder<>();
-      while (rs.next()) {
-        tableNamesBuilder.add('"' + rs.getString(1) + '"');
-      }
-      ImmutableList<String> tableNames = tableNamesBuilder.build();
-      if (!tableNames.isEmpty()) {
-        String sql =
-            String.format("TRUNCATE %s RESTART IDENTITY CASCADE", Joiner.on(',').join(tableNames));
-        executeSql(sql);
+    try (Statement statement = getReusableConnection().createStatement()) {
+      if (cachedResetSql == null) {
+        // We build up a list of tables/sequences to reset by querying SQL. We only need to do this
+        // once per schema.
+        StringBuilder sqlBuilder = new StringBuilder("BEGIN; SET CONSTRAINTS ALL DEFERRED; ");
+        try (ResultSet rs =
+            statement.executeQuery(
+                "SELECT table_name FROM information_schema.tables"
+                    + " WHERE table_schema = 'public' AND table_type = 'BASE TABLE';")) {
+          while (rs.next()) {
+            sqlBuilder.append(String.format("DELETE FROM \"%s\";", rs.getString(1)));
+          }
+        }
+        try (ResultSet rs =
+            statement.executeQuery(
+                "SELECT sequence_name FROM information_schema.sequences"
+                    + " WHERE sequence_schema = 'public';")) {
+          while (rs.next()) {
+            sqlBuilder.append(
+                String.format("ALTER SEQUENCE \"%s\" RESTART WITH 1;", rs.getString(1)));
+          }
+        }
+        sqlBuilder.append("COMMIT;");
+        cachedResetSql = sqlBuilder.toString();
       }
+      statement.execute(cachedResetSql);
     } catch (Exception e) {
+      closeReusableConnection();
       throw new RuntimeException(e);
     }
   }

db/src/main/resources/sql/schema/db-schema.sql.generated

@@ -951,49 +951,49 @@
     alter table if exists "ConsoleUpdateHistory" 
        add constraint FKnhl1eolgix64u90xv3pj6xa3x 
        foreign key (acting_user) 
-       references "User";
+       references "User" deferrable initially deferred;
 
     alter table if exists "DelegationSignerData" 
        add constraint FKtr24j9v14ph2mfuw2gsmt12kq 
        foreign key (domain_repo_id) 
-       references "Domain";
+       references "Domain" deferrable initially deferred;
 
     alter table if exists "DomainDsDataHistory" 
        add constraint FKo4ilgyyfnvppbpuivus565i0j 
        foreign key (domain_repo_id, domain_history_revision_id) 
-       references "DomainHistory";
+       references "DomainHistory" deferrable initially deferred;
 
     alter table if exists "DomainHistoryHost" 
        add constraint FKa9woh3hu8gx5x0vly6bai327n 
        foreign key (domain_history_domain_repo_id, domain_history_history_revision_id) 
-       references "DomainHistory";
+       references "DomainHistory" deferrable initially deferred;
 
     alter table if exists "DomainHost" 
        add constraint FKeq1guccbre1yk3oosgp2io554 
        foreign key (domain_repo_id) 
-       references "Domain";
+       references "Domain" deferrable initially deferred;
 
     alter table if exists "DomainTransactionRecord" 
        add constraint FKcjqe54u72kha71vkibvxhjye7 
        foreign key (domain_repo_id, history_revision_id) 
-       references "DomainHistory";
+       references "DomainHistory" deferrable initially deferred;
 
     alter table if exists "GracePeriod" 
        add constraint FK2mys4hojm6ev2g9tmy5aq6m7g 
        foreign key (domain_repo_id) 
-       references "Domain";
+       references "Domain" deferrable initially deferred;
 
     alter table if exists "GracePeriodHistory" 
        add constraint FK7w3cx8d55q8bln80e716tr7b8 
        foreign key (domain_repo_id, domain_history_revision_id) 
-       references "DomainHistory";
+       references "DomainHistory" deferrable initially deferred;
 
     alter table if exists "RegistryLock" 
        add constraint FK2lhcwpxlnqijr96irylrh1707 
        foreign key (relock_revision_id) 
-       references "RegistryLock";
+       references "RegistryLock" deferrable initially deferred;
 
     alter table if exists "SignedMarkRevocationEntry" 
        add constraint FK5ivlhvs3121yx2li5tqh54u4 
        foreign key (revision_id) 
-       references "SignedMarkRevocationList";
+       references "SignedMarkRevocationList" deferrable initially deferred;

db/src/test/java/google/registry/sql/flyway/SchemaTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.truth.TextDiffSubject.assertThat;
 
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.io.Resources;
 import google.registry.persistence.NomulusPostgreSql;
@@ -81,7 +82,9 @@ class SchemaTest {
   private final PostgreSQLContainer<?> sqlContainer =
       new PostgreSQLContainer<>(NomulusPostgreSql.getDockerImageName())
           .withClasspathResourceMapping(
-              MOUNTED_RESOURCE_PATH, CONTAINER_MOUNT_POINT, BindMode.READ_WRITE);
+              MOUNTED_RESOURCE_PATH, CONTAINER_MOUNT_POINT, BindMode.READ_WRITE)
+          // Use tmpFs so that we keep everything in memory. No point in writing to disk
+          .withTmpFs(ImmutableMap.of("/var/lib/postgresql/data", "rw"));
 
   @Test
   @DisabledIfSystemProperty(named = "deploy_to_existing_db", matches = ".*")