wip

Author: rkrage

Diff for: lib/pg_ha_migrations/safe_statements.rb

@@ -3,12 +3,12 @@ def safe_added_columns_without_default_value
     @safe_added_columns_without_default_value ||= []
   end
 
-  def safe_create_table(table, options={}, &block)
+  def safe_create_table(table, **options, &block)
     if options[:force]
       raise PgHaMigrations::UnsafeMigrationError.new(":force is NOT SAFE! Explicitly call unsafe_drop_table first if you want to recreate an existing table")
     end
 
-    unsafe_create_table(table, options, &block)
+    unsafe_create_table(table, **options, &block)
   end
 
   def safe_create_enum_type(name, values=nil)
@@ -29,15 +29,7 @@ def safe_add_enum_value(name, value)
     raw_execute("ALTER TYPE #{PG::Connection.quote_ident(name.to_s)} ADD VALUE '#{PG::Connection.escape_string(value)}'")
   end
 
-  def unsafe_rename_enum_value(name, old_value, new_value)
-    if ActiveRecord::Base.connection.postgresql_version < 10_00_00
-      raise PgHaMigrations::InvalidMigrationError, "Renaming an enum value is not supported on Postgres databases before version 10"
-    end
-
-    raw_execute("ALTER TYPE #{PG::Connection.quote_ident(name.to_s)} RENAME VALUE '#{PG::Connection.escape_string(old_value)}' TO '#{PG::Connection.escape_string(new_value)}'")
-  end
-
-  def safe_add_column(table, column, type, options = {})
+  def safe_add_column(table, column, type, **options)
     # Note: we don't believe we need to consider the odd case where
     # `:default => nil` or `:default => -> { null }` (or similar) is
     # passed because:
@@ -60,13 +52,7 @@ def safe_add_column(table, column, type, options = {})
       self.safe_added_columns_without_default_value << [table.to_s, column.to_s]
     end
 
-    unsafe_add_column(table, column, type, options)
-  end
-
-  def unsafe_add_column(table, column, type, options = {})
-    safely_acquire_lock_for_table(table) do
-      super(table, column, type, **options)
-    end
+    unsafe_add_column(table, column, type, **options)
   end
 
   def safe_change_column_default(table_name, column_name, default_value)
@@ -177,7 +163,7 @@ def unsafe_make_column_not_nullable(table, column, options={}) # options arg is
     end
   end
 
-  def safe_add_index_on_empty_table(table, columns, options={})
+  def safe_add_index_on_empty_table(table, columns, **options)
     if options[:algorithm] == :concurrently
       raise ArgumentError, "Cannot call safe_add_index_on_empty_table with :algorithm => :concurrently"
     end
@@ -193,11 +179,11 @@ def safe_add_index_on_empty_table(table, columns, options={})
     end
   end
 
-  def safe_add_concurrent_index(table, columns, options={})
+  def safe_add_concurrent_index(table, columns, **options)
     unsafe_add_index(table, columns, **options.merge(:algorithm => :concurrently))
   end
 
-  def safe_remove_concurrent_index(table, options={})
+  def safe_remove_concurrent_index(table, **options)
     unless options.is_a?(Hash) && options.key?(:name)
       raise ArgumentError, "Expected safe_remove_concurrent_index to be called with arguments (table_name, :name => ...)"
     end
@@ -329,20 +315,6 @@ def safe_rename_constraint(table, from:, to:)
     end
   end
 
-  def unsafe_remove_constraint(table, name:)
-    raise ArgumentError, "Expected <name> to be present" unless name.present?
-
-    quoted_table_name = connection.quote_table_name(table)
-    quoted_constraint_name = connection.quote_table_name(name)
-    sql = "ALTER TABLE #{quoted_table_name} DROP CONSTRAINT #{quoted_constraint_name}"
-
-    safely_acquire_lock_for_table(table) do
-      say_with_time "remove_constraint(#{table.inspect}, name: #{name.inspect})" do
-        connection.execute(sql)
-      end
-    end
-  end
-
   def safe_create_partitioned_table(table, partition_key:, type:, infer_primary_key: nil, **options, &block)
     raise ArgumentError, "Expected <partition_key> to be present" unless partition_key.present?
 
@@ -384,7 +356,7 @@ def safe_create_partitioned_table(table, partition_key:, type:, infer_primary_ke
 
     options[:options] = "PARTITION BY #{type.upcase} (#{quoted_partition_key})"
 
-    safe_create_table(table, options) do |td|
+    safe_create_table(table, **options) do |td|
       block.call(td) if block
 
       next unless options[:id]

Diff for: lib/pg_ha_migrations/unsafe_statements.rb

@@ -29,30 +29,47 @@ def self.delegate_unsafe_method_to_migration_base_class(method_name)
     ruby2_keywords "unsafe_#{method_name}"
   end
 
+  def self.delegate_unsafe_method_to_migration_base_class_with_lock(method_name, mode: :access_exclusive)
+    define_method("unsafe_#{method_name}") do |*args, &block|
+      if PgHaMigrations.config.check_for_dependent_objects
+        disallow_migration_method_if_dependent_objects!(method_name, arguments: args)
+      end
+
+      safely_acquire_lock_for_table(args.first, mode: mode) do
+        execute_ancestor_statement(method_name, *args, &block)
+      end
+    end
+    ruby2_keywords "unsafe_#{method_name}"
+  end
+
   def self.delegate_raw_method_to_migration_base_class(method_name)
     define_method("raw_#{method_name}") do |*args, &block|
       execute_ancestor_statement(method_name, *args, &block)
     end
     ruby2_keywords "raw_#{method_name}"
   end
 
-  # Direct dispatch to underlying Rails method as unsafe_<method_name> with dependent object check
-  delegate_unsafe_method_to_migration_base_class :add_check_constraint
-  delegate_unsafe_method_to_migration_base_class :add_column
-  delegate_unsafe_method_to_migration_base_class :add_foreign_key
-  delegate_unsafe_method_to_migration_base_class :change_column
-  delegate_unsafe_method_to_migration_base_class :change_column_default
-  delegate_unsafe_method_to_migration_base_class :change_table
+  # Direct dispatch to underlying Rails method as unsafe_<method_name> with dependent object check / safe lock acquisition
+  delegate_unsafe_method_to_migration_base_class_with_lock :add_check_constraint
+  delegate_unsafe_method_to_migration_base_class_with_lock :add_column
+  delegate_unsafe_method_to_migration_base_class_with_lock :add_foreign_key, mode: :share_row_exclusive # TODO: lock referenced table as well
+  delegate_unsafe_method_to_migration_base_class_with_lock :change_column
+  delegate_unsafe_method_to_migration_base_class_with_lock :change_column_default
+  delegate_unsafe_method_to_migration_base_class_with_lock :remove_check_constraint
+  delegate_unsafe_method_to_migration_base_class_with_lock :remove_column
+  delegate_unsafe_method_to_migration_base_class_with_lock :rename_column
+  delegate_unsafe_method_to_migration_base_class_with_lock :remove_foreign_key # TODO: lock referenced table as well
+
+  # Direct dispatch to underlying Rails method as unsafe_<method_name> with dependent object check / no locking
+  #
+  # We don't do locking on execute because it's so generic.
+  # For drop_table and rename_table, it's implied that the table is unused.
+  # If the table is still being used... that's a bigger problem that locking won't solve.
   delegate_unsafe_method_to_migration_base_class :drop_table
   delegate_unsafe_method_to_migration_base_class :execute
-  delegate_unsafe_method_to_migration_base_class :remove_check_constraint
-  delegate_unsafe_method_to_migration_base_class :remove_column
-  delegate_unsafe_method_to_migration_base_class :remove_foreign_key
-  delegate_unsafe_method_to_migration_base_class :remove_index
-  delegate_unsafe_method_to_migration_base_class :rename_column
   delegate_unsafe_method_to_migration_base_class :rename_table
 
-  # Direct dispatch to underlying Rails method as raw_<method_name> without dependent object check
+  # Direct dispatch to underlying Rails method as raw_<method_name> without dependent object check / locking
   delegate_raw_method_to_migration_base_class :add_check_constraint
   delegate_raw_method_to_migration_base_class :add_column
   delegate_raw_method_to_migration_base_class :add_foreign_key
@@ -72,7 +89,7 @@ def self.delegate_raw_method_to_migration_base_class(method_name)
   delegate_raw_method_to_migration_base_class :rename_table
 
   # Raises error if disable_default_migration_methods is true
-  # Otherwise, direct dispatch to underlying Rails method without dependent object check
+  # Otherwise, direct dispatch to underlying Rails method without dependent object check / locking
   disable_or_delegate_default_method :add_check_constraint, ":add_check_constraint is NOT SAFE! Use :safe_add_unvalidated_check_constraint and then :safe_validate_check_constraint instead"
   disable_or_delegate_default_method :add_column, ":add_column is NOT SAFE! Use safe_add_column instead"
   disable_or_delegate_default_method :add_foreign_key, ":add_foreign_key is NOT SAFE! Explicitly call :unsafe_add_foreign_key"
@@ -91,15 +108,21 @@ def self.delegate_raw_method_to_migration_base_class(method_name)
   disable_or_delegate_default_method :rename_column, ":rename_column is NOT SAFE! Explicitly call :unsafe_rename_column to proceed"
   disable_or_delegate_default_method :rename_table, ":rename_table is NOT SAFE! Explicitly call :unsafe_rename_table to proceed"
 
-  def unsafe_create_table(table, options={}, &block)
+  # Note that unsafe_* methods defined below do not run dependent object checks
+
+  def unsafe_change_table(*args, &block)
+    raise PgHaMigrations::UnsafeMigrationError.new(":change_table is too generic to even allow an unsafe variant. Instead use explicit methods for adding / removing / modifying columns.")
+  end
+
+  def unsafe_create_table(table, **options, &block)
     if options[:force] && !PgHaMigrations.config.allow_force_create_table
       raise PgHaMigrations::UnsafeMigrationError.new(":force is NOT SAFE! Explicitly call unsafe_drop_table first if you want to recreate an existing table")
     end
 
     execute_ancestor_statement(:create_table, table, **options, &block)
   end
 
-  def unsafe_add_index(table, column_names, options = {})
+  def unsafe_add_index(table, column_names, **options)
     if ((ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR >= 2) || ActiveRecord::VERSION::MAJOR > 5) &&
         column_names.is_a?(String) && /\W/.match?(column_names) && options.key?(:opclass)
       raise PgHaMigrations::InvalidMigrationError, "ActiveRecord drops the :opclass option when supplying a string containing an expression or list of columns; instead either supply an array of columns or include the opclass in the string for each column"
@@ -115,7 +138,71 @@ def unsafe_add_index(table, column_names, options = {})
 
     options[:name] = validated_index.name
 
-    execute_ancestor_statement(:add_index, table, column_names, **options)
+    if options[:algorithm] == :concurrently
+      execute_ancestor_statement(:add_index, table, column_names, **options)
+    else
+      safely_acquire_lock_for_table(table, mode: :share) do
+        execute_ancestor_statement(:add_index, table, column_names, **options)
+      end
+    end
+  end
+
+  def unsafe_remove_index(table, column = nil, **options)
+    if options[:algorithm]== :concurrently
+      execute_ancestor_statement(:remove_index, table, column, **options)
+    else
+      safely_acquire_lock_for_table(table) do
+        execute_ancestor_statement(:remove_index, table, column, **options)
+      end
+    end
+  end
+
+  def unsafe_rename_enum_value(name, old_value, new_value)
+    if ActiveRecord::Base.connection.postgresql_version < 10_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Renaming an enum value is not supported on Postgres databases before version 10"
+    end
+
+    raw_execute("ALTER TYPE #{PG::Connection.quote_ident(name.to_s)} RENAME VALUE '#{PG::Connection.escape_string(old_value)}' TO '#{PG::Connection.escape_string(new_value)}'")
+  end
+
+  def unsafe_make_column_not_nullable(table, column, **options) # options arg is only present for backwards compatiblity
+    safely_acquire_lock_for_table(table) do
+      raw_execute("ALTER TABLE #{table} ALTER COLUMN #{column} SET NOT NULL") # TODO: quoting
+    end
+  end
+
+  def unsafe_remove_constraint(table, name:)
+    raise ArgumentError, "Expected <name> to be present" unless name.present?
+
+    quoted_table_name = connection.quote_table_name(table)
+    quoted_constraint_name = connection.quote_table_name(name)
+    sql = "ALTER TABLE #{quoted_table_name} DROP CONSTRAINT #{quoted_constraint_name}"
+
+    safely_acquire_lock_for_table(table) do
+      say_with_time "remove_constraint(#{table.inspect}, name: #{name.inspect})" do
+        connection.execute(sql)
+      end
+    end
+  end
+
+  def unsafe_partman_update_config(table, **options)
+    invalid_options = options.keys - PgHaMigrations::PARTMAN_UPDATE_CONFIG_OPTIONS
+
+    raise ArgumentError, "Unrecognized argument(s): #{invalid_options}" unless invalid_options.empty?
+
+    PgHaMigrations::PartmanConfig.schema = _quoted_partman_schema
+
+    config = PgHaMigrations::PartmanConfig.find(_fully_qualified_table_name_for_partman(table))
+
+    config.assign_attributes(**options)
+
+    inherit_privileges_changed = config.inherit_privileges_changed?
+
+    say_with_time "partman_update_config(#{table.inspect}, #{options.map { |k,v| "#{k}: #{v.inspect}" }.join(", ")})" do
+      config.save!
+    end
+
+    safe_partman_reapply_privileges(table) if inherit_privileges_changed
   end
 
   def unsafe_partman_update_config(table, **options)

Diff for: spec/safe_statements_spec.rb

@@ -857,7 +857,7 @@ def up
           it "calls safely_acquire_lock_for_table" do
             migration = Class.new(migration_klass).new
 
-            expect(migration).to receive(:safely_acquire_lock_for_table).with(:foos)
+            expect(migration).to receive(:safely_acquire_lock_for_table).with(:foos, mode: :access_exclusive)
             migration.unsafe_add_column(:foos, :bar, :text)
           end
         end
@@ -1624,7 +1624,7 @@ def up
 
             expect do
               test_migration.suppress_messages { test_migration.migrate(:up) }
-            end.to make_database_queries(matching: /LOCK "public"\."foos" IN SHARE MODE/, count: 1)
+            end.to make_database_queries(matching: /LOCK "public"\."foos" IN SHARE MODE/, count: 2)
               .and(make_database_queries(matching: /CREATE INDEX "index_foos_on_bar"/, count: 1))
 
             indexes = ActiveRecord::Base.connection.indexes("foos")
@@ -2520,7 +2520,7 @@ def up
           end
         end
 
-        describe  "unsafe_add_index" do
+        describe "unsafe_add_index" do
           it "raises a helper warning when ActiveRecord is going to swallow per-column options" do
             migration = Class.new(migration_klass) do
               def up
@@ -2558,6 +2558,52 @@ def up
             end.not_to make_database_queries(matching: /int4_ops/)
           end
 
+          it "safely acquires SHARE lock when not creating indexes concurrently" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                unsafe_create_table :foos do |t|
+                  t.integer :bar
+                end
+              end
+            end
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                unsafe_add_index :foos, :bar
+              end
+            end
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to make_database_queries(matching: /LOCK "public"\."foos" IN SHARE MODE/, count: 1)
+              .and(make_database_queries(matching: /CREATE INDEX "index_foos_on_bar"/, count: 1))
+          end
+
+          it "skips lock acquisition when creating indexes concurrently" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                unsafe_create_table :foos do |t|
+                  t.integer :bar
+                end
+              end
+            end
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                unsafe_add_index :foos, :bar, algorithm: :concurrently
+              end
+            end
+
+            allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+            expect(ActiveRecord::Base.connection).to_not receive(:execute).with(/LOCK/)
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to make_database_queries(matching: /CREATE INDEX CONCURRENTLY "index_foos_on_bar"/, count: 1)
+          end
+
           it "generates index name with hashed identifier when default index name is too large" do
             setup_migration = Class.new(migration_klass) do
               def up
@@ -2667,19 +2713,7 @@ def up
             end
           end
 
-          it "raises a nice error if options isn't a hash" do
-            test_migration = Class.new(migration_klass) do
-              def up
-                safe_remove_concurrent_index :foos, [:column]
-              end
-            end
-
-            expect do
-              test_migration.suppress_messages { test_migration.migrate(:up) }
-            end.to raise_error(ArgumentError, "Expected safe_remove_concurrent_index to be called with arguments (table_name, :name => ...)")
-          end
-
-          it "raises a nice error if options is a hash with a :name key" do
+          it "raises a nice error if options does not have :name key" do
             test_migration = Class.new(migration_klass) do
               def up
                 safe_remove_concurrent_index :foos, :blah => :foo
@@ -4056,16 +4090,16 @@ def up
           it "renames change_table to unsafe_change_table" do
             migration = Class.new(migration_klass) do
               def up
-                unsafe_create_table :foos
-                unsafe_change_table :foos do |t|
-                  t.string :bar
-                end
+                unsafe_change_table :foos
               end
             end
 
-            migration.suppress_messages { migration.migrate(:up) }
-
-            expect(ActiveRecord::Base.connection.columns("foos").map(&:name)).to include("bar")
+            expect do
+              migration.suppress_messages { migration.migrate(:up) }
+            end.to raise_error(
+              PgHaMigrations::UnsafeMigrationError,
+              ":change_table is too generic to even allow an unsafe variant. Instead use explicit methods for adding / removing / modifying columns."
+            )
           end
 
           it "renames rename_table to unsafe_rename_table" do