From 17839aa473f3d2bb7b0749b047053f0a52afd3ee Mon Sep 17 00:00:00 2001
From: andres-portainer <91705312+andres-portainer@users.noreply.github.com>
Date: Thu, 30 Mar 2023 23:16:56 -0300
Subject: [PATCH] fix(endpointrelation): change a callback so it is
 transactional EE-5312 (#8729)

---
 api/dataservices/edgestack/edgestack.go       |  5 +++++
 api/dataservices/edgestack/tx.go              | 21 +++++++++++++++++--
 .../endpointrelation/endpointrelation.go      | 11 +++++++---
 api/dataservices/endpointrelation/tx.go       |  2 +-
 api/datastore/services.go                     |  2 +-
 5 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/api/dataservices/edgestack/edgestack.go b/api/dataservices/edgestack/edgestack.go
index c4a66c51b..d754ac025 100644
--- a/api/dataservices/edgestack/edgestack.go
+++ b/api/dataservices/edgestack/edgestack.go
@@ -159,6 +159,11 @@ func (service *Service) UpdateEdgeStackFunc(ID portainer.EdgeStackID, updateFunc
 	})
 }
 
+// UpdateEdgeStackFuncTx is a helper function used to call UpdateEdgeStackFunc inside a transaction.
+func (service *Service) UpdateEdgeStackFuncTx(tx portainer.Transaction, ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error {
+	return service.Tx(tx).UpdateEdgeStackFunc(ID, updateFunc)
+}
+
 // DeleteEdgeStack deletes an Edge stack.
 func (service *Service) DeleteEdgeStack(ID portainer.EdgeStackID) error {
 	service.mu.Lock()
diff --git a/api/dataservices/edgestack/tx.go b/api/dataservices/edgestack/tx.go
index 53640eea1..35301406d 100644
--- a/api/dataservices/edgestack/tx.go
+++ b/api/dataservices/edgestack/tx.go
@@ -1,7 +1,6 @@
 package edgestack
 
 import (
-	"errors"
 	"fmt"
 
 	portainer "github.com/portainer/portainer/api"
@@ -103,7 +102,25 @@ func (service ServiceTx) UpdateEdgeStack(ID portainer.EdgeStackID, edgeStack *po
 
 // UpdateEdgeStackFunc is a no-op inside a transaction.
 func (service ServiceTx) UpdateEdgeStackFunc(ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error {
-	return errors.New("cannot be called inside a transaction")
+	service.service.mu.Lock()
+	defer service.service.mu.Unlock()
+
+	edgeStack, err := service.EdgeStack(ID)
+	if err != nil {
+		return err
+	}
+
+	updateFunc(edgeStack)
+
+	err = service.UpdateEdgeStack(ID, edgeStack)
+	if err != nil {
+		return err
+	}
+
+	service.service.idxVersion[ID] = edgeStack.Version
+	service.service.cacheInvalidationFn(ID)
+
+	return nil
 }
 
 // DeleteEdgeStack deletes an Edge stack.
diff --git a/api/dataservices/endpointrelation/endpointrelation.go b/api/dataservices/endpointrelation/endpointrelation.go
index b8ecfaa4b..99cc14414 100644
--- a/api/dataservices/endpointrelation/endpointrelation.go
+++ b/api/dataservices/endpointrelation/endpointrelation.go
@@ -14,16 +14,21 @@ const BucketName = "endpoint_relations"
 
 // Service represents a service for managing environment(endpoint) relation data.
 type Service struct {
-	connection    portainer.Connection
-	updateStackFn func(ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error
+	connection      portainer.Connection
+	updateStackFn   func(ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error
+	updateStackFnTx func(tx portainer.Transaction, ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error
 }
 
 func (service *Service) BucketName() string {
 	return BucketName
 }
 
-func (service *Service) RegisterUpdateStackFunction(updateFunc func(ID portainer.EdgeStackID, updateFunc func(edgeStack *portainer.EdgeStack)) error) {
+func (service *Service) RegisterUpdateStackFunction(
+	updateFunc func(portainer.EdgeStackID, func(*portainer.EdgeStack)) error,
+	updateFuncTx func(portainer.Transaction, portainer.EdgeStackID, func(*portainer.EdgeStack)) error,
+) {
 	service.updateStackFn = updateFunc
+	service.updateStackFnTx = updateFuncTx
 }
 
 // NewService creates a new instance of a service.
diff --git a/api/dataservices/endpointrelation/tx.go b/api/dataservices/endpointrelation/tx.go
index 1da9f6b92..95e75f036 100644
--- a/api/dataservices/endpointrelation/tx.go
+++ b/api/dataservices/endpointrelation/tx.go
@@ -151,7 +151,7 @@ func (service ServiceTx) updateEdgeStacksAfterRelationChange(previousRelationSta
 				}
 			}
 
-			service.service.updateStackFn(refStackId, func(edgeStack *portainer.EdgeStack) {
+			service.service.updateStackFnTx(service.tx, refStackId, func(edgeStack *portainer.EdgeStack) {
 				edgeStack.NumDeployments = numDeployments
 			})
 		}
diff --git a/api/datastore/services.go b/api/datastore/services.go
index d186493fe..680885b7c 100644
--- a/api/datastore/services.go
+++ b/api/datastore/services.go
@@ -104,7 +104,7 @@ func (store *Store) initServices() error {
 		return err
 	}
 	store.EdgeStackService = edgeStackService
-	endpointRelationService.RegisterUpdateStackFunction(edgeStackService.UpdateEdgeStackFunc)
+	endpointRelationService.RegisterUpdateStackFunction(edgeStackService.UpdateEdgeStackFunc, edgeStackService.UpdateEdgeStackFuncTx)
 
 	edgeGroupService, err := edgegroup.NewService(store.connection)
 	if err != nil {