Use VKeys instead of Ofy keys in mutating command (#1682)

* Use VKeys instead of Ofy keys in mutating command

* Add createVKey to ImmutableObject

* Use SQL only VKeys
This commit is contained in:
sarahcaseybot 2022-06-27 17:49:24 -04:00 committed by GitHub
parent 89925f9ff2
commit 2c3279ba95
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 24 additions and 365 deletions

View file

@ -27,6 +27,7 @@ import com.google.common.base.Joiner;
import com.google.common.collect.Maps;
import com.googlecode.objectify.Key;
import com.googlecode.objectify.annotation.Ignore;
import google.registry.persistence.VKey;
import java.lang.annotation.Documented;
import java.lang.annotation.Retention;
import java.lang.annotation.Target;
@ -255,4 +256,8 @@ public abstract class ImmutableObject implements Cloneable {
public Map<String, Object> toDiffableFieldMap() {
return (Map<String, Object>) toMapRecursive(this);
}
public VKey createVKey() {
throw new UnsupportedOperationException("VKey creation is not supported for this entity");
}
}

View file

@ -293,6 +293,11 @@ public class Registry extends ImmutableObject implements Buildable, UnsafeSerial
return createVKey(key.getName());
}
@Override
public VKey<Registry> createVKey() {
return VKey.createSql(Registry.class, this.tldStrId);
}
/**
* The name of the pricing engine that this TLD uses.
*

View file

@ -22,7 +22,6 @@ import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Strings.emptyToNull;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.util.DatastoreServiceUtils.getNameOrId;
import static google.registry.util.DiffUtils.prettyPrintEntityDeepDiff;
import static java.util.stream.Collectors.joining;
@ -30,7 +29,6 @@ import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.googlecode.objectify.Key;
import google.registry.model.ImmutableObject;
import google.registry.persistence.VKey;
import java.util.ArrayList;
@ -77,65 +75,21 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
final ImmutableObject newEntity;
/** The key that points to the entity being changed. */
final VKey<?> key;
final VKey<?> vKey;
private EntityChange(ImmutableObject oldEntity, ImmutableObject newEntity) {
type = ChangeType.get(oldEntity != null, newEntity != null);
checkArgument(
type != ChangeType.UPDATE || Key.create(oldEntity).equals(Key.create(newEntity)),
type != ChangeType.UPDATE || oldEntity.createVKey().equals(newEntity.createVKey()),
"Both entity versions in an update must have the same Key.");
this.oldEntity = oldEntity;
this.newEntity = newEntity;
ImmutableObject entity = MoreObjects.firstNonNull(oldEntity, newEntity);
// This is one of the few cases where it is acceptable to create an asymmetric VKey (using
// createOfy()). We can use this code on datastore-only entities where we can't construct a
// SQL key.
VKey<?> createdKey;
try {
createdKey = VKey.from(Key.create(entity));
} catch (RuntimeException e) {
createdKey = VKey.createOfy(entity.getClass(), Key.create(entity));
}
key = createdKey;
}
/**
* EntityChange constructor that supports Vkey override. A Vkey is a key of an entity. This is a
* workaround to handle cases when a SqlEntity instance does not have a primary key before being
* persisted.
*/
private EntityChange(
@Nullable ImmutableObject oldEntity, @Nullable ImmutableObject newEntity, VKey<?> vkey) {
type = ChangeType.get(oldEntity != null, newEntity != null);
if (type == ChangeType.UPDATE) {
checkArgument(
Key.create(oldEntity).equals(Key.create(newEntity)),
"Both entity versions in an update must have the same Key.");
checkArgument(
Key.create(oldEntity).equals(vkey.getOfyKey()),
"The Key of the entity must be the same as the OfyKey of the vkey");
} else if (type == ChangeType.CREATE) {
checkArgument(
Key.create(newEntity).equals(vkey.getOfyKey()),
"Both entity versions in an update must have the same Key.");
} else if (type == ChangeType.DELETE) {
checkArgument(
Key.create(oldEntity).equals(vkey.getOfyKey()),
"The Key of the entity must be the same as the OfyKey of the vkey");
}
this.oldEntity = oldEntity;
this.newEntity = newEntity;
key = vkey;
vKey = MoreObjects.firstNonNull(oldEntity, newEntity).createVKey();
}
/** Returns a human-readable ID string for the entity being changed. */
String getEntityId() {
return String.format(
"%s@%s",
key.getOfyKey().getKind(),
// NB: try name before id, since name defaults to null, whereas id defaults to 0.
getNameOrId(key.getOfyKey().getRaw()));
return String.format("%s@%s", vKey.getKind().getSimpleName(), vKey.getSqlKey().toString());
}
/** Returns a string representation of this entity change. */
@ -195,7 +149,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
private void executeChange(EntityChange change) {
// Load the key of the entity to mutate and double-check that it hasn't been
// modified from the version that existed when the change was prepared.
Optional<?> existingEntity = tm().loadByKeyIfPresent(change.key);
Optional<?> existingEntity = tm().loadByKeyIfPresent(change.vKey);
checkState(
Objects.equals(change.oldEntity, existingEntity.orElse(null)),
"Entity changed since init() was called.\n%s",
@ -212,7 +166,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
tm().update(change.newEntity);
return;
case DELETE:
tm().delete(change.key);
tm().delete(change.vKey);
return;
}
throw new UnsupportedOperationException("Unknown entity change type: " + change.type);
@ -227,7 +181,7 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
ArrayList<EntityChange> nextBatch = new ArrayList<>();
for (EntityChange change : changedEntitiesMap.values()) {
nextBatch.add(change);
if (transactionBoundaries.contains(change.key)) {
if (transactionBoundaries.contains(change.vKey)) {
batches.add(ImmutableList.copyOf(nextBatch));
nextBatch.clear();
}
@ -249,30 +203,11 @@ public abstract class MutatingCommand extends ConfirmingCommand implements Comma
@Nullable ImmutableObject oldEntity, @Nullable ImmutableObject newEntity) {
EntityChange change = new EntityChange(oldEntity, newEntity);
checkArgument(
!changedEntitiesMap.containsKey(change.key),
!changedEntitiesMap.containsKey(change.vKey),
"Cannot apply multiple changes for the same entity: %s",
change.getEntityId());
changedEntitiesMap.put(change.key, change);
lastAddedKey = change.key;
}
/**
* Stages an entity change which will be applied by execute(), with the support of Vkey override.
* It supports cases of SqlEntity instances that do not have primary keys before being persisted.
*
* @param oldEntity the existing version of the entity, or null to create a new entity
* @param newEntity the new version of the entity to save, or null to delete the entity
* @param vkey the key of the entity
*/
protected void stageEntityChange(
@Nullable ImmutableObject oldEntity, @Nullable ImmutableObject newEntity, VKey vkey) {
EntityChange change = new EntityChange(oldEntity, newEntity, vkey);
checkArgument(
!changedEntitiesMap.containsKey(change.key),
"Cannot apply multiple changes for the same entity: %s",
change.getEntityId());
changedEntitiesMap.put(change.key, change);
lastAddedKey = change.key;
changedEntitiesMap.put(change.vKey, change);
lastAddedKey = change.vKey;
}
/**

View file

@ -1,166 +0,0 @@
// 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.checkState;
import static google.registry.model.ofy.ObjectifyService.auditedOfy;
import static java.nio.charset.StandardCharsets.UTF_8;
import com.beust.jcommander.Parameter;
import com.google.appengine.api.datastore.KeyFactory;
import com.google.common.base.Splitter;
import com.google.common.io.CharStreams;
import com.google.common.io.Files;
import com.googlecode.objectify.Key;
import google.registry.model.ImmutableObject;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.model.domain.DomainBase;
import google.registry.persistence.VKey;
import google.registry.util.NonFinalForTesting;
import google.registry.util.TypeUtils.TypeInstantiator;
import java.io.File;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.List;
/**
* 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.
*/
@DeleteAfterMigration
abstract class ReadEntityFromKeyPathCommand<T> extends MutatingCommand {
@Parameter(
names = "--key_paths_file",
description =
"Key paths file name, each line in the file should be a key literal. An example key"
+ " literal is: \"DomainBase\", \"111111-TEST\", \"HistoryEntry\", 2222222,"
+ " \"OneTime\", 3333333")
File keyPathsFile;
@NonFinalForTesting private static InputStream stdin = System.in;
private StringBuilder changeMessage = new StringBuilder();
abstract void process(T entity);
@Override
protected void init() throws Exception {
List<String> keyPaths =
keyPathsFile == null
? CharStreams.readLines(new InputStreamReader(stdin, UTF_8))
: Files.readLines(keyPathsFile, UTF_8);
for (String keyPath : keyPaths) {
Key<?> untypedKey = parseKeyPath(keyPath);
Object entity = auditedOfy().load().key(untypedKey).now();
if (entity == null) {
System.err.printf(
"Entity %s read from %s doesn't exist in Datastore! Skipping.%n",
untypedKey, keyPathsFile == null ? "STDIN" : "File " + keyPathsFile.getAbsolutePath());
continue;
}
Class<T> clazz = new TypeInstantiator<T>(getClass()) {}.getExactType();
if (clazz.isInstance(entity)) {
process((T) entity);
} else {
throw new IllegalArgumentException("Unsupported entity key: " + untypedKey);
}
flushTransaction();
}
}
@Override
protected void postBatchExecute() {
System.out.println(changeMessage);
}
void stageEntityKeyChange(ImmutableObject oldEntity, ImmutableObject newEntity) {
stageEntityChange(oldEntity, null);
stageEntityChange(null, newEntity);
appendChangeMessage(
String.format(
"Changed entity key from: %s to: %s", Key.create(oldEntity), Key.create(newEntity)));
}
void appendChangeMessage(String message) {
changeMessage.append(message);
}
private static boolean isKind(Key<?> key, Class<?> clazz) {
return key.getKind().equals(Key.getKind(clazz));
}
static Key<?> parseKeyPath(String keyPath) {
List<String> keyComponents = Splitter.on(',').splitToList(keyPath);
checkState(
keyComponents.size() > 0 && keyComponents.size() % 2 == 0,
"Invalid number of key components");
com.google.appengine.api.datastore.Key rawKey = null;
for (int i = 0, j = 1; j < keyComponents.size(); i += 2, j += 2) {
String kindLiteral = keyComponents.get(i).trim();
String idOrNameLiteral = keyComponents.get(j).trim();
rawKey = createDatastoreKey(rawKey, kindLiteral, idOrNameLiteral);
}
return Key.create(rawKey);
}
private static com.google.appengine.api.datastore.Key createDatastoreKey(
com.google.appengine.api.datastore.Key parent, String kindLiteral, String idOrNameLiteral) {
if (isLiteralString(idOrNameLiteral)) {
return KeyFactory.createKey(parent, removeQuotes(kindLiteral), removeQuotes(idOrNameLiteral));
} else {
return KeyFactory.createKey(
parent, removeQuotes(kindLiteral), Long.parseLong(idOrNameLiteral));
}
}
private static boolean isLiteralString(String raw) {
return raw.charAt(0) == '"' && raw.charAt(raw.length() - 1) == '"';
}
private static String removeQuotes(String literal) {
return literal.substring(1, literal.length() - 1);
}
static Key<DomainBase> getGrandParentAsDomain(Key<?> key) {
Key<?> grandParent;
try {
grandParent = key.getParent().getParent();
} catch (Throwable e) {
throw new IllegalArgumentException("Error retrieving grand parent key", e);
}
if (!isKind(grandParent, DomainBase.class)) {
throw new IllegalArgumentException(
String.format("Expected a Key<DomainBase> but got %s", grandParent));
}
return (Key<DomainBase>) grandParent;
}
static VKey<DomainBase> getGrandParentAsDomain(VKey<?> key) {
Key<DomainBase> grandParent;
try {
grandParent = key.getOfyKey().getParent().getParent();
} catch (Throwable e) {
throw new IllegalArgumentException("Error retrieving grand parent key", e);
}
if (!isKind(grandParent, DomainBase.class)) {
throw new IllegalArgumentException(
String.format("Expected a Key<DomainBase> but got %s", grandParent));
}
return VKey.create(DomainBase.class, grandParent.getName(), grandParent);
}
}

View file

@ -99,9 +99,7 @@ public final class RegistryTool {
.put("pending_escrow", PendingEscrowCommand.class)
.put("registrar_contact", RegistrarContactCommand.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)

View file

@ -1,55 +0,0 @@
// Copyright 2017 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.collect.Lists.partition;
import static google.registry.model.ofy.ObjectifyService.auditedOfy;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import google.registry.model.ImmutableObject;
import google.registry.persistence.VKey;
import java.util.List;
/**
* A command to load and resave an entity by websafe key.
*
* <p>This triggers @OnSave changes. If the entity was directly edited in the Datastore viewer, this
* can be used to make sure that the commit logs reflect the new state.
*/
@Parameters(
separators = " =",
commandDescription = "Load and resave entities by websafe key")
public final class ResaveEntitiesCommand extends MutatingCommand {
/** The number of resaves to do in a single transaction. */
private static final int BATCH_SIZE = 10;
// TODO(b/207376744): figure out if there's a guide that shows how a websafe key should look like
@Parameter(description = "Websafe keys", required = true)
List<String> mainParameters;
@Override
protected void init() {
for (List<String> batch : partition(mainParameters, BATCH_SIZE)) {
for (String websafeKey : batch) {
ImmutableObject entity =
(ImmutableObject) auditedOfy().load().key(VKey.create(websafeKey).getOfyKey()).now();
stageEntityChange(entity, entity);
}
flushTransaction();
}
}
}

View file

@ -1,62 +0,0 @@
// Copyright 2017 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 google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
import static org.joda.time.DateTimeZone.UTC;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import google.registry.model.EppResource;
import google.registry.persistence.VKey;
import google.registry.tools.CommandUtilities.ResourceType;
import org.joda.time.DateTime;
/**
* A command to load and resave an {@link EppResource} by foreign key.
*
* <p>This triggers @OnSave changes. If the entity was directly edited in the Datastore viewer, this
* can be used to make sure that the commit logs reflect the new state.
*/
@Parameters(
separators = " =",
commandDescription = "Load and resave EPP resources by foreign key")
public final class ResaveEppResourceCommand extends MutatingCommand {
@Parameter(
names = "--type",
description = "Resource type.")
protected ResourceType type;
@Parameter(
names = "--id",
description = "Foreign key of the resource.")
protected String uniqueId;
@Override
protected void init() {
VKey<? extends EppResource> resourceKey =
checkArgumentNotNull(
type.getKey(uniqueId, DateTime.now(UTC)),
"Could not find active resource of type %s: %s",
type,
uniqueId);
// Load the resource directly to bypass running cloneProjectedAtTime() automatically, which can
// cause stageEntityChange() to fail due to implicit projection changes.
EppResource resource = tm().loadByKey(resourceKey);
stageEntityChange(resource, resource);
}
}

View file

@ -106,9 +106,8 @@ class BillingVKeyTest {
return billingRecurrenceVKey.createVKey();
}
VKey<BillingVKeyTestEntity> createVKey() {
return VKey.create(
BillingVKeyTestEntity.class, id, Key.create(parent, BillingVKeyTestEntity.class, id));
public VKey<BillingVKeyTestEntity> createVKey() {
return VKey.createSql(BillingVKeyTestEntity.class, id);
}
}
}

View file

@ -96,8 +96,8 @@ class DomainHistoryVKeyTest {
this.domainHistoryVKey = domainHistoryVKey;
}
VKey<TestEntity> createVKey() {
return VKey.create(TestEntity.class, id, Key.create(parent, TestEntity.class, id));
public VKey<TestEntity> createVKey() {
return VKey.createSql(TestEntity.class, id);
}
}
}