try using regexes

3 files changed, 21 insertions, 60 deletions

diff --git a/activerecord/lib/active_record/attribute_methods.rb b/activerecord/lib/active_record/attribute_methods.rb
index ff381b4e0b..64f81ca582 100644
--- a/activerecord/lib/active_record/attribute_methods.rb
+++ b/activerecord/lib/active_record/attribute_methods.rb
@@ -167,12 +167,24 @@ module ActiveRecord
         end
       end
 
-      def enforce_raw_sql_whitelist(args, whitelist: attribute_names_and_aliases) # :nodoc:
+      # Regexp whitelist. Matches the following:
+      #   "#{table_name}.#{column_name}"
+      #   "#{column_name}"
+      COLUMN_NAME_WHITELIST = /\A(?:\w+\.)?\w+\z/i
+
+      # Regexp whitelist. Matches the following:
+      #   "#{table_name}.#{column_name}"
+      #   "#{table_name}.#{column_name} #{direction}"
+      #   "#{column_name}"
+      #   "#{column_name} #{direction}"
+      COLUMN_NAME_ORDER_WHITELIST = /\A(?:\w+\.)?\w+(?:\s+asc|\s+desc)?\z/i
+
+      def enforce_raw_sql_whitelist(args, whitelist: COLUMN_NAME_WHITELIST) # :nodoc:
         unexpected = args.reject do |arg|
-          whitelist.include?(arg.to_s) ||
-            arg.kind_of?(Arel::Node) ||
+          arg.kind_of?(Arel::Node) ||
             arg.is_a?(Arel::Nodes::SqlLiteral) ||
-            arg.is_a?(Arel::Attributes::Attribute)
+            arg.is_a?(Arel::Attributes::Attribute) ||
+            arg.to_s.split(/\s*,\s*/).all? { |part| whitelist.match?(part) }
         end
 
         return if unexpected.none?
@@ -195,22 +207,6 @@ module ActiveRecord
         end
       end
 
-      # Can the given name be treated as a column name? Returns true if name
-      # is attribute or attribute alias.
-      #
-      #   class Person < ActiveRecord::Base
-      #   end
-      #
-      #   Person.respond_to_attribute?(:name)
-      #   # => true
-      #
-      #   Person.respond_to_attribute?("foo")
-      #   # => false
-      def respond_to_attribute?(name) # :nodoc:
-        name = name.to_s
-        attribute_names.include?(name) || attribute_aliases.include?(name)
-      end
-
       # Returns true if the given attribute exists, otherwise false.
       #
       #   class Person < ActiveRecord::Base
@@ -242,18 +238,6 @@ module ActiveRecord
           ConnectionAdapters::NullColumn.new(name)
         end
       end
-
-      # An Array of String attribute names and aliases for accessing those
-      # attributes.
-      #
-      #   class Person < ActiveRecord::Base
-      #   end
-      #
-      #   Person.attribute_names_and_aliases
-      #   # => ["id", "created_at", "updated_at", "name", "age"]
-      def attribute_names_and_aliases # :nodoc:
-        attribute_names + attribute_aliases.keys
-      end
     end
 
     # A Person object with a name attribute can ask <tt>person.respond_to?(:name)</tt>,
diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 75795fe493..d49472fc70 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -175,7 +175,7 @@ module ActiveRecord
     # See also #ids.
     #
     def pluck(*column_names)
-      if loaded? && (column_names.map(&:to_s) - @klass.attribute_names_and_aliases).empty?
+      if loaded? && (column_names.map(&:to_s) - @klass.attribute_names - @klass.attribute_aliases.keys).empty?
         return records.pluck(*column_names)
       end
 
@@ -183,10 +183,10 @@ module ActiveRecord
         relation = apply_join_dependency
         relation.pluck(*column_names)
       else
-        enforce_raw_sql_whitelist(column_names, whitelist: allowed_pluck_columns)
+        enforce_raw_sql_whitelist(column_names)
         relation = spawn
         relation.select_values = column_names.map { |cn|
-          @klass.respond_to_attribute?(cn) ? arel_attribute(cn) : cn
+          @klass.has_attribute?(cn) || @klass.attribute_alias?(cn) ? arel_attribute(cn) : cn
         }
         result = skip_query_cache_if_necessary { klass.connection.select_all(relation.arel, nil) }
         result.cast_values(klass.attribute_types)
@@ -203,12 +203,6 @@ module ActiveRecord
 
     private
 
-      def allowed_pluck_columns
-        @klass.attribute_names_and_aliases.map do |name|
-          [name, "#{table_name}.#{name}"]
-        end.flatten
-      end
-
       def has_include?(column_name)
         eager_loading? || (includes_values.present? && column_name && column_name != :all)
       end
diff --git a/activerecord/lib/active_record/relation/query_methods.rb b/activerecord/lib/active_record/relation/query_methods.rb
index 094e5aa733..59a732168c 100644
--- a/activerecord/lib/active_record/relation/query_methods.rb
+++ b/activerecord/lib/active_record/relation/query_methods.rb
@@ -299,7 +299,7 @@ module ActiveRecord
     def order!(*args) # :nodoc:
       @klass.enforce_raw_sql_whitelist(
         column_names_from_order_arguments(args),
-        whitelist: allowed_order_columns
+        whitelist: AttributeMethods::ClassMethods::COLUMN_NAME_ORDER_WHITELIST
       )
 
       preprocess_order_args(args)
@@ -326,7 +326,7 @@ module ActiveRecord
     def reorder!(*args) # :nodoc:
       @klass.enforce_raw_sql_whitelist(
         column_names_from_order_arguments(args),
-        whitelist: allowed_order_columns
+        whitelist: AttributeMethods::ClassMethods::COLUMN_NAME_ORDER_WHITELIST
       )
 
       preprocess_order_args(args)
@@ -928,20 +928,6 @@ module ActiveRecord
 
     private
 
-      def allowed_order_columns
-        @klass.attribute_names_and_aliases.map do |name|
-          [name, "#{table_name}.#{name}"].map do |name|
-            [
-              name,
-              "#{name} asc",
-              "#{name} ASC",
-              "#{name} desc",
-              "#{name} DESC"
-            ]
-          end
-        end.flatten
-      end
-
       # Extract column names from arguments passed to #order or #reorder.
       def column_names_from_order_arguments(args)
         args.flat_map { |arg| arg.is_a?(Hash) ? arg.keys : arg }
@@ -1097,9 +1083,6 @@ module ActiveRecord
           when Arel::Nodes::Ordering
             o.reverse
           when String
-            # ensure we're not dealing with string subclass (Eg. Arel::Nodes::SqlLiteral)
-            o = String.new(o)
-
             if does_not_support_reverse?(o)
               raise IrreversibleOrderError, "Order #{o.inspect} can not be reversed automatically"
             end