Set up deployment of the Spec11 pipeline with JPA TM (#716)

* Set up deployment of the Spec11 pipeline with JPA TM

* Remove unnecessarily pipeline options setting

* Use enviroment name in BeamJpaModuleTest

* Fix checkstyle error

config/nom_build.py

@@ -19,6 +19,7 @@ import argparse
 import attr
 import io
 import os
+import shutil
 import subprocess
 import sys
 from typing import List, Union
@@ -58,6 +59,21 @@ PROPERTIES_HEADER = """\
 org.gradle.jvmargs=-Xmx1024m
 """
 
+# Help text to be displayed (in addition to the synopsis and flag help, which
+# are displayed automatically).
+HELP_TEXT = """\
+A wrapper around the gradle build that provides the following features:
+
+-   Converts properties into flags to guard against property name spelling errors
+    and to provide help descriptions for all properties.
+-   Provides pseudo-commands (with the ":nom:" prefix) that encapsulate common
+    actions that are difficult to implement in gradle.
+
+Pseudo-commands:
+    :nom:generate_golden_file - regenerates the golden file from the current
+        set of flyway files.
+"""
+
 # Define all of our special gradle properties here.
 PROPERTIES = [
     Property('mavenUrl',
@@ -256,8 +272,42 @@ def get_root() -> str:
     return cur_dir
 
 
-def main(args):
-    parser = argparse.ArgumentParser('nom_build')
+class Abort(Exception):
+    """Raised to terminate the process with a non-zero error code.
+
+    Parameters are ignored.
+    """
+
+
+def do_pseudo_task(task: str) -> None:
+    root = get_root()
+    if task == ':nom:generate_golden_file':
+        if not subprocess.call([f'{root}/gradlew', ':db:test']):
+            print('\033[33mWARNING:\033[0m Golden schema appears to be '
+                  'up-to-date.  If you are making schema changes,  be sure to '
+                  'add a flyway file for them.')
+            return
+        print('\033[33mWARNING:\033[0m Ignore the above failure, it is '
+              'expected.')
+
+        # Copy the new schema into place.
+        shutil.copy(f'{root}/db/build/resources/test/testcontainer/'
+                    'mount/dump.txt',
+                    f'{root}/db/src/main/resources/sql/schema/'
+                    'nomulus.golden.sql')
+
+        if subprocess.call([f'{root}/gradlew', ':db:test']):
+            print('\033[31mERROR:\033[0m Golden file test failed after '
+                  'copying schema.  Please check your flyway files.')
+            raise Abort()
+    else:
+        print(f'\033[31mERROR:\033[0m Unknown task {task}')
+        raise Abort()
+
+
+def main(args) -> int:
+    parser = argparse.ArgumentParser('nom_build', description=HELP_TEXT,
+                                     formatter_class=argparse.RawTextHelpFormatter)
     for prop in PROPERTIES:
         parser.add_argument('--' + prop.name, default=prop.default,
                             help=prop.desc)
@@ -296,7 +346,7 @@ def main(args):
     if args.generate_gradle_properties:
         with open(f'{root}/gradle.properties', 'w') as dst:
             dst.write(gradle_properties)
-        return
+        return 0
 
     # Verify that the gradle properties file is what we expect it to be.
     with open(f'{root}/gradle.properties') as src:
@@ -321,12 +371,39 @@ def main(args):
             if flag.has_arg:
                 gradle_command.append(arg_val)
 
+    # See if there are any special ":nom:" pseudo-tasks specified.
+    got_non_pseudo_tasks = False
+    for arg in args.non_flag_args[1:]:
+        if arg.startswith(':nom:'):
+            if got_non_pseudo_tasks:
+                # We can't currently deal with the situation of gradle tasks
+                # before pseudo-tasks.  This could be implemented by invoking
+                # gradle for only the set of gradle tasks before the pseudo
+                # task, but that's overkill for now.
+                print(f'\033[31mERROR:\033[0m Pseudo task ({arg}) must be '
+                      'specified prior to all actual gradle tasks.  Aborting.')
+                return 1
+            do_pseudo_task(arg)
+        else:
+            got_non_pseudo_tasks = True
+    non_flag_args = [
+        arg for arg in args.non_flag_args[1:] if not arg.startswith(':nom:')]
+
+    if not non_flag_args:
+        if not got_non_pseudo_tasks:
+            print('\033[33mWARNING:\033[0m No tasks specified.  Not '
+                  'doing anything')
+        return 0
+
     # Add the non-flag args (we exclude the first, which is the command name
     # itself) and run.
-    gradle_command.extend(args.non_flag_args[1:])
-    subprocess.call(gradle_command)
+    gradle_command.extend(non_flag_args)
+    return subprocess.call(gradle_command)
 
 
 if __name__ == '__main__':
-    main(sys.argv)
+    try:
+        sys.exit(main(sys.argv))
+    except Abort as ex:
+        sys.exit(1)
 

config/nom_build_test.py

@@ -14,6 +14,7 @@
 
 import io
 import os
+import shutil
 import unittest
 from unittest import mock
 import nom_build
@@ -67,6 +68,7 @@ class MyTest(unittest.TestCase):
             mock.patch.object(nom_build, 'print', self.print_fake).start())
 
         self.call_mock = mock.patch.object(subprocess, 'call').start()
+        self.copy_mock = mock.patch.object(shutil, 'copy').start()
 
         self.file_contents = {
             # Prefil with the actual file contents.
@@ -92,17 +94,32 @@ class MyTest(unittest.TestCase):
 
     def test_no_args(self):
         nom_build.main(['nom_build'])
-        self.assertEqual(self.printed, [])
-        self.call_mock.assert_called_with([GRADLEW])
+        self.assertEqual(self.printed,
+                         ['\x1b[33mWARNING:\x1b[0m No tasks specified.  Not '
+                         'doing anything'])
 
     def test_property_calls(self):
-        nom_build.main(['nom_build', '--testFilter=foo'])
-        self.call_mock.assert_called_with([GRADLEW, '-P', 'testFilter=foo'])
+        nom_build.main(['nom_build', 'task-name', '--testFilter=foo'])
+        self.call_mock.assert_called_with([GRADLEW, '-P', 'testFilter=foo',
+                                          'task-name'])
 
     def test_gradle_flags(self):
-        nom_build.main(['nom_build', '-d', '-b', 'foo'])
+        nom_build.main(['nom_build', 'task-name', '-d', '-b', 'foo'])
         self.call_mock.assert_called_with([GRADLEW, '--build-file', 'foo',
-                                          '--debug'])
+                                          '--debug', 'task-name'])
+
+    def test_generate_golden_file(self):
+        self.call_mock.side_effect = [1, 0]
+        nom_build.main(['nom_build', ':nom:generate_golden_file'])
+        self.call_mock.assert_has_calls([
+            mock.call([GRADLEW, ':db:test']),
+            mock.call([GRADLEW, ':db:test'])
+        ])
+
+    def test_generate_golden_file_nofail(self):
+        self.call_mock.return_value = 0
+        nom_build.main(['nom_build', ':nom:generate_golden_file'])
+        self.call_mock.assert_has_calls([mock.call([GRADLEW, ':db:test'])])
 
 unittest.main()
 

core/src/main/java/google/registry/beam/initsql/BeamJpaModule.java

@@ -54,37 +54,39 @@ public class BeamJpaModule {
 
   private static final String GCS_SCHEME = "gs://";
 
-  @Nullable private final String credentialFilePath;
+  @Nullable private final String sqlAccessInfoFile;
+  @Nullable private final String cloudKmsProjectId;
 
   /**
    * Constructs a new instance of {@link BeamJpaModule}.
    *
    * <p>Note: it is an unfortunately necessary antipattern to check for the validity of
-   * credentialFilePath in {@link #provideCloudSqlAccessInfo} rather than in the constructor.
+   * sqlAccessInfoFile in {@link #provideCloudSqlAccessInfo} rather than in the constructor.
    * Unfortunately, this is a restriction imposed upon us by Dagger. Specifically, because we use
    * this in at least one 1 {@link google.registry.tools.RegistryTool} command(s), it must be
    * instantiated in {@code google.registry.tools.RegistryToolComponent} for all possible commands;
    * Dagger doesn't permit it to ever be null. For the vast majority of commands, it will never be
    * used (so a null credential file path is fine in those cases).
    *
-   * @param credentialFilePath the path to a Cloud SQL credential file. This must refer to either a
+   * @param sqlAccessInfoFile the path to a Cloud SQL credential file. This must refer to either a
    *     real encrypted file on GCS as returned by {@link
    *     BackupPaths#getCloudSQLCredentialFilePatterns} or an unencrypted file on local filesystem
    *     with credentials to a test database.
    */
-  public BeamJpaModule(@Nullable String credentialFilePath) {
-    this.credentialFilePath = credentialFilePath;
+  public BeamJpaModule(@Nullable String sqlAccessInfoFile, @Nullable String cloudKmsProjectId) {
+    this.sqlAccessInfoFile = sqlAccessInfoFile;
+    this.cloudKmsProjectId = cloudKmsProjectId;
   }
 
   /** Returns true if the credential file is on GCS (and therefore expected to be encrypted). */
   private boolean isCloudSqlCredential() {
-    return credentialFilePath.startsWith(GCS_SCHEME);
+    return sqlAccessInfoFile.startsWith(GCS_SCHEME);
   }
 
   @Provides
   @Singleton
   SqlAccessInfo provideCloudSqlAccessInfo(Lazy<CloudSqlCredentialDecryptor> lazyDecryptor) {
-    checkArgument(!isNullOrEmpty(credentialFilePath), "Null or empty credentialFilePath");
+    checkArgument(!isNullOrEmpty(sqlAccessInfoFile), "Null or empty credentialFilePath");
     String line = readOnlyLineFromCredentialFile();
     if (isCloudSqlCredential()) {
       line = lazyDecryptor.get().decrypt(line);
@@ -101,7 +103,7 @@ public class BeamJpaModule {
 
   String readOnlyLineFromCredentialFile() {
     try {
-      ResourceId resourceId = FileSystems.matchSingleFileSpec(credentialFilePath).resourceId();
+      ResourceId resourceId = FileSystems.matchSingleFileSpec(sqlAccessInfoFile).resourceId();
       try (BufferedReader reader =
           new BufferedReader(
               new InputStreamReader(
@@ -141,8 +143,8 @@ public class BeamJpaModule {
 
   @Provides
   @Config("beamCloudKmsProjectId")
-  static String kmsProjectId() {
-    return "domain-registry-dev";
+  String kmsProjectId() {
+    return cloudKmsProjectId;
   }
 
   @Provides

core/src/main/java/google/registry/beam/initsql/InitSqlPipeline.java

@@ -24,7 +24,6 @@ import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.VersionedEntity;
 import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
 import google.registry.beam.initsql.Transforms.RemoveDomainBaseForeignKeys;
-import google.registry.beam.initsql.Transforms.SerializableSupplier;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainBase;
@@ -227,7 +226,7 @@ public class InitSqlPipeline implements Serializable {
             transformId,
             options.getMaxConcurrentSqlWriters(),
             options.getSqlWriteBatchSize(),
-            new JpaSupplierFactory(credentialFileUrl, jpaGetter)));
+            new JpaSupplierFactory(credentialFileUrl, options.getCloudKmsProjectId(), jpaGetter)));
   }
 
   private static ImmutableList<String> toKindStrings(Collection<Class<?>> entityClasses) {
@@ -235,26 +234,4 @@ public class InitSqlPipeline implements Serializable {
       return entityClasses.stream().map(Key::getKind).collect(ImmutableList.toImmutableList());
     }
   }
-
-  static class JpaSupplierFactory implements SerializableSupplier<JpaTransactionManager> {
-    private static final long serialVersionUID = 1L;
-
-    private String credentialFileUrl;
-    private SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter;
-
-    JpaSupplierFactory(
-        String credentialFileUrl,
-        SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter) {
-      this.credentialFileUrl = credentialFileUrl;
-      this.jpaGetter = jpaGetter;
-    }
-
-    @Override
-    public JpaTransactionManager get() {
-      return jpaGetter.apply(
-          DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
-              .beamJpaModule(new BeamJpaModule(credentialFileUrl))
-              .build());
-    }
-  }
 }

core/src/main/java/google/registry/beam/initsql/InitSqlPipelineOptions.java

@@ -60,6 +60,13 @@ public interface InitSqlPipelineOptions extends GcpOptions {
 
   void setEnvironment(String environment);
 
+  @Description("The GCP project that contains the keyring used for decrypting the "
+      + "SQL credential file.")
+  @Nullable
+  String getCloudKmsProjectId();
+
+  void setCloudKmsProjectId(String cloudKmsProjectId);
+
   @Description(
       "The maximum JDBC connection pool size on a VM. "
           + "This value should be equal to or greater than the number of cores on the VM.")

core/src/main/java/google/registry/beam/initsql/JpaSupplierFactory.java

@@ -0,0 +1,48 @@
+// Copyright 2020 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.beam.initsql;
+
+import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
+import google.registry.beam.initsql.Transforms.SerializableSupplier;
+import google.registry.persistence.transaction.JpaTransactionManager;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+
+public class JpaSupplierFactory implements SerializableSupplier<JpaTransactionManager> {
+
+  private static final long serialVersionUID = 1L;
+
+  private final String credentialFileUrl;
+  @Nullable private final String cloudKmsProjectId;
+  private final SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager>
+      jpaGetter;
+
+  public JpaSupplierFactory(
+      String credentialFileUrl,
+      @Nullable String cloudKmsProjectId,
+      SerializableFunction<JpaTransactionManagerComponent, JpaTransactionManager> jpaGetter) {
+    this.credentialFileUrl = credentialFileUrl;
+    this.cloudKmsProjectId = cloudKmsProjectId;
+    this.jpaGetter = jpaGetter;
+  }
+
+  @Override
+  public JpaTransactionManager get() {
+    return jpaGetter.apply(
+        DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
+            .beamJpaModule(new BeamJpaModule(credentialFileUrl, cloudKmsProjectId))
+            .build());
+  }
+}

core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java

@@ -347,8 +347,8 @@ public class DomainCreateFlow implements TransactionalFlow {
             .setRepoId(repoId)
             .setIdnTableName(validateDomainNameWithIdnTables(domainName))
             .setRegistrationExpirationTime(registrationExpirationTime)
-            .setAutorenewBillingEvent(Key.create(autorenewBillingEvent))
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .setLaunchNotice(hasClaimsNotice ? launchCreate.get().getNotice() : null)
             .setSmdId(signedMarkId)
             .setDsData(secDnsCreate.isPresent() ? secDnsCreate.get().getDsData() : null)

core/src/main/java/google/registry/flows/domain/DomainDeleteFlow.java

@@ -31,7 +31,6 @@ import static google.registry.model.ResourceTransferUtils.handlePendingTransferO
 import static google.registry.model.ResourceTransferUtils.updateForeignKeyIndexDeletionTime;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.ADD_FIELDS;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.RENEW_FIELDS;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -209,7 +208,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
       PollMessage.OneTime deletePollMessage =
           createDeletePollMessage(existingDomain, historyEntry, deletionTime);
       entitiesToSave.add(deletePollMessage);
-      builder.setDeletePollMessage(Key.create(deletePollMessage));
+      builder.setDeletePollMessage(deletePollMessage.createVKey());
     }
 
     // Cancel any grace periods that were still active, and set the expiration time accordingly.
@@ -222,8 +221,7 @@ public final class DomainDeleteFlow implements TransactionalFlow {
         if (gracePeriod.getOneTimeBillingEvent() != null) {
           // Take the amount of amount of registration time being refunded off the expiration time.
           // This can be either add grace periods or renew grace periods.
-          BillingEvent.OneTime oneTime =
-              ofy().load().key(gracePeriod.getOneTimeBillingEvent()).now();
+          BillingEvent.OneTime oneTime = tm().load(gracePeriod.getOneTimeBillingEvent());
           newExpirationTime = newExpirationTime.minusYears(oneTime.getPeriodYears());
         } else if (gracePeriod.getRecurringBillingEvent() != null) {
           // Take 1 year off the registration if in the autorenew grace period (no need to load the
@@ -370,12 +368,12 @@ public final class DomainDeleteFlow implements TransactionalFlow {
   private Money getGracePeriodCost(GracePeriod gracePeriod, DateTime now) {
     if (gracePeriod.getType() == GracePeriodStatus.AUTO_RENEW) {
       DateTime autoRenewTime =
-          ofy().load().key(checkNotNull(gracePeriod.getRecurringBillingEvent())).now()
+          tm().load(checkNotNull(gracePeriod.getRecurringBillingEvent()))
               .getRecurrenceTimeOfYear()
-                  .getLastInstanceBeforeOrAt(now);
+              .getLastInstanceBeforeOrAt(now);
       return getDomainRenewCost(targetId, autoRenewTime, 1);
     }
-    return ofy().load().key(checkNotNull(gracePeriod.getOneTimeBillingEvent())).now().getCost();
+    return tm().load(checkNotNull(gracePeriod.getOneTimeBillingEvent())).getCost();
   }
 
   @Nullable

core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java

@@ -517,14 +517,14 @@ public class DomainFlowUtils {
    */
   public static void updateAutorenewRecurrenceEndTime(DomainBase domain, DateTime newEndTime) {
     Optional<PollMessage.Autorenew> autorenewPollMessage =
-        Optional.ofNullable(ofy().load().key(domain.getAutorenewPollMessage()).now());
+        tm().maybeLoad(domain.getAutorenewPollMessage());
 
     // Construct an updated autorenew poll message. If the autorenew poll message no longer exists,
     // create a new one at the same id. This can happen if a transfer was requested on a domain
     // where all autorenew poll messages had already been delivered (this would cause the poll
     // message to be deleted), and then subsequently the transfer was canceled, rejected, or deleted
     // (which would cause the poll message to be recreated here).
-    Key<PollMessage.Autorenew> existingAutorenewKey = domain.getAutorenewPollMessage();
+    Key<PollMessage.Autorenew> existingAutorenewKey = domain.getAutorenewPollMessage().getOfyKey();
     PollMessage.Autorenew updatedAutorenewPollMessage =
         autorenewPollMessage.isPresent()
             ? autorenewPollMessage.get().asBuilder().setAutorenewEndTime(newEndTime).build()
@@ -542,7 +542,7 @@ public class DomainFlowUtils {
       ofy().save().entity(updatedAutorenewPollMessage);
     }
 
-    Recurring recurring = ofy().load().key(domain.getAutorenewBillingEvent()).now();
+    Recurring recurring = tm().load(domain.getAutorenewBillingEvent());
     ofy().save().entity(recurring.asBuilder().setRecurrenceEndTime(newEndTime).build());
   }
 

core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java

@@ -33,7 +33,6 @@ import static google.registry.util.DateTimeUtils.leapSafeAddYears;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.ParameterValueRangeErrorException;
 import google.registry.flows.ExtensionManager;
@@ -181,8 +180,8 @@ public final class DomainRenewFlow implements TransactionalFlow {
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(clientId)
             .setRegistrationExpirationTime(newExpirationTime)
-            .setAutorenewBillingEvent(Key.create(newAutorenewEvent))
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
+            .setAutorenewPollMessage(newAutorenewPollMessage.createVKey())
             .addGracePeriod(
                 GracePeriod.forBillingEvent(GracePeriodStatus.RENEW, explicitRenewEvent))
             .build();

core/src/main/java/google/registry/flows/domain/DomainRestoreRequestFlow.java

@@ -26,7 +26,6 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyNotReserved;
 import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNotBlocked;
 import static google.registry.flows.domain.DomainFlowUtils.verifyRegistrarIsActive;
 import static google.registry.model.ResourceTransferUtils.updateForeignKeyIndexDeletionTime;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
@@ -174,8 +173,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
             existingDomain, newExpirationTime, autorenewEvent, autorenewPollMessage, now, clientId);
     updateForeignKeyIndexDeletionTime(newDomain);
     entitiesToSave.add(newDomain, historyEntry, autorenewEvent, autorenewPollMessage);
-    ofy().save().entities(entitiesToSave.build());
-    ofy().delete().key(existingDomain.getDeletePollMessage());
+    tm().saveNewOrUpdateAll(entitiesToSave.build());
+    tm().delete(existingDomain.getDeletePollMessage());
     dnsQueue.addDomainRefreshTask(existingDomain.getDomainName());
     return responseBuilder
         .setExtensions(createResponseExtensions(feesAndCredits, feeUpdate, isExpired))
@@ -232,8 +231,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
         .setStatusValues(null)
         .setGracePeriods(null)
         .setDeletePollMessage(null)
-        .setAutorenewBillingEvent(Key.create(autorenewEvent))
-        .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+        .setAutorenewBillingEvent(autorenewEvent.createVKey())
+        .setAutorenewPollMessage(autorenewPollMessage.createVKey())
         .setLastEppUpdateTime(now)
         .setLastEppUpdateClientId(clientId)
         .build();

core/src/main/java/google/registry/flows/domain/DomainTransferApproveFlow.java

@@ -186,8 +186,8 @@ public final class DomainTransferApproveFlow implements TransactionalFlow {
                     .setTransferredRegistrationExpirationTime(newExpirationTime)
                     .build())
             .setRegistrationExpirationTime(newExpirationTime)
-            .setAutorenewBillingEvent(Key.create(autorenewEvent))
-            .setAutorenewPollMessage(Key.create(gainingClientAutorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
+            .setAutorenewPollMessage(gainingClientAutorenewPollMessage.createVKey())
             // Remove all the old grace periods and add a new one for the transfer.
             .setGracePeriods(
                 billingEvent.isPresent()

core/src/main/java/google/registry/model/billing/BillingEvent.java

@@ -630,9 +630,9 @@ public abstract class BillingEvent extends ImmutableObject
           .setParent(historyEntry);
       // Set the grace period's billing event using the appropriate Cancellation builder method.
       if (gracePeriod.getOneTimeBillingEvent() != null) {
-        builder.setOneTimeEventKey(VKey.from(gracePeriod.getOneTimeBillingEvent()));
+        builder.setOneTimeEventKey(gracePeriod.getOneTimeBillingEvent());
       } else if (gracePeriod.getRecurringBillingEvent() != null) {
-        builder.setRecurringEventKey(VKey.from(gracePeriod.getRecurringBillingEvent()));
+        builder.setRecurringEventKey(gracePeriod.getRecurringBillingEvent());
       }
       return builder.build();
     }

core/src/main/java/google/registry/model/domain/DomainBase.java

@@ -227,7 +227,8 @@ public class DomainBase extends EppResource
    * refer to a {@link PollMessage} timed to when the domain is fully deleted. If the domain is
    * restored, the message should be deleted.
    */
-  @Transient Key<PollMessage.OneTime> deletePollMessage;
+  @Column(name = "deletion_poll_message_id")
+  VKey<PollMessage.OneTime> deletePollMessage;
 
   /**
    * The recurring billing event associated with this domain's autorenewals.
@@ -237,7 +238,8 @@ public class DomainBase extends EppResource
    * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
    * should be created, and this field should be updated to point to the new one.
    */
-  @Transient Key<BillingEvent.Recurring> autorenewBillingEvent;
+  @Column(name = "billing_recurrence_id")
+  VKey<BillingEvent.Recurring> autorenewBillingEvent;
 
   /**
    * The recurring poll message associated with this domain's autorenewals.
@@ -247,7 +249,8 @@ public class DomainBase extends EppResource
    * {@link #registrationExpirationTime} is changed the recurrence should be closed, a new one
    * should be created, and this field should be updated to point to the new one.
    */
-  @Transient Key<PollMessage.Autorenew> autorenewPollMessage;
+  @Column(name = "autorenew_poll_message_id")
+  VKey<PollMessage.Autorenew> autorenewPollMessage;
 
   /** The unexpired grace periods for this domain (some of which may not be active yet). */
   @Transient @ElementCollection Set<GracePeriod> gracePeriods;
@@ -316,15 +319,15 @@ public class DomainBase extends EppResource
     return registrationExpirationTime;
   }
 
-  public Key<PollMessage.OneTime> getDeletePollMessage() {
+  public VKey<PollMessage.OneTime> getDeletePollMessage() {
     return deletePollMessage;
   }
 
-  public Key<BillingEvent.Recurring> getAutorenewBillingEvent() {
+  public VKey<BillingEvent.Recurring> getAutorenewBillingEvent() {
     return autorenewBillingEvent;
   }
 
-  public Key<PollMessage.Autorenew> getAutorenewPollMessage() {
+  public VKey<PollMessage.Autorenew> getAutorenewPollMessage() {
     return autorenewPollMessage;
   }
 
@@ -453,14 +456,8 @@ public class DomainBase extends EppResource
               .setRegistrationExpirationTime(expirationDate)
               // Set the speculatively-written new autorenew events as the domain's autorenew
               // events.
-              .setAutorenewBillingEvent(
-                  transferData.getServerApproveAutorenewEvent() == null
-                      ? null
-                      : transferData.getServerApproveAutorenewEvent().getOfyKey())
-              .setAutorenewPollMessage(
-                  transferData.getServerApproveAutorenewPollMessage() == null
-                      ? null
-                      : transferData.getServerApproveAutorenewPollMessage().getOfyKey());
+              .setAutorenewBillingEvent(transferData.getServerApproveAutorenewEvent())
+              .setAutorenewPollMessage(transferData.getServerApproveAutorenewPollMessage());
       if (transferData.getTransferPeriod().getValue() == 1) {
         // Set the grace period using a key to the prescheduled transfer billing event.  Not using
         // GracePeriod.forBillingEvent() here in order to avoid the actual Datastore fetch.
@@ -471,9 +468,7 @@ public class DomainBase extends EppResource
                     transferExpirationTime.plus(
                         Registry.get(getTld()).getTransferGracePeriodLength()),
                     transferData.getGainingClientId(),
-                    transferData.getServerApproveBillingEvent() == null
-                        ? null
-                        : transferData.getServerApproveBillingEvent().getOfyKey())));
+                    transferData.getServerApproveBillingEvent())));
       } else {
         // There won't be a billing event, so we don't need a grace period
         builder.setGracePeriods(ImmutableSet.of());
@@ -801,19 +796,17 @@ public class DomainBase extends EppResource
       return this;
     }
 
-    public Builder setDeletePollMessage(Key<PollMessage.OneTime> deletePollMessage) {
+    public Builder setDeletePollMessage(VKey<PollMessage.OneTime> deletePollMessage) {
       getInstance().deletePollMessage = deletePollMessage;
       return this;
     }
 
-    public Builder setAutorenewBillingEvent(
-        Key<BillingEvent.Recurring> autorenewBillingEvent) {
+    public Builder setAutorenewBillingEvent(VKey<BillingEvent.Recurring> autorenewBillingEvent) {
       getInstance().autorenewBillingEvent = autorenewBillingEvent;
       return this;
     }
 
-    public Builder setAutorenewPollMessage(
-        Key<PollMessage.Autorenew> autorenewPollMessage) {
+    public Builder setAutorenewPollMessage(VKey<PollMessage.Autorenew> autorenewPollMessage) {
       getInstance().autorenewPollMessage = autorenewPollMessage;
       return this;
     }

core/src/main/java/google/registry/model/domain/GracePeriod.java

@@ -17,12 +17,13 @@ package google.registry.model.domain;
 import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
-import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Ignore;
 import google.registry.model.ImmutableObject;
 import google.registry.model.billing.BillingEvent;
+import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
+import google.registry.persistence.VKey;
 import javax.annotation.Nullable;
 import javax.persistence.Column;
 import javax.persistence.GeneratedValue;
@@ -57,18 +58,18 @@ public class GracePeriod extends ImmutableObject {
 
   /**
    * The one-time billing event corresponding to the action that triggered this grace period, or
-   * null if not applicable.  Not set for autorenew grace periods (which instead use the field
-   * {@code billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
+   * null if not applicable. Not set for autorenew grace periods (which instead use the field {@code
+   * billingEventRecurring}) or for redemption grace periods (since deletes have no cost).
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  Key<BillingEvent.OneTime> billingEventOneTime = null;
+  VKey<BillingEvent.OneTime> billingEventOneTime = null;
 
   /**
    * The recurring billing event corresponding to the action that triggered this grace period, if
    * applicable - i.e. if the action was an autorenew - or null in all other cases.
    */
   // NB: Would @IgnoreSave(IfNull.class), but not allowed for @Embed collections.
-  Key<BillingEvent.Recurring> billingEventRecurring = null;
+  VKey<BillingEvent.Recurring> billingEventRecurring = null;
 
   public GracePeriodStatus getType() {
     return type;
@@ -91,8 +92,7 @@ public class GracePeriod extends ImmutableObject {
    * Returns the one time billing event. The value will only be non-null if the type of this grace
    * period is not AUTO_RENEW.
    */
-
-  public Key<BillingEvent.OneTime> getOneTimeBillingEvent() {
+  public VKey<BillingEvent.OneTime> getOneTimeBillingEvent() {
     return billingEventOneTime;
   }
 
@@ -100,16 +100,16 @@ public class GracePeriod extends ImmutableObject {
    * Returns the recurring billing event. The value will only be non-null if the type of this grace
    * period is AUTO_RENEW.
    */
-  public Key<BillingEvent.Recurring> getRecurringBillingEvent() {
+  public VKey<BillingEvent.Recurring> getRecurringBillingEvent() {
     return billingEventRecurring;
   }
 
   private static GracePeriod createInternal(
-       GracePeriodStatus type,
-       DateTime expirationTime,
-       String clientId,
-       @Nullable Key<BillingEvent.OneTime> billingEventOneTime,
-       @Nullable Key<BillingEvent.Recurring> billingEventRecurring) {
+      GracePeriodStatus type,
+      DateTime expirationTime,
+      String clientId,
+      @Nullable VKey<BillingEvent.OneTime> billingEventOneTime,
+      @Nullable VKey<BillingEvent.Recurring> billingEventRecurring) {
     checkArgument((billingEventOneTime == null) || (billingEventRecurring == null),
         "A grace period can have at most one billing event");
     checkArgument(
@@ -127,15 +127,15 @@ public class GracePeriod extends ImmutableObject {
   /**
    * Creates a GracePeriod for an (optional) OneTime billing event.
    *
-   * <p>Normal callers should always use {@link #forBillingEvent} instead, assuming they do not
-   * need to avoid loading the BillingEvent from Datastore.  This method should typically be
-   * called only from test code to explicitly construct GracePeriods.
+   * <p>Normal callers should always use {@link #forBillingEvent} instead, assuming they do not need
+   * to avoid loading the BillingEvent from Datastore. This method should typically be called only
+   * from test code to explicitly construct GracePeriods.
    */
   public static GracePeriod create(
       GracePeriodStatus type,
       DateTime expirationTime,
       String clientId,
-      @Nullable Key<BillingEvent.OneTime> billingEventOneTime) {
+      @Nullable VKey<BillingEvent.OneTime> billingEventOneTime) {
     return createInternal(type, expirationTime, clientId, billingEventOneTime, null);
   }
 
@@ -144,7 +144,7 @@ public class GracePeriod extends ImmutableObject {
       GracePeriodStatus type,
       DateTime expirationTime,
       String clientId,
-      Key<BillingEvent.Recurring> billingEventRecurring) {
+      VKey<Recurring> billingEventRecurring) {
     checkArgumentNotNull(billingEventRecurring, "billingEventRecurring cannot be null");
     return createInternal(type, expirationTime, clientId, null, billingEventRecurring);
   }
@@ -159,6 +159,6 @@ public class GracePeriod extends ImmutableObject {
   public static GracePeriod forBillingEvent(
       GracePeriodStatus type, BillingEvent.OneTime billingEvent) {
     return create(
-        type, billingEvent.getBillingTime(), billingEvent.getClientId(), Key.create(billingEvent));
+        type, billingEvent.getBillingTime(), billingEvent.getClientId(), billingEvent.createVKey());
   }
 }

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

@@ -16,17 +16,29 @@ package google.registry.tools;
 
 import com.beust.jcommander.Parameters;
 import google.registry.beam.spec11.Spec11Pipeline;
+import google.registry.config.CredentialModule.LocalCredential;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.util.GoogleCredentialsBundle;
+import google.registry.util.Retrier;
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 /** Nomulus command that deploys the {@link Spec11Pipeline} template. */
 @Parameters(commandDescription = "Deploy the Spec11 pipeline to GCS.")
 public class DeploySpec11PipelineCommand implements Command {
 
-  @Inject Spec11Pipeline spec11Pipeline;
+  @Inject @Config("projectId") String projectId;
+  @Inject @Config("beamStagingUrl") String beamStagingUrl;
+  @Inject @Config("spec11TemplateUrl")String spec11TemplateUrl;
+  @Inject @Config("reportingBucketUrl")String reportingBucketUrl;
+  @Inject @LocalCredential GoogleCredentialsBundle googleCredentialsBundle;
+  @Inject Retrier retrier;
+  @Inject @Nullable @Config("sqlAccessInfoFile") String sqlAccessInfoFile;
 
   @Override
   public void run() {
-    spec11Pipeline.deploy();
+    Spec11Pipeline pipeline = new Spec11Pipeline(projectId, beamStagingUrl, spec11TemplateUrl,
+        reportingBucketUrl, googleCredentialsBundle, retrier);
+    pipeline.deploy();
   }
 }
-

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

@@ -117,7 +117,7 @@ public final class DomainLockUtils {
               RegistryLock newLock =
                   RegistryLockDao.save(lock.asBuilder().setLockCompletionTimestamp(now).build());
               setAsRelock(newLock);
-              tm().transact(() -> applyLockStatuses(newLock, now));
+              tm().transact(() -> applyLockStatuses(newLock, now, isAdmin));
               return newLock;
             });
   }
@@ -171,7 +171,7 @@ public final class DomainLockUtils {
                       createLockBuilder(domainName, registrarId, registrarPocId, isAdmin)
                           .setLockCompletionTimestamp(now)
                           .build());
-              tm().transact(() -> applyLockStatuses(newLock, now));
+              tm().transact(() -> applyLockStatuses(newLock, now, isAdmin));
               setAsRelock(newLock);
               return newLock;
             });
@@ -222,18 +222,18 @@ public final class DomainLockUtils {
       String domainName, String registrarId, @Nullable String registrarPocId, boolean isAdmin) {
     DateTime now = jpaTm().getTransactionTime();
     DomainBase domainBase = getDomain(domainName, registrarId, now);
-    verifyDomainNotLocked(domainBase);
+    verifyDomainNotLocked(domainBase, isAdmin);
 
-    // Multiple pending actions are not allowed
+    // Multiple pending actions are not allowed for non-admins
     RegistryLockDao.getMostRecentByRepoId(domainBase.getRepoId())
         .ifPresent(
             previousLock ->
                 checkArgument(
                     previousLock.isLockRequestExpired(now)
-                        || previousLock.getUnlockCompletionTimestamp().isPresent(),
+                        || previousLock.getUnlockCompletionTimestamp().isPresent()
+                        || isAdmin,
                     "A pending or completed lock action already exists for %s",
                     previousLock.getDomainName()));
-
     return new RegistryLock.Builder()
         .setVerificationCode(stringGenerator.createString(VERIFICATION_CODE_LENGTH))
         .setDomainName(domainName)
@@ -250,6 +250,8 @@ public final class DomainLockUtils {
     Optional<RegistryLock> lockOptional =
         RegistryLockDao.getMostRecentVerifiedLockByRepoId(domainBase.getRepoId());
 
+    verifyDomainLocked(domainBase, isAdmin);
+
     RegistryLock.Builder newLockBuilder;
     if (isAdmin) {
       // Admins should always be able to unlock domains in case we get in a bad state
@@ -265,7 +267,6 @@ public final class DomainLockUtils {
                       .setLockCompletionTimestamp(now)
                       .setRegistrarId(registrarId));
     } else {
-      verifyDomainLocked(domainBase);
       RegistryLock lock =
           lockOptional.orElseThrow(
               () ->
@@ -293,16 +294,17 @@ public final class DomainLockUtils {
         .setRegistrarId(registrarId);
   }
 
-  private static void verifyDomainNotLocked(DomainBase domainBase) {
+  private static void verifyDomainNotLocked(DomainBase domainBase, boolean isAdmin) {
     checkArgument(
-        !domainBase.getStatusValues().containsAll(REGISTRY_LOCK_STATUSES),
+        isAdmin || !domainBase.getStatusValues().containsAll(REGISTRY_LOCK_STATUSES),
         "Domain %s is already locked",
         domainBase.getDomainName());
   }
 
-  private static void verifyDomainLocked(DomainBase domainBase) {
+  private static void verifyDomainLocked(DomainBase domainBase, boolean isAdmin) {
     checkArgument(
-        !Sets.intersection(domainBase.getStatusValues(), REGISTRY_LOCK_STATUSES).isEmpty(),
+        isAdmin || !Sets.intersection(domainBase.getStatusValues(), REGISTRY_LOCK_STATUSES)
+            .isEmpty(),
         "Domain %s is already unlocked",
         domainBase.getDomainName());
   }
@@ -311,7 +313,7 @@ public final class DomainLockUtils {
     DomainBase domain =
         loadByForeignKeyCached(DomainBase.class, domainName, now)
             .orElseThrow(
-                () -> new IllegalArgumentException(String.format("Unknown domain %s", domainName)));
+                () -> new IllegalArgumentException("Domain doesn't exist"));
     // The user must have specified either the correct registrar ID or the admin registrar ID
     checkArgument(
         registryAdminRegistrarId.equals(registrarId)
@@ -330,9 +332,9 @@ public final class DomainLockUtils {
                     String.format("Invalid verification code %s", verificationCode)));
   }
 
-  private void applyLockStatuses(RegistryLock lock, DateTime lockTime) {
+  private void applyLockStatuses(RegistryLock lock, DateTime lockTime, boolean isAdmin) {
     DomainBase domain = getDomain(lock.getDomainName(), lock.getRegistrarId(), lockTime);
-    verifyDomainNotLocked(domain);
+    verifyDomainNotLocked(domain, isAdmin);
 
     DomainBase newDomain =
         domain
@@ -345,9 +347,7 @@ public final class DomainLockUtils {
 
   private void removeLockStatuses(RegistryLock lock, boolean isAdmin, DateTime unlockTime) {
     DomainBase domain = getDomain(lock.getDomainName(), lock.getRegistrarId(), unlockTime);
-    if (!isAdmin) {
-      verifyDomainLocked(domain);
-    }
+    verifyDomainLocked(domain, isAdmin);
 
     DomainBase newDomain =
         domain

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

@@ -14,15 +14,7 @@
 
 package google.registry.tools;
 
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
-
 import com.beust.jcommander.Parameters;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Sets;
-import com.google.common.flogger.FluentLogger;
-import google.registry.model.domain.DomainBase;
-import google.registry.model.eppcommon.StatusValue;
-import org.joda.time.DateTime;
 
 /**
  * A command to registry lock domain names.
@@ -32,25 +24,6 @@ import org.joda.time.DateTime;
 @Parameters(separators = " =", commandDescription = "Registry lock a domain via EPP.")
 public class LockDomainCommand extends LockOrUnlockDomainCommand {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
-  @Override
-  protected boolean shouldApplyToDomain(String domain, DateTime now) {
-    DomainBase domainBase =
-        loadByForeignKey(DomainBase.class, domain, now)
-            .orElseThrow(
-                () ->
-                    new IllegalArgumentException(
-                        String.format("Domain '%s' does not exist or is deleted", domain)));
-    ImmutableSet<StatusValue> statusesToAdd =
-        Sets.difference(REGISTRY_LOCK_STATUSES, domainBase.getStatusValues()).immutableCopy();
-    if (statusesToAdd.isEmpty()) {
-      logger.atInfo().log("Domain '%s' is already locked and needs no updates.", domain);
-      return false;
-    }
-    return true;
-  }
-
   @Override
   protected void createAndApplyRequest(String domain) {
     domainLockUtils.administrativelyApplyLock(domain, clientId, null, true);

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

@@ -15,24 +15,28 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Iterables.partition;
 import static google.registry.model.eppcommon.StatusValue.SERVER_DELETE_PROHIBITED;
 import static google.registry.model.eppcommon.StatusValue.SERVER_TRANSFER_PROHIBITED;
 import static google.registry.model.eppcommon.StatusValue.SERVER_UPDATE_PROHIBITED;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.findDuplicates;
 
 import com.beust.jcommander.Parameter;
 import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.model.eppcommon.StatusValue;
 import java.util.List;
 import javax.inject.Inject;
-import org.joda.time.DateTime;
 
-/** Shared base class for commands to registry lock or unlock a domain via EPP. */
+/**
+ * Shared base class for commands to registry lock or unlock a domain via EPP.
+ */
 public abstract class LockOrUnlockDomainCommand extends ConfirmingCommand
     implements CommandWithRemoteApi {
 
@@ -57,7 +61,8 @@ public abstract class LockOrUnlockDomainCommand extends ConfirmingCommand
   @Config("registryAdminClientId")
   String registryAdminClientId;
 
-  @Inject DomainLockUtils domainLockUtils;
+  @Inject
+  DomainLockUtils domainLockUtils;
 
   protected ImmutableSet<String> getDomains() {
     return ImmutableSet.copyOf(mainParameters);
@@ -78,37 +83,34 @@ public abstract class LockOrUnlockDomainCommand extends ConfirmingCommand
   @Override
   protected String execute() {
     ImmutableSet.Builder<String> successfulDomainsBuilder = new ImmutableSet.Builder<>();
-    ImmutableSet.Builder<String> skippedDomainsBuilder = new ImmutableSet.Builder<>();
-    ImmutableSet.Builder<String> failedDomainsBuilder = new ImmutableSet.Builder<>();
+    ImmutableMap.Builder<String, String> failedDomainsToReasons = new ImmutableMap.Builder<>();
     partition(getDomains(), BATCH_SIZE)
         .forEach(
             batch ->
-                tm().transact(
-                        () -> {
-                          for (String domain : batch) {
-                            if (shouldApplyToDomain(domain, tm().getTransactionTime())) {
-                              try {
-                                createAndApplyRequest(domain);
-                              } catch (Throwable t) {
-                                logger.atSevere().withCause(t).log(
-                                    "Error when (un)locking domain %s.", domain);
-                                failedDomainsBuilder.add(domain);
-                              }
-                              successfulDomainsBuilder.add(domain);
-                            } else {
-                              skippedDomainsBuilder.add(domain);
-                            }
-                          }
-                        }));
+                // we require that the jpaTm is the outer transaction in DomainLockUtils
+                jpaTm().transact(() -> tm().transact(
+                    () -> {
+                      for (String domain : batch) {
+                        try {
+                          createAndApplyRequest(domain);
+                        } catch (Throwable t) {
+                          logger.atSevere().withCause(t).log(
+                              "Error when (un)locking domain %s.", domain);
+                          failedDomainsToReasons.put(domain, t.getMessage());
+                          continue;
+                        }
+                        successfulDomainsBuilder.add(domain);
+                      }
+                    })));
     ImmutableSet<String> successfulDomains = successfulDomainsBuilder.build();
-    ImmutableSet<String> skippedDomains = skippedDomainsBuilder.build();
-    ImmutableSet<String> failedDomains = failedDomainsBuilder.build();
+    ImmutableSet<String> failedDomains = failedDomainsToReasons.build().entrySet()
+        .stream()
+        .map(entry -> String.format("%s (%s)", entry.getKey(), entry.getValue()))
+        .collect(toImmutableSet());
     return String.format(
-        "Successfully locked/unlocked domains:\n%s\nSkipped domains:\n%s\nFailed domains:\n%s",
-        successfulDomains, skippedDomains, failedDomains);
+        "Successfully locked/unlocked domains:\n%s\nFailed domains:\n%s",
+        successfulDomains, failedDomains);
   }
 
-  protected abstract boolean shouldApplyToDomain(String domain, DateTime now);
-
   protected abstract void createAndApplyRequest(String domain);
 }

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

@@ -29,7 +29,7 @@ import com.google.appengine.tools.remoteapi.RemoteApiOptions;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
-import google.registry.beam.initsql.BeamJpaModule;
+import google.registry.backup.AppEngineEnvironment;
 import google.registry.config.RegistryConfig;
 import google.registry.model.ofy.ObjectifyService;
 import google.registry.persistence.transaction.TransactionManagerFactory;
@@ -68,9 +68,8 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
 
   @Parameter(
       names = {"--sql_access_info"},
-      description =
-          "Name of a file containing space-separated SQL access info used when deploying "
-              + "Beam pipelines")
+      description = "Name of a file containing space-separated SQL access info used when deploying "
+          + "Beam pipelines")
   private String sqlAccessInfoFile = null;
 
   // Do not make this final - compile-time constant inlining may interfere with JCommander.
@@ -168,7 +167,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
     component =
         DaggerRegistryToolComponent.builder()
             .credentialFilePath(credentialJson)
-            .beamJpaModule(new BeamJpaModule(sqlAccessInfoFile))
+            .sqlAccessInfoFile(sqlAccessInfoFile)
             .build();
 
     // JCommander stores sub-commands as nested JCommander objects containing a list of user objects
@@ -179,7 +178,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
             Iterables.getOnlyElement(jcommander.getCommands().get(parsedCommand).getObjects());
     loggingParams.configureLogging();  // Must be called after parameters are parsed.
 
-    try {
+    try (AppEngineEnvironment env = new AppEngineEnvironment()) {
       runCommand(command);
     } catch (RuntimeException ex) {
       if (Throwables.getRootCause(ex) instanceof LoginRequiredException) {

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

@@ -134,6 +134,9 @@ interface RegistryToolComponent {
     @BindsInstance
     Builder credentialFilePath(@Nullable @Config("credentialFilePath") String credentialFilePath);
 
+    @BindsInstance
+    Builder sqlAccessInfoFile(@Nullable @Config("sqlAccessInfoFile") String sqlAccessInfoFile);
+
     Builder beamJpaModule(BeamJpaModule beamJpaModule);
 
     RegistryToolComponent build();

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

@@ -14,16 +14,8 @@
 
 package google.registry.tools;
 
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
-
 import com.beust.jcommander.Parameters;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Sets;
-import com.google.common.flogger.FluentLogger;
-import google.registry.model.domain.DomainBase;
-import google.registry.model.eppcommon.StatusValue;
 import java.util.Optional;
-import org.joda.time.DateTime;
 
 /**
  * A command to registry unlock domain names.
@@ -33,25 +25,6 @@ import org.joda.time.DateTime;
 @Parameters(separators = " =", commandDescription = "Registry unlock a domain via EPP.")
 public class UnlockDomainCommand extends LockOrUnlockDomainCommand {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
-  @Override
-  protected boolean shouldApplyToDomain(String domain, DateTime now) {
-    DomainBase domainBase =
-        loadByForeignKey(DomainBase.class, domain, now)
-            .orElseThrow(
-                () ->
-                    new IllegalArgumentException(
-                        String.format("Domain '%s' does not exist or is deleted", domain)));
-    ImmutableSet<StatusValue> statusesToRemove =
-        Sets.intersection(domainBase.getStatusValues(), REGISTRY_LOCK_STATUSES).immutableCopy();
-    if (statusesToRemove.isEmpty()) {
-      logger.atInfo().log("Domain '%s' is already unlocked and needs no updates.", domain);
-      return false;
-    }
-    return true;
-  }
-
   @Override
   protected void createAndApplyRequest(String domain) {
     domainLockUtils.administrativelyApplyUnlock(domain, clientId, true, Optional.empty());

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

@@ -31,7 +31,6 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.domain.DomainBase;
 import google.registry.model.domain.Period;
@@ -224,8 +223,8 @@ class UnrenewDomainCommand extends ConfirmingCommand implements CommandWithRemot
             .setRegistrationExpirationTime(newExpirationTime)
             .setLastEppUpdateTime(now)
             .setLastEppUpdateClientId(domain.getCurrentSponsorClientId())
-            .setAutorenewBillingEvent(Key.create(newAutorenewEvent))
-            .setAutorenewPollMessage(Key.create(newAutorenewPollMessage))
+            .setAutorenewBillingEvent(newAutorenewEvent.createVKey())
+            .setAutorenewPollMessage(newAutorenewPollMessage.createVKey())
             .build();
     // In order to do it'll need to write out a new HistoryEntry (likely of type SYNTHETIC), a new
     // autorenew billing event and poll message, and a new one time poll message at the present time

core/src/test/java/google/registry/beam/initsql/BeamJpaExtension.java

@@ -54,7 +54,7 @@ public final class BeamJpaExtension implements BeforeEachCallback, AfterEachCall
     if (beamJpaModule != null) {
       return beamJpaModule;
     }
-    return beamJpaModule = new BeamJpaModule(credentialFile.getAbsolutePath());
+    return beamJpaModule = new BeamJpaModule(credentialFile.getAbsolutePath(), null);
   }
 
   @Override

core/src/test/java/google/registry/beam/initsql/BeamJpaModuleTest.java

@@ -82,7 +82,8 @@ class BeamJpaModuleTest {
         DaggerBeamJpaModule_JpaTransactionManagerComponent.builder()
             .beamJpaModule(
                 new BeamJpaModule(
-                    BackupPaths.getCloudSQLCredentialFilePatterns(environmentName).get(0)))
+                    BackupPaths.getCloudSQLCredentialFilePatterns(environmentName).get(0),
+                    String.format("domain-registry-%s", environmentName)))
             .build()
             .cloudSqlJpaTransactionManager();
     assertThat(

core/src/test/java/google/registry/beam/initsql/DomainBaseUtilTest.java

@@ -63,7 +63,7 @@ public class DomainBaseUtilTest {
   private DomainBase domain;
   private Entity domainEntity;
   private Key<OneTime> oneTimeBillKey;
-  private Key<BillingEvent.Recurring> recurringBillKey;
+  private VKey<BillingEvent.Recurring> recurringBillKey;
   private Key<DomainBase> domainKey;
 
   @RegisterExtension
@@ -102,11 +102,11 @@ public class DomainBaseUtilTest {
     Key<HistoryEntry> historyEntryKey =
         Key.create(persistResource(new HistoryEntry.Builder().setParent(domainKey).build()));
     oneTimeBillKey = Key.create(historyEntryKey, BillingEvent.OneTime.class, 1);
-    recurringBillKey = Key.create(historyEntryKey, BillingEvent.Recurring.class, 2);
-    Key<PollMessage.Autorenew> autorenewPollKey =
-        Key.create(historyEntryKey, PollMessage.Autorenew.class, 3);
-    Key<PollMessage.OneTime> onetimePollKey =
-        Key.create(historyEntryKey, PollMessage.OneTime.class, 1);
+    recurringBillKey = VKey.from(Key.create(historyEntryKey, BillingEvent.Recurring.class, 2));
+    VKey<PollMessage.Autorenew> autorenewPollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.Autorenew.class, 3));
+    VKey<PollMessage.OneTime> onetimePollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.OneTime.class, 1));
     // Set up a new persisted domain entity.
     domain =
         persistResource(
@@ -146,12 +146,10 @@ public class DomainBaseUtilTest {
                             .setPendingTransferExpirationTime(fakeClock.nowUtc())
                             .setServerApproveEntities(
                                 ImmutableSet.of(
-                                    VKey.from(oneTimeBillKey),
-                                    VKey.from(recurringBillKey),
-                                    VKey.from(autorenewPollKey)))
+                                    VKey.from(oneTimeBillKey), recurringBillKey, autorenewPollKey))
                             .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
-                            .setServerApproveAutorenewEvent(VKey.from(recurringBillKey))
-                            .setServerApproveAutorenewPollMessage(VKey.from(autorenewPollKey))
+                            .setServerApproveAutorenewEvent(recurringBillKey)
+                            .setServerApproveAutorenewPollMessage(autorenewPollKey)
                             .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
                             .setTransferStatus(TransferStatus.SERVER_APPROVED)
                             .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))

core/src/test/java/google/registry/beam/initsql/InitSqlPipelineTest.java

@@ -156,12 +156,12 @@ class InitSqlPipelineTest {
       Key<HistoryEntry> historyEntryKey = Key.create(historyEntry);
       Key<BillingEvent.OneTime> oneTimeBillKey =
           Key.create(historyEntryKey, BillingEvent.OneTime.class, 1);
-      Key<BillingEvent.Recurring> recurringBillKey =
-          Key.create(historyEntryKey, BillingEvent.Recurring.class, 2);
-      Key<PollMessage.Autorenew> autorenewPollKey =
-          Key.create(historyEntryKey, PollMessage.Autorenew.class, 3);
-      Key<PollMessage.OneTime> onetimePollKey =
-          Key.create(historyEntryKey, PollMessage.OneTime.class, 1);
+      VKey<BillingEvent.Recurring> recurringBillKey =
+          VKey.from(Key.create(historyEntryKey, BillingEvent.Recurring.class, 2));
+      VKey<PollMessage.Autorenew> autorenewPollKey =
+          VKey.from(Key.create(historyEntryKey, PollMessage.Autorenew.class, 3));
+      VKey<PollMessage.OneTime> onetimePollKey =
+          VKey.from(Key.create(historyEntryKey, PollMessage.OneTime.class, 1));
       domain =
           persistResource(
               new DomainBase.Builder()
@@ -200,12 +200,10 @@ class InitSqlPipelineTest {
                           .setPendingTransferExpirationTime(fakeClock.nowUtc())
                           .setServerApproveEntities(
                               ImmutableSet.of(
-                                  VKey.from(oneTimeBillKey),
-                                  VKey.from(recurringBillKey),
-                                  VKey.from(autorenewPollKey)))
+                                  VKey.from(oneTimeBillKey), recurringBillKey, autorenewPollKey))
                           .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
-                          .setServerApproveAutorenewEvent(VKey.from(recurringBillKey))
-                          .setServerApproveAutorenewPollMessage(VKey.from(autorenewPollKey))
+                          .setServerApproveAutorenewEvent(recurringBillKey)
+                          .setServerApproveAutorenewPollMessage(autorenewPollKey)
                           .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
                           .setTransferStatus(TransferStatus.SERVER_APPROVED)
                           .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))

core/src/test/java/google/registry/flows/FlowTestCase.java

@@ -22,6 +22,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static google.registry.model.eppcommon.EppXmlTransformer.marshal;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.POLL_MESSAGE_ID_STRIPPER;
 import static google.registry.testing.DatastoreHelper.getPollMessages;
 import static google.registry.testing.DatastoreHelper.stripBillingEventId;
@@ -196,12 +197,9 @@ public abstract class FlowTestCase<F extends Flow> {
     assertWithMessage("Billing event is present for grace period: " + gracePeriod)
         .that(gracePeriod.hasBillingEvent())
         .isTrue();
-    return ofy()
-        .load()
-        .key(
+    return tm().load(
             firstNonNull(
-                gracePeriod.getOneTimeBillingEvent(), gracePeriod.getRecurringBillingEvent()))
-        .now();
+                gracePeriod.getOneTimeBillingEvent(), gracePeriod.getRecurringBillingEvent()));
   }
 
   /**

core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -31,6 +31,7 @@ import static google.registry.model.registry.Registry.TldState.GENERAL_AVAILABIL
 import static google.registry.model.registry.Registry.TldState.PREDELEGATION;
 import static google.registry.model.registry.Registry.TldState.QUIET_PERIOD;
 import static google.registry.model.registry.Registry.TldState.START_DATE_SUNRISE;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.pricing.PricingEngineProxy.isDomainPremium;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.assertPollMessagesForResource;
@@ -257,8 +258,7 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
     HistoryEntry historyEntry = getHistoryEntries(domain).get(0);
     assertAboutDomains()
         .that(domain)
-        .hasRegistrationExpirationTime(
-            ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+        .hasRegistrationExpirationTime(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .and()
         .hasOnlyOneHistoryEntryWhich()
         .hasType(HistoryEntry.Type.DOMAIN_CREATE)

core/src/test/java/google/registry/flows/domain/DomainDeleteFlowTest.java

@@ -34,6 +34,7 @@ import static google.registry.model.reporting.DomainTransactionRecord.Transactio
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_DELETE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
@@ -139,8 +140,8 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
         persistResource(
             domain
                 .asBuilder()
-                .setAutorenewBillingEvent(Key.create(autorenewBillingEvent))
-                .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+                .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
+                .setAutorenewPollMessage(autorenewPollMessage.createVKey())
                 .build());
     assertTransactionalFlow(true);
   }
@@ -196,9 +197,9 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
                             GracePeriodStatus.AUTO_RENEW,
                             A_MONTH_AGO.plusDays(45),
                             "TheRegistrar",
-                            Key.create(autorenewBillingEvent))))
-                .setAutorenewBillingEvent(Key.create(autorenewBillingEvent))
-                .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+                            autorenewBillingEvent.createVKey())))
+                .setAutorenewBillingEvent(autorenewBillingEvent.createVKey())
+                .setAutorenewPollMessage(autorenewPollMessage.createVKey())
                 .build());
     assertTransactionalFlow(true);
   }
@@ -436,9 +437,9 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     DateTime deletionTime = domain.getDeletionTime();
     assertThat(getPollMessages("TheRegistrar", deletionTime.minusMinutes(1))).isEmpty();
     assertThat(getPollMessages("TheRegistrar", deletionTime)).hasSize(1);
-    assertThat(domain.getDeletePollMessage())
-        .isEqualTo(Key.create(getOnlyPollMessage("TheRegistrar")));
-    PollMessage.OneTime deletePollMessage = ofy().load().key(domain.getDeletePollMessage()).now();
+    assertThat(domain.getDeletePollMessage().getOfyKey())
+        .isEqualTo(getOnlyPollMessage("TheRegistrar").createVKey().getOfyKey());
+    PollMessage.OneTime deletePollMessage = tm().load(domain.getDeletePollMessage());
     assertThat(deletePollMessage.getMsg()).isEqualTo(expectedMessage);
   }
 
@@ -472,10 +473,7 @@ class DomainDeleteFlowTest extends ResourceFlowTestCase<DomainDeleteFlow, Domain
     // Modify the autorenew poll message so that it has unacked messages in the past. This should
     // prevent it from being deleted when the domain is deleted.
     persistResource(
-        ofy()
-            .load()
-            .key(reloadResourceByForeignKey().getAutorenewPollMessage())
-            .now()
+        tm().load(reloadResourceByForeignKey().getAutorenewPollMessage())
             .asBuilder()
             .setEventTime(A_MONTH_FROM_NOW.minusYears(3))
             .build());

core/src/test/java/google/registry/flows/domain/DomainInfoFlowTest.java

@@ -44,6 +44,7 @@ import google.registry.flows.domain.DomainFlowUtils.CurrencyUnitMismatchExceptio
 import google.registry.flows.domain.DomainFlowUtils.FeeChecksDontSupportPhasesException;
 import google.registry.flows.domain.DomainFlowUtils.RestoresAreAlwaysForOneYearException;
 import google.registry.flows.domain.DomainFlowUtils.TransfersAreAlwaysForOneYearException;
+import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.contact.ContactAuthInfo;
 import google.registry.model.contact.ContactResource;
@@ -59,6 +60,8 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.HostResource;
 import google.registry.model.ofy.RequestCapturingAsyncDatastoreService;
 import google.registry.model.registry.Registry;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineRule;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
@@ -324,16 +327,17 @@ class DomainInfoFlowTest extends ResourceFlowTestCase<DomainInfoFlow, DomainBase
   @Test
   void testSuccess_autoRenewGracePeriod() throws Exception {
     persistTestEntities(false);
+    Key<HistoryEntry> historyEntry =
+        Key.create(domain.createVKey().getOfyKey(), HistoryEntry.class, 67890);
+    VKey<BillingEvent.Recurring> recurringVKey =
+        VKey.from(Key.create(historyEntry, Recurring.class, 12345));
     // Add an AUTO_RENEW grace period to the saved resource.
     persistResource(
         domain
             .asBuilder()
             .addGracePeriod(
                 GracePeriod.createForRecurring(
-                    GracePeriodStatus.AUTO_RENEW,
-                    clock.nowUtc().plusDays(1),
-                    "foo",
-                    Key.create(Recurring.class, 12345)))
+                    GracePeriodStatus.AUTO_RENEW, clock.nowUtc().plusDays(1), "foo", recurringVKey))
             .build());
     doSuccessfulTest("domain_info_response_autorenewperiod.xml", false);
   }

core/src/test/java/google/registry/flows/domain/DomainRenewFlowTest.java

@@ -16,7 +16,7 @@ package google.registry.flows.domain;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.flows.domain.DomainTransferFlowTestCase.persistWithPendingTransfer;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
@@ -38,7 +38,6 @@ import static org.junit.Assert.assertThrows;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.FlowUtils.UnknownCurrencyEppException;
 import google.registry.flows.ResourceFlowTestCase;
@@ -137,8 +136,8 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
             .asBuilder()
             .setRegistrationExpirationTime(expirationTime)
             .setStatusValues(ImmutableSet.copyOf(statusValues))
-            .setAutorenewBillingEvent(Key.create(autorenewEvent))
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .build());
     clock.advanceOneMilli();
   }
@@ -175,7 +174,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRenew =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RENEW);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(newExpiration);
     assertAboutDomains()
         .that(domain)
@@ -470,10 +469,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
     persistDomain();
     // Modify the autorenew poll message so that it has an undelivered message in the past.
     persistResource(
-        ofy()
-            .load()
-            .key(reloadResourceByForeignKey().getAutorenewPollMessage())
-            .now()
+        tm().load(reloadResourceByForeignKey().getAutorenewPollMessage())
             .asBuilder()
             .setEventTime(expirationTime.minusYears(1))
             .build());

core/src/test/java/google/registry/flows/domain/DomainRestoreRequestFlowTest.java

@@ -15,7 +15,7 @@
 package google.registry.flows.domain;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEvents;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.getOnlyHistoryEntryOfType;
@@ -38,7 +38,6 @@ import static org.junit.Assert.assertThrows;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.EppException.UnimplementedExtensionException;
 import google.registry.flows.FlowUtils.UnknownCurrencyEppException;
@@ -116,13 +115,13 @@ class DomainRestoreRequestFlowTest
                     GracePeriodStatus.REDEMPTION, clock.nowUtc().plusDays(1), "foo", null))
             .setStatusValues(ImmutableSet.of(StatusValue.PENDING_DELETE))
             .setDeletePollMessage(
-                Key.create(
-                    persistResource(
+                persistResource(
                         new PollMessage.OneTime.Builder()
                             .setClientId("TheRegistrar")
                             .setEventTime(clock.nowUtc().plusDays(5))
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .build());
     clock.advanceOneMilli();
   }
@@ -146,7 +145,7 @@ class DomainRestoreRequestFlowTest
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRestore =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RESTORE);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expirationTime);
     assertAboutDomains()
         .that(domain)
@@ -214,7 +213,7 @@ class DomainRestoreRequestFlowTest
     DomainBase domain = reloadResourceByForeignKey();
     HistoryEntry historyEntryDomainRestore =
         getOnlyHistoryEntryOfType(domain, HistoryEntry.Type.DOMAIN_RESTORE);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(newExpirationTime);
     assertAboutDomains()
         .that(domain)

core/src/test/java/google/registry/flows/domain/DomainTransferApproveFlowTest.java

@@ -16,12 +16,12 @@ package google.registry.flows.domain;
 
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.NET_ADDS_4_YR;
 import static google.registry.model.reporting.DomainTransactionRecord.TransactionReportField.TRANSFER_SUCCESSFUL;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_APPROVE;
 import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_TRANSFER_REQUEST;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEventsForResource;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.deleteResource;
@@ -42,7 +42,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.Streams;
-import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
 import google.registry.flows.ResourceFlowUtils.BadAuthInfoForResourceException;
 import google.registry.flows.ResourceFlowUtils.ResourceDoesNotExistException;
@@ -192,7 +191,7 @@ class DomainTransferApproveFlowTest
     assertAboutHistoryEntries().that(historyEntryTransferApproved).hasOtherClientId("NewRegistrar");
     assertTransferApproved(domain, originalTransferData);
     assertAboutDomains().that(domain).hasRegistrationExpirationTime(expectedExpirationTime);
-    assertThat(ofy().load().key(domain.getAutorenewBillingEvent()).now().getEventTime())
+    assertThat(tm().load(domain.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expectedExpirationTime);
     // The poll message (in the future) to the losing registrar for implicit ack should be gone.
     assertThat(getPollMessages(domain, "TheRegistrar", clock.nowUtc().plusMonths(1))).isEmpty();
@@ -403,7 +402,7 @@ class DomainTransferApproveFlowTest
             .setEventTime(clock.nowUtc()) // The cancellation happens at the moment of transfer.
             .setBillingTime(
                 oldExpirationTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
-            .setRecurringEventKey(VKey.from(domain.getAutorenewBillingEvent())));
+            .setRecurringEventKey(domain.getAutorenewBillingEvent()));
   }
 
   @Test
@@ -634,7 +633,7 @@ class DomainTransferApproveFlowTest
   @Test
   void testSuccess_superuserExtension_transferPeriodZero_autorenewGraceActive() throws Exception {
     DomainBase domain = reloadResourceByForeignKey();
-    Key<Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to have auto-renewed just before the transfer request, so that it will have an
     // active autorenew grace period spanning the entire transfer window.
     DateTime autorenewTime = clock.nowUtc().minusDays(1);

core/src/test/java/google/registry/flows/domain/DomainTransferRequestFlowTest.java

@@ -318,8 +318,7 @@ class DomainTransferRequestFlowTest
             BillingEvent.class),
         Sets.union(expectedServeApproveBillingEvents, extraBillingEvents));
     // The domain's autorenew billing event should still point to the losing client's event.
-    BillingEvent.Recurring domainAutorenewEvent =
-        ofy().load().key(domain.getAutorenewBillingEvent()).now();
+    BillingEvent.Recurring domainAutorenewEvent = tm().load(domain.getAutorenewBillingEvent());
     assertThat(domainAutorenewEvent.getClientId()).isEqualTo("TheRegistrar");
     assertThat(domainAutorenewEvent.getRecurrenceEndTime()).isEqualTo(implicitTransferTime);
     // The original grace periods should remain untouched.
@@ -446,12 +445,7 @@ class DomainTransferRequestFlowTest
         .hasLastEppUpdateTime(implicitTransferTime)
         .and()
         .hasLastEppUpdateClientId("NewRegistrar");
-    assertThat(
-            ofy()
-                .load()
-                .key(domainAfterAutomaticTransfer.getAutorenewBillingEvent())
-                .now()
-                .getEventTime())
+    assertThat(tm().load(domainAfterAutomaticTransfer.getAutorenewBillingEvent()).getEventTime())
         .isEqualTo(expectedExpirationTime);
     // And after the expected grace time, the grace period should be gone.
     DomainBase afterGracePeriod =
@@ -955,7 +949,7 @@ class DomainTransferRequestFlowTest
   void testSuccess_superuserExtension_zeroPeriod_autorenewGraceActive() throws Exception {
     eppRequestSource = EppRequestSource.TOOL;
     setupDomain("example", "tld");
-    Key<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to have auto-renewed just before the transfer request, so that it will have an
     // active autorenew grace period spanning the entire transfer window.
     DateTime autorenewTime = clock.nowUtc().minusDays(1);
@@ -1110,7 +1104,7 @@ class DomainTransferRequestFlowTest
   @Test
   void testSuccess_autorenewGraceActive_throughoutTransferWindow() throws Exception {
     setupDomain("example", "tld");
-    Key<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to have auto-renewed just before the transfer request, so that it will have an
     // active autorenew grace period spanning the entire transfer window.
     DateTime autorenewTime = clock.nowUtc().minusDays(1);
@@ -1142,13 +1136,13 @@ class DomainTransferRequestFlowTest
             .setEventTime(clock.nowUtc().plus(Registry.get("tld").getAutomaticTransferLength()))
             .setBillingTime(autorenewTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
             // The cancellation should refer to the old autorenew billing event.
-            .setRecurringEventKey(VKey.from(existingAutorenewEvent)));
+            .setRecurringEventKey(existingAutorenewEvent));
   }
 
   @Test
   void testSuccess_autorenewGraceActive_onlyAtAutomaticTransferTime() throws Exception {
     setupDomain("example", "tld");
-    Key<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
+    VKey<BillingEvent.Recurring> existingAutorenewEvent = domain.getAutorenewBillingEvent();
     // Set domain to expire in 1 day, so that it will be in the autorenew grace period by the
     // automatic transfer time, even though it isn't yet.
     DateTime expirationTime = clock.nowUtc().plusDays(1);
@@ -1170,7 +1164,7 @@ class DomainTransferRequestFlowTest
             .setBillingTime(
                 expirationTime.plus(Registry.get("tld").getAutoRenewGracePeriodLength()))
             // The cancellation should refer to the old autorenew billing event.
-            .setRecurringEventKey(VKey.from(existingAutorenewEvent)));
+            .setRecurringEventKey(existingAutorenewEvent));
   }
 
   @Test

core/src/test/java/google/registry/model/billing/BillingEventTest.java

@@ -396,14 +396,15 @@ public class BillingEventTest extends EntityTestCase {
 
   @Test
   void testSuccess_cancellation_forGracePeriod_withRecurring() {
-    BillingEvent.Cancellation newCancellation = BillingEvent.Cancellation.forGracePeriod(
-        GracePeriod.createForRecurring(
-            GracePeriodStatus.AUTO_RENEW,
-            now.plusYears(1).plusDays(45),
-            "a registrar",
-            Key.create(recurring)),
-        historyEntry2,
-        "foo.tld");
+    BillingEvent.Cancellation newCancellation =
+        BillingEvent.Cancellation.forGracePeriod(
+            GracePeriod.createForRecurring(
+                GracePeriodStatus.AUTO_RENEW,
+                now.plusYears(1).plusDays(45),
+                "a registrar",
+                recurring.createVKey()),
+            historyEntry2,
+            "foo.tld");
     // Set ID to be the same to ignore for the purposes of comparison.
     newCancellation = newCancellation.asBuilder().setId(cancellationRecurring.getId()).build();
     assertThat(newCancellation).isEqualTo(cancellationRecurring);

core/src/test/java/google/registry/model/domain/DomainBaseTest.java

@@ -64,19 +64,19 @@ import org.junit.jupiter.api.Test;
 public class DomainBaseTest extends EntityTestCase {
 
   private DomainBase domain;
-  private Key<BillingEvent.OneTime> oneTimeBillKey;
-  private Key<BillingEvent.Recurring> recurringBillKey;
-  private Key<DomainBase> domainKey;
+  private VKey<BillingEvent.OneTime> oneTimeBillKey;
+  private VKey<BillingEvent.Recurring> recurringBillKey;
+  private VKey<DomainBase> domainKey;
 
   @BeforeEach
   void setUp() {
     createTld("com");
-    domainKey = Key.create(null, DomainBase.class, "4-COM");
+    domainKey = VKey.from(Key.create(null, DomainBase.class, "4-COM"));
     VKey<HostResource> hostKey =
         persistResource(
                 new HostResource.Builder()
                     .setHostName("ns1.example.com")
-                    .setSuperordinateDomain(VKey.from(domainKey))
+                    .setSuperordinateDomain(domainKey)
                     .setRepoId("1-COM")
                     .build())
             .createVKey();
@@ -95,13 +95,14 @@ public class DomainBaseTest extends EntityTestCase {
                     .build())
             .createVKey();
     Key<HistoryEntry> historyEntryKey =
-        Key.create(persistResource(new HistoryEntry.Builder().setParent(domainKey).build()));
-    oneTimeBillKey = Key.create(historyEntryKey, BillingEvent.OneTime.class, 1);
-    recurringBillKey = Key.create(historyEntryKey, BillingEvent.Recurring.class, 2);
-    Key<PollMessage.Autorenew> autorenewPollKey =
-        Key.create(historyEntryKey, PollMessage.Autorenew.class, 3);
-    Key<PollMessage.OneTime> onetimePollKey =
-        Key.create(historyEntryKey, PollMessage.OneTime.class, 1);
+        Key.create(
+            persistResource(new HistoryEntry.Builder().setParent(domainKey.getOfyKey()).build()));
+    oneTimeBillKey = VKey.from(Key.create(historyEntryKey, BillingEvent.OneTime.class, 1));
+    recurringBillKey = VKey.from(Key.create(historyEntryKey, BillingEvent.Recurring.class, 2));
+    VKey<PollMessage.Autorenew> autorenewPollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.Autorenew.class, 3));
+    VKey<PollMessage.OneTime> onetimePollKey =
+        VKey.from(Key.create(historyEntryKey, PollMessage.OneTime.class, 1));
     // Set up a new persisted domain entity.
     domain =
         persistResource(
@@ -138,13 +139,10 @@ public class DomainBaseTest extends EntityTestCase {
                             .setLosingClientId("losing")
                             .setPendingTransferExpirationTime(fakeClock.nowUtc())
                             .setServerApproveEntities(
-                                ImmutableSet.of(
-                                    VKey.from(oneTimeBillKey),
-                                    VKey.from(recurringBillKey),
-                                    VKey.from(autorenewPollKey)))
-                            .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
-                            .setServerApproveAutorenewEvent(VKey.from(recurringBillKey))
-                            .setServerApproveAutorenewPollMessage(VKey.from(autorenewPollKey))
+                                ImmutableSet.of(oneTimeBillKey, recurringBillKey, autorenewPollKey))
+                            .setServerApproveBillingEvent(oneTimeBillKey)
+                            .setServerApproveAutorenewEvent(recurringBillKey)
+                            .setServerApproveAutorenewPollMessage(autorenewPollKey)
                             .setTransferRequestTime(fakeClock.nowUtc().plusDays(1))
                             .setTransferStatus(TransferStatus.SERVER_APPROVED)
                             .setTransferRequestTrid(Trid.create("client-trid", "server-trid"))
@@ -327,7 +325,7 @@ public class DomainBaseTest extends EntityTestCase {
   private void assertTransferred(
       DomainBase domain,
       DateTime newExpirationTime,
-      Key<BillingEvent.Recurring> newAutorenewEvent) {
+      VKey<BillingEvent.Recurring> newAutorenewEvent) {
     assertThat(domain.getTransferData().getTransferStatus())
         .isEqualTo(TransferStatus.SERVER_APPROVED);
     assertThat(domain.getCurrentSponsorClientId()).isEqualTo("winner");
@@ -377,8 +375,8 @@ public class DomainBaseTest extends EntityTestCase {
             .build();
     DomainBase afterTransfer = domain.cloneProjectedAtTime(fakeClock.nowUtc().plusDays(1));
     DateTime newExpirationTime = oldExpirationTime.plusYears(1);
-    Key<BillingEvent.Recurring> serverApproveAutorenewEvent =
-        domain.getTransferData().getServerApproveAutorenewEvent().getOfyKey();
+    VKey<BillingEvent.Recurring> serverApproveAutorenewEvent =
+        domain.getTransferData().getServerApproveAutorenewEvent();
     assertTransferred(afterTransfer, newExpirationTime, serverApproveAutorenewEvent);
     assertThat(afterTransfer.getGracePeriods())
         .containsExactly(
@@ -389,7 +387,7 @@ public class DomainBaseTest extends EntityTestCase {
                     .plusDays(1)
                     .plus(Registry.get("com").getTransferGracePeriodLength()),
                 "winner",
-                Key.create(transferBillingEvent)));
+                transferBillingEvent.createVKey()));
     // If we project after the grace period expires all should be the same except the grace period.
     DomainBase afterGracePeriod =
         domain.cloneProjectedAtTime(
@@ -747,8 +745,8 @@ public class DomainBaseTest extends EntityTestCase {
             .setPendingTransferExpirationTime(transferExpirationTime)
             .setTransferStatus(TransferStatus.PENDING)
             .setGainingClientId("TheRegistrar")
-            .setServerApproveAutorenewEvent(VKey.from(recurringBillKey))
-            .setServerApproveBillingEvent(VKey.from(oneTimeBillKey))
+            .setServerApproveAutorenewEvent(recurringBillKey)
+            .setServerApproveBillingEvent(oneTimeBillKey)
             .build();
     domain =
         persistResource(

core/src/test/java/google/registry/model/domain/GracePeriodTest.java

@@ -24,6 +24,7 @@ import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
 import google.registry.testing.AppEngineRule;
 import org.joda.money.CurrencyUnit;
 import org.joda.money.Money;
@@ -46,23 +47,25 @@ public class GracePeriodTest {
 
   @BeforeEach
   void before() {
-    onetime = new BillingEvent.OneTime.Builder()
-      .setEventTime(now)
-      .setBillingTime(now.plusDays(1))
-      .setClientId("TheRegistrar")
-      .setCost(Money.of(CurrencyUnit.USD, 42))
-      .setParent(Key.create(HistoryEntry.class, 12345))
-      .setReason(Reason.CREATE)
-      .setPeriodYears(1)
-      .setTargetId("foo.google")
-      .build();
+    onetime =
+        new BillingEvent.OneTime.Builder()
+            .setEventTime(now)
+            .setBillingTime(now.plusDays(1))
+            .setClientId("TheRegistrar")
+            .setCost(Money.of(CurrencyUnit.USD, 42))
+            .setParent(
+                Key.create(Key.create(DomainBase.class, "domain"), HistoryEntry.class, 12345))
+            .setReason(Reason.CREATE)
+            .setPeriodYears(1)
+            .setTargetId("foo.google")
+            .build();
   }
 
   @Test
   void testSuccess_forBillingEvent() {
     GracePeriod gracePeriod = GracePeriod.forBillingEvent(GracePeriodStatus.ADD, onetime);
     assertThat(gracePeriod.getType()).isEqualTo(GracePeriodStatus.ADD);
-    assertThat(gracePeriod.getOneTimeBillingEvent()).isEqualTo(Key.create(onetime));
+    assertThat(gracePeriod.getOneTimeBillingEvent()).isEqualTo(onetime.createVKey());
     assertThat(gracePeriod.getRecurringBillingEvent()).isNull();
     assertThat(gracePeriod.getClientId()).isEqualTo("TheRegistrar");
     assertThat(gracePeriod.getExpirationTime()).isEqualTo(now.plusDays(1));
@@ -100,7 +103,7 @@ public class GracePeriodTest {
                     GracePeriodStatus.RENEW,
                     now.plusDays(1),
                     "TheRegistrar",
-                    Key.create(Recurring.class, 12345)));
+                    VKey.create(Recurring.class, 12345)));
     assertThat(thrown).hasMessageThat().contains("autorenew");
   }
 }

core/src/test/java/google/registry/rde/DomainBaseToXjcConverterTest.java

@@ -31,7 +31,6 @@ import static org.joda.money.CurrencyUnit.USD;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
@@ -302,8 +301,7 @@ public class DomainBaseToXjcConverterTest {
                     StatusValue.CLIENT_TRANSFER_PROHIBITED,
                     StatusValue.SERVER_UPDATE_PROHIBITED))
             .setAutorenewBillingEvent(
-                Key.create(
-                    persistResource(
+                persistResource(
                         new BillingEvent.Recurring.Builder()
                             .setReason(Reason.RENEW)
                             .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
@@ -312,10 +310,10 @@ public class DomainBaseToXjcConverterTest {
                             .setEventTime(END_OF_TIME)
                             .setRecurrenceEndTime(END_OF_TIME)
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setAutorenewPollMessage(
-                Key.create(
-                    persistResource(
+                persistResource(
                         new PollMessage.Autorenew.Builder()
                             .setTargetId("lol")
                             .setClientId("TheRegistrar")
@@ -323,7 +321,8 @@ public class DomainBaseToXjcConverterTest {
                             .setAutorenewEndTime(END_OF_TIME)
                             .setMsg("Domain was auto-renewed.")
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setTransferData(
                 new DomainTransferData.Builder()
                     .setGainingClientId("gaining")

core/src/test/java/google/registry/rde/RdeFixtures.java

@@ -26,7 +26,6 @@ import static org.joda.money.CurrencyUnit.USD;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
-import com.googlecode.objectify.Key;
 import google.registry.model.billing.BillingEvent;
 import google.registry.model.billing.BillingEvent.Flag;
 import google.registry.model.billing.BillingEvent.Reason;
@@ -147,8 +146,7 @@ final class RdeFixtures {
                     StatusValue.CLIENT_TRANSFER_PROHIBITED,
                     StatusValue.SERVER_UPDATE_PROHIBITED))
             .setAutorenewBillingEvent(
-                Key.create(
-                    persistResource(
+                persistResource(
                         new BillingEvent.Recurring.Builder()
                             .setReason(Reason.RENEW)
                             .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
@@ -157,10 +155,10 @@ final class RdeFixtures {
                             .setEventTime(END_OF_TIME)
                             .setRecurrenceEndTime(END_OF_TIME)
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setAutorenewPollMessage(
-                Key.create(
-                    persistSimpleResource(
+                persistSimpleResource(
                         new PollMessage.Autorenew.Builder()
                             .setTargetId(tld)
                             .setClientId("TheRegistrar")
@@ -168,7 +166,8 @@ final class RdeFixtures {
                             .setAutorenewEndTime(END_OF_TIME)
                             .setMsg("Domain was auto-renewed.")
                             .setParent(historyEntry)
-                            .build())))
+                            .build())
+                    .createVKey())
             .setTransferData(
                 new DomainTransferData.Builder()
                     .setGainingClientId("gaining")

core/src/test/java/google/registry/testing/DatastoreHelper.java

@@ -545,8 +545,8 @@ public class DatastoreHelper {
     return persistResource(
         domain
             .asBuilder()
-            .setAutorenewBillingEvent(Key.create(autorenewEvent))
-            .setAutorenewPollMessage(Key.create(autorenewPollMessage))
+            .setAutorenewBillingEvent(autorenewEvent.createVKey())
+            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
             .build());
   }
 
@@ -588,13 +588,13 @@ public class DatastoreHelper {
                 .build());
     // Modify the existing autorenew event to reflect the pending transfer.
     persistResource(
-        ofy().load().key(domain.getAutorenewBillingEvent()).now().asBuilder()
+        tm().load(domain.getAutorenewBillingEvent())
+            .asBuilder()
             .setRecurrenceEndTime(expirationTime)
             .build());
     // Update the end time of the existing autorenew poll message. We must delete it if it has no
     // events left in it.
-    PollMessage.Autorenew autorenewPollMessage =
-        ofy().load().key(domain.getAutorenewPollMessage()).now();
+    PollMessage.Autorenew autorenewPollMessage = tm().load(domain.getAutorenewPollMessage());
     if (autorenewPollMessage.getEventTime().isBefore(expirationTime)) {
       persistResource(
           autorenewPollMessage.asBuilder()

core/src/test/java/google/registry/tools/DomainLockUtilsTest.java

@@ -49,6 +49,7 @@ import google.registry.testing.AppEngineRule;
 import google.registry.testing.DatastoreHelper;
 import google.registry.testing.DeterministicStringGenerator;
 import google.registry.testing.FakeClock;
+import google.registry.testing.SqlHelper;
 import google.registry.testing.TaskQueueHelper.TaskMatcher;
 import google.registry.testing.UserInfo;
 import google.registry.util.AppEngineServiceUtils;
@@ -275,6 +276,37 @@ public final class DomainLockUtilsTest {
                 standardDays(6).plus(standardSeconds(30))));
   }
 
+  @Test
+  void testSuccess_adminCanLockLockedDomain_withNoSavedLock() {
+    // in the case of inconsistencies / errors, admins should have the ability to override
+    // whatever statuses exist on the domain
+    persistResource(domain.asBuilder().setStatusValues(REGISTRY_LOCK_STATUSES).build());
+    RegistryLock resultLock = domainLockUtils
+        .administrativelyApplyLock(DOMAIN_NAME, "TheRegistrar", POC_ID, true);
+    verifyProperlyLockedDomain(true);
+    assertThat(resultLock.getLockCompletionTimestamp()).isEqualTo(Optional.of(clock.nowUtc()));
+  }
+
+  @Test
+  void testSuccess_adminCanLockUnlockedDomain_withSavedLock() {
+    // in the case of inconsistencies / errors, admins should have the ability to override
+    // what the RegistryLock table says
+    SqlHelper.saveRegistryLock(new RegistryLock.Builder()
+        .setLockCompletionTimestamp(clock.nowUtc())
+        .setDomainName(DOMAIN_NAME)
+        .setVerificationCode("hi")
+        .setRegistrarId("TheRegistrar")
+        .setRepoId(domain.getRepoId())
+        .isSuperuser(false)
+        .setRegistrarPocId(POC_ID)
+        .build());
+    clock.advanceOneMilli();
+    RegistryLock resultLock = domainLockUtils
+        .administrativelyApplyLock(DOMAIN_NAME, "TheRegistrar", POC_ID, true);
+    verifyProperlyLockedDomain(true);
+    assertThat(resultLock.getLockCompletionTimestamp()).isEqualTo(Optional.of(clock.nowUtc()));
+  }
+
   @Test
   void testFailure_createUnlock_alreadyPendingUnlock() {
     RegistryLock lock =
@@ -317,7 +349,7 @@ public final class DomainLockUtilsTest {
                     domainLockUtils.saveNewRegistryLockRequest(
                         "asdf.tld", "TheRegistrar", POC_ID, false)))
         .hasMessageThat()
-        .isEqualTo("Unknown domain asdf.tld");
+        .isEqualTo("Domain doesn't exist");
   }
 
   @Test

core/src/test/java/google/registry/tools/LockDomainCommandTest.java

@@ -98,12 +98,9 @@ class LockDomainCommandTest extends CommandTestCase<LockDomainCommand> {
   }
 
   @Test
-  void testFailure_domainDoesntExist() {
-    IllegalArgumentException e =
-        assertThrows(
-            IllegalArgumentException.class,
-            () -> runCommandForced("--client=NewRegistrar", "missing.tld"));
-    assertThat(e).hasMessageThat().isEqualTo("Domain 'missing.tld' does not exist or is deleted");
+  void testFailure_domainDoesntExist() throws Exception {
+    runCommandForced("--client=NewRegistrar", "missing.tld");
+    assertInStdout("Failed domains:\n[missing.tld (Domain doesn't exist)]");
   }
 
   @Test

core/src/test/java/google/registry/tools/UnlockDomainCommandTest.java

@@ -108,19 +108,16 @@ class UnlockDomainCommandTest extends CommandTestCase<UnlockDomainCommand> {
   }
 
   @Test
-  void testFailure_domainDoesntExist() {
-    IllegalArgumentException e =
-        assertThrows(
-            IllegalArgumentException.class,
-            () -> runCommandForced("--client=TheRegistrar", "missing.tld"));
-    assertThat(e).hasMessageThat().isEqualTo("Domain 'missing.tld' does not exist or is deleted");
+  void testFailure_domainDoesntExist() throws Exception {
+    runCommandForced("--client=NewRegistrar", "missing.tld");
+    assertInStdout("Failed domains:\n[missing.tld (Domain doesn't exist)]");
   }
 
   @Test
-  void testSuccess_alreadyUnlockedDomain_performsNoAction() throws Exception {
+  void testSuccess_alreadyUnlockedDomain_staysUnlocked() throws Exception {
     DomainBase domain = persistActiveDomain("example.tld");
     runCommandForced("--client=TheRegistrar", "example.tld");
-    assertThat(reloadResource(domain)).isEqualTo(domain);
+    assertThat(reloadResource(domain).getStatusValues()).containsNoneIn(REGISTRY_LOCK_STATUSES);
   }
 
   @Test

core/src/test/java/google/registry/tools/UnrenewDomainCommandTest.java

@@ -20,6 +20,7 @@ import static google.registry.model.eppcommon.StatusValue.PENDING_DELETE;
 import static google.registry.model.eppcommon.StatusValue.PENDING_TRANSFER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.reporting.HistoryEntry.Type.SYNTHETIC;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatastoreHelper.assertBillingEventsEqual;
 import static google.registry.testing.DatastoreHelper.assertPollMessagesEqual;
 import static google.registry.testing.DatastoreHelper.createTld;
@@ -119,7 +120,7 @@ public class UnrenewDomainCommandTest extends CommandTestCase<UnrenewDomainComma
     HistoryEntry synthetic = getOnlyHistoryEntryOfType(domain, SYNTHETIC);
 
     assertBillingEventsEqual(
-        ofy().load().key(domain.getAutorenewBillingEvent()).now(),
+        tm().load(domain.getAutorenewBillingEvent()),
         new BillingEvent.Recurring.Builder()
             .setParent(synthetic)
             .setReason(Reason.RENEW)
@@ -148,7 +149,8 @@ public class UnrenewDomainCommandTest extends CommandTestCase<UnrenewDomainComma
                 .build()));
 
     // Check that fields on domain were updated correctly.
-    assertThat(domain.getAutorenewPollMessage().getParent()).isEqualTo(Key.create(synthetic));
+    assertThat(domain.getAutorenewPollMessage().getOfyKey().getParent())
+        .isEqualTo(Key.create(synthetic));
     assertThat(domain.getRegistrationExpirationTime()).isEqualTo(newExpirationTime);
     assertThat(domain.getLastEppUpdateTime()).isEqualTo(unrenewTime);
     assertThat(domain.getLastEppUpdateClientId()).isEqualTo("TheRegistrar");

core/src/test/java/google/registry/ui/server/registrar/RegistryLockPostActionTest.java

@@ -358,7 +358,7 @@ public final class RegistryLockPostActionTest {
                 "domainName", "bad.tld",
                 "isLock", true,
                 "password", "hi"));
-    assertFailureWithMessage(response, "Unknown domain bad.tld");
+    assertFailureWithMessage(response, "Domain doesn't exist");
   }
 
   @Test

core/src/test/resources/google/registry/model/schema.txt

@@ -164,14 +164,14 @@ class google.registry.model.domain.DomainAuthInfo {
 class google.registry.model.domain.DomainBase {
   @Id java.lang.String repoId;
   com.google.common.collect.ImmutableSortedMap<org.joda.time.DateTime, com.googlecode.objectify.Key<google.registry.model.ofy.CommitLogManifest>> revisions;
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$Recurring> autorenewBillingEvent;
-  com.googlecode.objectify.Key<google.registry.model.poll.PollMessage$Autorenew> autorenewPollMessage;
-  com.googlecode.objectify.Key<google.registry.model.poll.PollMessage$OneTime> deletePollMessage;
   google.registry.model.CreateAutoTimestamp creationTime;
   google.registry.model.UpdateAutoTimestamp updateTimestamp;
   google.registry.model.domain.DomainAuthInfo authInfo;
   google.registry.model.domain.launch.LaunchNotice launchNotice;
   google.registry.model.transfer.DomainTransferData transferData;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$Recurring> autorenewBillingEvent;
+  google.registry.persistence.VKey<google.registry.model.poll.PollMessage$Autorenew> autorenewPollMessage;
+  google.registry.persistence.VKey<google.registry.model.poll.PollMessage$OneTime> deletePollMessage;
   java.lang.String creationClientId;
   java.lang.String currentSponsorClientId;
   java.lang.String fullyQualifiedDomainName;
@@ -191,9 +191,9 @@ class google.registry.model.domain.DomainBase {
   org.joda.time.DateTime registrationExpirationTime;
 }
 class google.registry.model.domain.GracePeriod {
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$OneTime> billingEventOneTime;
-  com.googlecode.objectify.Key<google.registry.model.billing.BillingEvent$Recurring> billingEventRecurring;
   google.registry.model.domain.rgp.GracePeriodStatus type;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$OneTime> billingEventOneTime;
+  google.registry.persistence.VKey<google.registry.model.billing.BillingEvent$Recurring> billingEventRecurring;
   java.lang.String clientId;
   org.joda.time.DateTime expirationTime;
 }

db/README.md

@@ -34,9 +34,7 @@ Below are the steps to submit a schema change:
 2.  Run the `devTool generate_sql_schema` command to generate a new version of
     `db-schema.sql.generated`. The full command line to do this is:
 
-    `./gradlew devTool --args="-e localhost generate_sql_schema
-    --start_postgresql -o
-    /path/to/nomulus/db/src/main/resources/sql/schema/db-schema.sql.generated"`
+    `./nom_build generateSqlSchema`
 
 3.  Write an incremental DDL script that changes the existing schema to your new
     one. The generated SQL file from the previous step should help. New create
@@ -49,14 +47,18 @@ Below are the steps to submit a schema change:
     following the existing scripts in that folder. Note the double underscore in
     the naming pattern.
 
-4.  Run the `:db:test` task from the Gradle root project. The SchemaTest will
-    fail because the new schema does not match the golden file.
+4.  Run `./nom_build :nom:generate_golden_file`.  This is a pseudo-task
+    implemented in the `nom_build` script that does the following:
+    -   Runs the `:db:test` task from the Gradle root project. The SchemaTest
+        will fail because the new schema does not match the golden file.
 
-5.  Copy `db/build/resources/test/testcontainer/mount/dump.txt` to the golden
-    file `db/src/main/resources/sql/schema/nomulus.golden.sql`. Diff it against
-    the old version and verify that all changes are expected.
+    -   Copies `db/build/resources/test/testcontainer/mount/dump.txt` to the golden
+        file `db/src/main/resources/sql/schema/nomulus.golden.sql`.
 
-6.  Re-run the `:db:test` task. This time all tests should pass.
+    -   Re-runs the `:db:test` task. This time all tests should pass.
+
+    You'll want to have a look at the diffs in the golden schema to verify
+    that all changes are intentional.
 
 Relevant files (under db/src/main/resources/sql/schema/):
 
@@ -97,7 +99,7 @@ gcloud builds submit --config=release/cloudbuild-schema-deploy.yaml \
     --substitutions=TAG_NAME=${SCHEMA_TAG},_ENV=${SQL_ENV} \
     --project domain-registry-dev
 # Verify by checking Flyway Schema History:
-./gradlew :db:flywayInfo -PdbServer=${SQL_ENV}
+./nom_build :db:flywayInfo --dbServer=${SQL_ENV}
 ```
 
 #### Glass Breaking
@@ -135,9 +137,9 @@ test instance. E.g.,
 
 ```shell
 # Deploy to a local instance at standard port as the super user.
-gradlew :db:flywayMigrate -PdbServer=192.168.9.2 -PdbPassword=domain-registry
+./nom_build :db:flywayMigrate --dbServer=192.168.9.2 --dbPassword=domain-registry
 
 # Full specification of all parameters
-gradlew :db:flywayMigrate -PdbServer=192.168.9.2:5432 -PdbUser=postgres \
-    -PdbPassword=domain-registry
+./nom_build :db:flywayMigrate --dbServer=192.168.9.2:5432 --dbUser=postgres \
+    --dbPassword=domain-registry
 ```

db/src/main/resources/sql/flyway/V40__spec11threatmatch_remove_registrar_foreign_key.sql

@@ -0,0 +1,16 @@
+-- Copyright 2020 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.
+
+alter table if exists "Spec11ThreatMatch"
+  drop constraint "fk_safebrowsing_threat_registrar_id";

db/src/main/resources/sql/flyway/V41__add_columns_to_domain.sql

@@ -0,0 +1,32 @@
+-- Copyright 2020 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.
+
+ALTER TABLE "Domain" ADD COLUMN billing_recurrence_id int8;
+ALTER TABLE "Domain" ADD COLUMN autorenew_poll_message_id int8;
+ALTER TABLE "Domain" ADD COLUMN deletion_poll_message_id int8;
+
+ALTER TABLE IF EXISTS "Domain"
+   ADD CONSTRAINT fk_domain_billing_recurrence_id
+   FOREIGN KEY (billing_recurrence_id)
+   REFERENCES "BillingEvent";
+
+ALTER TABLE IF EXISTS "Domain"
+   ADD CONSTRAINT fk_domain_autorenew_poll_message_id
+   FOREIGN KEY (autorenew_poll_message_id)
+   REFERENCES "PollMessage";
+
+ALTER TABLE IF EXISTS "Domain"
+   ADD CONSTRAINT fk_domain_deletion_poll_message_id
+   FOREIGN KEY (deletion_poll_message_id)
+   REFERENCES "PollMessage";

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

@@ -232,7 +232,10 @@ create sequence history_id_sequence start 1 increment 1;
         admin_contact text,
         auth_info_repo_id text,
         auth_info_value text,
+        billing_recurrence_id int8,
+        autorenew_poll_message_id int8,
         billing_contact text,
+        deletion_poll_message_id int8,
         domain_name text,
         idn_table_name text,
         last_transfer_time timestamptz,
@@ -272,8 +275,8 @@ create sequence history_id_sequence start 1 increment 1;
 
     create table "GracePeriod" (
        id  bigserial not null,
-        billing_event_one_time bytea,
-        billing_event_recurring bytea,
+        billing_event_one_time int8,
+        billing_event_recurring int8,
         registrar_id text,
         expiration_time timestamptz,
         type int4,

db/src/main/resources/sql/schema/nomulus.golden.sql

@@ -398,7 +398,10 @@ CREATE TABLE public."Domain" (
     transfer_pending_expiration_time timestamp with time zone,
     transfer_request_time timestamp with time zone,
     transfer_status text,
-    update_timestamp timestamp with time zone
+    update_timestamp timestamp with time zone,
+    billing_recurrence_id bigint,
+    autorenew_poll_message_id bigint,
+    deletion_poll_message_id bigint
 );
 
 
@@ -1452,6 +1455,14 @@ ALTER TABLE ONLY public."Domain"
     ADD CONSTRAINT fk_domain_admin_contact FOREIGN KEY (admin_contact) REFERENCES public."Contact"(repo_id);
 
 
+--
+-- Name: Domain fk_domain_autorenew_poll_message_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."Domain"
+    ADD CONSTRAINT fk_domain_autorenew_poll_message_id FOREIGN KEY (autorenew_poll_message_id) REFERENCES public."PollMessage"(poll_message_id);
+
+
 --
 -- Name: Domain fk_domain_billing_contact; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1460,6 +1471,22 @@ ALTER TABLE ONLY public."Domain"
     ADD CONSTRAINT fk_domain_billing_contact FOREIGN KEY (billing_contact) REFERENCES public."Contact"(repo_id);
 
 
+--
+-- Name: Domain fk_domain_billing_recurrence_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."Domain"
+    ADD CONSTRAINT fk_domain_billing_recurrence_id FOREIGN KEY (billing_recurrence_id) REFERENCES public."BillingEvent"(billing_event_id);
+
+
+--
+-- Name: Domain fk_domain_deletion_poll_message_id; Type: FK CONSTRAINT; Schema: public; Owner: -
+--
+
+ALTER TABLE ONLY public."Domain"
+    ADD CONSTRAINT fk_domain_deletion_poll_message_id FOREIGN KEY (deletion_poll_message_id) REFERENCES public."PollMessage"(poll_message_id);
+
+
 --
 -- Name: Domain fk_domain_registrant_contact; Type: FK CONSTRAINT; Schema: public; Owner: -
 --
@@ -1596,14 +1623,6 @@ ALTER TABLE ONLY public."Spec11ThreatMatch"
     ADD CONSTRAINT fk_safebrowsing_threat_domain_repo_id FOREIGN KEY (domain_repo_id) REFERENCES public."Domain"(repo_id);
 
 
---
--- Name: Spec11ThreatMatch fk_safebrowsing_threat_registrar_id; Type: FK CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public."Spec11ThreatMatch"
-    ADD CONSTRAINT fk_safebrowsing_threat_registrar_id FOREIGN KEY (registrar_id) REFERENCES public."Registrar"(registrar_id);
-
-
 --
 -- Name: DomainHost fkfmi7bdink53swivs390m2btxg; Type: FK CONSTRAINT; Schema: public; Owner: -
 --