feat(orm): extend associable with collect_through and as_array

lib/prism/orm/associable.rb

@@ -1,40 +1,115 @@
 module Prism
   module Orm
     module Associable
+      DEPENDENT_OPTIONS = [nil, :destroy, :nullify].freeze
+
       def self.included(base)
         base.extend(ClassMethods)
       end
 
+      # Normalizes a has_many/has_one/belongs_to result into a flat, deduped
+      # (by id) array of records, walking each `through` record's `source`
+      # association. Shared by the lazy reader in ClassMethods and by
+      # Relation's eager-load preloader.
+      def self.collect_through(through_records, source_name)
+        seen = {}
+        as_array(through_records).each_with_object([]) do |through_record, list|
+          as_array(through_record.send(source_name)).each do |target|
+            next if target.nil? || seen[target.id]
+            seen[target.id] = true
+            list << target
+          end
+        end
+      end
+
+      def self.as_array(value)
+        case value
+        when nil then []
+        when Enumerable then value.to_a
+        else [value]
+        end
+      end
+
+      # Resolves the counter_cache column name against `klass.name` lazily,
+      # at hook-execution time rather than when `belongs_to` is declared -
+      # an anonymous `Class.new(Model) { belongs_to ..., counter_cache: true }`
+      # (the pattern this test suite uses throughout, constant-set only
+      # after the class body runs) would otherwise capture `klass.name ==
+      # nil` forever, silently producing a bogus "_count" column.
+      def self.counter_cache_column(klass, counter_cache)
+        counter_cache == true ? "#{Inflector.tableize(klass.name)}_count" : counter_cache.to_s
+      end
+
       module ClassMethods
         def associations
           @associations ||= {}
         end
 
-        def belongs_to(name, class_name: nil, foreign_key: nil)
+        def belongs_to(name, class_name: nil, foreign_key: nil, polymorphic: false, counter_cache: false, touch: false)
+          if polymorphic
+            return define_belongs_to_polymorphic(name, foreign_key: foreign_key, counter_cache: counter_cache,
+                                                        touch: touch)
+          end
+
           target = class_name || Inflector.camelize(name.to_s)
           key = (foreign_key || "#{name}_id").to_s
           associations[name.to_sym] = { type: :belongs_to, class_name: target, foreign_key: key }
 
           define_method(name) do
+            cache = association_cache
+            return cache[name] if cache.key?(name)
+            NPlusOneDetector.record_lazy_load(self, name)
             value = self[key]
-            next nil if value.nil?
-            @association_cache ||= {}
-            @association_cache[name] ||= Object.const_get(target).find(value)
+            cache[name] = value.nil? ? nil : Object.const_get(target).find(value)
           end
 
           define_method("#{name}=") do |record|
             self[key] = record&.id
-            @association_cache ||= {}
-            @association_cache[name] = record
-          end
+            association_cache[name] = record
           end
 
-        def has_many(name, class_name: nil, foreign_key: nil)
-          target = class_name || Inflector.classify(name.to_s)
+          define_counter_cache(target, key, counter_cache) if counter_cache
+          define_touch(target, key, touch) if touch
+        end
+
+        def has_one(name, class_name: nil, foreign_key: nil, dependent: nil, through: nil, source: nil, as: nil)
+          return define_has_one_through(name, through: through, source: source, dependent: dependent) if through
+          return define_has_one_polymorphic(name, class_name: class_name, as: as, dependent: dependent) if as
+
+          validate_dependent!(dependent)
+          target = class_name || Inflector.camelize(name.to_s)
           key = (foreign_key || "#{Inflector.underscore(self.name)}_id").to_s
-          associations[name.to_sym] = { type: :has_many, class_name: target, foreign_key: key }
+          associations[name.to_sym] = { type: :has_one, class_name: target, foreign_key: key, dependent: dependent }
 
           define_method(name) do
+            cache = association_cache
+            return cache[name] if cache.key?(name)
+            NPlusOneDetector.record_lazy_load(self, name)
+            cache[name] = Object.const_get(target).find_by(key => id)
+          end
+
+          define_method("#{name}=") do |record|
+            if record
+              record[key] = id
+              record.save
+            end
+            association_cache[name] = record
+          end
+        end
+
+        def has_many(name, class_name: nil, foreign_key: nil, dependent: nil, through: nil, source: nil, as: nil)
+          return define_has_many_through(name, through: through, source: source, dependent: dependent) if through
+          return define_has_many_polymorphic(name, class_name: class_name, as: as, dependent: dependent) if as
+
+          validate_dependent!(dependent)
+          target = class_name || Inflector.classify(name.to_s)
+          key = (foreign_key || "#{Inflector.underscore(self.name)}_id").to_s
+          associations[name.to_sym] = { type: :has_many, class_name: target, foreign_key: key, dependent: dependent }
+
+          define_method(name) do
+            cache = association_cache
+            return cache[name] if cache.key?(name)
+            NPlusOneDetector.record_lazy_load(self, name)
             Object.const_get(target).where(key => id)
           end
         end