From c7f1f22a25026e9b405a01d4342d07f0a71866e5 Mon Sep 17 00:00:00 2001
From: Alex Ghiculescu <alex@tanda.co>
Date: Wed, 10 Mar 2021 10:22:51 -0600
Subject: [PATCH] Handle custom `session_class` in `trim` task

Better approach to https://github.com/rails/activerecord-session_store/pull/178
---
 README.md                        | 7 +++++++
 lib/tasks/database.rake          | 4 ++--
 test/tasks/database_rake_test.rb | 8 ++++++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 2b9e829..c1b79f9 100644
--- a/README.md
+++ b/README.md
@@ -95,6 +95,13 @@ You must implement these methods:
 * `save`
 * `destroy`
 
+Note that some of the provided [rake tasks](https://github.com/rails/activerecord-session_store/blob/master/lib/tasks/database.rake)
+require additional methods to be implemented:
+
+* `db:sessions:clear` depends on `table_name`
+* `db:sessions:trim` depends on `where` and `delete_all`
+* `db:sessions:upgrade` depends on `secure!`
+
 The example SqlBypass class is a generic SQL session store. You may
 use it as a basis for high-performance database-specific stores.
 
diff --git a/lib/tasks/database.rake b/lib/tasks/database.rake
index 377a0b6..f7c55da 100644
--- a/lib/tasks/database.rake
+++ b/lib/tasks/database.rake
@@ -8,13 +8,13 @@ namespace 'db:sessions' do
 
   desc "Clear the sessions table"
   task :clear => [:environment, 'db:load_config'] do
-    ActiveRecord::Base.connection.execute "TRUNCATE TABLE #{ActiveRecord::SessionStore::Session.table_name}"
+    ActiveRecord::Base.connection.truncate(ActionDispatch::Session::ActiveRecordStore.session_class.table_name)
   end
 
   desc "Trim old sessions from the table (default: > 30 days)"
   task :trim => [:environment, 'db:load_config'] do
     cutoff_period = (ENV['SESSION_DAYS_TRIM_THRESHOLD'] || 30).to_i.days.ago
-    ActiveRecord::SessionStore::Session.
+    ActionDispatch::Session::ActiveRecordStore.session_class.
       where("updated_at < ?", cutoff_period).
       delete_all
   end
diff --git a/test/tasks/database_rake_test.rb b/test/tasks/database_rake_test.rb
index 85fd7be..a76455e 100644
--- a/test/tasks/database_rake_test.rb
+++ b/test/tasks/database_rake_test.rb
@@ -63,6 +63,14 @@ def test_upgrade_task
         assert Session.find_by_session_id(Rack::Session::SessionId.new("original_session_id").private_id)
         assert Session.find_by_session_id("2::secure_session_id")
       end
+
+      def test_clear_task
+        Session.create!(data: "obsolete")
+
+        Rake.application.invoke_task 'db:sessions:clear'
+
+        assert_equal 0, Session.count
+      end
     end
   end
 end