database_context.rb

# This class encapsulates a unit of work done for a particular tenant, connected to that tenant's database.
# ActiveRecord makes it _very_ hard to do in a simple manner and clever stuff is required, but it is knowable.
#
# What this class provides is a "misuse" of the database "roles" of ActiveRecord to have a role per tenant.
# If all the tenants are predefined, it can be done roughly so:
#
#     ActiveRecord::Base.legacy_connection_handling = false if ActiveRecord::Base.respond_to?(:legacy_connection_handling)
#     $databases.each_pair do |n, db_path|
#       config_hash = {
#         "adapter" => 'sqlite3',
#         "database" => db_path,
#         "pool" => 4
#       }
#       ActiveRecord::Base.connection_handler.establish_connection(config_hash, role: "database_#{n}")
#     end
#     
#     def named_databases_as_roles_using_connected_to(n, from_database_paths)
#       ActiveRecord::Base.connected_to(role: "database_#{n}") do
#         query_and_compare!(n)
#       end
#     end
#
# So what we do is this:
#
# * We want one connection pool per tenant (per database, thus)
# * We want to grab a connection from that pool and make sure our queries use that connection
# * Once we are done with our unit of work we want to return the connection to the pool
#
# This also uses a stack of Fibers because `connected_to` in ActiveRecord _wants_ to have a block, but for us
# "leaving" the context of a unit of work can happen in a Rack body close() call.
class DatabaseContext
  POOL_CHECK_MUTEX = Mutex.new
  
  def initialize(single_connection_config_hash)
    if ActiveRecord::Base.respond_to?(:legacy_connection_handling) && ActiveRecord::Base.legacy_connection_handling
      raise "ActiveRecord::Base.legacy_connection_handling is enabled (set to `true`) and we can't use roles that way."
    end

    @config_hash = single_connection_config_hash.with_indifferent_access
    @role_name = "tenant_db_#{Digest::SHA1.hexdigest(@config_hash.fetch(:database))}"
    @context_fiber = nil
  end

  def enter
    create_pool_for_database_if_none_available!
    @context_fiber = Fiber.new do
      ActiveRecord::Base.connected_to(role: @role_name) { Fiber.yield }
    end
    @context_fiber.resume
    true
  end

  def leave
    fiber, @context_fiber = @context_fiber, nil
    return unless fiber
    fiber.resume
  end

  def with(&blk)
    create_pool_for_database_if_none_available!
    ActiveRecord::Base.connected_to(role: @role_name, &blk)
  end

  # Unlike Rails, we connect not when the app boots - but at the start of the first request to a particular database.
  # If we don't connect - there will be a NoConnectionPool error.
  def create_pool_for_database_if_none_available!
    POOL_CHECK_MUTEX.synchronize do
      # Maybe there is a way to connect to a particular role using just ActiveRecord::Base.establish_connection, but I could not find it.
      # There is a way to connect to a particular _role_ - we have to connect using the connection handler instead of AR::Base though.
      # We pass it the config hash and specify the role. This is originally made for read replicas, but honey badger don't give a s_t.
      #
      # AR does not check whether we are trying to connect to the same DB under the same role, so we need to check whether there are
      # any pools for that role already. If there are - they will be used and there won't be any errors. If there are not - we need
      # to establish_connection.
      #
      # We do this under a mutex because I am too lazy to try and figure out whether `connection_pool_list` and `establish_connection`
      # are thread-safe or not, and under which versions. For 99% of requests it will be just an array size check anyway.
      if ActiveRecord::Base.connection_handler.connection_pool_list(@role_name).none?
        ActiveRecord::Base.connection_handler.establish_connection(@config_hash, role: @role_name)
      end
    end
  end
end