From 5faf3d283c6cc5450d64fe5fd2fec620bce11821 Mon Sep 17 00:00:00 2001 From: gbrodman Date: Mon, 13 Oct 2025 11:43:18 -0400 Subject: [PATCH] Differentiate between inserts and updates in flows (#2846) Updates (AKA merges) run an extra SELECT statement to figure out if the resource exists so that it can merge the entity into the existing object in Hibernate's schema. When we're inserting new rows (such as new poll messages or resource creates), we know that we don't need to do that merge. Doing this should save us some SELECT statements (this has borne out to be the truth in alpha) --- .../java/google/registry/flows/FlowUtils.java | 5 ++- .../registry/flows/custom/EntityChanges.java | 40 +++++++++++++------ .../flows/domain/DomainCreateFlow.java | 18 +++++---- .../flows/domain/DomainDeleteFlow.java | 15 ++++--- .../flows/domain/DomainFlowUtils.java | 2 +- .../flows/domain/DomainRenewFlow.java | 15 ++++--- .../domain/DomainRestoreRequestFlow.java | 13 +++--- .../domain/DomainTransferApproveFlow.java | 17 ++++---- .../domain/DomainTransferCancelFlow.java | 4 +- .../domain/DomainTransferRejectFlow.java | 4 +- .../domain/DomainTransferRequestFlow.java | 8 ++-- .../flows/domain/DomainUpdateFlow.java | 15 ++++--- .../registry/flows/host/HostCreateFlow.java | 5 ++- .../registry/flows/host/HostUpdateFlow.java | 15 +++---- .../TestDomainCreateFlowCustomLogic.java | 5 ++- 15 files changed, 104 insertions(+), 77 deletions(-) diff --git a/core/src/main/java/google/registry/flows/FlowUtils.java b/core/src/main/java/google/registry/flows/FlowUtils.java index c3c2b2425..8318b27ae 100644 --- a/core/src/main/java/google/registry/flows/FlowUtils.java +++ b/core/src/main/java/google/registry/flows/FlowUtils.java @@ -56,9 +56,10 @@ public final class FlowUtils { } } - /** Persists the saves and deletes in an {@link EntityChanges} to the DB. */ + /** Persists the inserts, updates, and deletes in an {@link EntityChanges} to the DB. */ public static void persistEntityChanges(EntityChanges entityChanges) { - tm().putAll(entityChanges.getSaves()); + tm().insertAll(entityChanges.getInserts()); + tm().updateAll(entityChanges.getUpdates()); tm().delete(entityChanges.getDeletes()); } diff --git a/core/src/main/java/google/registry/flows/custom/EntityChanges.java b/core/src/main/java/google/registry/flows/custom/EntityChanges.java index 00cac1567..87a9400c2 100644 --- a/core/src/main/java/google/registry/flows/custom/EntityChanges.java +++ b/core/src/main/java/google/registry/flows/custom/EntityChanges.java @@ -19,25 +19,30 @@ import com.google.common.collect.ImmutableSet; import google.registry.model.ImmutableObject; import google.registry.persistence.VKey; -/** A record that encapsulates database entities to both save and delete. */ +/** A record that encapsulates database entities to insert, update, and delete. */ public record EntityChanges( - ImmutableSet saves, ImmutableSet> deletes) { + ImmutableSet inserts, + ImmutableSet updates, + ImmutableSet> deletes) { - public ImmutableSet getSaves() { - return saves; + public ImmutableSet getInserts() { + return inserts; + } + + public ImmutableSet getUpdates() { + return updates; } - ; public ImmutableSet> getDeletes() { return deletes; } - ; public static Builder newBuilder() { - // Default both entities to save and entities to delete to empty sets, so that the build() - // method won't subsequently throw an exception if one doesn't end up being applicable. + // Default inserts, updates, and deletes to empty sets, so that the build() method won't + // subsequently throw an exception if one doesn't end up being applicable. return new AutoBuilder_EntityChanges_Builder() - .setSaves(ImmutableSet.of()) + .setInserts(ImmutableSet.of()) + .setUpdates(ImmutableSet.of()) .setDeletes(ImmutableSet.of()); } @@ -45,12 +50,21 @@ public record EntityChanges( @AutoBuilder public interface Builder { - Builder setSaves(ImmutableSet entitiesToSave); + Builder setInserts(ImmutableSet entitiesToInsert); - ImmutableSet.Builder savesBuilder(); + ImmutableSet.Builder insertsBuilder(); - default Builder addSave(ImmutableObject entityToSave) { - savesBuilder().add(entityToSave); + default Builder addInsert(ImmutableObject entityToInsert) { + insertsBuilder().add(entityToInsert); + return this; + } + + Builder setUpdates(ImmutableSet entitiesToUpdate); + + ImmutableSet.Builder updatesBuilder(); + + default Builder addUpdate(ImmutableObject entityToUpdate) { + updatesBuilder().add(entityToUpdate); return this; } diff --git a/core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java b/core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java index e06e6f4ed..e8cc745cf 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java @@ -357,11 +357,11 @@ public final class DomainCreateFlow implements MutatingFlow { domainHistoryId, registrationExpirationTime, isAnchorTenant, allocationToken); PollMessage.Autorenew autorenewPollMessage = createAutorenewPollMessage(domainHistoryId, registrationExpirationTime); - ImmutableSet.Builder entitiesToSave = new ImmutableSet.Builder<>(); - entitiesToSave.add(createBillingEvent, autorenewBillingEvent, autorenewPollMessage); + ImmutableSet.Builder entitiesToInsert = new ImmutableSet.Builder<>(); + entitiesToInsert.add(createBillingEvent, autorenewBillingEvent, autorenewPollMessage); // Bill for EAP cost, if any. if (!feesAndCredits.getEapCost().isZero()) { - entitiesToSave.add(createEapBillingEvent(feesAndCredits, createBillingEvent)); + entitiesToInsert.add(createEapBillingEvent(feesAndCredits, createBillingEvent)); } ImmutableSet reservationTypes = getReservationTypes(domainName); @@ -404,12 +404,13 @@ public final class DomainCreateFlow implements MutatingFlow { DomainHistory domainHistory = buildDomainHistory(domain, tld, now, period, tld.getAddGracePeriodLength()); if (reservationTypes.contains(NAME_COLLISION)) { - entitiesToSave.add( + entitiesToInsert.add( createNameCollisionOneTimePollMessage(targetId, domainHistory, registrarId, now)); } - entitiesToSave.add(domain, domainHistory); + entitiesToInsert.add(domain, domainHistory); + ImmutableSet.Builder entitiesToUpdate = new ImmutableSet.Builder<>(); if (allocationToken.isPresent() && allocationToken.get().getTokenType().isOneTimeUse()) { - entitiesToSave.add( + entitiesToUpdate.add( AllocationTokenFlowUtils.redeemToken( allocationToken.get(), domainHistory.getHistoryEntryId())); } @@ -422,7 +423,10 @@ public final class DomainCreateFlow implements MutatingFlow { .setNewDomain(domain) .setHistoryEntry(domainHistory) .setEntityChanges( - EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build()) + EntityChanges.newBuilder() + .setInserts(entitiesToInsert.build()) + .setUpdates(entitiesToUpdate.build()) + .build()) .setYears(years) .build()); persistEntityChanges(entityChanges); diff --git a/core/src/main/java/google/registry/flows/domain/DomainDeleteFlow.java b/core/src/main/java/google/registry/flows/domain/DomainDeleteFlow.java index e398fb2fb..1b1285c65 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainDeleteFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainDeleteFlow.java @@ -151,7 +151,7 @@ public final class DomainDeleteFlow implements MutatingFlow, SqlStatementLogging verifyDeleteAllowed(existingDomain, tld, now); flowCustomLogic.afterValidation( AfterValidationParameters.newBuilder().setExistingDomain(existingDomain).build()); - ImmutableSet.Builder entitiesToSave = new ImmutableSet.Builder<>(); + ImmutableSet.Builder entitiesToInsert = new ImmutableSet.Builder<>(); Domain.Builder builder; if (existingDomain.getStatusValues().contains(StatusValue.PENDING_TRANSFER)) { builder = @@ -221,7 +221,7 @@ public final class DomainDeleteFlow implements MutatingFlow, SqlStatementLogging } else { PollMessage.OneTime deletePollMessage = createDeletePollMessage(existingDomain, domainHistoryId, deletionTime); - entitiesToSave.add(deletePollMessage); + entitiesToInsert.add(deletePollMessage); builder.setDeletePollMessage(deletePollMessage.createVKey()); } } @@ -230,7 +230,7 @@ public final class DomainDeleteFlow implements MutatingFlow, SqlStatementLogging // registrar other than the sponsoring registrar (which will necessarily be a superuser). if (durationUntilDelete.isLongerThan(Duration.ZERO) && !registrarId.equals(existingDomain.getPersistedCurrentSponsorRegistrarId())) { - entitiesToSave.add( + entitiesToInsert.add( createImmediateDeletePollMessage(existingDomain, domainHistoryId, now, deletionTime)); } @@ -239,7 +239,7 @@ public final class DomainDeleteFlow implements MutatingFlow, SqlStatementLogging for (GracePeriod gracePeriod : existingDomain.getGracePeriods()) { // No cancellation is written if the grace period was not for a billable event. if (gracePeriod.hasBillingEvent()) { - entitiesToSave.add( + entitiesToInsert.add( BillingCancellation.forGracePeriod(gracePeriod, now, domainHistoryId, targetId)); if (gracePeriod.getBillingEvent() != null) { // Take the amount of registration time being refunded off the expiration time. @@ -271,7 +271,7 @@ public final class DomainDeleteFlow implements MutatingFlow, SqlStatementLogging // ResourceDeleteFlow since it's listed in serverApproveEntities. requestDomainDnsRefresh(existingDomain.getDomainName()); - entitiesToSave.add(newDomain, domainHistory); + entitiesToInsert.add(domainHistory); EntityChanges entityChanges = flowCustomLogic.beforeSave( BeforeSaveParameters.newBuilder() @@ -279,7 +279,10 @@ public final class DomainDeleteFlow implements MutatingFlow, SqlStatementLogging .setNewDomain(newDomain) .setHistoryEntry(domainHistory) .setEntityChanges( - EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build()) + EntityChanges.newBuilder() + .setInserts(entitiesToInsert.build()) + .addUpdate(newDomain) + .build()) .build()); BeforeResponseReturnData responseData = flowCustomLogic.beforeResponse( diff --git a/core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java b/core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java index f92b4d372..0f75637f3 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java +++ b/core/src/main/java/google/registry/flows/domain/DomainFlowUtils.java @@ -711,7 +711,7 @@ public class DomainFlowUtils { BillingRecurrence newBillingRecurrence = existingBillingRecurrence.asBuilder().setRecurrenceEndTime(newEndTime).build(); - tm().put(newBillingRecurrence); + tm().update(newBillingRecurrence); return newBillingRecurrence; } diff --git a/core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java b/core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java index 04bcdff97..120316c00 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainRenewFlow.java @@ -245,11 +245,13 @@ public final class DomainRenewFlow implements MutatingFlow { .build(); DomainHistory domainHistory = buildDomainHistory(newDomain, now, command.getPeriod(), tld.getRenewGracePeriodLength()); - ImmutableSet.Builder entitiesToSave = new ImmutableSet.Builder<>(); - entitiesToSave.add( - newDomain, domainHistory, explicitRenewEvent, newAutorenewEvent, newAutorenewPollMessage); + ImmutableSet.Builder entitiesToInsert = new ImmutableSet.Builder<>(); + ImmutableSet.Builder entitiesToUpdate = new ImmutableSet.Builder<>(); + entitiesToInsert.add( + domainHistory, explicitRenewEvent, newAutorenewEvent, newAutorenewPollMessage); + entitiesToUpdate.add(newDomain); if (allocationToken.isPresent() && allocationToken.get().getTokenType().isOneTimeUse()) { - entitiesToSave.add( + entitiesToUpdate.add( AllocationTokenFlowUtils.redeemToken( allocationToken.get(), domainHistory.getHistoryEntryId())); } @@ -262,7 +264,10 @@ public final class DomainRenewFlow implements MutatingFlow { .setYears(years) .setHistoryEntry(domainHistory) .setEntityChanges( - EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build()) + EntityChanges.newBuilder() + .setInserts(entitiesToInsert.build()) + .setUpdates(entitiesToUpdate.build()) + .build()) .build()); BeforeResponseReturnData responseData = flowCustomLogic.beforeResponse( diff --git a/core/src/main/java/google/registry/flows/domain/DomainRestoreRequestFlow.java b/core/src/main/java/google/registry/flows/domain/DomainRestoreRequestFlow.java index cb36d82b6..29fc413d3 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainRestoreRequestFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainRestoreRequestFlow.java @@ -146,18 +146,18 @@ public final class DomainRestoreRequestFlow implements MutatingFlow { verifyRestoreAllowed(command, existingDomain, feeUpdate, feesAndCredits, now); HistoryEntryId domainHistoryId = createHistoryEntryId(existingDomain); historyBuilder.setRevisionId(domainHistoryId.getRevisionId()); - ImmutableSet.Builder entitiesToSave = new ImmutableSet.Builder<>(); + ImmutableSet.Builder entitiesToInsert = new ImmutableSet.Builder<>(); DateTime newExpirationTime = existingDomain.getRegistrationExpirationTime().plusYears(isExpired ? 1 : 0); // Restore the expiration time on the deleted domain, except if that's already passed, then add // a year and bill for it immediately, with no grace period. if (isExpired) { - entitiesToSave.add( + entitiesToInsert.add( createRenewBillingEvent(domainHistoryId, feesAndCredits.getRenewCost(), now)); } // Always bill for the restore itself. - entitiesToSave.add( + entitiesToInsert.add( createRestoreBillingEvent(domainHistoryId, feesAndCredits.getRestoreCost(), now)); BillingRecurrence autorenewEvent = @@ -166,12 +166,14 @@ public final class DomainRestoreRequestFlow implements MutatingFlow { .setRecurrenceEndTime(END_OF_TIME) .setDomainHistoryId(domainHistoryId) .build(); + entitiesToInsert.add(autorenewEvent); PollMessage.Autorenew autorenewPollMessage = newAutorenewPollMessage(existingDomain) .setEventTime(newExpirationTime) .setAutorenewEndTime(END_OF_TIME) .setDomainHistoryId(domainHistoryId) .build(); + entitiesToInsert.add(autorenewPollMessage); Domain newDomain = performRestore( existingDomain, @@ -181,8 +183,9 @@ public final class DomainRestoreRequestFlow implements MutatingFlow { now, registrarId); DomainHistory domainHistory = buildDomainHistory(newDomain, now); - entitiesToSave.add(newDomain, domainHistory, autorenewEvent, autorenewPollMessage); - tm().putAll(entitiesToSave.build()); + entitiesToInsert.add(domainHistory); + tm().update(newDomain); + tm().insertAll(entitiesToInsert.build()); if (existingDomain.getDeletePollMessage() != null) { tm().delete(existingDomain.getDeletePollMessage()); } diff --git a/core/src/main/java/google/registry/flows/domain/DomainTransferApproveFlow.java b/core/src/main/java/google/registry/flows/domain/DomainTransferApproveFlow.java index da8ef3e87..cea6f6568 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainTransferApproveFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainTransferApproveFlow.java @@ -172,7 +172,7 @@ public final class DomainTransferApproveFlow implements MutatingFlow { .setDomainHistoryId(domainHistoryId) .build()); - ImmutableList.Builder entitiesToSave = new ImmutableList.Builder<>(); + ImmutableList.Builder entitiesToInsert = new ImmutableList.Builder<>(); // If we are within an autorenew grace period, cancel the autorenew billing event and don't // increase the registration time, since the transfer subsumes the autorenew's extra year. GracePeriod autorenewGrace = @@ -184,7 +184,7 @@ public final class DomainTransferApproveFlow implements MutatingFlow { // then the gaining registrar is not charged for the one-year renewal and the losing registrar // still needs to be charged for the auto-renew. if (billingEvent.isPresent()) { - entitiesToSave.add( + entitiesToInsert.add( BillingCancellation.forGracePeriod(autorenewGrace, now, domainHistoryId, targetId)); } } @@ -259,14 +259,11 @@ public final class DomainTransferApproveFlow implements MutatingFlow { PollMessage gainingClientPollMessage = createGainingTransferPollMessage( targetId, newDomain.getTransferData(), newExpirationTime, now, domainHistoryId); - billingEvent.ifPresent(entitiesToSave::add); - entitiesToSave.add( - autorenewEvent, - gainingClientPollMessage, - gainingClientAutorenewPollMessage, - newDomain, - domainHistory); - tm().putAll(entitiesToSave.build()); + billingEvent.ifPresent(entitiesToInsert::add); + entitiesToInsert.add( + autorenewEvent, gainingClientPollMessage, gainingClientAutorenewPollMessage, domainHistory); + tm().update(newDomain); + tm().insertAll(entitiesToInsert.build()); // Delete the billing event and poll messages that were written in case the transfer would have // been implicitly server approved. tm().delete(existingDomain.getTransferData().getServerApproveEntities()); diff --git a/core/src/main/java/google/registry/flows/domain/DomainTransferCancelFlow.java b/core/src/main/java/google/registry/flows/domain/DomainTransferCancelFlow.java index 0dc29671c..01969dca9 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainTransferCancelFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainTransferCancelFlow.java @@ -110,8 +110,8 @@ public final class DomainTransferCancelFlow implements MutatingFlow { Domain newDomain = denyPendingTransfer(existingDomain, TransferStatus.CLIENT_CANCELLED, now, registrarId); DomainHistory domainHistory = buildDomainHistory(newDomain, tld, now); - tm().putAll( - newDomain, + tm().update(newDomain); + tm().insertAll( domainHistory, createLosingTransferPollMessage( targetId, newDomain.getTransferData(), null, domainHistoryId)); diff --git a/core/src/main/java/google/registry/flows/domain/DomainTransferRejectFlow.java b/core/src/main/java/google/registry/flows/domain/DomainTransferRejectFlow.java index b9c6d39ed..351ee5e18 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainTransferRejectFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainTransferRejectFlow.java @@ -109,8 +109,8 @@ public final class DomainTransferRejectFlow implements MutatingFlow { Domain newDomain = denyPendingTransfer(existingDomain, TransferStatus.CLIENT_REJECTED, now, registrarId); DomainHistory domainHistory = buildDomainHistory(newDomain, tld, now); - tm().putAll( - newDomain, + tm().update(newDomain); + tm().insertAll( domainHistory, createGainingTransferPollMessage( targetId, newDomain.getTransferData(), null, now, domainHistoryId)); diff --git a/core/src/main/java/google/registry/flows/domain/DomainTransferRequestFlow.java b/core/src/main/java/google/registry/flows/domain/DomainTransferRequestFlow.java index 71b05d99a..071d590d2 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainTransferRequestFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainTransferRequestFlow.java @@ -283,11 +283,9 @@ public final class DomainTransferRequestFlow implements MutatingFlow { asyncTaskEnqueuer.enqueueAsyncResave( newDomain.createVKey(), now, ImmutableSortedSet.of(automaticTransferTime)); - tm().putAll( - new ImmutableSet.Builder<>() - .add(newDomain, domainHistory, requestPollMessage) - .addAll(serverApproveEntities) - .build()); + tm().put(newDomain); + tm().putAll(serverApproveEntities); + tm().insertAll(domainHistory, requestPollMessage); return responseBuilder .setResultFromCode(SUCCESS_WITH_ACTION_PENDING) .setResData(createResponse(period, existingDomain, newDomain, now)) diff --git a/core/src/main/java/google/registry/flows/domain/DomainUpdateFlow.java b/core/src/main/java/google/registry/flows/domain/DomainUpdateFlow.java index 2c475bdae..f43ad5274 100644 --- a/core/src/main/java/google/registry/flows/domain/DomainUpdateFlow.java +++ b/core/src/main/java/google/registry/flows/domain/DomainUpdateFlow.java @@ -190,14 +190,16 @@ public final class DomainUpdateFlow implements MutatingFlow { if (requiresDnsUpdate(existingDomain, newDomain)) { requestDomainDnsRefresh(targetId); } - ImmutableSet.Builder entitiesToSave = new ImmutableSet.Builder<>(); - entitiesToSave.add(newDomain, domainHistory); + ImmutableSet.Builder entitiesToInsert = new ImmutableSet.Builder<>(); + ImmutableSet.Builder entitiesToUpdate = new ImmutableSet.Builder<>(); + entitiesToUpdate.add(newDomain); + entitiesToInsert.add(domainHistory); Optional statusUpdateBillingEvent = createBillingEventForStatusUpdates(existingDomain, newDomain, domainHistory, now); - statusUpdateBillingEvent.ifPresent(entitiesToSave::add); + statusUpdateBillingEvent.ifPresent(entitiesToInsert::add); Optional serverStatusUpdatePollMessage = createPollMessageForServerStatusUpdates(existingDomain, newDomain, domainHistory, now); - serverStatusUpdatePollMessage.ifPresent(entitiesToSave::add); + serverStatusUpdatePollMessage.ifPresent(entitiesToInsert::add); EntityChanges entityChanges = flowCustomLogic.beforeSave( BeforeSaveParameters.newBuilder() @@ -205,7 +207,10 @@ public final class DomainUpdateFlow implements MutatingFlow { .setNewDomain(newDomain) .setExistingDomain(existingDomain) .setEntityChanges( - EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build()) + EntityChanges.newBuilder() + .setInserts(entitiesToInsert.build()) + .setUpdates(entitiesToUpdate.build()) + .build()) .build()); persistEntityChanges(entityChanges); return responseBuilder.build(); diff --git a/core/src/main/java/google/registry/flows/host/HostCreateFlow.java b/core/src/main/java/google/registry/flows/host/HostCreateFlow.java index 8f1fbfb65..fc99a8287 100644 --- a/core/src/main/java/google/registry/flows/host/HostCreateFlow.java +++ b/core/src/main/java/google/registry/flows/host/HostCreateFlow.java @@ -126,7 +126,8 @@ public final class HostCreateFlow implements MutatingFlow { .setSuperordinateDomain(superordinateDomain.map(Domain::createVKey).orElse(null)) .build(); historyBuilder.setType(HOST_CREATE).setHost(newHost); - ImmutableSet entitiesToSave = ImmutableSet.of(newHost, historyBuilder.build()); + ImmutableSet entitiesToInsert = + ImmutableSet.of(newHost, historyBuilder.build()); if (superordinateDomain.isPresent()) { tm().update( superordinateDomain @@ -138,7 +139,7 @@ public final class HostCreateFlow implements MutatingFlow { // they are only written as NS records from the referencing domain. requestHostDnsRefresh(targetId); } - tm().insertAll(entitiesToSave); + tm().insertAll(entitiesToInsert); return responseBuilder.setResData(HostCreateData.create(targetId, now)).build(); } diff --git a/core/src/main/java/google/registry/flows/host/HostUpdateFlow.java b/core/src/main/java/google/registry/flows/host/HostUpdateFlow.java index 8e4433f5d..63088ccc4 100644 --- a/core/src/main/java/google/registry/flows/host/HostUpdateFlow.java +++ b/core/src/main/java/google/registry/flows/host/HostUpdateFlow.java @@ -52,7 +52,6 @@ import google.registry.flows.annotations.ReportingSpec; import google.registry.flows.exceptions.ResourceHasClientUpdateProhibitedException; import google.registry.model.EppResource; import google.registry.model.ForeignKeyUtils; -import google.registry.model.ImmutableObject; import google.registry.model.domain.Domain; import google.registry.model.domain.metadata.MetadataExtension; import google.registry.model.eppcommon.StatusValue; @@ -198,16 +197,12 @@ public final class HostUpdateFlow implements MutatingFlow { .setPersistedCurrentSponsorRegistrarId(newPersistedRegistrarId) .build(); verifyHasIpsIffIsExternal(command, existingHost, newHost); - ImmutableSet.Builder entitiesToInsert = new ImmutableSet.Builder<>(); - ImmutableSet.Builder entitiesToUpdate = new ImmutableSet.Builder<>(); - entitiesToUpdate.add(newHost); if (isHostRename) { updateSuperordinateDomains(existingHost, newHost); } enqueueTasks(existingHost, newHost); - entitiesToInsert.add(historyBuilder.setType(HOST_UPDATE).setHost(newHost).build()); - tm().updateAll(entitiesToUpdate.build()); - tm().insertAll(entitiesToInsert.build()); + tm().update(newHost); + tm().insert(historyBuilder.setType(HOST_UPDATE).setHost(newHost).build()); return responseBuilder.build(); } @@ -290,7 +285,7 @@ public final class HostUpdateFlow implements MutatingFlow { && newHost.isSubordinate() && Objects.equals( existingHost.getSuperordinateDomain(), newHost.getSuperordinateDomain())) { - tm().put( + tm().update( tm().loadByKey(existingHost.getSuperordinateDomain()) .asBuilder() .removeSubordinateHost(existingHost.getHostName()) @@ -299,14 +294,14 @@ public final class HostUpdateFlow implements MutatingFlow { return; } if (existingHost.isSubordinate()) { - tm().put( + tm().update( tm().loadByKey(existingHost.getSuperordinateDomain()) .asBuilder() .removeSubordinateHost(existingHost.getHostName()) .build()); } if (newHost.isSubordinate()) { - tm().put( + tm().update( tm().loadByKey(newHost.getSuperordinateDomain()) .asBuilder() .addSubordinateHost(newHost.getHostName()) diff --git a/core/src/test/java/google/registry/flows/custom/TestDomainCreateFlowCustomLogic.java b/core/src/test/java/google/registry/flows/custom/TestDomainCreateFlowCustomLogic.java index bbd184b10..fac1f64d9 100644 --- a/core/src/test/java/google/registry/flows/custom/TestDomainCreateFlowCustomLogic.java +++ b/core/src/test/java/google/registry/flows/custom/TestDomainCreateFlowCustomLogic.java @@ -40,8 +40,9 @@ public class TestDomainCreateFlowCustomLogic extends DomainCreateFlowCustomLogic .setMsg("Custom logic was triggered") .build(); return EntityChanges.newBuilder() - .setSaves(parameters.entityChanges().getSaves()) - .addSave(extraPollMessage) + .setInserts(parameters.entityChanges().getInserts()) + .addInsert(extraPollMessage) + .setUpdates(parameters.entityChanges().getUpdates()) .setDeletes(parameters.entityChanges().getDeletes()) .build(); }