From cfcb80f64235e665fd82b66847e63caa1d0d496b Mon Sep 17 00:00:00 2001
From: Martin Artnik <martin@scinote.net>
Date: Mon, 6 Mar 2023 16:06:38 +0100
Subject: [PATCH] Fix assignment sync to new protocol drafts and deletion
 [SCI-8060]

---
 .../generate_user_assignments_job.rb          | 10 ++++
 app/models/concerns/assignable.rb             | 10 ++--
 app/models/protocol.rb                        | 60 ++++++++++++-------
 app/models/user_assignment.rb                 |  7 ++-
 4 files changed, 56 insertions(+), 31 deletions(-)

diff --git a/app/jobs/user_assignments/generate_user_assignments_job.rb b/app/jobs/user_assignments/generate_user_assignments_job.rb
index 884a940c2..dd49489e2 100644
--- a/app/jobs/user_assignments/generate_user_assignments_job.rb
+++ b/app/jobs/user_assignments/generate_user_assignments_job.rb
@@ -46,6 +46,16 @@ module UserAssignments
     end
 
     def assign_users_to_protocol(protocol)
+      if protocol.parent_id && protocol.in_repository_draft?
+        Protocol.transaction(requires_new: true) do
+          protocol.parent.user_assignments.find_each do |user_assignment|
+            protocol.parent.sync_child_protocol_user_assignment(user_assignment, protocol.id)
+          end
+        end
+
+        return
+      end
+
       return unless protocol.visible?
 
       protocol.create_public_user_assignments!(@assigned_by)
diff --git a/app/models/concerns/assignable.rb b/app/models/concerns/assignable.rb
index d091c5a43..7f70e4020 100644
--- a/app/models/concerns/assignable.rb
+++ b/app/models/concerns/assignable.rb
@@ -58,13 +58,13 @@ module Assignable
       self.class.name.in?(Extends::TOP_LEVEL_ASSIGNABLES)
     end
 
-    def after_user_assignment_save
-      # Optional, redefine in the assignable model.
-      # Will be called when an assignment is saved for the assignable model.
-    end
-
     private
 
+    def after_user_assignment_changed(user_assignment = nil)
+      # Optional, redefine in the assignable model.
+      # Will be called when an assignment is changed (save/destroy) for the assignable model.
+    end
+
     def create_users_assignments
       return if skip_user_assignments
 
diff --git a/app/models/protocol.rb b/app/models/protocol.rb
index e269bd525..4dbbb9c56 100644
--- a/app/models/protocol.rb
+++ b/app/models/protocol.rb
@@ -17,7 +17,6 @@ class Protocol < ApplicationRecord
   include TinyMceImages
 
   after_create :auto_assign_protocol_members, if: :visible?
-  after_create :sync_child_protocol_user_assignments, unless: -> { parent_id }
   after_destroy :decrement_linked_children
   after_save :update_user_assignments, if: -> { saved_change_to_visibility? && in_repository? }
   after_save :update_linked_children
@@ -685,33 +684,48 @@ class Protocol < ApplicationRecord
     end
   end
 
-  def after_user_assignment_save
-    sync_child_protocol_user_assignments
+  def child_version_protocols
+    published_versions.or(Protocol.where(id: draft&.id))
+  end
+
+  def sync_child_protocol_user_assignment(user_assignment, child_protocol_id = nil)
+    # Copy user assignments to child protocol(s)
+
+    Protocol.transaction(requires_new: true) do
+      # Reload to ensure a potential new draft is also included in child versions
+      reload
+
+      (
+        # all or single child version protocol
+        child_protocol_id ? child_version_protocols.where(id: child_protocol_id) : child_version_protocols
+      ).find_each do |child_protocol|
+        child_assignment = child_protocol.user_assignments.find_or_initialize_by(
+          user: user_assignment.user
+        )
+
+        if user_assignment.destroyed?
+          child_assignment.destroy! if child_assignment.persisted?
+          next
+        end
+
+        child_assignment.update!(
+          user_assignment.attributes.slice(
+            'user_role_id',
+            'assigned',
+            'assigned_by_id',
+            'team_id'
+          )
+        )
+      end
+    end
   end
 
   private
 
-  def sync_child_protocol_user_assignments
-    # Copy user assignments to child protocols
+  def after_user_assignment_changed(user_assignment)
+    return unless in_repository_published_original?
 
-    return if parent_id
-
-    Protocol.transaction do
-      user_assignments.find_each do |user_assignment|
-        linked_children.find_each do |child_protocol|
-          child_protocol.user_assignments.find_or_initialize_by(
-            user: user_assignment.user
-          ).update!(
-            user_assignment.attributes.slice(
-              'user_role_id',
-              'assigned',
-              'assigned_by_id',
-              'team_id'
-            )
-          )
-        end
-      end
-    end
+    sync_child_protocol_user_assignment(user_assignment)
   end
 
   def auto_assign_protocol_members
diff --git a/app/models/user_assignment.rb b/app/models/user_assignment.rb
index 921669ba5..9b54a5479 100644
--- a/app/models/user_assignment.rb
+++ b/app/models/user_assignment.rb
@@ -7,7 +7,8 @@ class UserAssignment < ApplicationRecord
   after_create :assign_team_child_objects, if: -> { assignable.is_a?(Team) }
   after_update :update_team_children_assignments, if: -> { assignable.is_a?(Team) && saved_change_to_user_role_id? }
   before_destroy :unassign_team_child_objects, if: -> { assignable.is_a?(Team) }
-  after_save :call_user_assignment_save_hook
+  after_destroy :call_user_assignment_changed_hook
+  after_save :call_user_assignment_changed_hook
 
   belongs_to :assignable, polymorphic: true, touch: true
   belongs_to :user_role
@@ -25,8 +26,8 @@ class UserAssignment < ApplicationRecord
 
   private
 
-  def call_user_assignment_save_hook
-    assignable.after_user_assignment_save
+  def call_user_assignment_changed_hook
+    assignable.__send__(:after_user_assignment_changed, self)
   end
 
   def assign_team_child_objects