Set statement timeout by default (#106)

* Set statement timeout by default

* Style

* Isolate statement insurer in its helper

* Post rebase

* review

lib/safe-pg-migrations/base.rb

@@ -5,6 +5,7 @@
 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'
@@ -40,7 +41,9 @@ def setup_and_teardown(migration, connection, &block)
         VerboseSqlLogger.new.setup if verbose?
         PLUGINS.each { |plugin| connection.extend(plugin) }
 
-        connection.with_setting :lock_timeout, SafePgMigrations.config.pg_lock_timeout, &block
+        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

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

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

@@ -1,20 +1,13 @@
 # frozen_string_literal: true
 
 module SafePgMigrations
-  module StatementInsurer # rubocop:disable Metrics/ModuleLength
+  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
-
     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)
@@ -24,37 +17,37 @@ def add_check_constraint(table_name, expression, **options)
       options = check_constraint_options(table_name, expression, options)
 
       Helpers::Logger.say_method_call :add_check_constraint, table_name, expression, **options,
-validate: false
+        validate: false
       super table_name, expression, **options, validate: false
 
       return unless options.fetch(:validate, true)
 
-      Helpers::Logger.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
@@ -69,84 +62,44 @@ def add_check_constraint(table_name, expression, **options)
       end
 
       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)
-      Helpers::Logger.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
+      return super if default || null || !Helpers::SatisfiedHelper.satisfies_change_column_null_requirements?
 
       add_check_constraint table_name, "#{column_name} IS NOT NULL"
 
       Helpers::Logger.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
+      super table_name, column_name, false
 
-      Helpers::Logger.say_method_call :remove_check_constraint, table_name,
-                                      "#{column_name} IS NOT NULL"
+      Helpers::Logger.say_method_call :remove_check_constraint, table_name, "#{column_name} IS NOT NULL"
       remove_check_constraint table_name, "#{column_name} IS NOT NULL"
     end
 
     def remove_column(table_name, column_name, *)
       foreign_key = foreign_key_for(table_name, column: column_name)
 
-      with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-        remove_foreign_key(table_name, name: foreign_key.name) if foreign_key
-        super
-      end
+      remove_foreign_key(table_name, name: foreign_key.name) if foreign_key
+      super
     end
 
     ruby2_keywords def drop_table(table_name, *args)
       foreign_keys(table_name).each do |foreign_key|
-        with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-          remove_foreign_key(table_name, name: foreign_key.name)
-        end
+        remove_foreign_key(table_name, name: foreign_key.name)
       end
 
       Helpers::Logger.say_method_call :drop_table, table_name, *args
 
-      with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-        super(table_name, *args)
-      end
-    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
-      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) }
+      super(table_name, *args)
     end
   end
 end

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

@@ -7,9 +7,7 @@ module AddColumn
         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.
@@ -23,18 +21,13 @@ module AddColumn
         default = options.delete(:default)
         null = options.delete(:null)
 
-        with_setting(:statement_timeout, SafePgMigrations.config.pg_statement_timeout) do
-          Helpers::Logger.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
 
         Helpers::Logger.say_method_call(:change_column_default, table_name, column_name, default)
         change_column_default(table_name, column_name, default)
 
-        Helpers::Logger.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
@@ -59,9 +52,11 @@ 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)
 
+        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")
+          batch.update_all("#{quoted_column_name} = DEFAULT")
+
           sleep SafePgMigrations.config.backfill_pause
         end
       end

test/IdempotentStatements/drop_table_test.rb

@@ -29,12 +29,8 @@ def test_first_fk_already_removed
       calls = record_calls(@connection, :execute) { run_migration }
 
       assert_calls <<~CALLS.strip.split("\n"), calls
-        SET statement_timeout TO '5s'
         ALTER TABLE "users" DROP CONSTRAINT "fk_rails_d15efa01b1"
-        SET statement_timeout TO '70s'
-        SET statement_timeout TO '5s'
         DROP TABLE "users"
-        SET statement_timeout TO '70s'
       CALLS
     end
 
@@ -44,11 +40,7 @@ def test_fks_already_removed
 
       calls = record_calls(@connection, :execute) { run_migration }
 
-      assert_calls <<~CALLS.strip.split("\n"), calls
-        SET statement_timeout TO '5s'
-        DROP TABLE "users"
-        SET statement_timeout TO '70s'
-      CALLS
+      assert_calls ['DROP TABLE "users"'], calls
     end
   end
 end