safe-pg-migrations 2.1.0 → 2.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2fafd98448f0ff697da386a337dc22753be2144feedba4793f87a351581a5607
-  data.tar.gz: 9653ae6395d4c9f27b69dba9dee9142a4535384c42e8d05fac9ca2e2ef1bdce6
+  metadata.gz: 624821b7134dee05186350edada6f27faaf8c820976837d32403dc5fffa4880a
+  data.tar.gz: bf41324b8d563de7bc893a302d45d8f175e890bb888a2f52e0b14be7676f942c
 SHA512:
-  metadata.gz: 64297c9ba4f040e7afa21fefa2acb2a26e27fb0bfee3d9f9cbf902ee4bb0ac087f4e22591a7b5470cf73dc7bc73345e59904649c02d7540bd4693372e7069549
-  data.tar.gz: 7a2640fd2494480d560cfa6f19059550ecc316eaf729df6d8f97a434c189f4565c49d47acbb02f437d9b1285a466a5615a79a84d6ba1dfb3d94bb54f87b9e788
+  metadata.gz: 76601dcf0b97013edc36de6ac351beb8e4f6c532a0e223cba95d1a29e4eb7f18b8a4e7b3e3f28e4e7bde02144a4de8c512c19e2725396578dcb322e040963168
+  data.tar.gz: 767876e75af3292f9c4d3fc82810ab37bb3d38dad7b8b87ee6ba7068ec5e4aa9967ac0e53170be77bda07dd3424f2d3b3b3a03d9b8dc6ebcda94cee21e61fcfc

data/README.md

@@ -119,7 +119,9 @@ PG will still needs to update every row of the table, and will most likely state
 **Safe PG Migrations** provides the extra option parameter `default_value_backfill:`. When your migration is adding a volatile default value, the option `:update_in_batches` can be set. It will automatically backfill the value in a safe manner.
 
 ```ruby
-add_column :users, :created_at, default: 'clock_timestamp()', default_value_backfill: :update_in_batches
+safety_assured do
+  add_column :users, :created_at, default: 'clock_timestamp()', default_value_backfill: :update_in_batches
+end
 ```
 
 More specifically, it will: 
@@ -142,11 +144,17 @@ Data backfill take time. If your table is big, your migrations will (safely) han
 
 2. manual data backfill (rake task, manual operation, ...)
 3. Second migration which change the column to null false (with **Safe PG Migrations**, `change_column_null` is safe and can be used; see section below)
-
 ---
 
 `default_value_backfill:` also accept the value `:auto` which is set by default. In this case, **Safe PG Migrations** will not backfill data and will let PostgreSQL handle it itself.
 
+### Preventing :update_in_batches when the table is too big
+
+`add_column` with `default_value_backfill: :update_in_batches` can be dangerous on big tables. To avoid unwanted long migrations, **Safe PG Migrations** does not automatically mark this usage as safe when used with `strong-migrations`, usage of `safety_assured` is required.
+
+It is also possible to set a threshold for the table size, above which the migration will fail. This can be done by setting the `default_value_backfill_threshold:` option in the configuration.
+
+
 </details>
 
 <details><summary id="safe_add_remove_index">Safe <code>add_index</code> and <code>remove_index</code></summary>
@@ -238,6 +246,35 @@ When a statement fails with a lock timeout, **Safe PG Migrations** retries it (5
 <details><summary>Blocking activity logging</summary>
 
 If a statement fails with a lock timeout, **Safe PG Migrations** will try to tell you what was the blocking statement.
+
+---
+**NOTE**
+
+Data logged by the Blocking activity logger can be sensitive (it will contain raw SQL queries, which can be hashes of password, user information, ...)
+
+If you cannot afford to log this type of data, you can either
+* Set `SafePgMigrations.config.blocking_activity_logger_verbose = false`. In this case, the logger will only log the pid of the blocking statement, which should be enough to investigate;
+* Provide a different logger through `SafePgMigrations.config.sensitive_logger = YourLogger.new`. Instead of using the default IO stream, SafePgMigrations will send sensitive data to the given logger which can be managed as you wish.
+
+---
+
+</details>
+
+<details><summary>Dropping a table</summary>
+
+Dropping a table can be difficult to achieve in a small amount of time if it holds several foreign keys to busy tables.
+To remove the table, PostgreSQL will have to acquire an access exclusive lock on all the tables referenced by the foreign keys.
+
+To solve this issue, **Safe Pg Migrations** will drop the foreign keys before dropping the table.
+
+---
+**NOTE**
+
+Dropping a table is a dangerous operation by nature. **Safe Pg Migrations** will not prevent the deletion of a table which
+would still be in use.
+
+---
+
 </details>
 
 <details><summary>Verbose SQL logging</summary>
@@ -288,33 +325,32 @@ So you can actually check that the `CREATE INDEX` statement will be performed co
 **Safe PG Migrations** can be customized, here is an example of a Rails initializer (the values are the default ones):
 
 ```ruby
-SafePgMigrations.config.safe_timeout = 5.seconds # Lock and statement timeout used for all DDL operations except from CREATE / DROP INDEX
+SafePgMigrations.config.safe_timeout = 5.seconds # Statement timeout used for all DDL operations except from CREATE / DROP INDEX
+
+SafePgMigrations.config.lock_timeout = nil # Lock timeout used for all DDL operations except from CREATE / DROP INDEX. If not set, safe_timeout will be used with a deduction of 1% to ensure that the lock timeout is raised in priority
 
 SafePgMigrations.config.blocking_activity_logger_verbose = true # Outputs the raw blocking queries on timeout. When false, outputs information about the lock instead
 
+SafePgMigrations.config.sensitive_logger = nil # When given, sensitive data will be sent to this logger instead of the standard output. Must implement method `info`.
+
 SafePgMigrations.config.blocking_activity_logger_margin = 1.second # Delay to output blocking queries before timeout. Must be shorter than safe_timeout
 
 SafePgMigrations.config.backfill_batch_size = 100_000 # Size of the batches used for backfilling when adding a column with a default value
 
 SafePgMigrations.config.backfill_pause = 0.5.second # Delay between each batch during a backfill. This ensure replication can happen safely. 
 
+SafePgMigrations.config.default_value_backfill_threshold = nil # When set, batch backfill will only be available if the table is under the given threshold. If the number of rows is higher (according to stats), the migration will fail. 
+
 SafePgMigrations.config.retry_delay = 1.minute # Delay between retries for retryable statements
 
 SafePgMigrations.config.max_tries = 5 # Number of retries before abortion of the migration
 ```
 
-## Running tests
-
-```bash
-bundle
-psql -h localhost -c 'CREATE DATABASE safe_pg_migrations_test'
-rake test
-```
-
 ## Authors
 
 - [Matthieu Prat](https://github.com/matthieuprat)
 - [Romain Choquet](https://github.com/rchoquet)
+- [Thomas Hareau](https://github.com/ThHareau)
 - [Paul-Etienne Coisne](https://github.com/coisnepe)
 
 ## License

data/lib/safe-pg-migrations/base.rb

@@ -1,16 +1,20 @@
 # frozen_string_literal: true
 
 require 'safe-pg-migrations/configuration'
+require 'safe-pg-migrations/helpers/logger'
 require 'safe-pg-migrations/helpers/satisfied_helper'
 require 'safe-pg-migrations/helpers/index_helper'
 require 'safe-pg-migrations/helpers/batch_over'
+require 'safe-pg-migrations/helpers/session_setting_management'
 require 'safe-pg-migrations/plugins/verbose_sql_logger'
 require 'safe-pg-migrations/plugins/blocking_activity_logger'
 require 'safe-pg-migrations/plugins/statement_insurer/add_column'
+require 'safe-pg-migrations/plugins/statement_insurer/change_column_null'
 require 'safe-pg-migrations/plugins/statement_insurer'
 require 'safe-pg-migrations/plugins/statement_retrier'
 require 'safe-pg-migrations/plugins/idempotent_statements'
 require 'safe-pg-migrations/plugins/useless_statements_logger'
+require 'safe-pg-migrations/plugins/strong_migrations_integration'
 require 'safe-pg-migrations/polyfills/index_definition_polyfill'
 require 'safe-pg-migrations/polyfills/verbose_query_logs_polyfill'
 
@@ -31,15 +35,29 @@ module SafePgMigrations
     def setup_and_teardown(migration, connection, &block)
       @pg_version_num = get_pg_version_num(connection)
       @alternate_connection = nil
-      @current_migration = migration
-      stdout_sql_logger = VerboseSqlLogger.new.setup if verbose?
-      PLUGINS.each { |plugin| connection.extend(plugin) }
 
-      connection.with_setting :lock_timeout, SafePgMigrations.config.pg_lock_timeout, &block
+      with_current_migration(migration) do
+        stdout_sql_logger = VerboseSqlLogger.new.setup if verbose?
+
+        VerboseSqlLogger.new.setup if verbose?
+        PLUGINS.each { |plugin| connection.extend(plugin) }
+
+        connection.with_setting :lock_timeout, SafePgMigrations.config.pg_lock_timeout do
+          connection.with_setting :statement_timeout, SafePgMigrations.config.pg_statement_timeout, &block
+        end
+      ensure
+        stdout_sql_logger&.teardown
+      end
     ensure
       close_alternate_connection
+    end
+
+    def with_current_migration(migration, &block)
+      @current_migration = migration
+
+      yield block
+    ensure
       @current_migration = nil
-      stdout_sql_logger&.teardown
     end
 
     def alternate_connection
@@ -53,16 +71,6 @@ module SafePgMigrations
       @alternate_connection = nil
     end
 
-    ruby2_keywords def say(*args)
-      return unless current_migration
-
-      current_migration.say(*args)
-    end
-
-    ruby2_keywords def say_method_call(method, *args)
-      say "#{method}(#{args.map(&:inspect) * ', '})", true
-    end
-
     def verbose?
       unless current_migration.class._safe_pg_migrations_verbose.nil?
         return current_migration.class._safe_pg_migrations_verbose
@@ -83,6 +91,8 @@ module SafePgMigrations
   end
 
   module Migration
+    include StrongMigrationsIntegration
+
     module ClassMethods
       attr_accessor :_safe_pg_migrations_verbose
 
@@ -98,28 +108,11 @@ module SafePgMigrations
     end
 
     def disable_ddl_transaction
-      UselessStatementsLogger.warn_useless '`disable_ddl_transaction`' if super
-      true
-    end
+      SafePgMigrations.with_current_migration(self) do
+        UselessStatementsLogger.warn_useless '`disable_ddl_transaction`' if super
 
-    SAFE_METHODS = %i[
-      execute
-      add_column
-      add_index
-      add_reference
-      add_belongs_to
-      change_column_null
-      add_foreign_key
-      add_check_constraint
-    ].freeze
-
-    SAFE_METHODS.each do |method|
-      define_method method do |*args|
-        return super(*args) unless respond_to?(:safety_assured)
-
-        safety_assured { super(*args) }
+        true
       end
-      ruby2_keywords method
     end
   end
 end

data/lib/safe-pg-migrations/configuration.rb

@@ -4,17 +4,49 @@ require 'active_support/core_ext/numeric/time'
 
 module SafePgMigrations
   class Configuration
-    attr_accessor :safe_timeout, :blocking_activity_logger_margin, :blocking_activity_logger_verbose,
-                  :backfill_batch_size, :backfill_pause, :retry_delay, :max_tries
+    attr_accessor(*%i[
+                    blocking_activity_logger_margin
+                    blocking_activity_logger_verbose
+                    default_value_backfill_threshold
+                    backfill_batch_size
+                    backfill_pause
+                    retry_delay
+                    max_tries
+                    sensitive_logger
+                  ])
+    attr_reader :lock_timeout, :safe_timeout
 
     def initialize
+      self.default_value_backfill_threshold = nil
       self.safe_timeout = 5.seconds
+      self.lock_timeout = nil
       self.blocking_activity_logger_margin = 1.second
       self.blocking_activity_logger_verbose = true
       self.backfill_batch_size = 100_000
       self.backfill_pause = 0.5.second
       self.retry_delay = 1.minute
       self.max_tries = 5
+      self.sensitive_logger = nil
+    end
+
+    def lock_timeout=(value)
+      raise 'Setting lock timeout to 0 disables the lock timeout and is dangerous' if value == 0.seconds
+
+      unless value.nil? || value < safe_timeout
+        raise ArgumentError, "Lock timeout (#{value}) cannot be greater than safe timeout (#{safe_timeout})"
+      end
+
+      @lock_timeout = value
+    end
+
+    def safe_timeout=(value)
+      raise 'Setting safe timeout to 0 disables the safe timeout and is dangerous' unless value
+
+      unless lock_timeout.nil? || value > lock_timeout
+        raise ArgumentError, "Safe timeout (#{value}) cannot be less than lock timeout (#{lock_timeout})"
+      end
+
+      @safe_timeout = value
     end
 
     def pg_statement_timeout
@@ -22,6 +54,8 @@ module SafePgMigrations
     end
 
     def pg_lock_timeout
+      return pg_duration lock_timeout if lock_timeout
+
       # if statement timeout and lock timeout have the same value, statement timeout will raise in priority. We actually
       # need the opposite for BlockingActivityLogger to detect lock timeouts correctly.
       # By reducing the lock timeout by a very small margin, we ensure that the lock timeout is raised in priority

data/lib/safe-pg-migrations/helpers/blocking_activity_formatter.rb

@@ -5,20 +5,20 @@ module SafePgMigrations
     module BlockingActivityFormatter
       def log_queries(queries)
         if queries.empty?
-          SafePgMigrations.say 'Could not find any blocking query.', true
+          Logger.say 'Could not find any blocking query.', sub_item: true
         else
-          SafePgMigrations.say(
-            "Statement was being blocked by the following #{'query'.pluralize(queries.size)}:",
-            true
-          )
-          SafePgMigrations.say '', true
+          Logger.say <<~MESSAGE.squish, sub_item: true
+            Statement was being blocked by the following #{'query'.pluralize(queries.size)}:
+          MESSAGE
+
+          Logger.say '', sub_item: true
           output_blocking_queries(queries)
-          SafePgMigrations.say(
-            'Beware, some of those queries might run in a transaction. In this case the locking query might be ' \
-            'located elsewhere in the transaction',
-            true
-          )
-          SafePgMigrations.say '', true
+          Logger.say <<~MESSAGE.squish, sub_item: true
+            Beware, some of those queries might run in a transaction. In this case the locking query might be located
+            elsewhere in the transaction
+          MESSAGE
+
+          Logger.say '', sub_item: true
         end
       end
 
@@ -27,8 +27,10 @@ module SafePgMigrations
       def output_blocking_queries(queries)
         if SafePgMigrations.config.blocking_activity_logger_verbose
           queries.each do |pid, query, start_time|
-            SafePgMigrations.say "Query with pid #{pid || 'null'} started #{format_start_time start_time}:  #{query}",
-                                 true
+            Logger.say(
+              "Query with pid #{pid || 'null'} started #{format_start_time start_time}: #{query}",
+              sub_item: true, sensitive: true
+            )
           end
         else
           output_confidentially_blocking_queries(queries)
@@ -37,14 +39,13 @@ module SafePgMigrations
 
       def output_confidentially_blocking_queries(queries)
         queries.each do |start_time, locktype, mode, pid, transactionid|
-          SafePgMigrations.say(
-            "Query with pid #{pid || 'null'} " \
-            "started #{format_start_time(start_time)}: " \
-            "lock type: #{locktype || 'null'}, " \
-            "lock mode: #{mode || 'null'}, " \
-            "lock transactionid: #{transactionid || 'null'}",
-            true
-          )
+          Logger.say <<~MESSAGE.squish, sub_item: true, sensitive: true
+            Query with pid #{pid || 'null'}
+            started #{format_start_time(start_time)}:
+            lock type: #{locktype || 'null'},
+            lock mode: #{mode || 'null'},
+            lock transactionid: #{transactionid || 'null'}",
+          MESSAGE
         end
       end
 

data/lib/safe-pg-migrations/helpers/logger.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module SafePgMigrations
+  module Helpers
+    module Logger
+      class << self
+        def say(message, sub_item: false, sensitive: false, warn_sensitive_logs: true)
+          return unless SafePgMigrations.current_migration
+
+          if sensitive
+            log_sensitive message, sub_item: sub_item
+            if warn_sensitive_logs && sensitive_logger?
+              log 'Sensitive data sent to sensitive logger', sub_item: sub_item
+            end
+          else
+            log message, sub_item: sub_item
+          end
+        end
+
+        def say_method_call(method, *args, sensitive: false, warn_sensitive_logs: true, **options)
+          args += [options] unless options.empty?
+
+          say "#{method}(#{args.map(&:inspect) * ', '})",
+              sub_item: true, sensitive: sensitive, warn_sensitive_logs: warn_sensitive_logs
+        end
+
+        private
+
+        def log(message, sub_item:)
+          SafePgMigrations.current_migration.say message, sub_item
+        end
+
+        def log_sensitive(message, sub_item:)
+          if sensitive_logger?
+            SafePgMigrations.config.sensitive_logger.info message
+          else
+            log message, sub_item: sub_item
+          end
+        end
+
+        def sensitive_logger?
+          SafePgMigrations.config.sensitive_logger
+        end
+      end
+    end
+  end
+end

data/lib/safe-pg-migrations/helpers/session_setting_management.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module SafePgMigrations
+  module Helpers
+    module SessionSettingManagement
+      def with_setting(key, value)
+        old_value = query_value("SHOW #{key}")
+        execute("SET #{key} TO #{quote(value)}")
+        begin
+          yield
+        ensure
+          begin
+            execute("SET #{key} TO #{quote(old_value)}")
+          rescue ActiveRecord::StatementInvalid => e
+            # Swallow `PG::InFailedSqlTransaction` exceptions so as to keep the
+            # original exception (if any).
+            raise unless e.cause.is_a?(PG::InFailedSqlTransaction)
+          end
+        end
+      end
+
+      def without_statement_timeout(&block)
+        with_setting(:statement_timeout, 0, &block)
+      end
+
+      def without_lock_timeout(&block)
+        with_setting(:lock_timeout, 0, &block)
+      end
+
+      def without_timeout(&block)
+        without_statement_timeout { without_lock_timeout(&block) }
+      end
+    end
+  end
+end

data/lib/safe-pg-migrations/plugins/blocking_activity_logger.rb

@@ -5,8 +5,8 @@ require_relative '../helpers/blocking_activity_selector'
 
 module SafePgMigrations
   module BlockingActivityLogger
-    include ::SafePgMigrations::Helpers::BlockingActivityFormatter
-    include ::SafePgMigrations::Helpers::BlockingActivitySelector
+    include Helpers::BlockingActivityFormatter
+    include Helpers::BlockingActivitySelector
 
     %i[
       add_column
@@ -18,7 +18,19 @@ module SafePgMigrations
       create_table
     ].each do |method|
       define_method method do |*args, &block|
-        log_blocking_queries_after_lock { super(*args, &block) }
+        log_context = lambda do
+          break unless SafePgMigrations.config.sensitive_logger
+
+          options = args.last.is_a?(Hash) ? args.last : {}
+
+          Helpers::Logger.say "Executing #{SafePgMigrations.current_migration.name}",
+                              sensitive: true, warn_sensitive_logs: false
+          Helpers::Logger.say_method_call method, *args, **options, sensitive: true, warn_sensitive_logs: false
+        end
+
+        log_blocking_queries_after_lock(log_context) do
+          super(*args, &block)
+        end
       end
       ruby2_keywords method
     end
@@ -50,7 +62,7 @@ module SafePgMigrations
       blocking_queries_retriever_thread.kill
     end
 
-    def log_blocking_queries_after_lock
+    def log_blocking_queries_after_lock(log_context)
       blocking_queries_retriever_thread =
         Thread.new do
           sleep delay_before_logging
@@ -63,14 +75,15 @@ module SafePgMigrations
 
       blocking_queries_retriever_thread.kill
     rescue ActiveRecord::LockWaitTimeout
-      SafePgMigrations.say 'Lock timeout.', true
+      Helpers::Logger.say 'Lock timeout.', sub_item: true
+      log_context.call
       queries =
         begin
           blocking_queries_retriever_thread.value
         rescue StandardError => e
-          SafePgMigrations.say(
+          Helpers::Logger.say(
             "Error while retrieving blocking queries: #{e}",
-            true
+            sub_item: true
           )
           nil
         end
@@ -81,12 +94,15 @@ module SafePgMigrations
     end
 
     def delay_before_logging
-      SafePgMigrations.config.safe_timeout -
-        SafePgMigrations.config.blocking_activity_logger_margin
+      timeout - SafePgMigrations.config.blocking_activity_logger_margin
     end
 
     def delay_before_retry
       SafePgMigrations.config.blocking_activity_logger_margin + SafePgMigrations.config.retry_delay
     end
+
+    def timeout
+      SafePgMigrations.config.lock_timeout || SafePgMigrations.config.safe_timeout
+    end
   end
 end

data/lib/safe-pg-migrations/plugins/idempotent_statements.rb

@@ -2,7 +2,7 @@
 
 module SafePgMigrations
   module IdempotentStatements
-    include SafePgMigrations::Helpers::IndexHelper
+    include Helpers::IndexHelper
 
     ruby2_keywords def add_index(table_name, column_name, *args)
       options = args.last.is_a?(Hash) ? args.last : {}
@@ -12,10 +12,9 @@ module SafePgMigrations
       return super unless index_name_exists?(index_definition.table, index_definition.name)
 
       if index_valid?(index_definition.name)
-        SafePgMigrations.say(
-          "/!\\ Index '#{index_definition.name}' already exists in '#{table_name}'. Skipping statement.",
-          true
-        )
+        Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+          /!\\ Index '#{index_definition.name}' already exists in '#{table_name}'. Skipping statement.
+        MESSAGE
         return
       end
 
@@ -26,13 +25,17 @@ module SafePgMigrations
     ruby2_keywords def add_column(table_name, column_name, type, *)
       return super unless column_exists?(table_name, column_name)
 
-      SafePgMigrations.say("/!\\ Column '#{column_name}' already exists in '#{table_name}'. Skipping statement.", true)
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Column '#{column_name}' already exists in '#{table_name}'. Skipping statement.
+      MESSAGE
     end
 
     ruby2_keywords def remove_column(table_name, column_name, type = nil, *)
       return super if column_exists?(table_name, column_name)
 
-      SafePgMigrations.say("/!\\ Column '#{column_name}' not found on table '#{table_name}'. Skipping statement.", true)
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Column '#{column_name}' not found on table '#{table_name}'. Skipping statement.
+      MESSAGE
     end
 
     ruby2_keywords def remove_index(table_name, *args)
@@ -41,41 +44,41 @@ module SafePgMigrations
 
       return super if index_name_exists?(table_name, index_name)
 
-      SafePgMigrations.say("/!\\ Index '#{index_name}' not found on table '#{table_name}'. Skipping statement.", true)
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Index '#{index_name}' not found on table '#{table_name}'. Skipping statement.
+      MESSAGE
     end
 
     ruby2_keywords def add_foreign_key(from_table, to_table, *args)
       options = args.last.is_a?(Hash) ? args.last : {}
-      suboptions = options.slice(:name, :column)
-      return super unless foreign_key_exists?(from_table, suboptions.present? ? nil : to_table, **suboptions)
+      sub_options = options.slice(:name, :column)
+      return super unless foreign_key_exists?(from_table, sub_options.present? ? nil : to_table, **sub_options)
 
-      SafePgMigrations.say(
-        "/!\\ Foreign key '#{from_table}' -> '#{to_table}' already exists. Skipping statement.",
-        true
-      )
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Foreign key '#{from_table}' -> '#{to_table}' already exists. Skipping statement.
+      MESSAGE
     end
 
     def remove_foreign_key(from_table, to_table = nil, **options)
       return super if foreign_key_exists?(from_table, to_table, **options)
 
       reference_name = to_table || options[:to_table] || options[:column] || options[:name]
-      SafePgMigrations.say(
-        "/!\\ Foreign key '#{from_table}' -> '#{reference_name}' does not exist. Skipping statement.",
-        true
-      )
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Foreign key '#{from_table}' -> '#{reference_name}' does not exist. Skipping statement.
+      MESSAGE
     end
 
     ruby2_keywords def create_table(table_name, *args)
       options = args.last.is_a?(Hash) ? args.last : {}
       return super if options[:force] || !table_exists?(table_name)
 
-      SafePgMigrations.say "/!\\ Table '#{table_name}' already exists.", true
+      Helpers::Logger.say "/!\\ Table '#{table_name}' already exists.", sub_item: true
 
       td = create_table_definition(table_name, *args)
 
       yield td if block_given?
 
-      SafePgMigrations.say(td.indexes.empty? ? '-- Skipping statement' : '-- Creating indexes', true)
+      Helpers::Logger.say td.indexes.empty? ? '-- Skipping statement' : '-- Creating indexes', sub_item: true
 
       td.indexes.each do |column_name, index_options|
         add_index(table_name, column_name, **index_options)
@@ -88,7 +91,9 @@ module SafePgMigrations
 
       return super if constraint_definition.nil?
 
-      SafePgMigrations.say "/!\\ Constraint '#{constraint_definition.name}' already exists. Skipping statement.", true
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Constraint '#{constraint_definition.name}' already exists. Skipping statement.
+      MESSAGE
     end
 
     def change_column_null(table_name, column_name, null, *)
@@ -96,10 +101,9 @@ module SafePgMigrations
 
       return super if column.null != null
 
-      SafePgMigrations.say(
-        "/!\\ Column '#{table_name}.#{column.name}' is already set to 'null: #{null}'. Skipping statement.",
-        true
-      )
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Column '#{table_name}.#{column.name}' is already set to 'null: #{null}'. Skipping statement.
+      MESSAGE
     end
 
     def validate_check_constraint(table_name, **options)
@@ -107,8 +111,9 @@ module SafePgMigrations
 
       return super unless constraint_definition.validated?
 
-      SafePgMigrations.say "/!\\ Constraint '#{constraint_definition.name}' already validated. Skipping statement.",
-                           true
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Constraint '#{constraint_definition.name}' already validated. Skipping statement.
+      MESSAGE
     end
 
     def change_column_default(table_name, column_name, default_or_changes)
@@ -122,11 +127,17 @@ module SafePgMigrations
 
       return super if new_alter_statement != previous_alter_statement
 
-      SafePgMigrations.say(
-        "/!\\ Column '#{table_name}.#{column.name}' is already set to 'default: #{column.default}'. " \
-        'Skipping statement.',
-        true
-      )
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Column '#{table_name}.#{column.name}' is already set to 'default: #{column.default}'. Skipping statement.
+      MESSAGE
+    end
+
+    ruby2_keywords def drop_table(table_name, *)
+      return super if table_exists?(table_name)
+
+      Helpers::Logger.say <<~MESSAGE.squish, sub_item: true
+        /!\\ Table '#{table_name} does not exist. Skipping statement.
+      MESSAGE
     end
   end
 end

data/lib/safe-pg-migrations/plugins/statement_insurer/add_column.rb

@@ -7,25 +7,27 @@ module SafePgMigrations
         options = args.last.is_a?(Hash) && args.last
         options ||= {}
 
-        if should_keep_default_implementation?(**options)
-          with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) { return super }
-        end
+        return super if should_keep_default_implementation?(**options)
+
+        raise <<~ERROR unless backfill_column_default_safe?(table_name)
+          Table #{table_name} has more than #{SafePgMigrations.config.default_value_backfill_threshold} rows.
+          Backfilling the default value for column #{column_name} on table #{table_name} would take too long.
+
+          Please revert this migration, and backfill the default value manually.
+
+          This check is configurable through the configuration "default_value_backfill_threshold".
+        ERROR
 
         default = options.delete(:default)
         null = options.delete(:null)
 
-        with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-          SafePgMigrations.say_method_call(:add_column, table_name, column_name, type, options)
-          super table_name, column_name, type, **options
-        end
+        Helpers::Logger.say_method_call(:add_column, table_name, column_name, type, options)
+        super table_name, column_name, type, **options
 
-        SafePgMigrations.say_method_call(:change_column_default, table_name, column_name, default)
+        Helpers::Logger.say_method_call(:change_column_default, table_name, column_name, default)
         change_column_default(table_name, column_name, default)
 
-        SafePgMigrations.say_method_call(:backfill_column_default, table_name, column_name)
-        without_statement_timeout do
-          backfill_column_default(table_name, column_name)
-        end
+        backfill_column_default(table_name, column_name)
 
         change_column_null(table_name, column_name, null) if null == false
       end
@@ -34,16 +36,27 @@ module SafePgMigrations
 
       def should_keep_default_implementation?(default: nil, default_value_backfill: :auto, **)
         default_value_backfill != :update_in_batches || !default ||
-          !SafePgMigrations::Helpers::SatisfiedHelper.satisfies_add_column_update_rows_backfill?
+          !Helpers::SatisfiedHelper.satisfies_add_column_update_rows_backfill?
+      end
+
+      def backfill_column_default_safe?(table_name)
+        return true if SafePgMigrations.config.default_value_backfill_threshold.nil?
+
+        row, = query("SELECT reltuples AS estimate FROM pg_class where relname = '#{table_name}';")
+        estimate, = row
+
+        estimate <= SafePgMigrations.config.default_value_backfill_threshold
       end
 
       def backfill_column_default(table_name, column_name)
         model = Class.new(ActiveRecord::Base) { self.table_name = table_name }
         quoted_column_name = quote_column_name(column_name)
 
-        SafePgMigrations::Helpers::BatchOver.new(model).each_batch do |batch|
-          batch
-            .update_all("#{quoted_column_name} = DEFAULT")
+        Helpers::Logger.say_method_call(:backfill_column_default, table_name, column_name)
+
+        Helpers::BatchOver.new(model).each_batch do |batch|
+          batch.update_all("#{quoted_column_name} = DEFAULT")
+
           sleep SafePgMigrations.config.backfill_pause
         end
       end

data/lib/safe-pg-migrations/plugins/statement_insurer/change_column_null.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module SafePgMigrations
+  module StatementInsurer
+    module ChangeColumnNull
+      def change_column_null(table_name, column_name, null, default = nil)
+        return super unless should_create_constraint? default, null
+
+        expression = "#{column_name} IS NOT NULL"
+        # constraint will be defined if the constraint was manually created in another migration
+        constraint = check_constraint_by_expression table_name, expression
+
+        default_name = check_constraint_name(table_name, expression: expression)
+        constraint_name = constraint&.name || default_name
+
+        add_check_constraint table_name, expression, name: constraint_name
+
+        Helpers::Logger.say_method_call :change_column_null, table_name, column_name, false
+        super table_name, column_name, false
+
+        return unless should_remove_constraint? default_name, constraint_name
+
+        Helpers::Logger.say_method_call :remove_check_constraint, table_name, expression, name: constraint_name
+        remove_check_constraint table_name, expression, name: constraint_name
+      end
+
+      private
+
+      def check_constraint_by_expression(table_name, expression)
+        check_constraints(table_name).detect { |check_constraint| check_constraint.expression = expression }
+      end
+
+      def should_create_constraint?(default, null)
+        !default && !null && Helpers::SatisfiedHelper.satisfies_change_column_null_requirements?
+      end
+
+      def should_remove_constraint?(default_name, constraint_name)
+        # we don't want to remove the constraint if it was created in another migration. The best guess we have here is
+        # that manually created constraint would likely have a name that is not the default name. This is not a perfect,
+        # a manually created constraint without a name would be removed. However, it is now replaced by the NOT NULL
+        # statement on the table, so this is not a big issue.
+        default_name == constraint_name
+      end
+    end
+  end
+end

data/lib/safe-pg-migrations/plugins/statement_insurer.rb

@@ -2,19 +2,13 @@
 
 module SafePgMigrations
   module StatementInsurer
+    include Helpers::SessionSettingManagement
     include AddColumn
-
-    %i[change_column].each do |method|
-      define_method method do |*args, &block|
-        with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) { super(*args, &block) }
-      end
-      ruby2_keywords method
-    end
+    include ChangeColumnNull
 
     def validate_check_constraint(table_name, **options)
-      without_statement_timeout do
-        super
-      end
+      Helpers::Logger.say_method_call :validate_check_constraint, table_name, **options
+      without_statement_timeout { super }
     end
 
     def add_check_constraint(table_name, expression, **options)
@@ -23,37 +17,38 @@ module SafePgMigrations
 
       options = check_constraint_options(table_name, expression, options)
 
-      SafePgMigrations.say_method_call :add_check_constraint, table_name, expression, **options, validate: false
+      Helpers::Logger.say_method_call :add_check_constraint, table_name, expression, **options,
+        validate: false
       super table_name, expression, **options, validate: false
 
       return unless options.fetch(:validate, true)
 
-      SafePgMigrations.say_method_call :validate_check_constraint, table_name, name: options[:name]
       validate_check_constraint table_name, name: options[:name]
     end
 
+    def validate_foreign_key(*, **)
+      without_statement_timeout { super }
+    end
+
     ruby2_keywords def add_foreign_key(from_table, to_table, *args)
       options = args.last.is_a?(Hash) ? args.last : {}
       validate_present = options.key?(:validate)
       options[:validate] = false unless validate_present
-      with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-        super(from_table, to_table, **options)
-      end
+
+      super(from_table, to_table, **options)
 
       return if validate_present
 
-      suboptions = options.slice(:name, :column)
-      without_statement_timeout { validate_foreign_key from_table, suboptions.present? ? nil : to_table, **suboptions }
+      sub_options = options.slice(:name, :column)
+      validate_foreign_key from_table, sub_options.present? ? nil : to_table, **sub_options
     end
 
     ruby2_keywords def create_table(*)
-      with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-        super do |td|
-          yield td if block_given?
-          td.indexes.map! do |key, index_options|
-            index_options[:algorithm] ||= :default
-            [key, index_options]
-          end
+      super do |td|
+        yield td if block_given?
+        td.indexes.map! do |key, index_options|
+          index_options[:algorithm] ||= :default
+          [key, index_options]
         end
       end
     end
@@ -67,61 +62,33 @@ module SafePgMigrations
         options[:algorithm] = :concurrently
       end
 
-      SafePgMigrations.say_method_call(:add_index, table_name, column_name, **options)
-
+      Helpers::Logger.say_method_call(:add_index, table_name, column_name, **options)
       without_timeout { super(table_name, column_name, **options) }
     end
 
     ruby2_keywords def remove_index(table_name, *args)
       options = args.last.is_a?(Hash) ? args.last : { column: args.last }
       options[:algorithm] = :concurrently unless options.key?(:algorithm)
-      SafePgMigrations.say_method_call(:remove_index, table_name, **options)
 
+      Helpers::Logger.say_method_call(:remove_index, table_name, **options)
       without_timeout { super(table_name, **options) }
     end
 
-    def change_column_null(table_name, column_name, null, default = nil)
-      if default || null || !Helpers::SatisfiedHelper.satisfies_change_column_null_requirements?
-        with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) { return super }
-      end
-
-      add_check_constraint table_name, "#{column_name} IS NOT NULL"
+    def remove_column(table_name, column_name, *)
+      foreign_key = foreign_key_for(table_name, column: column_name)
 
-      SafePgMigrations.say_method_call :change_column_null, table_name, column_name, false
-      with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-        super table_name, column_name, false
-      end
-
-      SafePgMigrations.say_method_call :remove_check_constraint, table_name, "#{column_name} IS NOT NULL"
-      remove_check_constraint table_name, "#{column_name} IS NOT NULL"
+      remove_foreign_key(table_name, name: foreign_key.name) if foreign_key
+      super
     end
 
-    def with_setting(key, value)
-      old_value = query_value("SHOW #{key}")
-      execute("SET #{key} TO #{quote(value)}")
-      begin
-        yield
-      ensure
-        begin
-          execute("SET #{key} TO #{quote(old_value)}")
-        rescue ActiveRecord::StatementInvalid => e
-          # Swallow `PG::InFailedSqlTransaction` exceptions so as to keep the
-          # original exception (if any).
-          raise unless e.cause.is_a?(PG::InFailedSqlTransaction)
-        end
+    ruby2_keywords def drop_table(table_name, *args)
+      foreign_keys(table_name).each do |foreign_key|
+        remove_foreign_key(table_name, name: foreign_key.name)
       end
-    end
-
-    def without_statement_timeout(&block)
-      with_setting(:statement_timeout, 0, &block)
-    end
 
-    def without_lock_timeout(&block)
-      with_setting(:lock_timeout, 0, &block)
-    end
+      Helpers::Logger.say_method_call :drop_table, table_name, *args
 
-    def without_timeout(&block)
-      without_statement_timeout { without_lock_timeout(&block) }
+      super(table_name, *args)
     end
   end
 end

data/lib/safe-pg-migrations/plugins/statement_retrier.rb

@@ -25,9 +25,9 @@ module SafePgMigrations
         raise unless remaining_tries > 0
 
         retry_delay = SafePgMigrations.config.retry_delay
-        SafePgMigrations.say "Retrying in #{retry_delay} seconds...", true
+        Helpers::Logger.say "Retrying in #{retry_delay} seconds...", sub_item: true
         sleep retry_delay
-        SafePgMigrations.say 'Retrying now.', true
+        Helpers::Logger.say 'Retrying now.', sub_item: true
         retry
       end
     end

data/lib/safe-pg-migrations/plugins/strong_migrations_integration.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module SafePgMigrations
+  module StrongMigrationsIntegration
+    class << self
+      def initialize
+        return unless strong_migration_available?
+
+        StrongMigrations.disable_check(:add_column_default)
+        StrongMigrations.disable_check(:add_column_default_callable)
+        StrongMigrations.add_check do |method, args|
+          next unless method == :add_column
+
+          options = args.last.is_a?(Hash) ? args.last : {}
+
+          default_value_backfill = options.fetch(:default_value_backfill, :auto)
+
+          if default_value_backfill == :update_in_batches
+            check_message = <<~CHECK
+              default_value_backfill: :update_in_batches will take time if the table is too big.
+
+              Your configuration sets a pause of #{SafePgMigrations.config.backfill_pause} seconds between batches of
+              #{SafePgMigrations.config.backfill_batch_size} rows. Each batch execution will take time as well. Please
+              check that the estimated duration of the migration is acceptable
+              before adding `safety_assured`.
+            CHECK
+
+            check_message += <<~CHECK if SafePgMigrations.config.default_value_backfill_threshold
+
+              Also, please note that SafePgMigrations is configured to raise if the table has more than
+              #{SafePgMigrations.config.default_value_backfill_threshold} rows.
+            CHECK
+
+            stop! check_message
+          end
+        end
+      end
+
+      private
+
+      def strong_migration_available?
+        Object.const_defined? :StrongMigrations
+      end
+    end
+
+    SAFE_METHODS = %i[
+      execute
+      add_index
+      add_reference
+      add_belongs_to
+      change_column_null
+      add_foreign_key
+      add_check_constraint
+    ].freeze
+
+    SAFE_METHODS.each do |method|
+      define_method method do |*args|
+        return super(*args) unless respond_to?(:safety_assured)
+
+        safety_assured { super(*args) }
+      end
+      ruby2_keywords method
+    end
+
+    ruby2_keywords def add_column(table_name, *args)
+      return super(table_name, *args) unless respond_to?(:safety_assured)
+
+      options = args.last.is_a?(Hash) ? args.last : {}
+
+      return safety_assured { super(table_name, *args) } if options.fetch(:default_value_backfill, :auto) == :auto
+
+      super(table_name, *args)
+    end
+  end
+end

data/lib/safe-pg-migrations/plugins/useless_statements_logger.rb

@@ -4,8 +4,10 @@ module SafePgMigrations
   module UselessStatementsLogger
     class << self
       ruby2_keywords def warn_useless(action, link = nil, *args)
-        SafePgMigrations.say "/!\\ No need to explicitly use #{action}, safe-pg-migrations does it for you", *args
-        SafePgMigrations.say "\t see #{link} for more details", *args if link
+        Helpers::Logger.say(
+          "/!\\ No need to explicitly use #{action}, safe-pg-migrations does it for you", *args
+        )
+        Helpers::Logger.say "\t see #{link} for more details", *args if link
       end
     end
 

data/lib/safe-pg-migrations/railtie.rb

@@ -9,6 +9,8 @@ module SafePgMigrations
         ActiveRecord::Migration.prepend(SafePgMigrations::Migration)
         ActiveRecord::Migration.singleton_class.prepend(SafePgMigrations::Migration::ClassMethods)
       end
+
+      SafePgMigrations::StrongMigrationsIntegration.initialize
     end
   end
 end

data/lib/safe-pg-migrations/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module SafePgMigrations
-  VERSION = '2.1.0'
+  VERSION = '2.2.1'
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: safe-pg-migrations
 version: !ruby/object:Gem::Version
-  version: 2.1.0
+  version: 2.2.1
 platform: ruby
 authors:
 - Matthieu Prat
@@ -10,7 +10,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-06-29 00:00:00.000000000 Z
+date: 2023-07-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -55,12 +55,16 @@ files:
 - lib/safe-pg-migrations/helpers/blocking_activity_formatter.rb
 - lib/safe-pg-migrations/helpers/blocking_activity_selector.rb
 - lib/safe-pg-migrations/helpers/index_helper.rb
+- lib/safe-pg-migrations/helpers/logger.rb
 - lib/safe-pg-migrations/helpers/satisfied_helper.rb
+- lib/safe-pg-migrations/helpers/session_setting_management.rb
 - lib/safe-pg-migrations/plugins/blocking_activity_logger.rb
 - lib/safe-pg-migrations/plugins/idempotent_statements.rb
 - lib/safe-pg-migrations/plugins/statement_insurer.rb
 - lib/safe-pg-migrations/plugins/statement_insurer/add_column.rb
+- lib/safe-pg-migrations/plugins/statement_insurer/change_column_null.rb
 - lib/safe-pg-migrations/plugins/statement_retrier.rb
+- lib/safe-pg-migrations/plugins/strong_migrations_integration.rb
 - lib/safe-pg-migrations/plugins/useless_statements_logger.rb
 - lib/safe-pg-migrations/plugins/verbose_sql_logger.rb
 - lib/safe-pg-migrations/polyfills/index_definition_polyfill.rb