lib/arel/extensions/update_manager.rb in arel_toolkit-0.3.0 vs lib/arel/extensions/update_manager.rb in arel_toolkit-0.4.0

@@ -1,9 +1,25 @@
-Arel::UpdateManager.class_eval do
-  def ==(other)
-    @ast == other.ast && @ctx == other.ctx
-  end
+# rubocop:disable Naming/MethodName
+# rubocop:disable Naming/UncommunicativeMethodParamName
 
-  protected
+module Arel
+  class UpdateManager < Arel::TreeManager
+    def ==(other)
+      other.is_a?(self.class) && @ast == other.ast && @ctx == other.ctx
+    end
 
-  attr_reader :ctx
+    protected
+
+    attr_reader :ctx
+  end
+
+  module Visitors
+    class Dot
+      def visit_Arel_UpdateManager(o)
+        visit_edge o, 'ast'
+      end
+    end
+  end
 end
+
+# rubocop:enable Naming/MethodName
+# rubocop:enable Naming/UncommunicativeMethodParamName