From e3a5456a3d29f766e86446dc07d4072ea202fc70 Mon Sep 17 00:00:00 2001
From: jslucas <jspencerlucas@gmail.com>
Date: Wed, 7 May 2025 00:34:26 -0400
Subject: [PATCH] Allow Safeguards to be added and/or overridden outside of
 core

---
 lib/database_cleaner/safeguard.rb       |  49 +++-
 spec/database_cleaner/safeguard_spec.rb | 293 ++++++++++++++----------
 2 files changed, 209 insertions(+), 133 deletions(-)

diff --git a/lib/database_cleaner/safeguard.rb b/lib/database_cleaner/safeguard.rb
index ec81afa7..1a5fbeee 100644
--- a/lib/database_cleaner/safeguard.rb
+++ b/lib/database_cleaner/safeguard.rb
@@ -22,7 +22,27 @@ def initialize
       end
     end
 
-    class AllowedUrl
+    # Base class for core-consumers to implement Safeguards
+    #
+    class Base
+      def run
+        raise NotImplementedError
+      end
+
+      def self.inherited(subclass)
+        DatabaseCleaner::Safeguard.registry << subclass
+        DatabaseCleaner::Safeguard.deprecated_registry.reject! do |const|
+          subclass.name.split("::").last == const.name.split("::").last
+        end
+      end
+    end
+
+    # Just a marker class for Safeguards implemented by core that are kept for backwards compatibility.
+    # Adapters should implement their own Safeguards using Safeguard::Base.
+    #
+    class Deprecated; end
+
+    class AllowedUrl < Deprecated
       def run
         return if skip?
         raise Error::UrlNotAllowed if database_url_not_allowed?
@@ -39,8 +59,7 @@ def skip?
         end
     end
 
-
-    class RemoteDatabaseUrl
+    class RemoteDatabaseUrl < Deprecated
       LOCAL = %w(localhost 127.0.0.1)
 
       def run
@@ -73,7 +92,7 @@ def skip?
         end
     end
 
-    class Production
+    class Production < Deprecated
       KEYS = %w(ENV APP_ENV RACK_ENV RAILS_ENV)
 
       def run
@@ -96,14 +115,30 @@ def skip?
         end
     end
 
-    CHECKS = [
+    def run
+      self.class.registry.each { |const| const.new.run }
+      self.class.deprecated_registry.each { |const| const.new.run }
+    end
+
+    @registry = []
+    @deprecated_registry = [
       RemoteDatabaseUrl,
       Production,
       AllowedUrl
     ]
 
-    def run
-      CHECKS.each { |const| const.new.run }
+    class << self
+      attr_reader :registry
+      attr_reader :deprecated_registry
+
+      def reset_registry!
+        @registry = []
+        @deprecated_registry = [
+          RemoteDatabaseUrl,
+          Production,
+          AllowedUrl
+        ]
+      end
     end
   end
 end
diff --git a/spec/database_cleaner/safeguard_spec.rb b/spec/database_cleaner/safeguard_spec.rb
index b4459a96..d6a2c2e0 100644
--- a/spec/database_cleaner/safeguard_spec.rb
+++ b/spec/database_cleaner/safeguard_spec.rb
@@ -2,200 +2,241 @@ module DatabaseCleaner
   RSpec.describe Safeguard do
     let(:cleaner)  { Cleaners.new }
 
-    describe 'DATABASE_URL is set' do
-      before { stub_const('ENV', 'DATABASE_URL' => database_url) }
+    describe 'Safeguards defined outside of core' do
+      after { described_class.reset_registry! }
+
+      describe 'when a subclass of Safeguard::Base is defined' do
+        before do
+          class AlwaysRaise < DatabaseCleaner::Safeguard::Base
+            def run
+              raise "Some error"
+            end
+          end
+        end
 
-      describe 'to any value' do
-        let(:database_url) { 'postgres://remote.host' }
+        after { DatabaseCleaner.send(:remove_const, :AlwaysRaise) if DatabaseCleaner.const_defined?(:AlwaysRaise) }
 
-        it 'raises' do
-          expect { cleaner.start }.to raise_error(Safeguard::Error::RemoteDatabaseUrl)
+        it 'is run' do
+          expect { cleaner.start }.to raise_error("Some error")
         end
       end
 
-      describe 'to a localhost url' do
-        let(:database_url) { 'postgres://localhost' }
+      describe 'when a subclass of Safeguard::Base shares the name of a Safeguard defined in core' do
+        before do
+          class RemoteDatabaseUrl < DatabaseCleaner::Safeguard::Base
+            def run
+              :ok
+            end
+          end
+        end
+
+        after { DatabaseCleaner.send(:remove_const, :RemoteDatabaseUrl) if DatabaseCleaner.const_defined?(:RemoteDatabaseUrl) }
+
+        it 'runs the Safeguard::Base subclass instead' do
+          expect_any_instance_of(described_class::RemoteDatabaseUrl).not_to receive(:run)
+          expect_any_instance_of(RemoteDatabaseUrl).to receive(:run)
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          cleaner.start
         end
       end
+    end
 
-      describe 'to a local, empty-host url' do
-        let(:database_url) { 'postgres:///' }
+    describe 'behavior of deprecated, core Safeguards' do
+      describe 'DATABASE_URL is set' do
+        before { stub_const('ENV', 'DATABASE_URL' => database_url) }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+        describe 'to any value' do
+          let(:database_url) { 'postgres://remote.host' }
+
+          it 'raises' do
+            expect { cleaner.start }.to raise_error(Safeguard::Error::RemoteDatabaseUrl)
+          end
         end
-      end
 
-      describe 'to a local tld url' do
-        let(:database_url) { 'postgres://postgres.local' }
+        describe 'to a localhost url' do
+          let(:database_url) { 'postgres://localhost' }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
         end
-      end
 
-      describe 'to a 127.0.0.1 url' do
-        let(:database_url) { 'postgres://127.0.0.1' }
+        describe 'to a local, empty-host url' do
+          let(:database_url) { 'postgres:///' }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
         end
-      end
 
-      describe 'to a sqlite url' do
-        let(:database_url) { 'sqlite://tmp/db.sqlite3' }
+        describe 'to a local tld url' do
+          let(:database_url) { 'postgres://postgres.local' }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
         end
-      end
 
-      describe 'to a sqlite3 url' do
-        let(:database_url) { 'sqlite3://tmp/db.sqlite3' }
+        describe 'to a 127.0.0.1 url' do
+          let(:database_url) { 'postgres://127.0.0.1' }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
         end
-      end
 
-      describe 'to a sqlite3 url with no slashes after the scheme' do
-        let(:database_url) { 'sqlite3:tmp/db.sqlite3' }
+        describe 'to a sqlite url' do
+          let(:database_url) { 'sqlite://tmp/db.sqlite3' }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
         end
-      end
 
-      describe 'DATABASE_CLEANER_ALLOW_REMOTE_DATABASE_URL is set' do
-        let(:database_url) { 'postgres://remote.host' }
-        before { stub_const('ENV', 'DATABASE_CLEANER_ALLOW_REMOTE_DATABASE_URL' => true) }
+        describe 'to a sqlite3 url' do
+          let(:database_url) { 'sqlite3://tmp/db.sqlite3' }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
         end
-      end
 
-      describe 'DatabaseCleaner.allow_remote_database_url is true' do
-        let(:database_url) { 'postgres://remote.host' }
-        before { DatabaseCleaner.allow_remote_database_url = true }
-        after  { DatabaseCleaner.allow_remote_database_url = nil }
+        describe 'to a sqlite3 url with no slashes after the scheme' do
+          let(:database_url) { 'sqlite3:tmp/db.sqlite3' }
 
-        it 'does not raise' do
-          expect { cleaner.start }.to_not raise_error
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
+        end
+
+        describe 'DATABASE_CLEANER_ALLOW_REMOTE_DATABASE_URL is set' do
+          let(:database_url) { 'postgres://remote.host' }
+          before { stub_const('ENV', 'DATABASE_CLEANER_ALLOW_REMOTE_DATABASE_URL' => true) }
+
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
+        end
+
+        describe 'DatabaseCleaner.allow_remote_database_url is true' do
+          let(:database_url) { 'postgres://remote.host' }
+          before { DatabaseCleaner.allow_remote_database_url = true }
+          after  { DatabaseCleaner.allow_remote_database_url = nil }
+
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
         end
-      end
 
-      describe 'DatabaseCleaner.url_allowlist is set' do
+        describe 'DatabaseCleaner.url_allowlist is set' do
 
-        shared_examples 'allowlist assigned' do
-          describe 'A remote url is on the allowlist' do
-            let(:database_url) { 'postgres://foo.bar' }
+          shared_examples 'allowlist assigned' do
+            describe 'A remote url is on the allowlist' do
+              let(:database_url) { 'postgres://foo.bar' }
 
-            it 'does not raise' do
-              expect { cleaner.start }.to_not raise_error
+              it 'does not raise' do
+                expect { cleaner.start }.to_not raise_error
+              end
             end
-          end
 
-          describe 'A remote url is not on the allowlist' do
-            let(:database_url) { 'postgress://bar.baz' }
+            describe 'A remote url is not on the allowlist' do
+              let(:database_url) { 'postgress://bar.baz' }
 
-            it 'raises a not allowed error' do
-              expect { cleaner.start }.to raise_error(Safeguard::Error::UrlNotAllowed)
+              it 'raises a not allowed error' do
+                expect { cleaner.start }.to raise_error(Safeguard::Error::UrlNotAllowed)
+              end
             end
-          end
 
-          describe 'A similar url not explicitly matched as a pattern' do
-            let(:database_url) { 'postgres://foo.bar?pool=8' }
+            describe 'A similar url not explicitly matched as a pattern' do
+              let(:database_url) { 'postgres://foo.bar?pool=8' }
 
-            it 'raises a not allowed error' do
-              expect { cleaner.start }.to raise_error(Safeguard::Error::UrlNotAllowed)
+              it 'raises a not allowed error' do
+                expect { cleaner.start }.to raise_error(Safeguard::Error::UrlNotAllowed)
+              end
             end
-          end
 
-          describe 'A remote url matches a pattern on the allowlist' do
-            let(:database_url) { 'postgres://bar.baz?pool=16' }
+            describe 'A remote url matches a pattern on the allowlist' do
+              let(:database_url) { 'postgres://bar.baz?pool=16' }
 
-            it 'does not raise' do
-              expect { cleaner.start }.to_not raise_error
+              it 'does not raise' do
+                expect { cleaner.start }.to_not raise_error
+              end
             end
-          end
 
-          describe 'A local url is on the allowlist' do
-            let(:database_url) { 'postgres://postgres@localhost' }
+            describe 'A local url is on the allowlist' do
+              let(:database_url) { 'postgres://postgres@localhost' }
 
-            it 'does not raise' do
-              expect { cleaner.start }.to_not raise_error
+              it 'does not raise' do
+                expect { cleaner.start }.to_not raise_error
+              end
             end
-          end
 
-          describe 'A local url is not on the allowlist' do
-            let(:database_url) { 'postgres://localhost' }
+            describe 'A local url is not on the allowlist' do
+              let(:database_url) { 'postgres://localhost' }
 
-            it 'raises a not allowed error' do
-              expect { cleaner.start }.to raise_error(Safeguard::Error::UrlNotAllowed)
+              it 'raises a not allowed error' do
+                expect { cleaner.start }.to raise_error(Safeguard::Error::UrlNotAllowed)
+              end
             end
-          end
 
-          describe 'A url that matches a proc' do
-            let(:database_url) { 'redis://test:test@foo.bar' }
+            describe 'A url that matches a proc' do
+              let(:database_url) { 'redis://test:test@foo.bar' }
 
-            it 'does not raise' do
-              expect { cleaner.start }.to_not raise_error
+              it 'does not raise' do
+                expect { cleaner.start }.to_not raise_error
+              end
             end
           end
-        end
 
-        let(:url_allowlist) do
-          [
-            'postgres://postgres@localhost',
-            'postgres://foo.bar',
-            %r{^postgres://bar.baz},
-            proc { |x| URI.parse(x).user == 'test' }
-          ]
-        end
+          let(:url_allowlist) do
+            [
+              'postgres://postgres@localhost',
+              'postgres://foo.bar',
+              %r{^postgres://bar.baz},
+              proc { |x| URI.parse(x).user == 'test' }
+            ]
+          end
 
-        describe 'url_allowlist' do
-          before { DatabaseCleaner.url_allowlist = url_allowlist }
-          after  { DatabaseCleaner.url_allowlist = nil }
-          it_behaves_like 'allowlist assigned'
-        end
+          describe 'url_allowlist' do
+            before { DatabaseCleaner.url_allowlist = url_allowlist }
+            after  { DatabaseCleaner.url_allowlist = nil }
+            it_behaves_like 'allowlist assigned'
+          end
 
-        describe 'url_whitelist' do
-          before { DatabaseCleaner.url_whitelist = url_allowlist }
-          after  { DatabaseCleaner.url_whitelist = nil }
-          it_behaves_like 'allowlist assigned'
-        end
+          describe 'url_whitelist' do
+            before { DatabaseCleaner.url_whitelist = url_allowlist }
+            after  { DatabaseCleaner.url_whitelist = nil }
+            it_behaves_like 'allowlist assigned'
+          end
 
+        end
       end
-    end
 
-    describe 'ENV is set to production' do
-      %w(ENV APP_ENV RACK_ENV RAILS_ENV).each do |key|
-        describe "on #{key}" do
-          before { stub_const('ENV', key => "production") }
+      describe 'ENV is set to production' do
+        %w(ENV APP_ENV RACK_ENV RAILS_ENV).each do |key|
+          describe "on #{key}" do
+            before { stub_const('ENV', key => "production") }
 
-          it 'raises' do
-            expect { cleaner.start }.to raise_error(Safeguard::Error::ProductionEnv)
+            it 'raises' do
+              expect { cleaner.start }.to raise_error(Safeguard::Error::ProductionEnv)
+            end
           end
-        end
 
-        describe 'DATABASE_CLEANER_ALLOW_PRODUCTION is set' do
-          before { stub_const('ENV', 'DATABASE_CLEANER_ALLOW_PRODUCTION' => true) }
+          describe 'DATABASE_CLEANER_ALLOW_PRODUCTION is set' do
+            before { stub_const('ENV', 'DATABASE_CLEANER_ALLOW_PRODUCTION' => true) }
 
-          it 'does not raise' do
-            expect { cleaner.start }.to_not raise_error
+            it 'does not raise' do
+              expect { cleaner.start }.to_not raise_error
+            end
           end
-        end
 
-        describe 'DatabaseCleaner.allow_production is true' do
-          before { DatabaseCleaner.allow_production = true }
-          after  { DatabaseCleaner.allow_production = nil }
+          describe 'DatabaseCleaner.allow_production is true' do
+            before { DatabaseCleaner.allow_production = true }
+            after  { DatabaseCleaner.allow_production = nil }
 
-          it 'does not raise' do
-            expect { cleaner.start }.to_not raise_error
+            it 'does not raise' do
+              expect { cleaner.start }.to_not raise_error
+            end
           end
         end
       end