Merge pull request rails#51801 from Shopify/refactor-find_by_token_for

Author: byroot

activerecord/lib/active_record/relation.rb

@@ -66,7 +66,7 @@ def exec_explain(&block)
 
     include Enumerable
     include FinderMethods, Calculations, SpawnMethods, QueryMethods, Batches, Explain, Delegation
-    include SignedId::RelationMethods
+    include SignedId::RelationMethods, TokenFor::RelationMethods
 
     attr_reader :table, :klass, :loaded, :predicate_builder
     attr_accessor :skip_preloading_value

activerecord/lib/active_record/token_for.rb

@@ -35,6 +35,24 @@ def resolve_token(token)
       end
     end
 
+    module RelationMethods
+      # Finds a record using a given +token+ for a predefined +purpose+. Returns
+      # +nil+ if the token is invalid or the record was not found.
+      def find_by_token_for(purpose, token)
+        raise UnknownPrimaryKey.new(self) unless model.primary_key
+        model.token_definitions.fetch(purpose).resolve_token(token) { |id| find_by(model.primary_key => id) }
+      end
+
+      # Finds a record using a given +token+ for a predefined +purpose+. Raises
+      # ActiveSupport::MessageVerifier::InvalidSignature if the token is invalid
+      # (e.g. expired, bad format, etc). Raises ActiveRecord::RecordNotFound if
+      # the token is valid but the record was not found.
+      def find_by_token_for!(purpose, token)
+        model.token_definitions.fetch(purpose).resolve_token(token) { |id| find(id) } ||
+          (raise ActiveSupport::MessageVerifier::InvalidSignature)
+      end
+    end
+
     module ClassMethods
       # Defines the behavior of tokens generated for a specific +purpose+.
       # A token can be generated by calling TokenFor#generate_token_for on a
@@ -85,20 +103,12 @@ def generates_token_for(purpose, expires_in: nil, &block)
         self.token_definitions = token_definitions.merge(purpose => TokenDefinition.new(self, purpose, expires_in, block))
       end
 
-      # Finds a record using a given +token+ for a predefined +purpose+. Returns
-      # +nil+ if the token is invalid or the record was not found.
-      def find_by_token_for(purpose, token)
-        raise UnknownPrimaryKey.new(self) unless primary_key
-        token_definitions.fetch(purpose).resolve_token(token) { |id| find_by(primary_key => id) }
+      def find_by_token_for(purpose, token) # :nodoc:
+        all.find_by_token_for(purpose, token)
       end
 
-      # Finds a record using a given +token+ for a predefined +purpose+. Raises
-      # ActiveSupport::MessageVerifier::InvalidSignature if the token is invalid
-      # (e.g. expired, bad format, etc). Raises ActiveRecord::RecordNotFound if
-      # the token is valid but the record was not found.
-      def find_by_token_for!(purpose, token)
-        token_definitions.fetch(purpose).resolve_token(token) { |id| find(id) } ||
-          (raise ActiveSupport::MessageVerifier::InvalidSignature)
+      def find_by_token_for!(purpose, token) # :nodoc:
+        all.find_by_token_for!(purpose, token)
       end
     end