feat(orm): extend associable with has_and_belongs_to_many

lib/prism/orm/associable.rb

@@ -113,6 +113,179 @@ module Prism
             Object.const_get(target).where(key => id)
           end
         end
+
+        def has_and_belongs_to_many(name, class_name: nil, join_table: nil, foreign_key: nil,
+                                    association_foreign_key: nil)
+          target = class_name || Inflector.classify(name.to_s)
+          jt = (join_table || habtm_join_table_name(self.name, target)).to_s
+          left_key = (foreign_key || "#{Inflector.underscore(self.name)}_id").to_s
+          right_key = (association_foreign_key || "#{Inflector.underscore(target)}_id").to_s
+          associations[name.to_sym] = {
+            type: :habtm, class_name: target, join_table: jt, left_key: left_key, right_key: right_key
+          }
+
+          define_method(name) do
+            HabtmCollection.new(self, target: Object.const_get(target), join_table: jt,
+                                      left_key: left_key, right_key: right_key)
+          end
+
+          define_method("#{name}=") do |records|
+            send(name).replace(records)
+          end
+        end
+
+        def dependent_associations
+          associations.select { |_, meta| meta[:dependent] }
+        end
+
+        private
+
+        # has_many :through - derives its collection by walking an existing
+        # association (`through`) and reading another association (`source`)
+        # off each of those records, e.g. `has_many :commenters, through:
+        # :comments, source: :commenter`. Unlike plain has_many this holds no
+        # foreign key of its own, so `dependent` is not supported.
+        def define_has_many_through(name, through:, source:, dependent: nil)
+          raise ArgumentError, 'has_many :through requires a through: association' unless through
+          if dependent
+            raise ArgumentError, "dependent is not supported with has_many :through (#{name.inspect})"
+          end
+          source_name = (source || Inflector.singularize(name.to_s)).to_sym
+          associations[name.to_sym] = { type: :has_many_through, through: through.to_sym, source: source_name }
+
+          define_method(name) do
+            cache = association_cache
+            return cache[name] if cache.key?(name)
+            NPlusOneDetector.record_lazy_load(self, name)
+            cache[name] = Associable.collect_through(send(through), source_name)
+          end
+        end
+
+        # has_one :through - single-record counterpart to has_many :through:
+        # walks an existing has_one/belongs_to association (`through`) and
+        # reads another single-record association (`source`) off it, e.g.
+        # `has_one :account_history, through: :account`. Holds no foreign key
+        # of its own, so `dependent` is not supported.
+        def define_has_one_through(name, through:, source:, dependent: nil)
+          raise ArgumentError, 'has_one :through requires a through: association' unless through
+          if dependent
+            raise ArgumentError, "dependent is not supported with has_one :through (#{name.inspect})"
+          end
+          source_name = (source || name).to_sym
+          associations[name.to_sym] = { type: :has_one_through, through: through.to_sym, source: source_name }
+
+          define_method(name) do
+            cache = association_cache
+            return cache[name] if cache.key?(name)
+            NPlusOneDetector.record_lazy_load(self, name)
+            through_record = send(through)
+            cache[name] = through_record && through_record.send(source_name)
+          end
+        end
+
+        # belongs_to :commentable, polymorphic: true - the owner holds both
+        # the foreign key (commentable_id) and a type column (commentable_type,
+        # the target's class name as a string), so the same foreign key can
+        # point at rows in different tables depending on the record.
+        #
+        # counter_cache:/touch: both work here too, exactly as on a plain
+        # belongs_to, except which owner CLASS to reach for is read off
+        # type_key per record/per saved-change instead of being the single
+        # fixed class_name a plain belongs_to always has - see the type_key:
+        # branches in define_counter_cache/define_touch below.
+        def define_belongs_to_polymorphic(name, foreign_key:, counter_cache: false, touch: false)
+          key = (foreign_key || "#{name}_id").to_s
+          type_key = "#{name}_type"
+          associations[name.to_sym] = { type: :belongs_to, polymorphic: true, foreign_key: key, type_key: type_key }
+
+          define_method(name) do
+            cache = association_cache
+            return cache[name] if cache.key?(name)
+            NPlusOneDetector.record_lazy_load(self, name)
+            target_id = self[key]
+            target_type = self[type_key]
+            cache[name] = (target_id.nil? || target_type.nil?) ? nil : Object.const_get(target_type).find(target_id)
+          end
+
+          define_method("#{name}=") do |record|
+            self[key] = record&.id
+            self[type_key] = record&.class&.name
+            association_cache[name] = record
+          end
+
+          define_counter_cache(nil, key, counter_cache, type_key: type_key) if counter_cache
+          define_touch(nil, key, touch, type_key: type_key) if touch
+        end
+
+        # has_many :comments, as: :commentable - the inverse side of a
+        # polymorphic belongs_to: scopes the child table by both the foreign
+        # key and the type column so rows belonging to a different owner
+        # class with the same id never leak in.
+        def define_has_many_polymorphic(name, class_name:, as:, dependent: nil)
+          validate_dependent!(dependent)
+          target = class_name || Inflector.classify(name.to_s)
+          key = "#{as}_id"
+          type_key = "#{as}_type"
+          associations[name.to_sym] = {
+            type: :has_many, class_name: target, foreign_key: key, type_key: type_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, type_key => self.class.name)
+          end
+        end
+
+        # has_one :preview_image, as: :imageable - single-record counterpart
+        # to define_has_many_polymorphic.
+        def define_has_one_polymorphic(name, class_name:, as:, dependent: nil)
+          validate_dependent!(dependent)
+          target = class_name || Inflector.camelize(name.to_s)
+          key = "#{as}_id"
+          type_key = "#{as}_type"
+          associations[name.to_sym] = {
+            type: :has_one, class_name: target, foreign_key: key, type_key: type_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, type_key => self.class.name)
+          end
+
+          define_method("#{name}=") do |record|
+            if record
+              record[key] = id
+              record[type_key] = self.class.name
+              record.save
+            end
+            association_cache[name] = record
+          end
+        end
+
+        # belongs_to ..., counter_cache: true - keeps a denormalized count of
+        # this class's own rows on the parent row pointed at by the foreign
+        # key, so `post.comments.count` becomes a plain column read
+        # (`post.comments_count`) instead of a query, mirroring Rails'
+        # `belongs_to :post, counter_cache: true` on the child (Comment)
+        # side. `counter_cache: true` defaults the column to this class's
+        # own table name pluralized-as-a-suffix (Comment -> comments_count);
+        # `counter_cache: :custom_column` names it explicitly. The three
+        # hooks run inside the same transaction as the triggering
+        # create/update/destroy (Store#increment_counter is reentrant), so a
+        # rolled-back child write rolls the counter bump back with it.
+        #
+        # type_key: is set by define_belongs_to_polymorphic - a polymorphic
+        # owner's class varies per row (read off that column), unlike a
+        # plain belongs_to's single fixed target_class_name, so id and type
+        # are tracked as a PAIR everywhere below instead of an id alone.
+        # Either half of that pair can in principle change independently of
+        # the other on update, so old/new are resolved separately for each.
+        def define_counter_cache(target_class_name, foreign_key, counter_cache, type_key: nil)
+        end
       end
     end
   end