pg_ha_migrations 1.8.0 → 2.1.0

data/lib/pg_ha_migrations/safe_statements.rb

@@ -3,12 +3,16 @@ module PgHaMigrations::SafeStatements
     @safe_added_columns_without_default_value ||= []
   end
 
-  def safe_create_table(table, options={}, &block)
+  def partman_extension
+    @partman_extension ||= PgHaMigrations::Extension.new("pg_partman")
+  end
+
+  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)
@@ -16,28 +20,20 @@ module PgHaMigrations::SafeStatements
     when nil
       raise ArgumentError, "safe_create_enum_type expects a set of values; if you want an enum with no values please pass an empty array"
     when []
-      unsafe_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM ()")
+      raw_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM ()")
     else
       escaped_values = values.map do |value|
         "'#{PG::Connection.escape_string(value.to_s)}'"
       end
-      unsafe_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM (#{escaped_values.join(',')})")
+      raw_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM (#{escaped_values.join(',')})")
     end
   end
 
   def safe_add_enum_value(name, value)
-    unsafe_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
-
-    unsafe_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)}'")
+    raw_execute("ALTER TYPE #{PG::Connection.quote_ident(name.to_s)} ADD VALUE '#{PG::Connection.escape_string(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:
@@ -53,20 +49,14 @@ module PgHaMigrations::SafeStatements
         raise PgHaMigrations::UnsafeMigrationError.new(":default is not safe if the default value is volatile. Use safe_change_column_default afterwards then backfill the data to prevent locking the table")
       end
     elsif options[:null] == false
-      raise PgHaMigrations::UnsafeMigrationError.new(":null => false is NOT SAFE if the table has data! If you _really_ want to do this, use unsafe_make_column_not_nullable")
+      raise PgHaMigrations::UnsafeMigrationError.new(":null => false is NOT SAFE if the table has data! If you want to do this, use safe_make_column_not_nullable")
     end
 
     unless options.has_key?(:default)
       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)
@@ -133,22 +123,126 @@ module PgHaMigrations::SafeStatements
   end
 
   def safe_make_column_nullable(table, column)
+    quoted_table_name = connection.quote_table_name(table)
+    quoted_column_name = connection.quote_column_name(column)
+
     safely_acquire_lock_for_table(table) do
-      unsafe_execute "ALTER TABLE #{table} ALTER COLUMN #{column} DROP NOT NULL"
+      raw_execute "ALTER TABLE #{quoted_table_name} ALTER COLUMN #{quoted_column_name} DROP NOT NULL"
+    end
+  end
+
+  # Postgres 12+ can use a valid CHECK constraint to validate that no values of a column are null, avoiding
+  # a full table scan while holding an exclusive lock on the table when altering a column to NOT NULL
+  #
+  # Source:
+  # https://dba.stackexchange.com/questions/267947/how-can-i-set-a-column-to-not-null-without-locking-the-table-during-a-table-scan/268128#268128
+  # (https://archive.is/X55up)
+  def safe_make_column_not_nullable(table, column)
+    if ActiveRecord::Base.connection.postgresql_version < 12_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Cannot safely make a column non-nullable before Postgres 12"
+    end
+
+    validated_table = PgHaMigrations::Table.from_table_name(table)
+    quoted_column_name = connection.quote_column_name(column)
+    column_str = column.to_s
+
+    # First, look for existing constraints that match the IS NOT NULL pattern for this column
+    existing_constraint = validated_table.check_constraints.select do |c|
+      c.definition =~ /\ACHECK \(*(#{Regexp.escape(column_str)}|#{Regexp.escape(quoted_column_name)}) IS NOT NULL\)*\Z/i
+    end.first
+
+    constraint_name = nil
+    if existing_constraint
+      if existing_constraint.validated
+        say "Found existing validated constraint #{existing_constraint.inspect} for column #{column_str}, using it directly"
+      else
+        say "Found existing unvalidated constraint #{existing_constraint.inspect} for column #{column_str}, validating it first"
+        safe_validate_check_constraint(table, name: existing_constraint.name)
+      end
+      constraint_name = existing_constraint.name
+    else
+      # Create a temporary constraint if no matching constraints exist
+      constraint_name = "tmp_not_null_constraint_#{OpenSSL::Digest::SHA256.hexdigest(column.to_s).first(7)}"
+
+      safe_add_unvalidated_check_constraint(table, "#{quoted_column_name} IS NOT NULL", name: constraint_name)
+      safe_validate_check_constraint(table, name: constraint_name)
     end
+
+    # "Ordinarily this is checked during the ALTER TABLE by scanning the entire table; however, if a
+    # valid CHECK constraint is found which proves no NULL can exist, then the table scan is
+    # skipped."
+    # See: https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DESC-SET-DROP-NOT-NULL
+    unsafe_make_column_not_nullable(table, column)
+
+    # Always drop the constraint at the end, whether it was existing or temporary
+    unsafe_remove_constraint(table, name: constraint_name)
   end
 
-  def unsafe_make_column_not_nullable(table, column, options={}) # options arg is only present for backwards compatiblity
+  # This method is a variant of `safe_make_column_not_nullable` that is expected to always be fast;
+  # i.e., it will not perform a full table scan to check for null values.
+  def safe_make_column_not_nullable_from_check_constraint(table, column, constraint_name:, drop_constraint: true)
+    unless ActiveRecord::Base.connection.postgresql_version >= 12_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Cannot safely make a column non-nullable before Postgres 12"
+    end
+
+    unless constraint_name
+      raise ArgumentError, "Expected <constraint_name> to be present"
+    end
+    constraint_name = constraint_name.to_s
+
+    quoted_table_name = connection.quote_table_name(table)
+    quoted_column_name = connection.quote_column_name(column)
+
+    validated_table = PgHaMigrations::Table.from_table_name(table)
+    constraint = validated_table.check_constraints.find do |c|
+      c.name == constraint_name
+    end
+
+    unless constraint
+      raise PgHaMigrations::InvalidMigrationError, "The provided constraint does not exist"
+    end
+
+    unless constraint.validated
+      raise PgHaMigrations::InvalidMigrationError, "The provided constraint is not validated"
+    end
+
+    # The constraint has to actually prove that no null values exist, so the
+    # constraint condition can't simply include the `IS NOT NULL` check. We
+    # don't try to handle all possible cases here. For example,
+    # `a IS NOT NULL AND b IS NOT NULL` would prove what we need, but it would
+    # be complicated to check. We must ensure, however, that we're not too
+    # loose. For example, `a IS NOT NULL OR b IS NOT NULL` would not prove that
+    # `a IS NOT NULL`.
+    unless constraint.definition =~ /\ACHECK \(*(#{Regexp.escape(column.to_s)}|#{Regexp.escape(quoted_column_name)}) IS NOT NULL\)*\Z/i
+      raise PgHaMigrations::InvalidMigrationError, "The provided constraint does not enforce non-null values for the column"
+    end
+
+    # We don't want to acquire an exclusive lock on the table twice, and we also don't want it to be
+    # posssible to have the NOT NULL constraint addition succeed while the constraint removal fails,
+    # so we acquire the lock once and do both operations in the same block.
     safely_acquire_lock_for_table(table) do
-      unsafe_execute "ALTER TABLE #{table} ALTER COLUMN #{column} SET NOT NULL"
+      # "Ordinarily this is checked during the ALTER TABLE by scanning the entire table; however, if a
+      # valid CHECK constraint is found which proves no NULL can exist, then the table scan is
+      # skipped."
+      # See: https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DESC-SET-DROP-NOT-NULL
+      unsafe_make_column_not_nullable(table, column)
+
+      if drop_constraint
+        unsafe_remove_constraint(table, name: constraint_name)
+      end
     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
 
+    # Check if nulls_not_distinct was provided but PostgreSQL version doesn't support it
+    if options[:nulls_not_distinct] && ActiveRecord::Base.connection.postgresql_version < 15_00_00
+      raise PgHaMigrations::InvalidMigrationError, "nulls_not_distinct option requires PostgreSQL 15 or higher"
+    end
+
     # Avoids taking out an unnecessary SHARE lock if the table does have data
     ensure_small_table!(table, empty: true)
 
@@ -160,11 +254,16 @@ module PgHaMigrations::SafeStatements
     end
   end
 
-  def safe_add_concurrent_index(table, columns, options={})
+  def safe_add_concurrent_index(table, columns, **options)
+    # Check if nulls_not_distinct was provided but PostgreSQL version doesn't support it
+    if options[:nulls_not_distinct] && ActiveRecord::Base.connection.postgresql_version < 15_00_00
+      raise PgHaMigrations::InvalidMigrationError, "nulls_not_distinct option requires PostgreSQL 15 or higher"
+    end
+
     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
@@ -184,8 +283,13 @@ module PgHaMigrations::SafeStatements
     using: nil,
     unique: nil,
     where: nil,
-    comment: nil
+    comment: nil,
+    nulls_not_distinct: nil
   )
+    # Check if nulls_not_distinct was provided but PostgreSQL version doesn't support it
+    if !nulls_not_distinct.nil? && ActiveRecord::Base.connection.postgresql_version < 15_00_00
+      raise PgHaMigrations::InvalidMigrationError, "nulls_not_distinct option requires PostgreSQL 15 or higher"
+    end
 
     if ActiveRecord::Base.connection.postgresql_version < 11_00_00
       raise PgHaMigrations::InvalidMigrationError, "Concurrent partitioned index creation not supported on Postgres databases before version 11"
@@ -208,21 +312,19 @@ module PgHaMigrations::SafeStatements
       PgHaMigrations::Index.from_table_and_columns(child_table, columns)
     end
 
-    # TODO: take out ShareLock after issue #39 is implemented
-    safely_acquire_lock_for_table(parent_table.fully_qualified_name) do
-      # CREATE INDEX ON ONLY parent_table
-      unsafe_add_index(
-        parent_table.fully_qualified_name,
-        columns,
-        name: parent_index.name,
-        if_not_exists: if_not_exists,
-        using: using,
-        unique: unique,
-        where: where,
-        comment: comment,
-        algorithm: :only, # see lib/pg_ha_migrations/hacks/add_index_on_only.rb
-      )
-    end
+    # CREATE INDEX ON ONLY parent_table
+    unsafe_add_index(
+      parent_table.fully_qualified_name,
+      columns,
+      name: parent_index.name,
+      if_not_exists: if_not_exists,
+      using: using,
+      unique: unique,
+      nulls_not_distinct: nulls_not_distinct,
+      where: where,
+      comment: comment,
+      algorithm: :only, # see lib/pg_ha_migrations/hacks/add_index_on_only.rb
+    )
 
     child_indexes.each do |child_index|
       add_index_method = if child_index.table.natively_partitioned?
@@ -239,6 +341,7 @@ module PgHaMigrations::SafeStatements
         if_not_exists: if_not_exists,
         using: using,
         unique: unique,
+        nulls_not_distinct: nulls_not_distinct,
         where: where,
       )
     end
@@ -261,7 +364,7 @@ module PgHaMigrations::SafeStatements
   end
 
   def safe_set_maintenance_work_mem_gb(gigabytes)
-    unsafe_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
+    raw_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
   end
 
   def safe_add_unvalidated_check_constraint(table, expression, name:)
@@ -296,20 +399,6 @@ module PgHaMigrations::SafeStatements
     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?
 
@@ -351,7 +440,7 @@ module PgHaMigrations::SafeStatements
 
     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]
@@ -370,15 +459,7 @@ module PgHaMigrations::SafeStatements
     end
   end
 
-  def safe_partman_create_parent(table, **options)
-    if options[:retention].present? || options[:retention_keep_table] == false
-      raise PgHaMigrations::UnsafeMigrationError.new(":retention and/or :retention_keep_table => false can potentially result in data loss if misconfigured. Please use unsafe_partman_create_parent if you want to set these options")
-    end
-
-    unsafe_partman_create_parent(table, **options)
-  end
-
-  def unsafe_partman_create_parent(
+  def safe_partman_create_parent(
     table,
     partition_key:,
     interval:,
@@ -397,6 +478,18 @@ module PgHaMigrations::SafeStatements
       raise PgHaMigrations::InvalidMigrationError, "Native partitioning with partman not supported on Postgres databases before version 11"
     end
 
+    raise PgHaMigrations::MissingExtensionError, "The pg_partman extension is not installed" unless partman_extension.installed?
+
+    if partman_extension.major_version >= 5 || PgHaMigrations.config.partman_5_compatibility_mode
+      if PgHaMigrations::PARTMAN_UNSUPPORTED_INTERVALS.include?(interval)
+        raise PgHaMigrations::InvalidMigrationError,
+          "Special partition interval values (#{interval}) are no longer supported. " \
+          "Please use a supported interval time value from core PostgreSQL " \
+          "#{(partman_extension.major_version < 5 ? "or turn partman 5 compatibility mode off " : "")}" \
+          "(https://www.postgresql.org/docs/current/datatype-datetime.html#DATATYPE-INTERVAL-INPUT)"
+      end
+    end
+
     formatted_start_partition = nil
 
     if start_partition.present?
@@ -411,16 +504,22 @@ module PgHaMigrations::SafeStatements
       end
     end
 
+    validated_table = PgHaMigrations::PartmanTable.from_table_name(table)
+    validated_template_table = template_table ? PgHaMigrations::PartmanTable.from_table_name(template_table) : nil
+
     create_parent_options = {
-      parent_table: _fully_qualified_table_name_for_partman(table),
-      template_table: template_table ? _fully_qualified_table_name_for_partman(template_table) : nil,
+      parent_table: validated_table.fully_qualified_name,
+      template_table: validated_template_table&.fully_qualified_name,
       control: partition_key,
-      type: "native",
       interval: interval,
       premake: premake,
       start_partition: formatted_start_partition,
     }.compact
 
+    if partman_extension.major_version < 5
+      create_parent_options[:type] = "native"
+    end
+
     create_parent_sql = create_parent_options.map { |k, v| "p_#{k} := #{connection.quote(v)}" }.join(", ")
 
     log_message = "partman_create_parent(#{table.inspect}, " \
@@ -431,7 +530,7 @@ module PgHaMigrations::SafeStatements
       "template_table: #{template_table.inspect})"
 
     say_with_time(log_message) do
-      connection.execute("SELECT #{_quoted_partman_schema}.create_parent(#{create_parent_sql})")
+      connection.execute("SELECT #{partman_extension.quoted_schema}.create_parent(#{create_parent_sql})")
     end
 
     update_config_options = {
@@ -442,6 +541,10 @@ module PgHaMigrations::SafeStatements
     }.compact
 
     unsafe_partman_update_config(table, **update_config_options)
+
+    if PgHaMigrations.config.partman_5_compatibility_mode && partman_extension.major_version < 5
+      unsafe_partman_standardize_partition_naming(table)
+    end
   end
 
   def safe_partman_update_config(table, **options)
@@ -452,55 +555,16 @@ module PgHaMigrations::SafeStatements
     unsafe_partman_update_config(table, **options)
   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
+  def safe_partman_reapply_privileges(table)
+    raise PgHaMigrations::MissingExtensionError, "The pg_partman extension is not installed" unless partman_extension.installed?
 
-    safe_partman_reapply_privileges(table) if inherit_privileges_changed
-  end
+    validated_table = PgHaMigrations::PartmanTable.from_table_name(table)
 
-  def safe_partman_reapply_privileges(table)
     say_with_time "partman_reapply_privileges(#{table.inspect})" do
-      connection.execute("SELECT #{_quoted_partman_schema}.reapply_privileges('#{_fully_qualified_table_name_for_partman(table)}')")
+      connection.execute("SELECT #{partman_extension.quoted_schema}.reapply_privileges('#{validated_table.fully_qualified_name}')")
     end
   end
 
-  def _quoted_partman_schema
-    schema = connection.select_value(<<~SQL)
-      SELECT nspname
-      FROM pg_namespace JOIN pg_extension
-        ON pg_namespace.oid = pg_extension.extnamespace
-      WHERE pg_extension.extname = 'pg_partman'
-    SQL
-
-    raise PgHaMigrations::InvalidMigrationError, "The pg_partman extension is not installed" unless schema.present?
-
-    connection.quote_schema_name(schema)
-  end
-
-  def _fully_qualified_table_name_for_partman(table)
-    table = PgHaMigrations::Table.from_table_name(table)
-
-    [table.schema, table.name].each do |identifier|
-      if identifier.to_s !~ /^[a-z_][a-z_\d]*$/
-        raise PgHaMigrations::InvalidMigrationError, "Partman requires schema / table names to be lowercase with underscores"
-      end
-    end.join(".")
-  end
-
   def _per_migration_caller
     @_per_migration_caller ||= Kernel.caller
   end
@@ -528,40 +592,50 @@ module PgHaMigrations::SafeStatements
     super(conn, direction)
   end
 
-  def safely_acquire_lock_for_table(table, mode: :access_exclusive, &block)
-    nested_target_table = Thread.current[__method__]
-
+  def safely_acquire_lock_for_table(*tables, mode: :access_exclusive, &block)
     _check_postgres_adapter!
 
-    target_table = PgHaMigrations::Table.from_table_name(table, mode)
+    target_tables = PgHaMigrations::TableCollection.from_table_names(tables, mode)
 
-    if nested_target_table
-      if nested_target_table != target_table
-        raise PgHaMigrations::InvalidMigrationError, "Nested lock detected! Cannot acquire lock on #{target_table.fully_qualified_name} while #{nested_target_table.fully_qualified_name} is locked."
-      elsif nested_target_table.mode < target_table.mode
-        raise PgHaMigrations::InvalidMigrationError, "Lock escalation detected! Cannot change lock level from :#{nested_target_table.mode} to :#{target_table.mode} for #{target_table.fully_qualified_name}."
+    if @parent_lock_tables
+      if !target_tables.subset?(@parent_lock_tables)
+        raise PgHaMigrations::InvalidMigrationError,
+          "Nested lock detected! Cannot acquire lock on #{target_tables.to_sql} " \
+          "while #{@parent_lock_tables.to_sql} is locked."
+      end
+
+      if @parent_lock_tables.mode < target_tables.mode
+        raise PgHaMigrations::InvalidMigrationError,
+          "Lock escalation detected! Cannot change lock level from :#{@parent_lock_tables.mode} " \
+          "to :#{target_tables.mode} for #{target_tables.to_sql}."
       end
-    else
-      Thread.current[__method__] = target_table
-    end
 
-    # Locking a partitioned table will also lock child tables (including sub-partitions),
-    # so we need to check for blocking queries on those tables as well
-    target_tables = target_table.partitions(include_sub_partitions: true, include_self: true)
+      # If in a nested context and all of the above checks have passed,
+      # we have already acquired the lock (or a lock at a higher level),
+      # and can simply execute the block and short-circuit.
+      block.call
+
+      return
+    end
 
     successfully_acquired_lock = false
 
     until successfully_acquired_lock
-      while (
+      loop do
         blocking_transactions = PgHaMigrations::BlockingDatabaseTransactions.find_blocking_transactions("#{PgHaMigrations::LOCK_TIMEOUT_SECONDS} seconds")
-        blocking_transactions.any? do |query|
+
+        # Locking a partitioned table will also lock child tables (including sub-partitions),
+        # so we need to check for blocking queries on those tables as well
+        target_tables_with_partitions = target_tables.with_partitions
+
+        break unless blocking_transactions.any? do |query|
           query.tables_with_locks.any? do |locked_table|
-            target_tables.any? do |target_table|
+            target_tables_with_partitions.any? do |target_table|
               target_table.conflicts_with?(locked_table)
             end
           end
         end
-      )
+
         say "Waiting on blocking transactions:"
         blocking_transactions.each do |blocking_transaction|
           say blocking_transaction.description
@@ -570,16 +644,21 @@ module PgHaMigrations::SafeStatements
       end
 
       connection.transaction do
-        adjust_timeout_method = connection.postgresql_version >= 9_03_00 ? :adjust_lock_timeout : :adjust_statement_timeout
         begin
-          method(adjust_timeout_method).call(PgHaMigrations::LOCK_TIMEOUT_SECONDS) do
-            connection.execute("LOCK #{target_table.fully_qualified_name} IN #{target_table.mode.to_sql} MODE;")
+          # A lock timeout would apply to each individual table in the query,
+          # so we made a conscious decision to use a statement timeout here
+          # to keep behavior consistent in a multi-table lock scenario.
+          adjust_statement_timeout(PgHaMigrations::LOCK_TIMEOUT_SECONDS) do
+            connection.execute("LOCK #{target_tables.to_sql} IN #{target_tables.mode.to_sql} MODE;")
           end
           successfully_acquired_lock = true
         rescue ActiveRecord::StatementInvalid => e
+          # It is still possible to hit a lock timeout if the session has
+          # that value set to something less than LOCK_TIMEOUT_SECONDS.
+          # We should retry when either of these exceptions are raised.
           if e.message =~ /PG::LockNotAvailable.+ lock timeout/ || e.message =~ /PG::QueryCanceled.+ statement timeout/
             sleep_seconds = PgHaMigrations::LOCK_FAILURE_RETRY_DELAY_MULTLIPLIER * PgHaMigrations::LOCK_TIMEOUT_SECONDS
-            say "Timed out trying to acquire #{target_table.mode.to_sql} lock on the #{target_table.fully_qualified_name} table."
+            say "Timed out trying to acquire #{target_tables.mode.to_sql} lock on #{target_tables.to_sql}."
             say "Sleeping for #{sleep_seconds}s to allow potentially queued up queries to finish before continuing."
             sleep(sleep_seconds)
 
@@ -590,12 +669,16 @@ module PgHaMigrations::SafeStatements
         end
 
         if successfully_acquired_lock
-          block.call
+          @parent_lock_tables = target_tables
+
+          begin
+            block.call
+          ensure
+            @parent_lock_tables = nil
+          end
         end
       end
     end
-  ensure
-    Thread.current[__method__] = nil unless nested_target_table
   end
 
   def adjust_lock_timeout(timeout_seconds = PgHaMigrations::LOCK_TIMEOUT_SECONDS, &block)