module Stagehand
  module Connection
    def self.with_production_writes(model, &block)
      model.connection.allow_writes(&block)
    end

    module AdapterExtensions
      def self.prepended(base)
        base.set_callback :checkout, :after, :update_readonly_state
        base.set_callback :checkin, :before, :clear_readonly_state
      end

      def exec_insert(*)
        handle_readonly_writes!
        super
      end

      def exec_update(*)
        handle_readonly_writes!
        super
      end

      def exec_delete(*)
        handle_readonly_writes!
        super
      end

      def allow_writes(&block)
        state = readonly?
        readonly!(false)
        return block.call
      ensure
        readonly!(state)
      end

      def readonly!(state = true)
        @readonly = state
      end

      def readonly?
        !!@readonly
      end

      private

      def update_readonly_state
        readonly! unless Configuration.single_connection? || @config[:database] == Database.staging_database_name
      end

      def clear_readonly_state
        readonly!(false)
      end

      def handle_readonly_writes!
        if !readonly?
          return
        elsif Configuration.allow_unsynced_production_writes?
          Rails.logger.warn "Writing directly to #{@config[:database]} database using readonly connection"
        else
          raise(UnsyncedProductionWrite, "Attempted to write directly to #{@config[:database]} database using readonly connection")
        end
      end
    end
  end


  # EXCEPTIONS

  class UnsyncedProductionWrite < StandardError; end
end

ActiveRecord::Base.connection.class.prepend(Stagehand::Connection::AdapterExtensions)

module Stagehand
  module Connection
    def self.with_production_writes(model, &block)
      state = allow_unsynced_production_writes?
      allow_unsynced_production_writes!(true)
      return block.call
    ensure
      allow_unsynced_production_writes!(state)
    end

    def self.allow_unsynced_production_writes!(state = true)
      Thread.current[:stagehand_allow_unsynced_production_writes] = state
    end

    def self.allow_unsynced_production_writes?
      !!Thread.current[:stagehand_allow_unsynced_production_writes]
    end

    module AdapterExtensions
      def exec_insert(*)
        handle_readonly_writes!
        super
      end

      def exec_update(*)
        handle_readonly_writes!
        super
      end

      def exec_delete(*)
        handle_readonly_writes!
        super
      end

      private

      def write_access?
        Configuration.single_connection? || @config[:database] == Database.staging_database_name || Connection.allow_unsynced_production_writes?
      end

      def handle_readonly_writes!
        if write_access?
          return
        elsif Configuration.allow_unsynced_production_writes?
          Rails.logger.warn "Writing directly to #{@config[:database]} database using readonly connection"
        else
          raise(UnsyncedProductionWrite, "Attempted to write directly to #{@config[:database]} database using readonly connection")
        end
      end
    end
  end


  # EXCEPTIONS

  class UnsyncedProductionWrite < StandardError; end
end

ActiveRecord::Base.connection.class.prepend(Stagehand::Connection::AdapterExtensions)

Make unsynced production write blocks more robust
Make unsynced production write blocks more robust

When we create a block that allows unsynced production writes, we expect that any connection opened within that block would be allowed to write to the production database. Unfortunately, only the current `ActiveRecord::Base.connection` instance is marked as being read/write. This has caused an issue when attempting to allow rake db:test:load_structure to rake both databases as the connection we mark as read/write is not the one that ends up being used to write to the production database.

Instead of storing the state on the connection instance, we store it on `Thread.current`. This ensures that any connection we open for the duration of the block will be allowed to write to the production database.
