From 0a7756c7461415d6b0fbda163f433313bc7e4010 Mon Sep 17 00:00:00 2001
From: Daria Mayorova <mayorova@users.noreply.github.com>
Date: Mon, 30 Sep 2024 13:23:41 +0200
Subject: [PATCH] Improvements for orphan objects removal

---
 lib/tasks/multitenant/tenants.rake | 31 +++++++++++++++++++++---------
 1 file changed, 22 insertions(+), 9 deletions(-)

diff --git a/lib/tasks/multitenant/tenants.rake b/lib/tasks/multitenant/tenants.rake
index 71bdd66913..5f47ae7d83 100644
--- a/lib/tasks/multitenant/tenants.rake
+++ b/lib/tasks/multitenant/tenants.rake
@@ -86,19 +86,21 @@ namespace :multitenant do
 
       provider_account_ids = Account.where(provider: true).pluck(:id) + [Account.master.id]
 
-      ApplicationRecord.descendants.each do |model|
-        next unless model.table_exists? && model.column_names.include?('tenant_id')
-
+      base_models.each do |model|
         orphaned_objects = model.where.not(tenant_id: provider_account_ids)
 
         if orphaned_objects.exists?
-          puts "Found orphaned objects for model #{model.name}:"
-          orphaned_objects.find_each { |obj| puts "- ID: #{obj.id}, Tenant ID: #{obj.tenant_id}" }
-
-          if destroy
-            puts "Destroying orphaned objects for model #{model.name}..."
-            orphaned_objects.in_batches(of: 100).destroy_all
+          puts "Found #{orphaned_objects.size} orphaned objects for model #{model.name}:"
+          orphaned_objects.find_in_batches(batch_size: 100).with_index do |batch, index|
+            puts "Processing batch #{index+1} of model #{model.name}..."
+            batch.each do |object|
+              puts "- ID: #{object.id}, Tenant ID: #{object.tenant_id}"
+              object.destroy if destroy
+            end
           end
+          orphaned_objects.find_each { |obj|  }
+
+
         else
           puts "No orphaned objects found for model #{model.name}."
         end
@@ -107,6 +109,17 @@ namespace :multitenant do
       puts 'Orphaned objects check completed.'
     end
 
+    def base_models
+      all_models = ApplicationRecord.descendants.select(&:arel_table).reject(&:abstract_class?)
+      all_models.select! { _1.attribute_names.include? "tenant_id"}
+      # we only want base STI classes, not the children
+      all_models.select do |model|
+        base_class = model.base_class
+        # either current model is the base_class or we can't find a base class amongst the discovered models (which would be very weird)
+        base_class == model || all_models.none? { |potential_parent| potential_parent == base_class }
+      end
+    end
+
     def update_tenant_ids(tenant_id_block, association_block, condition, **args)
       query = args[:table_name].constantize.joining(&association_block).where.has(&condition)
       puts "------ Updating #{args[:table_name]} ------"