From 1771c305655fb2900904d281be071fdd6d730dc2 Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Fri, 10 Oct 2025 14:59:27 -0400 Subject: [PATCH] Differentiate between inserts and updates in flows 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 c3c2b242589..8318b27aee1 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 static void validateRegistrarIsLoggedIn(String registrarId) throws EppExc } } - /** 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 00cac1567d3..87a9400c2bb 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 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 static Builder newBuilder() { @AutoBuilder public interface Builder { - Builder setSaves(ImmutableSet entitiesToSave); + Builder setInserts(ImmutableSet entitiesToInsert); + + ImmutableSet.Builder insertsBuilder(); + + default Builder addInsert(ImmutableObject entityToInsert) { + insertsBuilder().add(entityToInsert); + return this; + } + + Builder setUpdates(ImmutableSet entitiesToUpdate); - ImmutableSet.Builder savesBuilder(); + ImmutableSet.Builder updatesBuilder(); - default Builder addSave(ImmutableObject entityToSave) { - savesBuilder().add(entityToSave); + 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 e06e6f4ed46..e8cc745cf53 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 EppResponse run() throws EppException { 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 EppResponse run() throws EppException { 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 EppResponse run() throws EppException { .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 e398fb2fb2c..1b1285c651c 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 EppResponse run() throws EppException { 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 EppResponse run() throws EppException { } else { PollMessage.OneTime deletePollMessage = createDeletePollMessage(existingDomain, domainHistoryId, deletionTime); - entitiesToSave.add(deletePollMessage); + entitiesToInsert.add(deletePollMessage); builder.setDeletePollMessage(deletePollMessage.createVKey()); } } @@ -230,7 +230,7 @@ public EppResponse run() throws EppException { // 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 EppResponse run() throws EppException { 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 EppResponse run() throws EppException { // 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 EppResponse run() throws EppException { .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 f92b4d372ae..0f75637f352 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 static BillingRecurrence updateAutorenewRecurrenceEndTime( 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 04bcdff97e3..120316c004f 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 EppResponse run() throws EppException { .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 EppResponse run() throws EppException { .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 cb36d82b62f..29fc413d351 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 EppResponse run() throws EppException { 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 EppResponse run() throws EppException { .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 EppResponse run() throws EppException { 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 da8ef3e8774..cea6f656825 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 EppResponse run() throws EppException { .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 EppResponse run() throws EppException { // 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 EppResponse run() throws EppException { 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 0dc29671c64..01969dca9d2 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 EppResponse run() throws EppException { 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 b9c6d39ed62..351ee5e18fd 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 EppResponse run() throws EppException { 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 71b05d99a5a..071d590d2e1 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 EppResponse run() throws EppException { 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 2c475bdaeae..f43ad52749c 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 EppResponse run() throws EppException { 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 EppResponse run() throws EppException { .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 8f1fbfb651f..fc99a8287a1 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 EppResponse run() throws EppException { .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 EppResponse run() throws EppException { // 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 8e4433f5d6b..63088ccc4c2 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.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 EppResponse run() throws EppException { .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 @@ private static void updateSuperordinateDomains(Host existingHost, Host newHost) && newHost.isSubordinate() && Objects.equals( existingHost.getSuperordinateDomain(), newHost.getSuperordinateDomain())) { - tm().put( + tm().update( tm().loadByKey(existingHost.getSuperordinateDomain()) .asBuilder() .removeSubordinateHost(existingHost.getHostName()) @@ -299,14 +294,14 @@ private static void updateSuperordinateDomains(Host existingHost, Host newHost) 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 bbd184b1052..fac1f64d995 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 EntityChanges beforeSave(BeforeSaveParameters parameters) { .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(); }