Add SetSqlReplayCheckpoint command for SQL replay (#895 )

* Add SetSqlReplayCheckpoint command for SQL replay We should set this to the same time that we initially populate the SQL database from Datastore.
Convert HostResourceTest to work with Cloud SQL (#905 )
2026-07-10 18:13:02 +00:00 · 2020-12-11 17:41:06 -05:00 · 2020-12-11 13:17:55 -05:00 · 2020-12-10 15:25:25 -05:00 · 2020-12-10 11:29:44 -05:00 · 2020-12-09 15:52:56 -05:00
40 changed files with 1371 additions and 219 deletions
@@ -50,11 +50,21 @@ public class ResaveAllEppResourcesAction implements Runnable {
  @Inject Response response;
  @Inject ResaveAllEppResourcesAction() {}

+  /**
+   * The number of shards to run the map-only mapreduce on.
+   *
+   * <p>This is less than the default of 100 because we only run this action monthly and can afford
+   * it being slower, but we don't want to write out lots of large commit logs in a short period of
+   * time because they make the Cloud SQL migration tougher.
+   */
+  private static final int NUM_SHARDS = 10;
+
  @Override
  public void run() {
    mrRunner
        .setJobName("Re-save all EPP resources")
        .setModuleName("backend")
+        .setDefaultMapShards(NUM_SHARDS)
        .runMapOnly(
            new ResaveAllEppResourcesActionMapper(),
            ImmutableList.of(EppResourceInputs.createKeyInput(EppResource.class)))
@@ -20,6 +20,7 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Throwables.throwIfUnchecked;
 import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
 import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
+import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.JpaRetries.isFailedTxnRetriable;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.setJpaTm;
@@ -68,6 +69,7 @@ import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.ProcessFunction;
+import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
@@ -264,9 +266,9 @@ public final class Transforms {
  }

  /**
-   * Returns a {@link PTransform} that writes a {@link PCollection} of entities to a SQL database.
-   * and outputs an empty {@code PCollection<Void>}. This allows other operations to {@link
-   * org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
+   * Returns a {@link PTransform} that writes a {@link PCollection} of {@link VersionedEntity}s to a
+   * SQL database. and outputs an empty {@code PCollection<Void>}. This allows other operations to
+   * {@link org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
   *
   * <p>Errors are handled according to the pipeline runner's default policy. As part of a one-time
   * job, we will not add features unless proven necessary.
@@ -282,16 +284,53 @@ public final class Transforms {
      int maxWriters,
      int batchSize,
      SerializableSupplier<JpaTransactionManager> jpaSupplier) {
-    return new PTransform<PCollection<VersionedEntity>, PCollection<Void>>() {
+    return writeToSql(
+        transformId,
+        maxWriters,
+        batchSize,
+        jpaSupplier,
+        (e) -> ofy().toPojo(e.getEntity().get()),
+        TypeDescriptor.of(VersionedEntity.class));
+  }
+
+  /**
+   * Returns a {@link PTransform} that writes a {@link PCollection} of entities to a SQL database.
+   * and outputs an empty {@code PCollection<Void>}. This allows other operations to {@link
+   * org.apache.beam.sdk.transforms.Wait wait} for the completion of this transform.
+   *
+   * <p>The converter and type descriptor are generics so that we can convert any type of entity to
+   * an object to be placed in SQL.
+   *
+   * <p>Errors are handled according to the pipeline runner's default policy. As part of a one-time
+   * job, we will not add features unless proven necessary.
+   *
+   * @param transformId a unique ID for an instance of the returned transform
+   * @param maxWriters the max number of concurrent writes to SQL, which also determines the max
+   *     number of connection pools created
+   * @param batchSize the number of entities to write in each operation
+   * @param jpaSupplier supplier of a {@link JpaTransactionManager}
+   * @param jpaConverter the function that converts the input object to a JPA entity
+   * @param objectDescriptor the type descriptor of the input object
+   */
+  public static <T> PTransform<PCollection<T>, PCollection<Void>> writeToSql(
+      String transformId,
+      int maxWriters,
+      int batchSize,
+      SerializableSupplier<JpaTransactionManager> jpaSupplier,
+      SerializableFunction<T, Object> jpaConverter,
+      TypeDescriptor<T> objectDescriptor) {
+    return new PTransform<PCollection<T>, PCollection<Void>>() {
      @Override
-      public PCollection<Void> expand(PCollection<VersionedEntity> input) {
+      public PCollection<Void> expand(PCollection<T> input) {
        return input
            .apply(
                "Shard data for " + transformId,
-                MapElements.into(kvs(integers(), TypeDescriptor.of(VersionedEntity.class)))
+                MapElements.into(kvs(integers(), objectDescriptor))
                    .via(ve -> KV.of(ThreadLocalRandom.current().nextInt(maxWriters), ve)))
            .apply("Batch output by shard " + transformId, GroupIntoBatches.ofSize(batchSize))
-            .apply("Write in batch for " + transformId, ParDo.of(new SqlBatchWriter(jpaSupplier)));
+            .apply(
+                "Write in batch for " + transformId,
+                ParDo.of(new SqlBatchWriter<T>(jpaSupplier, jpaConverter)));
      }
    };
  }
@@ -385,18 +424,22 @@ public final class Transforms {
   * to hold the {@code JpaTransactionManager} instance, we must ensure that JpaTransactionManager
   * is not changed or torn down while being used by some instance.
   */
-  private static class SqlBatchWriter extends DoFn<KV<Integer, Iterable<VersionedEntity>>, Void> {
+  private static class SqlBatchWriter<T> extends DoFn<KV<Integer, Iterable<T>>, Void> {

    private static int instanceCount = 0;
    private static JpaTransactionManager originalJpa;

    private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
+    private final SerializableFunction<T, Object> jpaConverter;

    private transient Ofy ofy;
    private transient SystemSleeper sleeper;

-    SqlBatchWriter(SerializableSupplier<JpaTransactionManager> jpaSupplier) {
+    SqlBatchWriter(
+        SerializableSupplier<JpaTransactionManager> jpaSupplier,
+        SerializableFunction<T, Object> jpaConverter) {
      this.jpaSupplier = jpaSupplier;
+      this.jpaConverter = jpaConverter;
    }

    @Setup
@@ -429,13 +472,11 @@ public final class Transforms {
    }

    @ProcessElement
-    public void processElement(@Element KV<Integer, Iterable<VersionedEntity>> kv) {
+    public void processElement(@Element KV<Integer, Iterable<T>> kv) {
      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
        ImmutableList<Object> ofyEntities =
            Streams.stream(kv.getValue())
-                .map(VersionedEntity::getEntity)
-                .map(Optional::get)
-                .map(ofy::toPojo)
+                .map(this.jpaConverter::apply)
                .collect(ImmutableList.toImmutableList());
        retry(() -> jpaTm().transact(() -> jpaTm().putAll(ofyEntities)));
      }
@@ -20,7 +20,7 @@ import static google.registry.beam.BeamUtils.getQueryFromFile;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableSet;
-import google.registry.backup.AppEngineEnvironment;
+import google.registry.beam.initsql.Transforms;
 import google.registry.beam.initsql.Transforms.SerializableSupplier;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
 import google.registry.config.CredentialModule.LocalCredential;
@@ -43,7 +43,6 @@ import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
-import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -191,34 +190,27 @@ public class Spec11Pipeline implements Serializable {
      PCollection<Subdomain> domains,
      EvaluateSafeBrowsingFn evaluateSafeBrowsingFn,
      ValueProvider<String> dateProvider) {
-
    PCollection<KV<Subdomain, ThreatMatch>> subdomainsSql =
        domains.apply("Run through SafeBrowsing API", ParDo.of(evaluateSafeBrowsingFn));
-    /* Store ThreatMatch objects in SQL. */
+    TypeDescriptor<KV<Subdomain, ThreatMatch>> descriptor =
+        new TypeDescriptor<KV<Subdomain, ThreatMatch>>() {};
    subdomainsSql.apply(
-        ParDo.of(
-            new DoFn<KV<Subdomain, ThreatMatch>, Void>() {
-              @ProcessElement
-              public void processElement(ProcessContext context) {
-                // create the Spec11ThreatMatch from Subdomain and ThreatMatch
-                try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-                  Subdomain subdomain = context.element().getKey();
-                  Spec11ThreatMatch threatMatch =
-                      new Spec11ThreatMatch.Builder()
-                          .setThreatTypes(
-                              ImmutableSet.of(
-                                  ThreatType.valueOf(context.element().getValue().threatType())))
-                          .setCheckDate(
-                              LocalDate.parse(dateProvider.get(), ISODateTimeFormat.date()))
-                          .setDomainName(subdomain.domainName())
-                          .setDomainRepoId(subdomain.domainRepoId())
-                          .setRegistrarId(subdomain.registrarId())
-                          .build();
-                  JpaTransactionManager jpaTransactionManager = jpaSupplierFactory.get();
-                  jpaTransactionManager.transact(() -> jpaTransactionManager.insert(threatMatch));
-                }
-              }
-            }));
+        Transforms.writeToSql(
+            "Spec11ThreatMatch",
+            4,
+            4,
+            jpaSupplierFactory,
+            (kv) -> {
+              Subdomain subdomain = kv.getKey();
+              return new Spec11ThreatMatch.Builder()
+                  .setThreatTypes(ImmutableSet.of(ThreatType.valueOf(kv.getValue().threatType())))
+                  .setCheckDate(LocalDate.parse(dateProvider.get(), ISODateTimeFormat.date()))
+                  .setDomainName(subdomain.domainName())
+                  .setDomainRepoId(subdomain.domainRepoId())
+                  .setRegistrarId(subdomain.registrarId())
+                  .build();
+            },
+            descriptor));

    /* Store ThreatMatch objects in JSON. */
    PCollection<KV<Subdomain, ThreatMatch>> subdomainsJson =
@@ -415,6 +415,14 @@ public final class RegistryConfig {
      return config.cloudSql.instanceConnectionName;
    }

+    @Provides
+    @Config("cloudSqlDbInstanceName")
+    public static String providesCloudSqlDbInstance(RegistryConfigSettings config) {
+      // Format of instanceConnectionName: project-id:region:instance-name
+      int lastColonIndex = config.cloudSql.instanceConnectionName.lastIndexOf(':');
+      return config.cloudSql.instanceConnectionName.substring(lastColonIndex + 1);
+    }
+
    @Provides
    @Config("cloudDnsRootUrl")
    public static Optional<String> getCloudDnsRootUrl(RegistryConfigSettings config) {
@@ -290,7 +290,7 @@ public abstract class BillingEvent extends ImmutableObject
        @javax.persistence.Index(columnList = "allocationToken")
      })
  @AttributeOverride(name = "id", column = @Column(name = "billing_event_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
  public static class OneTime extends BillingEvent implements DatastoreAndSqlEntity {

    /** The billable value. */
@@ -464,7 +464,7 @@ public abstract class BillingEvent extends ImmutableObject
        @javax.persistence.Index(columnList = "recurrence_time_of_year")
      })
  @AttributeOverride(name = "id", column = @Column(name = "billing_recurrence_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
  public static class Recurring extends BillingEvent implements DatastoreAndSqlEntity {

    /**
@@ -559,7 +559,7 @@ public abstract class BillingEvent extends ImmutableObject
        @javax.persistence.Index(columnList = "billingTime")
      })
  @AttributeOverride(name = "id", column = @Column(name = "billing_cancellation_id"))
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
  public static class Cancellation extends BillingEvent implements DatastoreAndSqlEntity {

    /** The billing time of the charge that is being cancelled. */
@@ -680,7 +680,7 @@ public abstract class BillingEvent extends ImmutableObject
  /** An event representing a modification of an existing one-time billing event. */
  @ReportedOn
  @Entity
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
  public static class Modification extends BillingEvent implements DatastoreOnlyEntity {

    /** The change in cost that should be applied to the original billing event. */
@@ -279,7 +279,7 @@ public abstract class PollMessage extends ImmutableObject
  @EntitySubclass(index = false)
  @javax.persistence.Entity
  @DiscriminatorValue("ONE_TIME")
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
  public static class OneTime extends PollMessage {

    // Response data. Objectify cannot persist a base class type, so we must have a separate field
@@ -432,7 +432,7 @@ public abstract class PollMessage extends ImmutableObject
  @EntitySubclass(index = false)
  @javax.persistence.Entity
  @DiscriminatorValue("AUTORENEW")
-  @WithLongVKey
+  @WithLongVKey(compositeKey = true)
  public static class Autorenew extends PollMessage {

    /** The target id of the autorenew event. */
@@ -31,4 +31,12 @@ public @interface WithLongVKey {
   * class name will be "VKeyConverter_" concatenated with the suffix.
   */
  String classNameSuffix() default "";
+
+  /**
+   * Set to true if this is a composite vkey.
+   *
+   * <p>For composite VKeys, we don't attempt to define an objectify key when loading from SQL: the
+   * enclosing class has to take care of that.
+   */
+  boolean compositeKey() default false;
 }
@@ -31,4 +31,12 @@ public @interface WithStringVKey {
   * class name will be "VKeyConverter_" concatenated with the suffix.
   */
  String classNameSuffix() default "";
+
+  /**
+   * Set to true if this is a composite vkey.
+   *
+   * <p>For composite VKeys, we don't attempt to define an objectify key when loading from SQL: the
+   * enclosing class has to take care of that.
+   */
+  boolean compositeKey() default false;
 }
@@ -14,6 +14,7 @@

 package google.registry.persistence.converter;

+import com.googlecode.objectify.Key;
 import google.registry.persistence.VKey;
 import javax.annotation.Nullable;
 import javax.persistence.AttributeConverter;
@@ -29,7 +30,28 @@ public abstract class VKeyConverter<T, C> implements AttributeConverter<VKey<? e
  @Override
  @Nullable
  public VKey<? extends T> convertToEntityAttribute(@Nullable C dbData) {
-    return dbData == null ? null : VKey.createSql(getAttributeClass(), dbData);
+    if (dbData == null) {
+      return null;
+    }
+    Class<? extends T> clazz = getAttributeClass();
+    Key ofyKey = null;
+    if (!hasCompositeOfyKey()) {
+      // If this isn't a composite key, we can create the Ofy key from the SQL key.
+      ofyKey =
+          dbData instanceof String
+              ? Key.create(clazz, (String) dbData)
+              : Key.create(clazz, (Long) dbData);
+      return VKey.create(clazz, dbData, ofyKey);
+    } else {
+      // We don't know how to create the Ofy key and probably don't have everything necessary to do
+      // it anyway, so just create an asymmetric key - the containing object will have to convert it
+      // into a symmetric key.
+      return VKey.createSql(clazz, dbData);
+    }
+  }
+
+  protected boolean hasCompositeOfyKey() {
+    return false;
  }

  /** Returns the class of the attribute. */
@@ -22,6 +22,9 @@ import java.util.Optional;
 /** A Cloud Secret Manager client for Nomulus, bound to a specific GCP project. */
 public interface SecretManagerClient {

+  /** Returns the project name with which this client is associated. */
+  String getProject();
+
  /**
   * Creates a new secret in the Cloud Secret Manager with no data.
   *
@@ -32,6 +35,9 @@ public interface SecretManagerClient {
   */
  void createSecret(String secretId);

+  /** Checks if a secret with the given {@code secretId} already exists. */
+  boolean secretExists(String secretId);
+
  /** Returns all secret IDs in the Cloud Secret Manager. */
  Iterable<String> listSecrets();

@@ -67,6 +73,24 @@ public interface SecretManagerClient {
   */
  String getSecretData(String secretId, Optional<String> version);

+  /**
+   * Enables a secret version.
+   *
+   * @param secretId The ID of the secret
+   * @param version The version of the secret to fetch. If not provided, the {@code latest} version
+   *     will be returned
+   */
+  void enableSecretVersion(String secretId, String version);
+
+  /**
+   * Disables a secret version.
+   *
+   * @param secretId The ID of the secret
+   * @param version The version of the secret to fetch. If not provided, the {@code latest} version
+   *     will be returned
+   */
+  void disableSecretVersion(String secretId, String version);
+
  /**
   * Destroys a secret version.
   *
@@ -29,12 +29,11 @@ import com.google.cloud.secretmanager.v1.SecretName;
 import com.google.cloud.secretmanager.v1.SecretPayload;
 import com.google.cloud.secretmanager.v1.SecretVersion;
 import com.google.cloud.secretmanager.v1.SecretVersionName;
-import com.google.common.collect.Iterables;
+import com.google.common.collect.Streams;
 import com.google.protobuf.ByteString;
 import google.registry.util.Retrier;
 import java.util.Optional;
 import java.util.concurrent.Callable;
-import javax.inject.Inject;

 /** Implements {@link SecretManagerClient} on Google Cloud Platform. */
 public class SecretManagerClientImpl implements SecretManagerClient {
@@ -42,13 +41,17 @@ public class SecretManagerClientImpl implements SecretManagerClient {
  private final SecretManagerServiceClient csmClient;
  private final Retrier retrier;

-  @Inject
  SecretManagerClientImpl(String project, SecretManagerServiceClient csmClient, Retrier retrier) {
    this.project = project;
    this.csmClient = csmClient;
    this.retrier = retrier;
  }

+  @Override
+  public String getProject() {
+    return project;
+  }
+
  @Override
  public void createSecret(String secretId) {
    checkNotNull(secretId, "secretId");
@@ -57,12 +60,25 @@ public class SecretManagerClientImpl implements SecretManagerClient {
        () -> csmClient.createSecret(ProjectName.of(project), secretId, secretSettings));
  }

+  @Override
+  public boolean secretExists(String secretId) {
+    checkNotNull(secretId, "secretId");
+    try {
+      callSecretManager(() -> csmClient.getSecret(SecretName.of(project, secretId)));
+      return true;
+    } catch (NoSuchSecretResourceException e) {
+      return false;
+    }
+  }
+
  @Override
  public Iterable<String> listSecrets() {
    ListSecretsPagedResponse response =
        callSecretManager(() -> csmClient.listSecrets(ProjectName.of(project)));
-    return Iterables.transform(
-        response.iterateAll(), secret -> SecretName.parse(secret.getName()).getSecret());
+    return () ->
+        Streams.stream(response.iterateAll())
+            .map(secret -> SecretName.parse(secret.getName()).getSecret())
+            .iterator();
  }

  @Override
@@ -70,8 +86,10 @@ public class SecretManagerClientImpl implements SecretManagerClient {
    checkNotNull(secretId, "secretId");
    ListSecretVersionsPagedResponse response =
        callSecretManager(() -> csmClient.listSecretVersions(SecretName.of(project, secretId)));
-    return Iterables.transform(
-        response.iterateAll(), SecretManagerClientImpl::toSecretVersionState);
+    return () ->
+        Streams.stream(response.iterateAll())
+            .map(SecretManagerClientImpl::toSecretVersionState)
+            .iterator();
  }

  private static SecretVersionState toSecretVersionState(SecretVersion secretVersion) {
@@ -108,6 +126,22 @@ public class SecretManagerClientImpl implements SecretManagerClient {
                .toStringUtf8());
  }

+  @Override
+  public void enableSecretVersion(String secretId, String version) {
+    checkNotNull(secretId, "secretId");
+    checkNotNull(version, "version");
+    callSecretManager(
+        () -> csmClient.enableSecretVersion(SecretVersionName.of(project, secretId, version)));
+  }
+
+  @Override
+  public void disableSecretVersion(String secretId, String version) {
+    checkNotNull(secretId, "secretId");
+    checkNotNull(version, "version");
+    callSecretManager(
+        () -> csmClient.disableSecretVersion(SecretVersionName.of(project, secretId, version)));
+  }
+
  @Override
  public void destroySecretVersion(String secretId, String version) {
    checkNotNull(secretId, "secretId");
@@ -14,12 +14,12 @@

 package google.registry.privileges.secretmanager;

-import static com.google.common.base.Preconditions.checkNotNull;

 import com.google.cloud.secretmanager.v1.SecretManagerServiceClient;
 import dagger.Component;
 import dagger.Module;
 import dagger.Provides;
+import google.registry.config.RegistryConfig.Config;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.util.Retrier;
 import google.registry.util.UtilsModule;
@@ -28,19 +28,16 @@ import javax.inject.Singleton;

 /** Provides bindings for {@link SecretManagerClient}. */
@Module
-public class SecretManagerModule {
-
-  private final String project;
-
-  public SecretManagerModule(String project) {
-    this.project = checkNotNull(project, "project");
-  }
+public abstract class SecretManagerModule {

  @Provides
  @Singleton
-  SecretManagerClient provideSecretManagerClient(Retrier retrier) {
+  static SecretManagerClient provideSecretManagerClient(
+      @Config("projectId") String project, Retrier retrier) {
    try {
-      return new SecretManagerClientImpl(project, SecretManagerServiceClient.create(), retrier);
+      SecretManagerServiceClient stub = SecretManagerServiceClient.create();
+      Runtime.getRuntime().addShutdownHook(new Thread(stub::close));
+      return new SecretManagerClientImpl(project, stub, retrier);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
@@ -0,0 +1,55 @@
+// 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.privileges.secretmanager;
+
+import static avro.shaded.com.google.common.base.Preconditions.checkState;
+
+import com.google.auto.value.AutoValue;
+import java.util.List;
+
+/**
+ * Contains the login name and password of a Cloud SQL user.
+ *
+ * <p>User must take care not to include the {@link #SEPARATOR} in property values.
+ */
+@AutoValue
+public abstract class SqlCredential {
+
+  public static final Character SEPARATOR = ' ';
+
+  public abstract String login();
+
+  public abstract String password();
+
+  @Override
+  public final String toString() {
+    // Use Object.toString(), which does not show object data.
+    return super.toString();
+  }
+
+  public final String toFormattedString() {
+    return String.format("%s%c%s", login(), SEPARATOR, password());
+  }
+
+  public static SqlCredential fromFormattedString(String sqlCredential) {
+    List<String> items = com.google.common.base.Splitter.on(SEPARATOR).splitToList(sqlCredential);
+    checkState(items.size() == 2, "Invalid SqlCredential string.");
+    return of(items.get(0), items.get(1));
+  }
+
+  public static SqlCredential of(String login, String password) {
+    return new AutoValue_SqlCredential(login, password);
+  }
+}
@@ -0,0 +1,118 @@
+// 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.privileges.secretmanager;
+
+import com.google.cloud.secretmanager.v1.SecretVersionName;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.privileges.secretmanager.SecretManagerClient.NoSuchSecretResourceException;
+import google.registry.privileges.secretmanager.SecretManagerClient.SecretAlreadyExistsException;
+import java.util.Optional;
+import javax.inject.Inject;
+
+/**
+ * Storage of SQL users' login credentials, backed by Cloud Secret Manager.
+ *
+ * <p>A user's credential is stored with one level of indirection using two secret IDs: Each version
+ * of the <em>credential data</em> is stored as follows: its secret ID is determined by {@link
+ * #getCredentialDataSecretId(SqlUser, String dbInstance)}, and the value of each version is a
+ * {@link SqlCredential}, serialized using {@link SqlCredential#toFormattedString}. The 'live'
+ * version of the credential is saved under the 'live pointer' secret explained below.
+ *
+ * <p>The pointer to the 'live' version of the credential data is stored as follows: its secret ID
+ * is determined by {@link #getLiveLabelSecretId(SqlUser, String dbInstance)}; and the value of each
+ * version is a {@link SecretVersionName} in String form, pointing to a version of the credential
+ * data. Only the 'latest' version of this secret should be used. It is guaranteed to be valid.
+ *
+ * <p>The indirection in credential storage makes it easy to handle failures in the credential
+ * change process.
+ */
+public class SqlCredentialStore {
+  private final SecretManagerClient csmClient;
+  private final String dbInstance;
+
+  @Inject
+  SqlCredentialStore(
+      SecretManagerClient csmClient, @Config("cloudSqlDbInstanceName") String dbInstance) {
+    this.csmClient = csmClient;
+    this.dbInstance = dbInstance;
+  }
+
+  public SqlCredential getCredential(SqlUser user) {
+    SecretVersionName credentialName = getLiveCredentialSecretVersion(user);
+    return SqlCredential.fromFormattedString(
+        csmClient.getSecretData(
+            credentialName.getSecret(), Optional.of(credentialName.getSecretVersion())));
+  }
+
+  public void createOrUpdateCredential(SqlUser user, String password) {
+    SecretVersionName dataName = saveCredentialData(user, password);
+    saveLiveLabel(user, dataName);
+  }
+
+  public void deleteCredential(SqlUser user) {
+    try {
+      csmClient.deleteSecret(getCredentialDataSecretId(user, dbInstance));
+    } catch (NoSuchSecretResourceException e) {
+      // ok
+    }
+    try {
+      csmClient.deleteSecret(getLiveLabelSecretId(user, dbInstance));
+    } catch (NoSuchSecretResourceException e) {
+      // ok.
+    }
+  }
+
+  private void createSecretIfAbsent(String secretId) {
+    try {
+      csmClient.createSecret(secretId);
+    } catch (SecretAlreadyExistsException ignore) {
+      // Not a problem.
+    }
+  }
+
+  private SecretVersionName saveCredentialData(SqlUser user, String password) {
+    String credentialDataSecretId = getCredentialDataSecretId(user, dbInstance);
+    createSecretIfAbsent(credentialDataSecretId);
+    String credentialVersion =
+        csmClient.addSecretVersion(
+            credentialDataSecretId,
+            SqlCredential.of(createDatabaseLoginName(user), password).toFormattedString());
+    return SecretVersionName.of(csmClient.getProject(), credentialDataSecretId, credentialVersion);
+  }
+
+  private void saveLiveLabel(SqlUser user, SecretVersionName dataVersionName) {
+    String liveLabelSecretId = getLiveLabelSecretId(user, dbInstance);
+    createSecretIfAbsent(liveLabelSecretId);
+    csmClient.addSecretVersion(liveLabelSecretId, dataVersionName.toString());
+  }
+
+  private SecretVersionName getLiveCredentialSecretVersion(SqlUser user) {
+    return SecretVersionName.parse(
+        csmClient.getSecretData(getLiveLabelSecretId(user, dbInstance), Optional.empty()));
+  }
+
+  private static String getLiveLabelSecretId(SqlUser user, String dbInstance) {
+    return String.format("sql-cred-live-label-%s-%s", user.geUserName(), dbInstance);
+  }
+
+  private static String getCredentialDataSecretId(SqlUser user, String dbInstance) {
+    return String.format("sql-cred-data-%s-%s", user.geUserName(), dbInstance);
+  }
+
+  // WIP: when b/170230882 is complete, login will be versioned.
+  private static String createDatabaseLoginName(SqlUser user) {
+    return user.geUserName();
+  }
+}
@@ -0,0 +1,62 @@
+// 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.privileges.secretmanager;
+
+import com.google.common.base.Ascii;
+
+/**
+ * SQL user information for privilege management purposes.
+ *
+ * <p>A {@link RobotUser} represents a software system accessing the database using its own
+ * credential. Robots are well known and enumerated in {@link RobotId}.
+ */
+public abstract class SqlUser {
+
+  private final UserType type;
+  private final String userName;
+
+  protected SqlUser(UserType type, String userName) {
+    this.type = type;
+    this.userName = userName;
+  }
+
+  public UserType getType() {
+    return type;
+  }
+
+  public String geUserName() {
+    return userName;
+  }
+
+  /** Cloud SQL user types. Please see class javadoc of {@link SqlUser} for more information. */
+  enum UserType {
+    // Work in progress. Human user will be added.
+    ROBOT
+  }
+
+  /** Enumerates the {@link RobotUser RobotUsers} in the system. */
+  public enum RobotId {
+    NOMULUS;
+  }
+
+  /** Information of a RobotUser for privilege management purposes. */
+  // Work in progress. Eventually will be provided based on configuration.
+  public static class RobotUser extends SqlUser {
+
+    public RobotUser(RobotId robot) {
+      super(UserType.ROBOT, Ascii.toLowerCase(robot.name()));
+    }
+  }
+}
@@ -45,10 +45,10 @@ import google.registry.model.domain.DomainBase;
@Parameters(
    separators = " =",
    commandDescription = "Dedupe BillingEvent.OneTime entities with duplicate IDs.")
-public class DedupeOneTimeBillingEventIdsCommand extends DedupeEntityIdsCommand<OneTime> {
+public class DedupeOneTimeBillingEventIdsCommand extends ReadEntityFromKeyPathCommand<OneTime> {

  @Override
-  void dedupe(OneTime entity) {
+  void process(OneTime entity) {
    Key<BillingEvent> key = Key.create(entity);
    Key<DomainBase> domainKey = getGrandParentAsDomain(key);
    DomainBase domain = ofy().load().key(domainKey).now();
@@ -54,11 +54,10 @@ import java.util.Set;
@Parameters(
    separators = " =",
    commandDescription = "Dedupe BillingEvent.Recurring entities with duplicate IDs.")
-public class DedupeRecurringBillingEventIdsCommand
-    extends DedupeEntityIdsCommand<BillingEvent.Recurring> {
+public class DedupeRecurringBillingEventIdsCommand extends ReadEntityFromKeyPathCommand<Recurring> {

  @Override
-  void dedupe(Recurring recurring) {
+  void process(Recurring recurring) {
    // Loads the associated DomainBase and BillingEvent.OneTime entities that
    // may have reference to this BillingEvent.Recurring entity.
    Key<DomainBase> domainKey = getGrandParentAsDomain(Key.create(recurring));
@@ -0,0 +1,73 @@
+// 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.tools;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.base.Ascii;
+import google.registry.privileges.secretmanager.SecretManagerClient.SecretManagerException;
+import google.registry.privileges.secretmanager.SqlCredential;
+import google.registry.privileges.secretmanager.SqlCredentialStore;
+import google.registry.privileges.secretmanager.SqlUser;
+import google.registry.privileges.secretmanager.SqlUser.RobotUser;
+import google.registry.tools.params.PathParameter;
+import java.io.FileOutputStream;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import javax.inject.Inject;
+
+/**
+ * Command to get a Cloud SQL credential in the Secret Manager.
+ *
+ * <p>This command is a short-term tool that will be deprecated by the planned privilege server.
+ */
+@Parameters(separators = " =", commandDescription = "Get the Cloud SQL Credential for a given user")
+public class GetSqlCredentialCommand implements Command {
+
+  @Inject SqlCredentialStore store;
+
+  @Parameter(names = "--user", description = "The Cloud SQL user.", required = true)
+  private String user;
+
+  @Parameter(
+      names = {"-o", "--output"},
+      description = "Name of output file for key data.",
+      validateWith = PathParameter.OutputFile.class)
+  private Path outputPath = null;
+
+  @Inject
+  GetSqlCredentialCommand() {}
+
+  @Override
+  public void run() throws Exception {
+    SqlUser sqlUser = new RobotUser(SqlUser.RobotId.valueOf(Ascii.toUpperCase(user)));
+
+    SqlCredential credential;
+    try {
+      credential = store.getCredential(sqlUser);
+    } catch (SecretManagerException e) {
+      System.out.println(e.getMessage());
+      return;
+    }
+
+    if (outputPath == null) {
+      System.out.printf("[%s]\n", credential.toFormattedString());
+      return;
+    }
+    try (FileOutputStream out = new FileOutputStream(outputPath.toFile())) {
+      out.write(credential.toFormattedString().getBytes(StandardCharsets.UTF_8));
+    }
+  }
+}
@@ -33,8 +33,13 @@ import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.util.List;

-/** Base Command to dedupe entities with duplicate IDs. */
-abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
+/**
+ * Base Command to read entities from Datastore by their key paths retrieved from BigQuery.
+ *
+ * <p>The key path is the value of column __key__.path of the entity's BigQuery table. Its value is
+ * converted from the entity's key.
+ */
+abstract class ReadEntityFromKeyPathCommand<T> extends MutatingCommand {

  @Parameter(
      names = "--key_paths_file",
@@ -48,7 +53,7 @@ abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {

  private StringBuilder changeMessage = new StringBuilder();

-  abstract void dedupe(T entity);
+  abstract void process(T entity);

  @Override
  protected void init() throws Exception {
@@ -69,7 +74,7 @@ abstract class DedupeEntityIdsCommand<T> extends MutatingCommand {
      }
      Class<T> clazz = new TypeInstantiator<T>(getClass()) {}.getExactType();
      if (clazz.isInstance(entity)) {
-        dedupe((T) entity);
+        process((T) entity);
      } else {
        throw new IllegalArgumentException("Unsupported entity key: " + untypedKey);
      }
@@ -78,6 +78,7 @@ public final class RegistryTool {
          .put("get_routing_map", GetRoutingMapCommand.class)
          .put("get_schema", GetSchemaCommand.class)
          .put("get_schema_tree", GetSchemaTreeCommand.class)
+          .put("get_sql_credential", GetSqlCredentialCommand.class)
          .put("get_tld", GetTldCommand.class)
          .put("ghostryde", GhostrydeCommand.class)
          .put("hash_certificate", HashCertificateCommand.class)
@@ -99,12 +100,15 @@ public final class RegistryTool {
          .put("populate_null_registrar_fields", PopulateNullRegistrarFieldsCommand.class)
          .put("registrar_contact", RegistrarContactCommand.class)
          .put("remove_ip_address", RemoveIpAddressCommand.class)
+          .put("remove_registry_one_key", RemoveRegistryOneKeyCommand.class)
          .put("renew_domain", RenewDomainCommand.class)
          .put("resave_entities", ResaveEntitiesCommand.class)
          .put("resave_environment_entities", ResaveEnvironmentEntitiesCommand.class)
          .put("resave_epp_resource", ResaveEppResourceCommand.class)
+          .put("save_sql_credential", SaveSqlCredentialCommand.class)
          .put("send_escrow_report_to_icann", SendEscrowReportToIcannCommand.class)
          .put("set_num_instances", SetNumInstancesCommand.class)
+          .put("set_sql_replay_checkpoint", SetSqlReplayCheckpointCommand.class)
          .put("setup_ote", SetupOteCommand.class)
          .put("uniform_rapid_suspension", UniformRapidSuspensionCommand.class)
          .put("unlock_domain", UnlockDomainCommand.class)
@@ -34,6 +34,7 @@ import google.registry.keyring.kms.KmsModule;
 import google.registry.persistence.PersistenceModule;
 import google.registry.persistence.PersistenceModule.NomulusToolJpaTm;
 import google.registry.persistence.transaction.JpaTransactionManager;
+import google.registry.privileges.secretmanager.SecretManagerModule;
 import google.registry.rde.RdeModule;
 import google.registry.request.Modules.DatastoreServiceModule;
 import google.registry.request.Modules.Jackson2Module;
@@ -74,6 +75,7 @@ import javax.inject.Singleton;
      PersistenceModule.class,
      RdeModule.class,
      RequestFactoryModule.class,
+      SecretManagerModule.class,
      URLFetchServiceModule.class,
      UrlFetchTransportModule.class,
      UserServiceModule.class,
@@ -118,6 +120,8 @@ interface RegistryToolComponent {

  void inject(GetOperationStatusCommand command);

+  void inject(GetSqlCredentialCommand command);
+
  void inject(GhostrydeCommand command);

  void inject(ImportDatastoreCommand command);
@@ -138,10 +142,14 @@ interface RegistryToolComponent {

  void inject(RenewDomainCommand command);

+  void inject(SaveSqlCredentialCommand command);
+
  void inject(SendEscrowReportToIcannCommand command);

  void inject(SetNumInstancesCommand command);

+  void inject(SetSqlReplayCheckpointCommand command);
+
  void inject(SetupOteCommand command);

  void inject(UnlockDomainCommand command);
@@ -0,0 +1,69 @@
+// 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.tools;
+
+import com.beust.jcommander.Parameters;
+import com.googlecode.objectify.Key;
+import google.registry.model.domain.DomainBase;
+import google.registry.persistence.VKey;
+import javax.annotation.Nullable;
+import org.joda.time.DateTime;
+
+/** Command to remove the Registry 1.0 key in {@link DomainBase} entity. */
+@Parameters(separators = " =", commandDescription = "Remove .")
+public class RemoveRegistryOneKeyCommand extends ReadEntityFromKeyPathCommand<DomainBase> {
+
+  @Override
+  void process(DomainBase entity) {
+    // Assert that the DomainBase entity must be deleted before 2017-08-01(most of the problematic
+    // entities were deleted before 2017, though there are still a few entities deleted in 2017-07).
+    // This is because we finished the Registry 2.0 migration in 2017 and should not generate any
+    // Registry 1.0 key after it.
+    if (!entity.getDeletionTime().isBefore(DateTime.parse("2017-08-01T00:00:00Z"))) {
+      throw new IllegalStateException(
+          String.format(
+              "Entity's deletion time %s is not before 2017-08-01T00:00:00Z",
+              entity.getDeletionTime()));
+    }
+    boolean hasChange = false;
+    DomainBase.Builder domainBuilder = entity.asBuilder();
+    // We only found the registry 1.0 key existed in fields autorenewBillingEvent,
+    // autorenewPollMessage and deletePollMessage so we just need to check these fields for each
+    // entity.
+    if (isRegistryOneKey(entity.getAutorenewBillingEvent())) {
+      domainBuilder.setAutorenewBillingEvent(null);
+      hasChange = true;
+    }
+    if (isRegistryOneKey(entity.getAutorenewPollMessage())) {
+      domainBuilder.setAutorenewPollMessage(null);
+      hasChange = true;
+    }
+    if (isRegistryOneKey(entity.getDeletePollMessage())) {
+      domainBuilder.setDeletePollMessage(null);
+      hasChange = true;
+    }
+    if (hasChange) {
+      stageEntityChange(entity, domainBuilder.build());
+    }
+  }
+
+  private static boolean isRegistryOneKey(@Nullable VKey<?> vKey) {
+    if (vKey == null || vKey.getOfyKey() == null || vKey.getOfyKey().getParent() == null) {
+      return false;
+    }
+    Key<?> parentKey = vKey.getOfyKey().getParent();
+    return parentKey.getKind().equals("EntityGroupRoot") && parentKey.getName().equals("per-tld");
+  }
+}
@@ -0,0 +1,69 @@
+// 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.tools;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.base.Ascii;
+import google.registry.privileges.secretmanager.SqlCredentialStore;
+import google.registry.privileges.secretmanager.SqlUser;
+import google.registry.privileges.secretmanager.SqlUser.RobotUser;
+import google.registry.tools.params.PathParameter;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import javax.inject.Inject;
+
+/**
+ * Command to create or update a Cloud SQL credential in the Secret Manager.
+ *
+ * <p>This command is a short-term tool that will be deprecated by the planned privilege server.
+ */
+@Parameters(
+    separators = " =",
+    commandDescription = "Create or update the Cloud SQL Credential for a given user")
+public class SaveSqlCredentialCommand implements Command {
+
+  @Inject SqlCredentialStore store;
+
+  @Parameter(names = "--user", description = "The Cloud SQL user.", required = true)
+  private String user;
+
+  @Parameter(
+      names = {"--input"},
+      description =
+          "Name of input file for the password. If absent, command will prompt for "
+              + "password in console.",
+      validateWith = PathParameter.InputFile.class)
+  private Path inputPath = null;
+
+  @Inject
+  SaveSqlCredentialCommand() {}
+
+  @Override
+  public void run() throws Exception {
+    String password = getPassword();
+    SqlUser sqlUser = new RobotUser(SqlUser.RobotId.valueOf(Ascii.toUpperCase(user)));
+    store.createOrUpdateCredential(sqlUser, password);
+    System.out.printf("\nDone:[%s]\n", password);
+  }
+
+  private String getPassword() throws Exception {
+    if (inputPath != null) {
+      return Files.readAllLines(inputPath, StandardCharsets.UTF_8).get(0);
+    }
+    return System.console().readLine("Please enter the password: ").trim();
+  }
+}
@@ -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.tools;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
+import com.beust.jcommander.Parameter;
+import com.beust.jcommander.Parameters;
+import com.google.common.collect.Iterables;
+import google.registry.schema.replay.SqlReplayCheckpoint;
+import java.util.List;
+import org.joda.time.DateTime;
+
+/** Command to set {@link SqlReplayCheckpoint} to a particular, post-initial-population time. */
+@Parameters(separators = " =", commandDescription = "Set SqlReplayCheckpoint to a particular time")
+public class SetSqlReplayCheckpointCommand extends ConfirmingCommand
+    implements CommandWithRemoteApi {
+
+  @Parameter(description = "Time to which SqlReplayCheckpoint will be set", required = true)
+  List<DateTime> mainParameters;
+
+  @Override
+  protected String prompt() {
+    checkArgument(mainParameters.size() == 1, "Must provide exactly one DateTime to set");
+    return String.format(
+        "Set SqlReplayCheckpoint to %s?", Iterables.getOnlyElement(mainParameters));
+  }
+
+  @Override
+  protected String execute() {
+    DateTime dateTime = Iterables.getOnlyElement(mainParameters);
+    jpaTm().transact(() -> SqlReplayCheckpoint.set(dateTime));
+    return String.format("Set SqlReplayCheckpoint time to %s", dateTime);
+  }
+}
@@ -20,7 +20,6 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.withSettings;

@@ -285,8 +284,7 @@ class Spec11PipelineTest {
            .build();

    verify(mockJpaTm).transact(any(Runnable.class));
-    verify(mockJpaTm).insert(expected);
-    verifyNoMoreInteractions(mockJpaTm);
+    verify(mockJpaTm).putAll(ImmutableList.of(expected));
  }

  /**
@@ -26,6 +26,7 @@ import static google.registry.testing.DatabaseHelper.newDomainBase;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistDeletedDomain;
 import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.testing.DatabaseHelper.persistResources;
 import static google.registry.testing.DomainBaseSubject.assertAboutDomains;
 import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions;
 import static google.registry.testing.HistoryEntrySubject.assertAboutHistoryEntries;
@@ -69,12 +70,17 @@ import google.registry.model.registry.Registry;
 import google.registry.model.reporting.DomainTransactionRecord;
 import google.registry.model.reporting.DomainTransactionRecord.TransactionReportField;
 import google.registry.model.reporting.HistoryEntry;
+import google.registry.testing.ReplayExtension;
 import java.util.Map;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.RegisterExtension;

 /** Unit tests for {@link DomainRenewFlow}. */
 class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBase> {
@@ -96,50 +102,66 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa

  private final DateTime expirationTime = DateTime.parse("2000-04-03T22:00:00.0Z");

+  @Order(value = Order.DEFAULT - 3)
+  @RegisterExtension
+  final SetClockExtension setClockExtension = new SetClockExtension();
+
+  @Order(value = Order.DEFAULT - 2)
+  @RegisterExtension
+  final ReplayExtension replayExtension = new ReplayExtension(clock);
+
  @BeforeEach
  void initDomainTest() {
    createTld("tld");
-    clock.setTo(expirationTime.minusMillis(2));
    setEppInput("domain_renew.xml", ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "5"));
  }

  private void persistDomain(StatusValue... statusValues) throws Exception {
    DomainBase domain = newDomainBase(getUniqueIdFromCommand());
-    HistoryEntry historyEntryDomainCreate =
-        persistResource(
-            new HistoryEntry.Builder()
-                .setParent(domain)
-                .setType(HistoryEntry.Type.DOMAIN_CREATE)
-                .build());
-    BillingEvent.Recurring autorenewEvent =
-        persistResource(
-            new BillingEvent.Recurring.Builder()
-                .setReason(Reason.RENEW)
-                .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
-                .setTargetId(getUniqueIdFromCommand())
-                .setClientId("TheRegistrar")
-                .setEventTime(expirationTime)
-                .setRecurrenceEndTime(END_OF_TIME)
-                .setParent(historyEntryDomainCreate)
-                .build());
-    PollMessage.Autorenew autorenewPollMessage =
-        persistResource(
-            new PollMessage.Autorenew.Builder()
-                .setTargetId(getUniqueIdFromCommand())
-                .setClientId("TheRegistrar")
-                .setEventTime(expirationTime)
-                .setAutorenewEndTime(END_OF_TIME)
-                .setMsg("Domain was auto-renewed.")
-                .setParent(historyEntryDomainCreate)
-                .build());
-    persistResource(
-        domain
-            .asBuilder()
-            .setRegistrationExpirationTime(expirationTime)
-            .setStatusValues(ImmutableSet.copyOf(statusValues))
-            .setAutorenewBillingEvent(autorenewEvent.createVKey())
-            .setAutorenewPollMessage(autorenewPollMessage.createVKey())
-            .build());
+    tm().transact(
+            () -> {
+              try {
+                HistoryEntry historyEntryDomainCreate =
+                    new HistoryEntry.Builder()
+                        .setParent(domain)
+                        .setType(HistoryEntry.Type.DOMAIN_CREATE)
+                        .setModificationTime(clock.nowUtc())
+                        .build();
+                BillingEvent.Recurring autorenewEvent =
+                    new BillingEvent.Recurring.Builder()
+                        .setReason(Reason.RENEW)
+                        .setFlags(ImmutableSet.of(Flag.AUTO_RENEW))
+                        .setTargetId(getUniqueIdFromCommand())
+                        .setClientId("TheRegistrar")
+                        .setEventTime(expirationTime)
+                        .setRecurrenceEndTime(END_OF_TIME)
+                        .setParent(historyEntryDomainCreate)
+                        .build();
+                PollMessage.Autorenew autorenewPollMessage =
+                    new PollMessage.Autorenew.Builder()
+                        .setTargetId(getUniqueIdFromCommand())
+                        .setClientId("TheRegistrar")
+                        .setEventTime(expirationTime)
+                        .setAutorenewEndTime(END_OF_TIME)
+                        .setMsg("Domain was auto-renewed.")
+                        .setParent(historyEntryDomainCreate)
+                        .build();
+                DomainBase newDomain =
+                    domain
+                        .asBuilder()
+                        .setRegistrationExpirationTime(expirationTime)
+                        .setStatusValues(ImmutableSet.copyOf(statusValues))
+                        .setAutorenewBillingEvent(autorenewEvent.createVKey())
+                        .setAutorenewPollMessage(autorenewPollMessage.createVKey())
+                        .build();
+                persistResources(
+                    ImmutableSet.of(
+                        historyEntryDomainCreate, autorenewEvent,
+                        autorenewPollMessage, newDomain));
+              } catch (Exception e) {
+                throw new RuntimeException("Error persisting domain", e);
+              }
+            });
    clock.advanceOneMilli();
  }

@@ -255,6 +277,7 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa

  @Test
  void testSuccess() throws Exception {
+    clock.advanceOneMilli();
    persistDomain();
    doSuccessfulTest(
        "domain_renew_response.xml",
@@ -789,4 +812,21 @@ class DomainRenewFlowTest extends ResourceFlowTestCase<DomainRenewFlow, DomainBa
                TransactionReportField.netRenewsFieldFromYears(5),
                1));
  }
+
+  /**
+   * Local extension so we can set the clock correctly prior to using it to define basic entities in
+   * AppEngineExtension.
+   *
+   * <p>This has to happen first, we'll get timestamp inversions if we try to set the clock after
+   * these objects are created.
+   */
+  class SetClockExtension implements BeforeEachCallback {
+    @Override
+    public void beforeEach(ExtensionContext context) {
+      // we have to go far enough back before the expiration time so that the clock advances we do
+      // after each persist don't move us into a grace period.  The current value is likely beyond
+      // what is necessary, but this doesn't hurt.
+      clock.setTo(expirationTime.minusMillis(20));
+    }
+  }
 }
@@ -90,11 +90,14 @@ import google.registry.model.host.HostResource;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.persistence.VKey;
+import google.registry.testing.ReplayExtension;
 import java.util.Optional;
 import org.joda.money.Money;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;

 /** Unit tests for {@link DomainUpdateFlow}. */
 class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, DomainBase> {
@@ -112,6 +115,10 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
  private ContactResource mak21Contact;
  private ContactResource unusedContact;

+  @Order(value = Order.DEFAULT - 2)
+  @RegisterExtension
+  final ReplayExtension replayExtension = new ReplayExtension(clock);
+
  @BeforeEach
  void initDomainTest() {
    createTld("tld");
@@ -146,6 +153,7 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
    persistResource(
        new HistoryEntry.Builder()
            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setModificationTime(clock.nowUtc())
            .setParent(domain)
            .build());
    clock.advanceOneMilli();
@@ -168,6 +176,7 @@ class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow, Domain
    persistResource(
        new HistoryEntry.Builder()
            .setType(HistoryEntry.Type.DOMAIN_CREATE)
+            .setModificationTime(clock.nowUtc())
            .setParent(domain)
            .build());
    clock.advanceOneMilli();
@@ -84,10 +84,10 @@ public class DomainBaseSqlTest {
    saveRegistrar("registrar1");
    saveRegistrar("registrar2");
    saveRegistrar("registrar3");
-    contactKey = VKey.createSql(ContactResource.class, "contact_id1");
-    contact2Key = VKey.createSql(ContactResource.class, "contact_id2");
+    contactKey = createKey(ContactResource.class, "contact_id1");
+    contact2Key = createKey(ContactResource.class, "contact_id2");

-    host1VKey = VKey.createSql(HostResource.class, "host1");
+    host1VKey = createKey(HostResource.class, "host1");

    domain =
        new DomainBase.Builder()
@@ -696,6 +696,10 @@ public class DomainBaseSqlTest {
    assertThat(domain.getGracePeriods()).isEqualTo(gracePeriods);
  }

+  private <T> VKey<T> createKey(Class<T> clazz, String name) {
+    return VKey.create(clazz, name, Key.create(clazz, name));
+  }
+
  private <T> VKey<T> createLegacyVKey(Class<T> clazz, long id) {
    return VKey.create(
        clazz, id, Key.create(Key.create(EntityGroupRoot.class, "per-tld"), clazz, id));
@@ -14,7 +14,6 @@

 package google.registry.model.history;

-import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableObjects;
 import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
@@ -92,10 +91,7 @@ public class DomainHistoryTest extends EntityTestCase {
              assertDomainHistoriesEqual(fromDatabase, domainHistory);
              assertThat(fromDatabase.getParentVKey()).isEqualTo(domainHistory.getParentVKey());
              assertThat(fromDatabase.getNsHosts())
-                  .containsExactlyElementsIn(
-                      domainHistory.getNsHosts().stream()
-                          .map(key -> VKey.createSql(HostResource.class, key.getSqlKey()))
-                          .collect(toImmutableSet()));
+                  .containsExactlyElementsIn(domainHistory.getNsHosts());
            });
  }

@@ -17,13 +17,17 @@ package google.registry.model.host;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth8.assertThat;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.cloneAndSetAutoTimestamps;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistNewRegistrars;
 import static google.registry.testing.DatabaseHelper.persistResource;
 import static google.registry.testing.HostResourceSubject.assertAboutHosts;
 import static org.junit.jupiter.api.Assertions.assertThrows;

+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
 import google.registry.model.EntityTestCase;
@@ -32,11 +36,14 @@ import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.eppcommon.Trid;
 import google.registry.model.transfer.DomainTransferData;
 import google.registry.model.transfer.TransferStatus;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
+import google.registry.testing.TestOfyOnly;
 import org.joda.time.DateTime;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;

 /** Unit tests for {@link HostResource}. */
+@DualDatabaseTest
 class HostResourceTest extends EntityTestCase {

  private final DateTime day3 = fakeClock.nowUtc();
@@ -49,6 +56,7 @@ class HostResourceTest extends EntityTestCase {
  @BeforeEach
  void setUp() {
    createTld("com");
+    persistNewRegistrars("gaining", "losing", "thisRegistrar", "thatRegistrar");
    // Set up a new persisted registrar entity.
    domain =
        persistResource(
@@ -71,9 +79,9 @@ class HostResourceTest extends EntityTestCase {
                new HostResource.Builder()
                    .setRepoId("DEADBEEF-COM")
                    .setHostName("ns1.example.com")
-                    .setCreationClientId("a registrar")
+                    .setCreationClientId("thisRegistrar")
                    .setLastEppUpdateTime(fakeClock.nowUtc())
-                    .setLastEppUpdateClientId("another registrar")
+                    .setLastEppUpdateClientId("thatRegistrar")
                    .setLastTransferTime(fakeClock.nowUtc())
                    .setInetAddresses(ImmutableSet.of(InetAddresses.forString("127.0.0.1")))
                    .setStatusValues(ImmutableSet.of(StatusValue.OK))
@@ -81,13 +89,22 @@ class HostResourceTest extends EntityTestCase {
                    .build()));
  }

-  @Test
+  @TestOfyAndSql
  void testPersistence() {
+    HostResource newHost = host.asBuilder().setRepoId("NEWHOST").build();
+    tm().transact(() -> tm().insert(newHost));
+    assertThat(ImmutableList.of(tm().transact(() -> tm().load(newHost.createVKey()))))
+        .comparingElementsUsing(immutableObjectCorrespondence("revisions"))
+        .containsExactly(newHost);
+  }
+
+  @TestOfyOnly
+  void testLoadingByForeignKey() {
    assertThat(loadByForeignKey(HostResource.class, host.getForeignKey(), fakeClock.nowUtc()))
        .hasValue(host);
  }

-  @Test
+  @TestOfyOnly
  void testIndexing() throws Exception {
    // Clone it and save it before running the indexing test so that its transferData fields are
    // populated from the superordinate domain.
@@ -100,7 +117,7 @@ class HostResourceTest extends EntityTestCase {
        "currentSponsorClientId");
  }

-  @Test
+  @TestOfyAndSql
  void testEmptyStringsBecomeNull() {
    assertThat(
            new HostResource.Builder()
@@ -122,7 +139,7 @@ class HostResourceTest extends EntityTestCase {
        .isNotNull();
  }

-  @Test
+  @TestOfyAndSql
  void testEmptySetsBecomeNull() {
    assertThat(new HostResource.Builder().setInetAddresses(null).build().inetAddresses).isNull();
    assertThat(new HostResource.Builder().setInetAddresses(ImmutableSet.of()).build().inetAddresses)
@@ -135,7 +152,7 @@ class HostResourceTest extends EntityTestCase {
        .isNotNull();
  }

-  @Test
+  @TestOfyAndSql
  void testImplicitStatusValues() {
    // OK is implicit if there's no other statuses.
    assertAboutHosts()
@@ -157,13 +174,13 @@ class HostResourceTest extends EntityTestCase {
        .hasExactlyStatusValues(StatusValue.CLIENT_HOLD);
  }

-  @Test
+  @TestOfyAndSql
  void testToHydratedString_notCircular() {
    // If there are circular references, this will overflow the stack.
    host.toHydratedString();
  }

-  @Test
+  @TestOfyAndSql
  void testFailure_uppercaseHostName() {
    IllegalArgumentException thrown =
        assertThrows(
@@ -173,7 +190,7 @@ class HostResourceTest extends EntityTestCase {
        .contains("Host name must be in puny-coded, lower-case form");
  }

-  @Test
+  @TestOfyAndSql
  void testFailure_utf8HostName() {
    IllegalArgumentException thrown =
        assertThrows(
@@ -183,14 +200,14 @@ class HostResourceTest extends EntityTestCase {
        .contains("Host name must be in puny-coded, lower-case form");
  }

-  @Test
+  @TestOfyAndSql
  void testComputeLastTransferTime_hostNeverSwitchedDomains_domainWasNeverTransferred() {
    domain = domain.asBuilder().setLastTransferTime(null).build();
    host = host.asBuilder().setLastTransferTime(null).setLastSuperordinateChange(null).build();
    assertThat(host.computeLastTransferTime(domain)).isNull();
  }

-  @Test
+  @TestOfyAndSql
  void testComputeLastTransferTime_hostNeverSwitchedDomains_domainWasTransferred() {
    // Host was created on Day 1.
    // Domain was transferred on Day 2.
@@ -205,7 +222,7 @@ class HostResourceTest extends EntityTestCase {
    assertThat(host.computeLastTransferTime(domain)).isEqualTo(day2);
  }

-  @Test
+  @TestOfyAndSql
  void testComputeLastTransferTime_hostCreatedAfterDomainWasTransferred() {
    // Domain was transferred on Day 1.
    // Host was created subordinate to domain on Day 2.
@@ -217,9 +234,9 @@ class HostResourceTest extends EntityTestCase {
                    .setCreationTime(day2)
                    .setRepoId("DEADBEEF-COM")
                    .setHostName("ns1.example.com")
-                    .setCreationClientId("a registrar")
+                    .setCreationClientId("thisRegistrar")
                    .setLastEppUpdateTime(fakeClock.nowUtc())
-                    .setLastEppUpdateClientId("another registrar")
+                    .setLastEppUpdateClientId("thatRegistrar")
                    .setInetAddresses(ImmutableSet.of(InetAddresses.forString("127.0.0.1")))
                    .setStatusValues(ImmutableSet.of(StatusValue.OK))
                    .setSuperordinateDomain(domain.createVKey())
@@ -227,7 +244,7 @@ class HostResourceTest extends EntityTestCase {
    assertThat(host.computeLastTransferTime(domain)).isNull();
  }

-  @Test
+  @TestOfyAndSql
  void testComputeLastTransferTime_hostWasTransferred_domainWasNeverTransferred() {
    // Host was transferred on Day 1.
    // Host was made subordinate to domain on Day 2.
@@ -237,7 +254,7 @@ class HostResourceTest extends EntityTestCase {
    assertThat(host.computeLastTransferTime(domain)).isEqualTo(day1);
  }

-  @Test
+  @TestOfyAndSql
  void testComputeLastTransferTime_domainWasTransferredBeforeHostBecameSubordinate() {
    // Host was transferred on Day 1.
    // Domain was transferred on Day 2.
@@ -247,7 +264,7 @@ class HostResourceTest extends EntityTestCase {
    assertThat(host.computeLastTransferTime(domain)).isEqualTo(day1);
  }

-  @Test
+  @TestOfyAndSql
  void testComputeLastTransferTime_domainWasTransferredAfterHostBecameSubordinate() {
    // Host was transferred on Day 1.
    // Host was made subordinate to domain on Day 2.
@@ -19,8 +19,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.

 import google.registry.persistence.VKey;
 import google.registry.persistence.WithLongVKey;
-import google.registry.persistence.transaction.JpaTestRules;
-import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension;
+import google.registry.testing.AppEngineExtension;
 import javax.persistence.Entity;
 import javax.persistence.Id;
 import org.junit.jupiter.api.Test;
@@ -30,33 +29,56 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 public class LongVKeyConverterTest {

  @RegisterExtension
-  public final JpaUnitTestExtension jpaExtension =
-      new JpaTestRules.Builder()
-          .withEntityClass(TestEntity.class, VKeyConverter_LongType.class)
-          .buildUnitTestRule();
+  public final AppEngineExtension appEngineExtension =
+      new AppEngineExtension.Builder()
+          .withDatastoreAndCloudSql()
+          .withoutCannedData()
+          .withJpaUnitTestEntities(
+              TestLongEntity.class,
+              VKeyConverter_LongType.class,
+              VKeyConverter_CompositeLongType.class)
+          .withOfyTestEntities(TestLongEntity.class, CompositeKeyTestLongEntity.class)
+          .build();

  @Test
  void testRoundTrip() {
-    TestEntity original = new TestEntity(VKey.createSql(TestEntity.class, 10L));
+    TestLongEntity original =
+        new TestLongEntity(
+            VKey.createSql(TestLongEntity.class, 10L),
+            VKey.createSql(CompositeKeyTestLongEntity.class, 20L));
    jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));

-    TestEntity retrieved =
-        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "id"));
+    TestLongEntity retrieved =
+        jpaTm().transact(() -> jpaTm().getEntityManager().find(TestLongEntity.class, "id"));
    assertThat(retrieved.number.getSqlKey()).isEqualTo(10L);
+    assertThat(retrieved.number.getOfyKey().getId()).isEqualTo(10L);
+
+    assertThat(retrieved.composite.getSqlKey()).isEqualTo(20L);
+    assertThat(retrieved.composite.maybeGetOfyKey().isPresent()).isFalse();
  }

-  @Entity(name = "TestEntity")
+  @Entity(name = "TestLongEntity")
+  @com.googlecode.objectify.annotation.Entity
  @WithLongVKey(classNameSuffix = "LongType")
-  static class TestEntity {
-    @Id String id = "id";
+  static class TestLongEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";

-    VKey<TestEntity> number;
+    VKey<TestLongEntity> number;
+    VKey<CompositeKeyTestLongEntity> composite;

-    TestEntity(VKey<TestEntity> number) {
+    TestLongEntity(VKey<TestLongEntity> number, VKey<CompositeKeyTestLongEntity> composite) {
      this.number = number;
+      this.composite = composite;
    }

    /** Default constructor, needed for hibernate. */
-    public TestEntity() {}
+    public TestLongEntity() {}
+  }
+
+  @Entity(name = "CompositeKeyTestLongEntity")
+  @com.googlecode.objectify.annotation.Entity
+  @WithLongVKey(classNameSuffix = "CompositeLongType", compositeKey = true)
+  static class CompositeKeyTestLongEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
  }
 }
@@ -19,8 +19,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.

 import google.registry.persistence.VKey;
 import google.registry.persistence.WithStringVKey;
-import google.registry.persistence.transaction.JpaTestRules;
-import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension;
+import google.registry.testing.AppEngineExtension;
 import javax.persistence.Entity;
 import javax.persistence.Id;
 import org.junit.jupiter.api.Test;
@@ -30,36 +29,61 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 public class StringVKeyConverterTest {

  @RegisterExtension
-  public final JpaUnitTestExtension jpaExtension =
-      new JpaTestRules.Builder()
-          .withEntityClass(TestEntity.class, VKeyConverter_StringType.class)
-          .buildUnitTestRule();
+  public final AppEngineExtension appEngineExtension =
+      new AppEngineExtension.Builder()
+          .withDatastoreAndCloudSql()
+          .withoutCannedData()
+          .withJpaUnitTestEntities(
+              TestStringEntity.class,
+              VKeyConverter_StringType.class,
+              VKeyConverter_CompositeStringType.class)
+          .withOfyTestEntities(TestStringEntity.class, CompositeKeyTestStringEntity.class)
+          .build();

  @Test
  void testRoundTrip() {
-    TestEntity original =
-        new TestEntity("TheRealSpartacus", VKey.createSql(TestEntity.class, "ImSpartacus!"));
+    TestStringEntity original =
+        new TestStringEntity(
+            "TheRealSpartacus",
+            VKey.createSql(TestStringEntity.class, "ImSpartacus!"),
+            VKey.createSql(CompositeKeyTestStringEntity.class, "NoImSpartacus!"));
    jpaTm().transact(() -> jpaTm().getEntityManager().persist(original));

-    TestEntity retrieved =
+    TestStringEntity retrieved =
        jpaTm()
-            .transact(() -> jpaTm().getEntityManager().find(TestEntity.class, "TheRealSpartacus"));
+            .transact(
+                () -> jpaTm().getEntityManager().find(TestStringEntity.class, "TheRealSpartacus"));
    assertThat(retrieved.other.getSqlKey()).isEqualTo("ImSpartacus!");
+    assertThat(retrieved.other.getOfyKey().getName()).isEqualTo("ImSpartacus!");
+
+    assertThat(retrieved.composite.getSqlKey()).isEqualTo("NoImSpartacus!");
+    assertThat(retrieved.composite.maybeGetOfyKey().isPresent()).isFalse();
  }

-  @Entity(name = "TestEntity")
+  @Entity(name = "TestStringEntity")
+  @com.googlecode.objectify.annotation.Entity
  @WithStringVKey(classNameSuffix = "StringType")
-  static class TestEntity {
-    @Id String id;
+  static class TestStringEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id;

-    VKey<TestEntity> other;
+    VKey<TestStringEntity> other;
+    VKey<CompositeKeyTestStringEntity> composite;

-    TestEntity(String id, VKey<TestEntity> other) {
+    TestStringEntity(
+        String id, VKey<TestStringEntity> other, VKey<CompositeKeyTestStringEntity> composite) {
      this.id = id;
      this.other = other;
+      this.composite = composite;
    }

    /** Default constructor, needed for hibernate. */
-    public TestEntity() {}
+    public TestStringEntity() {}
+  }
+
+  @Entity(name = "CompositeKeyTestStringEntity")
+  @com.googlecode.objectify.annotation.Entity
+  @WithStringVKey(classNameSuffix = "CompositeStringType", compositeKey = true)
+  static class CompositeKeyTestStringEntity {
+    @com.googlecode.objectify.annotation.Id @Id String id = "id";
  }
 }
@@ -32,6 +32,11 @@ public class FakeSecretManagerClient implements SecretManagerClient {
  @Inject
  FakeSecretManagerClient() {}

+  @Override
+  public String getProject() {
+    return "fake_project";
+  }
+
  @Override
  public void createSecret(String secretId) {
    checkNotNull(secretId, "secretId");
@@ -41,6 +46,12 @@ public class FakeSecretManagerClient implements SecretManagerClient {
    secrets.put(secretId, new SecretEntry(secretId));
  }

+  @Override
+  public boolean secretExists(String secretId) {
+    checkNotNull(secretId, "secretId");
+    return secrets.containsKey(secretId);
+  }
+
  @Override
  public Iterable<String> listSecrets() {
    return ImmutableSet.copyOf(secrets.keySet());
@@ -78,6 +89,28 @@ public class FakeSecretManagerClient implements SecretManagerClient {
    return secretEntry.getVersion(version).getData();
  }

+  @Override
+  public void enableSecretVersion(String secretId, String version) {
+    checkNotNull(secretId, "secretId");
+    checkNotNull(version, "version");
+    SecretEntry secretEntry = secrets.get(secretId);
+    if (secretEntry == null) {
+      throw new NoSuchSecretResourceException(null);
+    }
+    secretEntry.enableVersion(version);
+  }
+
+  @Override
+  public void disableSecretVersion(String secretId, String version) {
+    checkNotNull(secretId, "secretId");
+    checkNotNull(version, "version");
+    SecretEntry secretEntry = secrets.get(secretId);
+    if (secretEntry == null) {
+      throw new NoSuchSecretResourceException(null);
+    }
+    secretEntry.disableVersion(version);
+  }
+
  @Override
  public void destroySecretVersion(String secretId, String version) {
    checkNotNull(secretId, "secretId");
@@ -118,6 +151,20 @@ public class FakeSecretManagerClient implements SecretManagerClient {
      return state;
    }

+    void enable() {
+      if (state.equals(State.DESTROYED)) {
+        throw new SecretManagerException(null);
+      }
+      state = State.ENABLED;
+    }
+
+    void disable() {
+      if (state.equals(State.DESTROYED)) {
+        throw new SecretManagerException(null);
+      }
+      state = State.DISABLED;
+    }
+
    void destroy() {
      data = null;
      state = State.DESTROYED;
@@ -145,6 +192,8 @@ public class FakeSecretManagerClient implements SecretManagerClient {
        return versions.get(index);
      } catch (NumberFormatException e) {
        throw new IllegalArgumentException("Invalid version " + version.get());
+      } catch (ArrayIndexOutOfBoundsException e) {
+        throw new NoSuchSecretResourceException(null);
      }
    }

@@ -156,15 +205,16 @@ public class FakeSecretManagerClient implements SecretManagerClient {
      return builder.build();
    }

+    void enableVersion(String version) {
+      getVersion(Optional.of(version)).enable();
+    }
+
+    void disableVersion(String version) {
+      getVersion(Optional.of(version)).disable();
+    }
+
    void destroyVersion(String version) {
-      try {
-        int index = Integer.valueOf(version);
-        versions.get(index).destroy();
-      } catch (NumberFormatException e) {
-        throw new IllegalArgumentException("Invalid version " + version);
-      } catch (ArrayIndexOutOfBoundsException e) {
-        throw new NoSuchSecretResourceException(null);
-      }
+      getVersion(Optional.of(version)).destroy();
    }
  }
 }
@@ -18,12 +18,17 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;

 import com.google.cloud.secretmanager.v1.SecretVersion.State;
+import google.registry.privileges.secretmanager.SecretManagerClient.NoSuchSecretResourceException;
 import google.registry.privileges.secretmanager.SecretManagerClient.SecretAlreadyExistsException;
 import google.registry.privileges.secretmanager.SecretManagerClient.SecretManagerException;
+import google.registry.util.Retrier;
+import google.registry.util.SystemSleeper;
+import java.io.IOException;
 import java.util.Optional;
 import java.util.UUID;
-import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;

 /**
@@ -43,52 +48,50 @@ public class SecretManagerClientTest {
  private static final String SECRET_ID_PREFIX = "TEST_" + UUID.randomUUID() + "_";
  // Used for unique secret id generation.
  private static int seqno = 0;
-
  private static SecretManagerClient secretManagerClient;
  private static boolean isUnitTest = true;

-  private static String nextSecretId() {
-    return SECRET_ID_PREFIX + seqno++;
-  }
+  private String secretId;

  @BeforeAll
  static void beforeAll() {
    String environmentName = System.getProperty("test.gcp_integration.env");
    if (environmentName != null) {
      secretManagerClient =
-          DaggerSecretManagerModule_SecretManagerComponent.builder()
-              .secretManagerModule(
-                  new SecretManagerModule(String.format("domain-registry-%s", environmentName)))
-              .build()
-              .secretManagerClient();
+          SecretManagerModule.provideSecretManagerClient(
+              String.format("domain-registry-%s", environmentName),
+              new Retrier(new SystemSleeper(), 1));
      isUnitTest = false;
    } else {
      secretManagerClient = new FakeSecretManagerClient();
    }
  }

-  @AfterAll
-  static void afterAll() {
+  @BeforeEach
+  void beforeEach() {
+    secretId = SECRET_ID_PREFIX + seqno++;
+  }
+
+  @AfterEach
+  void afterEach() throws IOException {
    if (isUnitTest) {
      return;
    }
-    for (String secretId : secretManagerClient.listSecrets()) {
-      if (secretId.startsWith(SECRET_ID_PREFIX)) {
-        secretManagerClient.deleteSecret(secretId);
-      }
+    try {
+      secretManagerClient.deleteSecret(secretId);
+    } catch (NoSuchSecretResourceException e) {
+      // deleteSecret() deleted it already.
    }
  }

  @Test
  void createSecret_success() {
-    String secretId = nextSecretId();
    secretManagerClient.createSecret(secretId);
    assertThat(secretManagerClient.listSecrets()).contains(secretId);
  }

  @Test
  void createSecret_duplicate() {
-    String secretId = nextSecretId();
    secretManagerClient.createSecret(secretId);
    assertThrows(
        SecretAlreadyExistsException.class, () -> secretManagerClient.createSecret(secretId));
@@ -96,16 +99,25 @@ public class SecretManagerClientTest {

  @Test
  void addSecretVersion() {
-    String secretId = nextSecretId();
    secretManagerClient.createSecret(secretId);
    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
    assertThat(secretManagerClient.listSecretVersions(secretId, State.ENABLED))
        .containsExactly(version);
  }

+  @Test
+  void secretExists_true() {
+    secretManagerClient.createSecret(secretId);
+    assertThat(secretManagerClient.secretExists(secretId)).isTrue();
+  }
+
+  @Test
+  void secretExists_False() {
+    assertThat(secretManagerClient.secretExists(secretId)).isFalse();
+  }
+
  @Test
  void getSecretData_byVersion() {
-    String secretId = nextSecretId();
    secretManagerClient.createSecret(secretId);
    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
    assertThat(secretManagerClient.getSecretData(secretId, Optional.of(version)))
@@ -114,15 +126,70 @@ public class SecretManagerClientTest {

  @Test
  void getSecretData_latestVersion() {
-    String secretId = nextSecretId();
    secretManagerClient.createSecret(secretId);
    secretManagerClient.addSecretVersion(secretId, "mydata");
    assertThat(secretManagerClient.getSecretData(secretId, Optional.empty())).isEqualTo("mydata");
  }

+  @Test
+  void disableSecretVersion() {
+    secretManagerClient.createSecret(secretId);
+    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
+    secretManagerClient.disableSecretVersion(secretId, version);
+    assertThat(secretManagerClient.listSecretVersions(secretId, State.DISABLED)).contains(version);
+  }
+
+  @Test
+  void disableSecretVersion_ignoreAlreadyDisabled() {
+    secretManagerClient.createSecret(secretId);
+    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
+    secretManagerClient.disableSecretVersion(secretId, version);
+    assertThat(secretManagerClient.listSecretVersions(secretId, State.DISABLED)).contains(version);
+    secretManagerClient.disableSecretVersion(secretId, version);
+  }
+
+  @Test
+  void disableSecretVersion_destroyed() {
+    secretManagerClient.createSecret(secretId);
+    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
+    secretManagerClient.destroySecretVersion(secretId, version);
+    assertThat(secretManagerClient.listSecretVersions(secretId, State.DESTROYED)).contains(version);
+    assertThrows(
+        SecretManagerException.class,
+        () -> secretManagerClient.disableSecretVersion(secretId, version));
+  }
+
+  @Test
+  void enableSecretVersion_ignoreAlreadyEnabled() {
+    secretManagerClient.createSecret(secretId);
+    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
+    assertThat(secretManagerClient.listSecretVersions(secretId, State.ENABLED)).contains(version);
+    secretManagerClient.enableSecretVersion(secretId, version);
+  }
+
+  @Test
+  void enableSecretVersion() {
+    secretManagerClient.createSecret(secretId);
+    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
+    secretManagerClient.disableSecretVersion(secretId, version);
+    assertThat(secretManagerClient.listSecretVersions(secretId, State.DISABLED)).contains(version);
+    secretManagerClient.enableSecretVersion(secretId, version);
+    assertThat(secretManagerClient.listSecretVersions(secretId, State.ENABLED)).contains(version);
+  }
+
+  @Test
+  void enableSecretVersion_destroyed() {
+    secretManagerClient.createSecret(secretId);
+    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
+    secretManagerClient.destroySecretVersion(secretId, version);
+    assertThat(secretManagerClient.listSecretVersions(secretId, State.DESTROYED)).contains(version);
+    assertThrows(
+        SecretManagerException.class,
+        () -> secretManagerClient.enableSecretVersion(secretId, version));
+  }
+
  @Test
  void destroySecretVersion() {
-    String secretId = nextSecretId();
    secretManagerClient.createSecret(secretId);
    String version = secretManagerClient.addSecretVersion(secretId, "mydata");
    secretManagerClient.destroySecretVersion(secretId, version);
@@ -134,7 +201,6 @@ public class SecretManagerClientTest {

  @Test
  void deleteSecret() {
-    String secretId = nextSecretId();
    secretManagerClient.createSecret(secretId);
    assertThat(secretManagerClient.listSecrets()).contains(secretId);
    secretManagerClient.deleteSecret(secretId);
@@ -0,0 +1,63 @@
+// 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.privileges.secretmanager;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.secretmanager.v1.SecretVersionName;
+import google.registry.privileges.secretmanager.SqlUser.RobotId;
+import google.registry.privileges.secretmanager.SqlUser.RobotUser;
+import java.util.Optional;
+import org.junit.jupiter.api.Test;
+
+/** Unit tests for {@link SqlCredentialStore}. */
+public class SqlCredentialStoreTest {
+
+  private final SecretManagerClient client = new FakeSecretManagerClient();
+  private final SqlCredentialStore credentialStore = new SqlCredentialStore(client, "db");
+  private SqlUser user = new RobotUser(RobotId.NOMULUS);
+
+  @Test
+  void createSecret() {
+    credentialStore.createOrUpdateCredential(user, "password");
+    assertThat(client.secretExists("sql-cred-live-label-nomulus-db")).isTrue();
+    assertThat(
+            SecretVersionName.parse(
+                    client.getSecretData("sql-cred-live-label-nomulus-db", Optional.empty()))
+                .getSecret())
+        .isEqualTo("sql-cred-data-nomulus-db");
+    assertThat(client.secretExists("sql-cred-data-nomulus-db")).isTrue();
+    assertThat(client.getSecretData("sql-cred-data-nomulus-db", Optional.empty()))
+        .isEqualTo("nomulus password");
+  }
+
+  @Test
+  void getCredential() {
+    credentialStore.createOrUpdateCredential(user, "password");
+    SqlCredential credential = credentialStore.getCredential(user);
+    assertThat(credential.login()).isEqualTo("nomulus");
+    assertThat(credential.password()).isEqualTo("password");
+  }
+
+  @Test
+  void deleteCredential() {
+    credentialStore.createOrUpdateCredential(user, "password");
+    assertThat(client.secretExists("sql-cred-live-label-nomulus-db")).isTrue();
+    assertThat(client.secretExists("sql-cred-data-nomulus-db")).isTrue();
+    credentialStore.deleteCredential(user);
+    assertThat(client.secretExists("sql-cred-live-label-nomulus-db")).isFalse();
+    assertThat(client.secretExists("sql-cred-data-nomulus-db")).isFalse();
+  }
+}
@@ -750,6 +750,20 @@ public class DatabaseHelper {
            .build());
  }

+  /** Persists and returns a {@link Registrar} with the specified registrarId. */
+  public static Registrar persistNewRegistrar(String registrarId) {
+    return persistNewRegistrar(registrarId, registrarId + " name", Registrar.Type.REAL, 100L);
+  }
+
+  /** Persists and returns a list of {@link Registrar}s with the specified registrarIds. */
+  public static ImmutableList<Registrar> persistNewRegistrars(String... registrarIds) {
+    ImmutableList.Builder<Registrar> newRegistrars = new ImmutableList.Builder<>();
+    for (String registrarId : registrarIds) {
+      newRegistrars.add(persistNewRegistrar(registrarId));
+    }
+    return newRegistrars.build();
+  }
+
  private static Iterable<BillingEvent> getBillingEvents() {
    return transactIfJpaTm(
        () ->
@@ -0,0 +1,122 @@
+// 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.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.ImmutableObjectSubject.immutableObjectCorrespondence;
+import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.newDomainBase;
+import static google.registry.testing.DatabaseHelper.persistResource;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.truth.Correspondence;
+import com.googlecode.objectify.Key;
+import google.registry.model.ImmutableObject;
+import google.registry.model.billing.BillingEvent;
+import google.registry.model.common.EntityGroupRoot;
+import google.registry.model.domain.DomainBase;
+import google.registry.model.poll.PollMessage;
+import google.registry.model.reporting.HistoryEntry;
+import google.registry.persistence.VKey;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Unit test for {@link RemoveRegistryOneKeyCommand}. */
+public class RemoveRegistryOneKeyCommandTest extends CommandTestCase<RemoveRegistryOneKeyCommand> {
+  DomainBase domain;
+  HistoryEntry historyEntry;
+
+  @BeforeEach
+  void beforeEach() {
+    createTld("foobar");
+    domain =
+        newDomainBase("foo.foobar")
+            .asBuilder()
+            .setDeletionTime(DateTime.parse("2016-01-01T00:00:00Z"))
+            .setAutorenewBillingEvent(createRegistryOneVKey(BillingEvent.Recurring.class, 100L))
+            .setAutorenewPollMessage(createRegistryOneVKey(PollMessage.Autorenew.class, 200L))
+            .setDeletePollMessage(createRegistryOneVKey(PollMessage.OneTime.class, 300L))
+            .build();
+  }
+
+  @Test
+  void removeRegistryOneKeyInDomainBase_succeeds() throws Exception {
+    DomainBase origin = persistResource(domain);
+
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(domain)));
+
+    DomainBase persisted = ofy().load().key(domain.createVKey().getOfyKey()).now();
+    assertThat(ImmutableList.of(persisted))
+        .comparingElementsUsing(getDomainBaseCorrespondence())
+        .containsExactly(origin);
+    assertThat(persisted.getAutorenewBillingEvent()).isNull();
+    assertThat(persisted.getAutorenewPollMessage()).isNull();
+    assertThat(persisted.getDeletePollMessage()).isNull();
+  }
+
+  @Test
+  void removeRegistryOneKeyInDomainBase_notModifyRegistryTwoKey() throws Exception {
+    DomainBase origin =
+        persistResource(
+            domain
+                .asBuilder()
+                .setAutorenewBillingEvent(
+                    createRegistryTwoVKey(BillingEvent.Recurring.class, domain, 300L))
+                .build());
+
+    runCommand(
+        "--force",
+        "--key_paths_file",
+        writeToNamedTmpFile("keypath.txt", getKeyPathLiteral(domain)));
+
+    DomainBase persisted = ofy().load().key(domain.createVKey().getOfyKey()).now();
+    assertThat(ImmutableList.of(persisted))
+        .comparingElementsUsing(getDomainBaseCorrespondence())
+        .containsExactly(origin);
+    assertThat(persisted.getAutorenewBillingEvent())
+        .isEqualTo(createRegistryTwoVKey(BillingEvent.Recurring.class, domain, 300L));
+    assertThat(persisted.getAutorenewPollMessage()).isNull();
+    assertThat(persisted.getDeletePollMessage()).isNull();
+  }
+
+  private static String getKeyPathLiteral(Object entity) {
+    Key<?> key = Key.create(entity);
+    return String.format("\"DomainBase\", \"%s\"", key.getName());
+  }
+
+  private static <T> VKey<T> createRegistryOneVKey(Class<T> clazz, long id) {
+    Key<?> parent = Key.create(EntityGroupRoot.class, "per-tld");
+    return VKey.create(clazz, id, Key.create(parent, clazz, id));
+  }
+
+  private static <T> VKey<T> createRegistryTwoVKey(Class<T> clazz, DomainBase domain, long id) {
+    Key<?> parent = Key.create(domain.createVKey().getOfyKey(), HistoryEntry.class, 1000L);
+    return VKey.create(clazz, id, Key.create(parent, clazz, id));
+  }
+
+  private static Correspondence<ImmutableObject, ImmutableObject> getDomainBaseCorrespondence() {
+    return immutableObjectCorrespondence(
+        "revisions",
+        "updateTimestamp",
+        "autorenewBillingEvent",
+        "autorenewPollMessage",
+        "deletePollMessage");
+  }
+}
@@ -0,0 +1,42 @@
+// 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.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.junit.Assert.assertThrows;
+
+import google.registry.schema.replay.SqlReplayCheckpoint;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+
+class SetSqlReplayCheckpointCommandTest extends CommandTestCase<SetSqlReplayCheckpointCommand> {
+
+  @Test
+  void testSuccess() throws Exception {
+    assertThat(jpaTm().transact(SqlReplayCheckpoint::get)).isEqualTo(START_OF_TIME);
+    DateTime timeToSet = DateTime.parse("2000-06-06T22:00:00.0Z");
+    runCommandForced(timeToSet.toString());
+    assertThat(jpaTm().transact(SqlReplayCheckpoint::get)).isEqualTo(timeToSet);
+  }
+
+  @Test
+  void testFailure_multipleParams() throws Exception {
+    DateTime one = DateTime.parse("2000-06-06T22:00:00.0Z");
+    DateTime two = DateTime.parse("2001-06-06T22:00:00.0Z");
+    assertThrows(IllegalArgumentException.class, () -> runCommand(one.toString(), two.toString()));
+  }
+}
@@ -18,6 +18,12 @@ plugins {

 dependencies {
    def deps = rootProject.dependencyMap
+
+    // Custom-built objectify jar at commit ecd5165, included in Nomulus
+    // release.
+    compile files(
+        "${rootDir}/third_party/objectify/v4_1/objectify-4.1.3.jar")
+
    compile deps['com.google.code.findbugs:jsr305']
    compile deps['com.google.guava:guava']
    compile deps['com.squareup:javapoet']
@@ -27,11 +27,14 @@ import com.squareup.javapoet.TypeSpec;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.annotation.processing.AbstractProcessor;
 import javax.annotation.processing.RoundEnvironment;
 import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.AnnotationValue;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
@@ -47,9 +50,12 @@ import javax.persistence.Converter;
 public abstract class AbstractVKeyProcessor extends AbstractProcessor {

  private static final String CONVERTER_CLASS_NAME_TEMP = "VKeyConverter_%s";
-  // The method with same name should be defined in StringVKey and LongVKey
+  // The method with same name should be defined in WithStringVKey and WithLongVKey
  private static final String CLASS_NAME_SUFFIX_KEY = "classNameSuffix";

+  // Method in WithStringVKey and WithLongVKey to indicate that this is a composite key.
+  private static final String COMPOSITE_KEY_KEY = "compositeKey";
+
  abstract Class<?> getSqlColumnType();

  abstract String getAnnotationSimpleName();
@@ -76,18 +82,18 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                        actualAnnotation.size() == 1,
                        String.format(
                            "type can have only 1 %s annotation", getAnnotationSimpleName()));
-                    String converterClassNameSuffix =
-                        actualAnnotation.get(0).getElementValues().entrySet().stream()
-                            .filter(
-                                entry ->
-                                    entry
-                                        .getKey()
-                                        .getSimpleName()
-                                        .toString()
-                                        .equals(CLASS_NAME_SUFFIX_KEY))
-                            .map(entry -> ((String) entry.getValue().getValue()).trim())
-                            .findFirst()
-                            .orElse("");
+                    String converterClassNameSuffix = "";
+                    boolean hasCompositeOfyKey = false;
+                    for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> entry :
+                        actualAnnotation.get(0).getElementValues().entrySet()) {
+                      String keyName = entry.getKey().getSimpleName().toString();
+                      Object value = entry.getValue().getValue();
+                      if (keyName.equals(CLASS_NAME_SUFFIX_KEY)) {
+                        converterClassNameSuffix = ((String) value).trim();
+                      } else if (keyName.equals(COMPOSITE_KEY_KEY)) {
+                        hasCompositeOfyKey = (Boolean) value;
+                      }
+                    }
                    if (converterClassNameSuffix.isEmpty()) {
                      converterClassNameSuffix =
                          getTypeUtils().asElement(entityType).getSimpleName().toString();
@@ -97,7 +103,8 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                      createJavaFile(
                              getPackageName(annotatedTypeElement),
                              String.format(CONVERTER_CLASS_NAME_TEMP, converterClassNameSuffix),
-                              entityType)
+                              entityType,
+                              hasCompositeOfyKey)
                          .writeTo(processingEnv.getFiler());
                    } catch (IOException e) {
                      throw new UncheckedIOException(e);
@@ -108,7 +115,10 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
  }

  private JavaFile createJavaFile(
-      String packageName, String converterClassName, TypeMirror entityTypeMirror) {
+      String packageName,
+      String converterClassName,
+      TypeMirror entityTypeMirror,
+      boolean hasCompositeOfyKey) {
    TypeName entityType = ClassName.get(entityTypeMirror);

    ParameterizedTypeName attributeConverter =
@@ -127,7 +137,7 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
            .addStatement("return $T.class", entityType)
            .build();

-    TypeSpec vKeyConverter =
+    TypeSpec.Builder classBuilder =
        TypeSpec.classBuilder(converterClassName)
            .addAnnotation(
                AnnotationSpec.builder(ClassName.get(Converter.class))
@@ -135,9 +145,24 @@ public abstract class AbstractVKeyProcessor extends AbstractProcessor {
                    .build())
            .addModifiers(Modifier.FINAL)
            .superclass(attributeConverter)
-            .addMethod(getAttributeClass)
-            .build();
+            .addMethod(getAttributeClass);

+    // If this is a converter for a composite vkey type, generate an override for the default
+    // {@link google.registry.persistence.VKeyConverter.hasCompositeOfyKey()} method, which returns
+    // false.
+    if (hasCompositeOfyKey) {
+      MethodSpec hasCompositeOfyKeyMethod =
+          MethodSpec.methodBuilder("hasCompositeOfyKey")
+              .addAnnotation(Override.class)
+              .addModifiers(Modifier.PROTECTED)
+              .returns(boolean.class)
+              .addStatement("return true", entityType)
+              .build();
+
+      classBuilder.addMethod(hasCompositeOfyKeyMethod);
+    }
+
+    TypeSpec vKeyConverter = classBuilder.build();
    return JavaFile.builder(packageName, vKeyConverter).build();
  }
Author	SHA1	Message	Date
gbrodman	7c364b4471	Add SetSqlReplayCheckpoint command for SQL replay (#895 ) * Add SetSqlReplayCheckpoint command for SQL replay We should set this to the same time that we initially populate the SQL database from Datastore.	2020-12-11 17:41:06 -05:00
Shicong Huang	b5137c3d05	Convert HostResourceTest to work with Cloud SQL (#905 )	2020-12-11 13:17:55 -05:00
Ben McIlwain	6a9929019a	Use 10 workers instead of the default 100 for re-save all EPP resources (#904 ) * Use 10 workers instead of the default 100 for re-save all EPP resources The intended/desired effect is to have a larger number of GCS commit log diffs spread out over a longer period of time, with each diff itself being significantly smaller. This should retain roughly the same amount of total work for the async Cloud SQL replication action to have to deal with, but spread across 10X as much time.	2020-12-10 15:25:25 -05:00
Weimin Yu	83ed448741	Add a credential store backed by Secret Manager (#901 ) * Add a credential store backed by Secret Manager Added a SqlCredentialStore that stores user credentials with one level of indirection: for each credential, an addtional secret is used to identify the 'live' version of the credential. This is a work in progress and the overall design is explained in go/dr-sql-security. Also added two nomulus commands for credential management. They are stop-gap measures that will be deprecated by the planned privilege management system.	2020-12-10 11:29:44 -05:00
gbrodman	2c6ee6dae9	Parameterize the serialization of objects being written to SQL (#892 ) * Parameterize the serialization of objects being written to SQL We shouldn't require that objects written to SQL during a Beam pipeline be VersionedEntity objects -- they may be non-Objectify entities. As a result, we should allow the user to specify what the objects are that should be written to SQL. Note: we will need to clean up the Spec11PipelineTest more but that can be out of the scope of this PR. * Overload the method and add a bit of javadoc * Actually use the overloaded function	2020-12-09 15:52:56 -05:00
Shicong Huang	a181d6a720	Add a command to remove Registry 1.0 key in DomainBase (#900 )	2020-12-09 10:30:23 -05:00
Michael Muller	db19f9ea4f	Resurrect symmetric vkeys when possible. (#899 ) * Resurrect symmetric vkeys when possible. AbstractVKeyConverter had never been updated to generate symmetric VKeys for simple (i.e. non-composite) VKey types. Update it to default to generating a VKey with a simple Ofy key. With this change, composite VKeys must specify the "compositeKey = true" field in the With*VKey annotations. Doing so creates an asymmetric (SQL only) VKey that can triggers higher-level logic to generate a corrected VKey containing the composite Key.	2020-12-08 15:13:44 -05:00
Michael Muller	c7e6192929	Add replay extension to more domain tests (#893 ) * Add replay extension to more domain tests Add replay to DomainRenewFlowTest and DomainUpdateFlowTest. * Formatting fix	2020-12-04 14:15:24 -05:00