diff --git a/lib/activerecord-multi-tenant/arel_visitors_depth_first.rb b/lib/activerecord-multi-tenant/arel_visitors_depth_first.rb
index a0c0f3b2..89fb875d 100644
--- a/lib/activerecord-multi-tenant/arel_visitors_depth_first.rb
+++ b/lib/activerecord-multi-tenant/arel_visitors_depth_first.rb
@@ -39,7 +39,7 @@ def unary(obj)
 
     def function(obj)
       visit obj.expressions
-      visit obj.alias
+      visit obj.alias if obj.respond_to?(:alias)
       visit obj.distinct
     end
     alias visit_Arel_Nodes_Avg    function
@@ -54,12 +54,12 @@ def visit_Arel_Nodes_NamedFunction(obj)
       visit obj.name
       visit obj.expressions
       visit obj.distinct
-      visit obj.alias
+      visit obj.alias if obj.respond_to?(:alias)
     end
 
     def visit_Arel_Nodes_Count(obj)
       visit obj.expressions
-      visit obj.alias
+      visit obj.alias if obj.respond_to?(:alias)
       visit obj.distinct
     end
 
diff --git a/lib/activerecord-multi-tenant/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb
index bc165218..598a1414 100644
--- a/lib/activerecord-multi-tenant/migrations.rb
+++ b/lib/activerecord-multi-tenant/migrations.rb
@@ -82,14 +82,12 @@ def create_table(table_name, options = {}, &block)
 end
 ActiveRecord::ConnectionAdapters::SchemaStatements.prepend(MultiTenant::SchemaStatementsExtensions)
 
-module ActiveRecord
-  class SchemaDumper
+module MultiTenant
+  module SchemaDumper
     private
 
-    alias initialize_without_citus initialize
-
     def initialize(connection, options = {})
-      initialize_without_citus(connection, options)
+      super
 
       citus_version =
         begin
@@ -114,10 +112,8 @@ def initialize(connection, options = {})
     end
 
     # Support for create_distributed_table & create_reference_table
-    alias table_without_citus table
-
     def table(table, stream)
-      table_without_citus(table, stream)
+      super
       table_name = remove_prefix_and_suffix(table)
       distribution_column = @distribution_columns[table_name]
       if distribution_column
@@ -130,3 +126,5 @@ def table(table, stream)
     end
   end
 end
+
+ActiveRecord::SchemaDumper.prepend(MultiTenant::SchemaDumper)
diff --git a/lib/activerecord-multi-tenant/relation_extension.rb b/lib/activerecord-multi-tenant/relation_extension.rb
index 2e3d5e3c..731bb263 100644
--- a/lib/activerecord-multi-tenant/relation_extension.rb
+++ b/lib/activerecord-multi-tenant/relation_extension.rb
@@ -5,7 +5,7 @@ module ActiveRecordRelationExtension
     # Overrides the delete_all method to include tenant scoping
     def delete_all
       # Call the original delete_all method if the current tenant is identified by an ID
-      return super if MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil?
+      return super if MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil? || primary_key.nil?
 
       stmt = Arel::DeleteManager.new.from(table)
       stmt.wheres = [generate_in_condition_subquery]
@@ -21,7 +21,23 @@ def update_all(updates)
 
       stmt = Arel::UpdateManager.new
       stmt.table(table)
-      stmt.set Arel.sql(klass.send(:sanitize_sql_for_assignment, updates))
+
+      # Handle updates differently based on whether it's a hash with Arel nodes or regular values
+      if updates.is_a?(Hash) && updates.values.any? { |value| value.is_a?(Arel::Nodes::Node) }
+        # For hash with Arel nodes, build assignments directly
+        assignments = updates.map do |column, value|
+          if value.is_a?(Arel::Nodes::Node)
+            [table[column], value]
+          else
+            [table[column], Arel::Nodes.build_quoted(value, table[column])]
+          end
+        end
+        stmt.set assignments
+      else
+        # For regular updates, use sanitize_sql_for_assignment
+        stmt.set Arel.sql(klass.send(:sanitize_sql_for_assignment, updates))
+      end
+
       stmt.wheres = [generate_in_condition_subquery]
 
       klass.connection.update(stmt, "#{klass} Update All").tap { reset }
diff --git a/lib/activerecord-multi-tenant/version.rb b/lib/activerecord-multi-tenant/version.rb
index 20299e68..dd2e281b 100644
--- a/lib/activerecord-multi-tenant/version.rb
+++ b/lib/activerecord-multi-tenant/version.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module MultiTenant
-  VERSION = '2.4.0'
+  VERSION = '2.4.2'.freeze
 end