with_recursive.rb 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. # frozen_string_literal: true
  2. # Add support for writing recursive CTEs in ActiveRecord
  3. # Initially from Lorin Thwaits (https://github.com/lorint) as per comment:
  4. # https://github.com/vlado/activerecord-cte/issues/16#issuecomment-1433043310
  5. # Modified from the above code to change the signature to
  6. # `with_recursive(hash)` and extending CTE hash values to also includes arrays
  7. # of values that get turned into UNION ALL expressions.
  8. # This implementation has been merged in Rails: https://github.com/rails/rails/pull/51601
  9. module ActiveRecord
  10. module QueryMethodsExtensions
  11. def with_recursive(*args)
  12. @with_is_recursive = true
  13. check_if_method_has_arguments!(__callee__, args)
  14. spawn.with_recursive!(*args)
  15. end
  16. # Like #with_recursive but modifies the relation in place.
  17. def with_recursive!(*args) # :nodoc:
  18. self.with_values += args
  19. @with_is_recursive = true
  20. self
  21. end
  22. private
  23. def build_with(arel)
  24. return if with_values.empty?
  25. with_statements = with_values.map do |with_value|
  26. raise ArgumentError, "Unsupported argument type: #{with_value} #{with_value.class}" unless with_value.is_a?(Hash)
  27. build_with_value_from_hash(with_value)
  28. end
  29. # Was: arel.with(with_statements)
  30. @with_is_recursive ? arel.with(:recursive, with_statements) : arel.with(with_statements)
  31. end
  32. def build_with_value_from_hash(hash)
  33. hash.map do |name, value|
  34. Arel::Nodes::TableAlias.new(build_with_expression_from_value(value), name)
  35. end
  36. end
  37. def build_with_expression_from_value(value)
  38. case value
  39. when Arel::Nodes::SqlLiteral then Arel::Nodes::Grouping.new(value)
  40. when ActiveRecord::Relation then value.arel
  41. when Arel::SelectManager then value
  42. when Array then value.map { |e| build_with_expression_from_value(e) }.reduce { |result, value| Arel::Nodes::UnionAll.new(result, value) }
  43. else
  44. raise ArgumentError, "Unsupported argument type: `#{value}` #{value.class}"
  45. end
  46. end
  47. end
  48. end
  49. ActiveSupport.on_load(:active_record) do
  50. ActiveRecord::QueryMethods.prepend(ActiveRecord::QueryMethodsExtensions)
  51. end