Convert AuthenticatedRegAccessor and OteStats to SQL (#1039)

This required adding a new HistoryEntryDao method but it's fairly similar to the ones we already have.
2026-08-15 11:46:08 +00:00 · 2021-04-02 11:41:26 -04:00
parent 4e7dd7a95a
commit c077aca433
9 changed files with 230 additions and 119 deletions
@@ -18,7 +18,6 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABILITY;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -41,10 +40,10 @@ import google.registry.model.registry.Registry;
 import google.registry.model.registry.Registry.TldState;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.model.registry.label.PremiumListDualDao;
+import google.registry.persistence.VKey;
 import google.registry.util.CidrAddressBlock;
 import java.util.Collection;
 import java.util.Optional;
-import java.util.Set;
 import java.util.function.Function;
 import java.util.regex.Pattern;
 import org.joda.money.CurrencyUnit;
@@ -225,15 +224,14 @@ public final class OteAccountBuilder {
  }

  /**
-   * Persists all the OT&amp;E entities to datastore.
+   * Persists all the OT&amp;E entities to the database.
   *
   * @return map from the new clientIds created to the new TLDs they have access to. Can be used to
   *     go over all the newly created Registrars / Registries / RegistrarContacts if any
   *     post-creation work is needed.
   */
  public ImmutableMap<String, String> buildAndPersist() {
-    // save all the entitiesl in a single transaction
-    tm().transact(this::saveAllEntities);
+    saveAllEntities();
    return clientIdToTld;
  }

@@ -246,30 +244,38 @@ public final class OteAccountBuilder {

  /** Saves all the OT&amp;E entities we created. */
  private void saveAllEntities() {
-    tm().assertInTransaction();
-
    // use ImmutableObject instead of Registry so that the Key generation doesn't break
-    ImmutableList<ImmutableObject> registries = ImmutableList.of(sunriseTld, gaTld, eapTld);
+    ImmutableList<Registry> registries = ImmutableList.of(sunriseTld, gaTld, eapTld);
    ImmutableList<RegistrarContact> contacts = contactsBuilder.build();

-    if (!replaceExisting) {
-      ImmutableList<Key<ImmutableObject>> keys =
-          Streams.concat(registries.stream(), registrars.stream(), contacts.stream())
-              .map(Key::create)
-              .collect(toImmutableList());
-      Set<Key<ImmutableObject>> existingKeys = ofy().load().keys(keys).keySet();
-      checkState(
-          existingKeys.isEmpty(),
-          "Found existing object(s) conflicting with OT&E objects: %s",
-          existingKeys);
-    }
-    // Save the Registries (TLDs) first
-    ofy().save().entities(registries).now();
-    // Now we can set the allowedTlds for the registrars
-    registrars = registrars.stream().map(this::addAllowedTld).collect(toImmutableList());
-    // and we can save the registrars and contacts!
-    ofy().save().entities(registrars);
-    ofy().save().entities(contacts);
+    tm().transact(
+            () -> {
+              if (!replaceExisting) {
+                ImmutableList<VKey<? extends ImmutableObject>> keys =
+                    Streams.concat(
+                            registries.stream()
+                                .map(registry -> Registry.createVKey(registry.getTldStr())),
+                            registrars.stream().map(Registrar::createVKey),
+                            contacts.stream().map(RegistrarContact::createVKey))
+                        .collect(toImmutableList());
+                ImmutableMap<VKey<? extends ImmutableObject>, ImmutableObject> existingObjects =
+                    tm().loadByKeysIfPresent(keys);
+                checkState(
+                    existingObjects.isEmpty(),
+                    "Found existing object(s) conflicting with OT&E objects: %s",
+                    existingObjects.keySet());
+              }
+              // Save the Registries (TLDs) first
+              tm().putAll(registries);
+            });
+    // Now we can set the allowedTlds for the registrars in a new transaction
+    tm().transact(
+            () -> {
+              registrars = registrars.stream().map(this::addAllowedTld).collect(toImmutableList());
+              // and we can save the registrars and contacts!
+              tm().putAll(registrars);
+              tm().putAll(contacts);
+            });
  }

  private Registrar addAllowedTld(Registrar registrar) {
@@ -17,7 +17,6 @@ package google.registry.model;
 import static com.google.common.base.Predicates.equalTo;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static google.registry.model.eppcommon.EppXmlTransformer.unmarshal;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.DomainNameUtils.ACE_PREFIX;

@@ -28,7 +27,6 @@ import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Multiset;
 import com.googlecode.objectify.Key;
-import com.googlecode.objectify.cmd.Query;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.fee.FeeCreateCommandExtension;
 import google.registry.model.domain.launch.LaunchCreateExtension;
@@ -39,6 +37,7 @@ import google.registry.model.eppinput.EppInput.ResourceCommandWrapper;
 import google.registry.model.host.HostCommand;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.model.reporting.HistoryEntry.Type;
+import google.registry.model.reporting.HistoryEntryDao;
 import google.registry.xml.XmlException;
 import java.util.Arrays;
 import java.util.EnumSet;
@@ -196,16 +195,10 @@ public class OteStats {
   * <p>Stops when it notices that all tests have passed.
   */
  private OteStats recordRegistrarHistory(String registrarName) {
-    ImmutableCollection<String> clientIds =
+    ImmutableCollection<String> registrarIds =
        OteAccountBuilder.createClientIdToTldMap(registrarName).keySet();

-    Query<HistoryEntry> query =
-        ofy()
-            .load()
-            .type(HistoryEntry.class)
-            .filter("clientId in", clientIds)
-            .order("modificationTime");
-    for (HistoryEntry historyEntry : query) {
+    for (HistoryEntry historyEntry : HistoryEntryDao.loadHistoryObjectsByRegistrars(registrarIds)) {
      try {
        record(historyEntry);
      } catch (XmlException e) {
@@ -21,7 +21,9 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;

+import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Streams;
 import google.registry.model.EppResource;
 import google.registry.model.contact.ContactHistory;
 import google.registry.model.contact.ContactResource;
@@ -30,7 +32,11 @@ import google.registry.model.domain.DomainHistory;
 import google.registry.model.host.HostHistory;
 import google.registry.model.host.HostResource;
 import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.CriteriaQueryBuilder;
 import java.util.Comparator;
+import java.util.stream.Stream;
+import javax.persistence.criteria.CriteriaBuilder;
+import javax.persistence.criteria.CriteriaQuery;
 import org.joda.time.DateTime;

 /**
@@ -85,21 +91,56 @@ public class HistoryEntryDao {
    }
  }

+  /** Loads all history objects from all time from the given registrars. */
+  public static Iterable<? extends HistoryEntry> loadHistoryObjectsByRegistrars(
+      ImmutableCollection<String> registrarIds) {
+    if (tm().isOfy()) {
+      return ofy()
+          .load()
+          .type(HistoryEntry.class)
+          .filter("clientId in", registrarIds)
+          .order("modificationTime");
+    } else {
+      return jpaTm()
+          .transact(
+              () ->
+                  Streams.concat(
+                          loadHistoryObjectFromSqlByRegistrars(ContactHistory.class, registrarIds),
+                          loadHistoryObjectFromSqlByRegistrars(DomainHistory.class, registrarIds),
+                          loadHistoryObjectFromSqlByRegistrars(HostHistory.class, registrarIds))
+                      .sorted(Comparator.comparing(HistoryEntry::getModificationTime))
+                      .collect(toImmutableList()));
+    }
+  }
+
+  private static Stream<? extends HistoryEntry> loadHistoryObjectFromSqlByRegistrars(
+      Class<? extends HistoryEntry> historyClass, ImmutableCollection<String> registrarIds) {
+    return jpaTm()
+        .getEntityManager()
+        .createQuery(
+            CriteriaQueryBuilder.create(historyClass)
+                .whereFieldIsIn("clientId", registrarIds)
+                .build())
+        .getResultStream();
+  }
+
  private static Iterable<? extends HistoryEntry> loadHistoryObjectsForResourceFromSql(
      VKey<? extends EppResource> parentKey, DateTime afterTime, DateTime beforeTime) {
+    // The class we're searching from is based on which parent type (e.g. Domain) we have
    Class<? extends HistoryEntry> historyClass = getHistoryClassFromParent(parentKey.getKind());
+    // The field representing repo ID unfortunately varies by history class
    String repoIdFieldName = getRepoIdFieldNameFromHistoryClass(historyClass);
-    String tableName = jpaTm().getEntityManager().getMetamodel().entity(historyClass).getName();
-    String queryString =
-        String.format(
-            "SELECT entry FROM %s entry WHERE entry.modificationTime >= :afterTime AND "
-                + "entry.modificationTime <= :beforeTime AND entry.%s = :parentKey",
-            tableName, repoIdFieldName);
+    CriteriaBuilder criteriaBuilder = jpaTm().getEntityManager().getCriteriaBuilder();
+    CriteriaQuery<? extends HistoryEntry> criteriaQuery =
+        CriteriaQueryBuilder.create(historyClass)
+            .where("modificationTime", criteriaBuilder::greaterThanOrEqualTo, afterTime)
+            .where("modificationTime", criteriaBuilder::lessThanOrEqualTo, beforeTime)
+            .where(repoIdFieldName, criteriaBuilder::equal, parentKey.getSqlKey().toString())
+            .build();
+
    return jpaTm()
-        .query(queryString, historyClass)
-        .setParameter("afterTime", afterTime)
-        .setParameter("beforeTime", beforeTime)
-        .setParameter("parentKey", parentKey.getSqlKey().toString())
+        .getEntityManager()
+        .createQuery(criteriaQuery)
        .getResultStream()
        .sorted(Comparator.comparing(HistoryEntry::getModificationTime))
        .collect(toImmutableList());
@@ -127,15 +168,14 @@ public class HistoryEntryDao {

  private static Iterable<? extends HistoryEntry> loadAllHistoryObjectsFromSql(
      Class<? extends HistoryEntry> historyClass, DateTime afterTime, DateTime beforeTime) {
+    CriteriaBuilder criteriaBuilder = jpaTm().getEntityManager().getCriteriaBuilder();
    return jpaTm()
-        .query(
-            String.format(
-                "SELECT entry FROM %s entry WHERE entry.modificationTime >= :afterTime AND "
-                    + "entry.modificationTime <= :beforeTime",
-                jpaTm().getEntityManager().getMetamodel().entity(historyClass).getName()),
-            historyClass)
-        .setParameter("afterTime", afterTime)
-        .setParameter("beforeTime", beforeTime)
+        .getEntityManager()
+        .createQuery(
+            CriteriaQueryBuilder.create(historyClass)
+                .where("modificationTime", criteriaBuilder::greaterThanOrEqualTo, afterTime)
+                .where("modificationTime", criteriaBuilder::lessThanOrEqualTo, beforeTime)
+                .build())
        .getResultList();
  }
 }
@@ -19,6 +19,9 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Streams.stream;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;

 import com.google.appengine.api.users.User;
 import com.google.common.annotations.VisibleForTesting;
@@ -316,30 +319,46 @@ public class AuthenticatedRegistrarAccessor {
    logger.atInfo().log("Checking registrar contacts for user ID %s", user.getUserId());

    // Find all registrars that have a registrar contact with this user's ID.
-    ImmutableList<Key<Registrar>> accessibleClientIds =
-        stream(ofy().load().type(RegistrarContact.class).filter("gaeUserId", user.getUserId()))
-            .map(RegistrarContact::getParent)
-            .collect(toImmutableList());
-    // Filter out disabled registrars (note that pending registrars still allow console login).
-    ofy().load().keys(accessibleClientIds).values().stream()
-        .filter(registrar -> registrar.getState() != State.DISABLED)
-        .forEach(registrar -> builder.put(registrar.getClientId(), Role.OWNER));
+    if (tm().isOfy()) {
+      ImmutableList<Key<Registrar>> accessibleClientIds =
+          stream(ofy().load().type(RegistrarContact.class).filter("gaeUserId", user.getUserId()))
+              .map(RegistrarContact::getParent)
+              .collect(toImmutableList());
+      // Filter out disabled registrars (note that pending registrars still allow console login).
+      ofy().load().keys(accessibleClientIds).values().stream()
+          .filter(registrar -> registrar.getState() != State.DISABLED)
+          .forEach(registrar -> builder.put(registrar.getClientId(), Role.OWNER));
+    } else {
+      jpaTm()
+          .transact(
+              () ->
+                  jpaTm()
+                      .query(
+                          "SELECT r FROM Registrar r INNER JOIN RegistrarPoc rp ON "
+                              + "r.clientIdentifier = rp.registrarId WHERE rp.gaeUserId = "
+                              + ":gaeUserId AND r.state != :state",
+                          Registrar.class)
+                      .setParameter("gaeUserId", user.getUserId())
+                      .setParameter("state", State.DISABLED)
+                      .getResultStream()
+                      .forEach(registrar -> builder.put(registrar.getClientId(), Role.OWNER)));
+    }

    // Admins have ADMIN access to all registrars, and also OWNER access to the registry registrar
    // and all non-REAL or non-live registrars.
    if (isAdmin) {
-      ofy()
-          .load()
-          .type(Registrar.class)
-          .forEach(
-              registrar -> {
-                if (registrar.getType() != Registrar.Type.REAL
-                    || !registrar.isLive()
-                    || registrar.getClientId().equals(registryAdminClientId)) {
-                  builder.put(registrar.getClientId(), Role.OWNER);
-                }
-                builder.put(registrar.getClientId(), Role.ADMIN);
-              });
+      transactIfJpaTm(
+          () ->
+              tm().loadAllOf(Registrar.class)
+                  .forEach(
+                      registrar -> {
+                        if (registrar.getType() != Registrar.Type.REAL
+                            || !registrar.isLive()
+                            || registrar.getClientId().equals(registryAdminClientId)) {
+                          builder.put(registrar.getClientId(), Role.OWNER);
+                        }
+                        builder.put(registrar.getClientId(), Role.ADMIN);
+                      }));
    }

    return builder.build();