Working version contained within template model

assets/javascripts/discourse/templates/components/wizard-custom-field.hbs

@@ -44,7 +44,7 @@
       imageUrl=field.image
       onUploadDone=(action "imageUploadDone")
       onUploadDeleted=(action "imageUploadDeleted")
-      type=(concat "wizard-" wizard.id)
+      type="wizard-field-image"
       class="no-repeat contain-image"
       id=(concat "wizard-field-" field.id "-image-upload")}}
   </div>

assets/javascripts/discourse/templates/components/wizard-custom-step.hbs

@@ -18,7 +18,7 @@
       imageUrl=step.banner
       onUploadDone=(action "bannerUploadDone")
       onUploadDeleted=(action "bannerUploadDeleted")
-      type=(concat "wizard-" wizard.id)
+      type="wizard-step-banner"
       class="no-repeat contain-image"
       id=(concat "wizard-step-" step.id "-banner-upload")}}
   </div>

lib/custom_wizard/template.rb

@@ -29,10 +29,10 @@ class CustomWizard::Template
     ActiveRecord::Base.transaction do
       schedule_save_jobs unless opts[:skip_jobs]
       PluginStore.set(CustomWizard::PLUGIN_NAME, @data[:id], @data)
+      ensure_wizard_upload_references!
     end
 
     self.class.clear_cache_keys
-    remove_unused_wizard_upload_references
 
     @data[:id]
   end
@@ -53,17 +53,21 @@ class CustomWizard::Template
     PluginStore.get(CustomWizard::PLUGIN_NAME, wizard_id)
   end
 
+  def self.find_record(wizard_id)
+    PluginStoreRow.find_by(plugin_name: CustomWizard::PLUGIN_NAME, key: wizard_id)
+  end
+
   def self.remove(wizard_id)
     wizard = CustomWizard::Wizard.create(wizard_id)
     return false if !wizard
 
     ActiveRecord::Base.transaction do
+      ensure_wizard_upload_references!(wizard_id)
       PluginStore.remove(CustomWizard::PLUGIN_NAME, wizard.id)
       clear_user_wizard_redirect(wizard_id, after_time: !!wizard.after_time)
     end
 
     clear_cache_keys
-    remove_wizard_upload_references
 
     true
   end
@@ -125,6 +129,18 @@ class CustomWizard::Template
     CustomWizard::Cache.new(AFTER_TIME_CACHE_KEY).delete
   end
 
+  def self.ensure_wizard_upload_references!(wizard_id, wizard_upload_ids = [])
+    wizard_record = find_record(wizard_id)
+
+    if wizard_record
+      UploadReference.ensure_exist!(
+        upload_ids: wizard_upload_ids,
+        target_type: "PluginStoreRow",
+        target_id: wizard_record.id
+      )
+    end
+  end
+
   private
 
   def normalize_data
@@ -179,7 +195,7 @@ class CustomWizard::Template
     end
   end
 
-  def collect_upload_ids
+  def ensure_wizard_upload_references!
     upload_ids = []
 
     @data[:steps].each do |step|
@@ -190,21 +206,7 @@ class CustomWizard::Template
       end
     end
 
-    upload_ids
+    upload_ids = upload_ids.select { |upload_id| Upload.exists?(upload_id) }
-  end
+    self.class.ensure_wizard_upload_references!(@data[:id], upload_ids)
-
-  def wizard_upload_references
-    @wizard_upload_references ||= begin
-      record = PluginStoreRow.find_by(plugin_name: CustomWizard::PLUGIN_NAME, key: @data[:id])
-      UploadReference.where(target_type: "CustomWizard", target_id: custom_wizard_record.id)
-    end
-  end
-
-  def remove_unused_wizard_upload_references
-    wizard_upload_references.where.not(upload_id: collect_upload_ids).delete_all
-  end
-
-  def remove_wizard_upload_references
-    wizard_upload_references.delete_all
   end
 end

plugin.rb

@@ -226,21 +226,5 @@ after_initialize do
     ::DiscourseTagging.singleton_class.prepend CustomWizardDiscourseTagging
   end
 
-  on(:after_upload_creation) do |upload, opts|
-    from_wizard = opts[:type].include?("wizard-")
-    wizard_id = opts[:type].split("wizard-").last
-    wizard_record = PluginStoreRow.find_by(plugin_name: CustomWizard::PLUGIN_NAME, key: wizard_id)
-
-    if wizard_record
-      UploadReference.create(
-        upload_id: upload.id,
-        target_type: "CustomWizard",
-        target_id: wizard_record.id,
-        created_at: Time.zone.now,
-        updated_at: Time.zone.now
-      )
-    end
-  end
-
   DiscourseEvent.trigger(:custom_wizard_ready)
 end

spec/components/custom_wizard/template_spec.rb

@@ -2,6 +2,7 @@
 
 describe CustomWizard::Template do
   fab!(:user) { Fabricate(:user) }
+  fab!(:upload) { Fabricate(:upload) }
 
   let(:template_json) {
     JSON.parse(File.open(
@@ -54,6 +55,74 @@ describe CustomWizard::Template do
     ).to eq(true)
   end
 
+  context "upload references" do
+    it "are added if a wizard has a step banner" do
+      template_json['steps'][0]['banner'] = upload.url
+      template_json['steps'][0]['banner_upload_id'] = upload.id
+      CustomWizard::Template.save(template_json, skip_jobs: true)
+      wizard_record = CustomWizard::Template.find_record(template_json["id"])
+      expect(
+        UploadReference.exists?(
+          upload_id: upload.id,
+          target_type: "PluginStoreRow",
+          target_id: wizard_record.id
+        )
+      ).to eq(true)
+    end
+
+    it "are added if a wizard has a field image" do
+      template_json['steps'][0]["fields"][0]['image'] = upload.url
+      template_json['steps'][0]["fields"][0]['image_upload_id'] = upload.id
+      CustomWizard::Template.save(template_json, skip_jobs: true)
+      wizard_record = CustomWizard::Template.find_record(template_json["id"])
+      expect(
+        UploadReference.exists?(
+          upload_id: upload.id,
+          target_type: "PluginStoreRow",
+          target_id: wizard_record.id
+        )
+      ).to eq(true)
+    end
+
+    it "are removed if a wizard step banner is removed" do
+      template_json['steps'][0]['banner'] = upload.url
+      template_json['steps'][0]['banner_upload_id'] = upload.id
+      CustomWizard::Template.save(template_json, skip_jobs: true)
+
+      template_json['steps'][0]['banner'] = nil
+      template_json['steps'][0]['banner_upload_id'] = nil
+      CustomWizard::Template.save(template_json, skip_jobs: true)
+      wizard_record = CustomWizard::Template.find_record(template_json["id"])
+      expect(
+        UploadReference.exists?(target_type: "PluginStoreRow")
+      ).to eq(false)
+    end
+
+    it "are removed if a wizard field image is removed" do
+      template_json['steps'][0]["fields"][0]['image'] = upload.url
+      template_json['steps'][0]["fields"][0]['image_upload_id'] = upload.id
+      CustomWizard::Template.save(template_json, skip_jobs: true)
+
+      template_json['steps'][0]["fields"][0]['image'] = nil
+      template_json['steps'][0]["fields"][0]['image_upload_id'] = nil
+      CustomWizard::Template.save(template_json, skip_jobs: true)
+      wizard_record = CustomWizard::Template.find_record(template_json["id"])
+      expect(
+        UploadReference.exists?(target_type: "PluginStoreRow")
+      ).to eq(false)
+    end
+
+    it "are removed if a wizard is removed" do
+      template_json['steps'][0]["fields"][0]['image'] = upload.url
+      template_json['steps'][0]["fields"][0]['image_upload_id'] = upload.id
+      CustomWizard::Template.save(template_json, skip_jobs: true)
+      CustomWizard::Template.remove(template_json["id"])
+      expect(
+        UploadReference.exists?(target_type: "PluginStoreRow")
+      ).to eq(false)
+    end
+  end
+
   context "wizard template list" do
     before do
       template_json_2 = template_json.dup