ActiveRecord Concern to mark a slice of a model's association tree as freezable, to freeze it at some point, and then to work with frozen versions of those objects moving forward

snapshotable.rb

# frozen_string_literal: true

require "bundler/inline"

gemfile(true) do
  source "https://rubygems.org"

  git_source(:github) { |repo| "https://github.com/#{repo}.git" }

  # Activate the gem you are reporting the issue against.
  gem "activerecord", "5.2.0"
  gem "sqlite3", "~> 1.3.6"
end

require "active_record"
# require "minitest/autorun"
require "logger"

# This connection will do for database-independent bug reports.
ActiveRecord::Base.establish_connection(adapter: "sqlite3", database: ":memory:")
ActiveRecord::Base.logger = Logger.new(STDOUT)

ActiveRecord::Schema.define do
  create_table :identities, force: true do |t|
    t.string :email
    t.timestamps
  end

  create_table :actors, force: true do |t|
    t.references :identity, foreign_key: true
    t.boolean :payable, default: true
    t.timestamps
  end

  create_table :accounts, force: true do |t|
    t.references :identity, foreign_key: true
    t.boolean :personal, default: true
    t.timestamps
  end

  create_table :addresses, force: true do |t|
    t.references :actor, foreign_key: true
    t.string :zip
    t.timestamps
  end

  create_table :payout_methods, force: true do |t|
    t.references :actor, foreign_key: true
    t.string :type
    t.json :payout_data
    t.timestamps
  end

  create_table :bills, force: true do |t|
    t.references :actor, foreign_key: true
    t.integer :amount_in_cents
    t.timestamps

    t.text :association_snapshot
  end

  create_table :bill_items, force: true do |t|
    t.references :actor, foreign_key: true
    t.references :bill, foreign_key: true, optional: true
    t.integer :amount_in_cents
    t.timestamps
  end
end

# ------------------------------------------------------------------------------

# TODO: we want to turn this concern into a gem
# Usage:
# 1. Add a new column `association_snapshot` of type text to your table
# 2. In the class e.g.:
# class Bill < ApplicationRecord
#   include Snapshotable
#   snapshotable actor: [:address, {identity: :account}]
# end
#
module Snapshotable
  extend ActiveSupport::Concern

  # this module will be included into each of our namespaced clone classes
  # it ensures the snapshotted instances are read-only
  module Clone
    # mark all instances as readonly, such that attempting to update will result in an `ActiveRecord::ReadOnlyRecord` error.
    def readonly?
      true
    end

    def snapshotted?
      true
    end

    # don't allow reloading from database
    def reload
      self
    end

    # don't even allow in-memory Ruby attribute updates
    def _write_attribute(attr_name, value)
      # if someone or something is trying to set a db attribute to its current value, let it
      # this happens, for example, when Paranoia has hooked into the initialization of a model
      if attributes.include?(attr_name) && attributes[attr_name] == value
        super
        # if someone or something is trying to set an accessor to its current value, let it
      elsif instance_variables.include?("@#{attr_name}".to_sym) && instance_variable_get("@#{attr_name}") == value
        super
        # if Rails or someone is trying to set the STI attribute for a model, let it
      elsif self.class.finder_needs_type_condition? && self.class.inheritance_column == attr_name && self.class.sti_name == value
        super
        # otherwise, we do not allow changing attributes of our snapshotted models
      else
        _raise_readonly_record_error
      end
    end

    # ensure that attributes are displayed, even if no longer present on the parent class schema
    def inspect
      inspection = if @attributes
                     attribute_names
                       .select { |name| has_attribute?(name) }
                       .map { |name| "#{name}: #{attribute_for_inspect(name)}" }
                       .join(', ')
                   else
                     'not initialized'
                   end

      "#<#{self.class} #{inspection}>"
    end

    # make this class a custom serializer class, so that we don't have to use the default `Coders::YAMLColumn`
    # which will try to introspect on the class and thus requires a db connection, which doesn't always exist then this code is being run
    # and which isn't necessary for serialization. Below is a simplified copy of the `Coders::YAMLColumn` load and dump logic tho.
    module ClassMethods
      def dump(obj)
        return if obj.nil?

        YAML.dump obj
      end

      def load(yaml)
        return new if yaml.nil?
        return yaml unless yaml.is_a?(String) && /^---/.match?(yaml)

        YAML.safe_load(yaml) || new
      end
    end

    def self.included(receiver)
      receiver.extend ClassMethods
    end
  end

  class_methods do
    attr_reader :snapshotable_associations

    # explain which "neighborhood" of the model's associations will need to be snapshotted at some point
    def snapshotable(associations_tree)
      # associations_tree will be something like:
      # { actor: [:address, { identity: :account }] }

      # set a class level instance variable such that the instance method will have access to this data
      @snapshotable_associations = associations_tree

      # tell Rails that we will store an Hash of model objects, so that Rails will marshall these objects
      # and allow us to have full objects to work with in the app
      serialize :association_snapshot, Hash
    end
  end

  included do
    # when fetched directly from the database we need to link the association methods to the clone classes
    after_initialize do
      _overwrite_association_methods_in_clone_classes(association_snapshot) if association_snapshot.present?
    end
  end

  # take a snapshot of the "neighborhood" of associations, marshall the read-only objects, and store in db
  def take_snapshot!
    # write this array of read-only clone instance objects into the db
    # where the `serialize` set at the class level will marshall these objects
    take_snapshot.tap { save! }
  end

  # Only assign the associations without writing into the db
  def take_snapshot
    self.association_snapshot = snapshotted_clones
    self
  end

  def snapshotable?
    attributes.include? 'association_snapshot'
  end

  private

  def snapshotted_clones
    # grab the hash of the associations tree declared by the `snapshotable` class macro
    # { actor: [:address, { identity: :account }] }
    associations_tree = self.class.snapshotable_associations

    # restructure the associations tree into a flat structure that will be much easier to work with
    # { [:actor]=>nil,
    #   [:actor, :address]=>nil,
    #   [:actor, :identity]=>nil,
    #   [:actor, :identity, :account]=>nil }
    flat_associations_tree = _flatten_associations_tree(associations_tree)

    # walk the associations paths to gather the current model instance objects for each association
    # { [:actor]=>#<Actor id>,
    #   [:actor, :address]=>#<Address id>,
    #   [:actor, :identity]=>#<Identity id>,
    #   [:actor, :identity, :account]=>#<Account id> }
    associated_instances_tree = _gather_associated_instances(flat_associations_tree)

    # for each model instance, convert it into a namespaced read-only clone
    # { [:actor]=>#<Snapshotable::Actor id>,
    #   [:actor, :address]=>#<Snapshotable::Address id>,
    #   [:actor, :identity]=>#<Snapshotable::Identity id>,
    #   [:actor, :identity, :account]=>#<Snapshotable::Account id> }
    associated_clones_tree = _snapshot_associated_instances(associated_instances_tree)

    # ensure that the association methods return the cloned objects, not the live objects from the db
    _overwrite_association_methods_in_clone_classes(associated_clones_tree)

    # return the hash to be serialized and stored in the `association_snapshot` field
    associated_clones_tree
  end

  def _flatten_associations_tree(tree, keys = [], output = {})
    # level 0
    #   { :tree=>{:actor=>[:address, {:identity=>:account}]},
    #     :keys=>[],
    #     :output=>{} }
    # level 1
    #   { :tree=>[:address, {:identity=>:account}],
    #     :keys=>[:actor],
    #     :output=>{[:actor]=>nil} }
    # level 2.1
    #   { :tree=>:address,
    #     :keys=>[:actor],
    #     :output=>{[:actor]=>nil} }
    # level 2.2
    #   { :tree=>{:identity=>:account},
    #     :keys=>[:actor],
    #     :output=>{[:actor]=>nil, [:actor, :address]=>nil} }
    # level 3
    #   { :tree=>:account,
    #     :keys=>[:actor, :identity],
    #     :output=>{[:actor]=>nil, [:actor, :address]=>nil, [:actor, :identity]=>nil}}
    # output
    #   { [:actor]=>nil,
    #     [:actor, :address]=>nil,
    #     [:actor, :identity]=>nil,
    #     [:actor, :identity, :account]=>nil }

    case tree
    when Hash
      tree.each do |association, sub_tree|
        output[[*keys, association]] = nil

        _flatten_associations_tree(
          sub_tree,
          [*keys, association],
          output
        )
      end
    when Array
      tree.each do |association|
        _flatten_associations_tree(
          association,
          keys,
          output
        )
      end
    else
      output[[*keys, tree]] = nil
    end

    output
  end

  def _gather_associated_instances(tree)
    # { [:actor]=>nil,
    #   [:actor, :address]=>nil,
    #   [:actor, :identity]=>nil,
    #   [:actor, :identity, :account]=>nil }
    tree.each_with_object({}) do |(associations_path, _), new_tree|
      new_tree[associations_path] = _walk_path_to_get_object(associations_path)
    end
    # { [:actor]=>#<Actor id>,
    #   [:actor, :address]=>#<Address id>,
    #   [:actor, :identity]=>#<Identity id>,
    #   [:actor, :identity, :account]=>#<Account id> }
  end

  def _snapshot_associated_instances(tree)
    # { [:actor]=>#<Actor id>,
    #   [:actor, :address]=>#<Address id>,
    #   [:actor, :identity]=>#<Identity id>,
    #   [:actor, :identity, :account]=>#<Account id> }
    tree.transform_values do |associated_instance|
      associated_instance&.becomes _prepare_namespaced_readonly_clone_class(associated_instance.class)
    end
    # { [:actor]=>#<Snapshotable::Actor id>,
    #   [:actor, :address]=>#<Snapshotable::Address id>,
    #   [:actor, :identity]=>#<Snapshotable::Identity id>,
    #   [:actor, :identity, :account]=>#<Snapshotable::Account id> }
  end

  def _overwrite_association_methods_in_clone_classes(tree)
    # { [:actor]=>#<Snapshotable::Actor id>,
    #   [:actor, :address]=>#<Snapshotable::Address id>,
    #   [:actor, :identity]=>#<Snapshotable::Identity id>,
    #   [:actor, :identity, :account]=>#<Snapshotable::Account id> }
    tree.each do |associations_path, associated_clone_instance|
      *parent_path, association_method = associations_path

      parent_instance = _walk_path_to_get_object(parent_path)

      parent_instance.define_singleton_method association_method do
        if try(:snapshotable?) || try(:snapshotted?)
          associated_clone_instance.freeze
        else
          super
        end
      end
    end

    true
  end

  def _prepare_namespaced_readonly_clone_class(base_class)
    associated_model_name = base_class.name.to_s
    namespaced_readonly_clone_class = _create_namespaced_readonly_clone_class_for(base_class)
    *parent_path, association_constant = associated_model_name.split('::')

    parent_constant = parent_path.reduce(Snapshotable) do |base, const_name|
      base.const_get(const_name, inherit = false) # rubocop:disable Lint/UselessAssignment
    rescue NameError
      base.const_set(const_name, Module.new)
    end

    # We have to ensure that we are checking *only* whether the constant is defined _under_
    # the `parent_constant`. `const_defined?` will walk up the scopes tree, so `Namespace.const_defined?(Const)` will return `true` is `::Const` is defined.
    # We need to check if and only if `Namespace::Const` is defined already.
    parent_constant.const_set(association_constant, namespaced_readonly_clone_class) unless parent_constant.constants.include? association_constant.to_sym

    parent_constant.const_get(association_constant)
  end

  def _create_namespaced_readonly_clone_class_for(base_class)
    Class.new(base_class) { include Clone }
  end

  def _walk_path_to_get_object(path)
    path.reduce(self) do |object, signal|
      break nil unless object.respond_to? signal

      object.public_send(signal)
    end
  end
end

# uncomment this if you want to test that snapshotable works without a db connection
# ActiveRecord::Base.remove_connection

class Identity < ActiveRecord::Base
  has_one :account
  has_one :actor
end

class Actor < ActiveRecord::Base
  belongs_to :identity
  has_one :address
  has_one :payout_method
  has_many :bill_items
  has_many :bills
end

class Account < ActiveRecord::Base
  belongs_to :identity
end

class Address < ActiveRecord::Base
  belongs_to :actor

  def zip_in_reverse
    zip.reverse
  end
end

class PayoutMethod < ActiveRecord::Base
  belongs_to :worker
end

class PayoutMethod::Bank < PayoutMethod
  store_accessor :payout_data, :iban
end

class PayoutMethod::Paypal < PayoutMethod
  store_accessor :payout_data, :email
end

class Bill < ActiveRecord::Base
  include Snapshotable

  belongs_to :actor
  has_many :bill_items

  snapshotable actor: [:address, :payout_method, {identity: :account}]
end

class BillItem < ActiveRecord::Base
  belongs_to :actor
  belongs_to :bill, optional: true
end

# ------------------------------------------------------------------------------

person_1_identity = Identity.create!(email: "person_1@email.address")
person_2_identity = Identity.create!(email: "person_2@email.address")

person_1_actor = person_1_identity.create_actor!(payable: true)
person_2_actor = person_2_identity.create_actor!(payable: true)

person_1_account = person_1_identity.create_account!(personal: true)
person_2_account = person_2_identity.create_account!(personal: true)

person_1_address = person_1_actor.create_address!(zip: 12345)
person_2_address = person_2_actor.create_address!(zip: 67890)

person_1_payout_method = person_1_actor.create_payout_method!(type: "PayoutMethod::Bank", payout_data: {iban: 'DE1234'})
person_2_payout_method = person_2_actor.create_payout_method!(type: "PayoutMethod::Paypal", payout_data: {email: 'payout@email.com'})

person_1_actor.bill_items.create!(amount_in_cents: 100)
person_1_actor.bill_items.create!(amount_in_cents: 200)
person_1_actor.bill_items.create!(amount_in_cents: 300)

person_2_actor.bill_items.create!(amount_in_cents: 100)
person_2_actor.bill_items.create!(amount_in_cents: 200)
person_2_actor.bill_items.create!(amount_in_cents: 300)

person_1_bill = person_1_actor.bills.create!(amount_in_cents: person_1_actor.bill_items.sum(:amount_in_cents))
person_2_bill = person_2_actor.bills.create!(amount_in_cents: person_2_actor.bill_items.sum(:amount_in_cents))

person_1_actor.bill_items.update_all(bill_id: person_1_bill.id)
person_2_actor.bill_items.update_all(bill_id: person_2_bill.id)

person_1_bill.take_snapshot!

# ------------------------------------------------------------------------------

# class TestFreezableBelongsTo < Minitest::Test
#   def setup
#     @original_price = 10
#     @updated_price = 200
#     @user = User.create
#     @widget = Widget.create(price: @original_price, category: :special)
#     @widget_purchase = WidgetPurchase.create(user: @user, widget: @widget)
#   end

#   def test_1
#     assert_equal @widget.attributes, @widget_purchase.live_widget.attributes
#     assert_equal WidgetPurchase::Widget.new.attributes, @widget_purchase.frozen_widget.attributes
#   end

#   def test_3
#     @widget.update(price: @updated_price)

#     assert_equal @updated_price, @widget_purchase.widget.price
#   end

#   def test_4
#     @widget_purchase.freeze_widget!
#     @widget.update(price: @updated_price)

#     assert_equal @original_price, @widget_purchase.widget.price
#   end
# end

# class TestFreezableHasMany < Minitest::Test
#   def setup
#     @widgets_purchase = WidgetsPurchase.create
#     Widget.create(widgets_purchase: @widgets_purchase)
#   end

#   def test_1
#     assert_equal Widget.where(widgets_purchase: @widgets_purchase).to_a, @widgets_purchase.live_widgets.to_a
#     assert_equal [], @widgets_purchase.frozen_widgets
#   end

#   def test_2
#     assert_equal @widgets_purchase.live_widgets, @widgets_purchase.widgets
#   end

#   def test_3
#     orignal_widgets = @widgets_purchase.widgets.map { |u| u.becomes(WidgetsPurchase::Widget) }
#     @widgets_purchase.freeze_widgets!
#     Widget.update_all(widgets_purchase_id: nil)

#     assert_equal orignal_widgets, @widgets_purchase.frozen_widgets
#     assert_equal @widgets_purchase.frozen_widgets, @widgets_purchase.widgets
#     assert_equal Widget.none, @widgets_purchase.live_widgets
#   end
# end