module RemoteDb
  module Concerns
    module RestrictedColumns
      extend ActiveSupport::Concern

      class ForbiddenColumnException < Exception; end

      included do
        # Rails 4.2 introduced a new attribute reader method `_read_attribute`.
        # However, older versions of Rails (<= 4.1) use `read_attribute`.
        [:_read_attribute, :read_attribute].each do |method|
          if self.method_defined?(method)
            define_method("#{method}_with_safety") do |name|
              original_db_column = self.class.original_columns.include?(name.to_sym)
              visible_db_column  = self.class.visible_columns.include? (name.to_sym)
              without_safety_method = "#{method}_without_safety"
              if (original_db_column && visible_db_column) || !original_db_column
                send(without_safety_method, name)
              else
                raise ForbiddenColumnException, "Column #{name} is not allowed for access."
              end
            end
            alias_method_chain method, :safety

            break
          end
        end
      end

      module ClassMethods

        def visible_columns
          @visible_columns ||= []
        end

        def original_columns
          @original_columns ||= []
        end

        # Note: This is a hack that relies on ActiveRecord's internals. Most of the
        # logic is originally from: ActiveRecord::ModelSchema#reset_column_information`

        def table_columns=(visible_columns)
          unless abstract_class?
            @visible_columns  = visible_columns
            @original_columns = columns.map(&:name).map(&:to_sym)
            @columns.reject! do |column|
              !visible_columns.include?(column.name.to_sym)
            end

            @column_names = @content_columns = @column_defaults = @columns_hash = nil
            @dynamic_methods_hash = nil
            @inheritance_column = nil unless defined?(@explicit_inheritance_column) && @explicit_inheritance_column
            @arel_engine = @relation = nil
          end
        end
      end
    end
  end
end