From eb0b67a14ad4b07522455645c5888c4b0c32157a Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Thu, 7 Nov 2024 03:18:36 -0800
Subject: [PATCH] Reset state for lock timeout retries

---
 lib/strong_migrations/checker.rb  |  5 +++++
 lib/strong_migrations/migrator.rb |  2 +-
 test/timeouts_test.rb             | 16 ++++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/lib/strong_migrations/checker.rb b/lib/strong_migrations/checker.rb
index 87f9163b..1c240969 100644
--- a/lib/strong_migrations/checker.rb
+++ b/lib/strong_migrations/checker.rb
@@ -11,10 +11,15 @@ class << self
 
     def initialize(migration)
       @migration = migration
+      reset
+    end
+
+    def reset
       @new_tables = []
       @new_columns = []
       @timeouts_set = false
       @committed = false
+      @transaction_disabled = false
     end
 
     def self.safety_assured
diff --git a/lib/strong_migrations/migrator.rb b/lib/strong_migrations/migrator.rb
index 2e8dca20..1168e7fc 100644
--- a/lib/strong_migrations/migrator.rb
+++ b/lib/strong_migrations/migrator.rb
@@ -12,7 +12,7 @@ def ddl_transaction(migration, ...)
       # retry migration since the entire transaction needs to be rerun
       checker.retry_lock_timeouts(check_committed: true) do
         # failed transaction reverts timeout, so need to re-apply
-        checker.timeouts_set = false
+        checker.reset
 
         super(migration, ...)
       end
diff --git a/test/timeouts_test.rb b/test/timeouts_test.rb
index 852e72cd..13b87284 100644
--- a/test/timeouts_test.rb
+++ b/test/timeouts_test.rb
@@ -170,6 +170,22 @@ def test_lock_timeout_retries_add_index
     migrate AddIndexConcurrently, direction: :down
   end
 
+  def test_lock_timeout_retries_add_index_safe_by_default
+    skip unless postgresql?
+
+    StrongMigrations.safe_by_default = true
+
+    assert_raises(ActiveRecord::LockWaitTimeout) do
+      with_lock_timeout_retries do
+        migrate AddIndex
+      end
+    end
+
+    migrate AddIndex, direction: :down
+  ensure
+    StrongMigrations.safe_by_default = false
+  end
+
   def reset_timeouts
     StrongMigrations.lock_timeout = nil
     StrongMigrations.statement_timeout = nil