Fix ownership association and Rake tasks

CHANGELOG.md

@@ -8,6 +8,7 @@ User-visible changes worth mentioning.
 ## master
 
 - [#PR ID] Add your PR description here.
+- [#1478] Fix ownership association and Rake tasks when custom models configured.
 - [#1477] Respect `ActiveRecord::Base.pluralize_table_names` for Doorkeeper table names.
 
 ## 5.5.0.rc2

lib/doorkeeper/orm/active_record.rb

@@ -37,15 +37,7 @@ def self.initialize_application_owner!
       end
 
       def self.lazy_load(&block)
-        # ActiveSupport has no public interface to check if something
-        # already lazy-loaded :(
-        loaded = ActiveSupport.instance_variable_get(:"@loaded") || {}
-
-        if loaded.key?(:active_record)
-          block.call
-        else
-          ActiveSupport.on_load(:active_record, {}, &block)
-        end
+        ActiveSupport.on_load(:active_record, {}, &block)
       end
 
       def self.models

lib/doorkeeper/rake/db.rake

@@ -13,7 +13,7 @@ namespace :doorkeeper do
     namespace :cleanup do
       desc "Removes stale access tokens"
       task revoked_tokens: "doorkeeper:setup" do
-        cleaner = Doorkeeper::StaleRecordsCleaner.new(Doorkeeper::AccessToken)
+        cleaner = Doorkeeper::StaleRecordsCleaner.new(Doorkeeper.config.access_token_model)
         cleaner.clean_revoked
       end
 
@@ -26,13 +26,13 @@ namespace :doorkeeper do
 
       desc "Removes stale access grants"
       task revoked_grants: "doorkeeper:setup" do
-        cleaner = Doorkeeper::StaleRecordsCleaner.new(Doorkeeper::AccessGrant)
+        cleaner = Doorkeeper::StaleRecordsCleaner.new(Doorkeeper.config.access_grant_model)
         cleaner.clean_revoked
       end
 
       desc "Removes expired (TTL passed) access grants"
       task expired_grants: "doorkeeper:setup" do
-        cleaner = Doorkeeper::StaleRecordsCleaner.new(Doorkeeper::AccessGrant)
+        cleaner = Doorkeeper::StaleRecordsCleaner.new(Doorkeeper.config.access_grant_model)
         cleaner.clean_expired(Doorkeeper.config.authorization_code_expires_in)
       end
     end

lib/doorkeeper/rake/setup.rake

@@ -2,5 +2,10 @@
 
 namespace :doorkeeper do
   task setup: :environment do
+    # Dirty hack to manually initialize AR because of lazy auto-loading,
+    # in other case we'll see NameError: uninitialized constant Doorkeeper::AccessToken
+    if Doorkeeper.config.orm == :active_record && defined?(::ActiveRecord::Base)
+      Object.const_get("::ActiveRecord::Base")
+    end
   end
 end