support non-native partitioned tables

Co-authored-by: Ryan Krage <[email protected]>
braintree · Sep 14, 2023 · c8d2ace · c8d2ace
1 parent eabfdcd
commit c8d2ace
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 5 deletions.
diff --git a/lib/pg_ha_migrations/safe_statements.rb b/lib/pg_ha_migrations/safe_statements.rb
@@ -187,7 +187,7 @@ def safe_add_concurrent_partitioned_index(
 
     parent_table = PgHaMigrations::Struct::Table.from_table_name(table)
 
-    raise PgHaMigrations::InvalidMigrationError, "Table #{parent_table.inspect} is not a partitioned table" unless parent_table.partitioned?
+    raise PgHaMigrations::InvalidMigrationError, "Table #{parent_table.inspect} is not a partitioned table" unless parent_table.natively_partitioned?
 
     parent_index = if name.present?
       PgHaMigrations::Struct::Index.new(name, parent_table)
@@ -199,7 +199,7 @@ def safe_add_concurrent_partitioned_index(
     return if if_not_exists && parent_index.valid?
 
     child_indexes = parent_table.partitions.map do |child_table|
-      raise PgHaMigrations::InvalidMigrationError, "Partitioned table #{parent_table.inspect} contains sub-partitions" if child_table.partitioned?
+      raise PgHaMigrations::InvalidMigrationError, "Partitioned table #{parent_table.inspect} contains sub-partitions" if child_table.natively_partitioned?
 
       PgHaMigrations::Struct::Index.from_table_and_columns(child_table, columns)
     end
@@ -537,7 +537,7 @@ def safely_acquire_lock_for_table(table, &block)
 
     # 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.concat(table.partitions(include_sub_partitions: true)) if table.partitioned?
+    target_tables.concat(table.partitions(include_sub_partitions: true))
 
     successfully_acquired_lock = false
 

diff --git a/lib/pg_ha_migrations/struct.rb b/lib/pg_ha_migrations/struct.rb
@@ -42,7 +42,7 @@ def self.from_table_name(table)
       new(pg_name.identifier, schema)
     end
 
-    def partitioned?
+    def natively_partitioned?
       connection.select_value(<<~SQL)
         SELECT EXISTS (
           SELECT 1
@@ -69,7 +69,7 @@ def partitions(include_sub_partitions: false)
 
       if include_sub_partitions
         sub_partitions = tables.each_with_object([]) do |table, arr|
-          arr.concat(table.partitions(include_sub_partitions: true)) if table.partitioned?
+          arr.concat(table.partitions(include_sub_partitions: true))
         end
 
         tables.concat(sub_partitions)

diff --git a/spec/safe_statements_spec.rb b/spec/safe_statements_spec.rb
@@ -3620,6 +3620,46 @@ def up
               end
             end
 
+            it "waits to acquire a lock if the table is non-natively partitioned and child table is blocked" do
+              stub_const("PgHaMigrations::LOCK_TIMEOUT_SECONDS", 1)
+
+              ActiveRecord::Base.connection.execute(<<~SQL)
+                CREATE TABLE bogus_table_child(pk SERIAL, i INTEGER) INHERITS (#{table_name})
+              SQL
+
+              begin
+                thread = Thread.new do
+                  ActiveRecord::Base.connection.execute(<<~SQL)
+                    LOCK bogus_table_child;
+                    SELECT pg_sleep(3);
+                  SQL
+                end
+
+                sleep 1.1
+
+                expect(PgHaMigrations::BlockingDatabaseTransactions).to receive(:find_blocking_transactions)
+                  .at_least(3)
+                  .times
+                  .and_call_original
+
+                migration.suppress_messages do
+                  migration.safely_acquire_lock_for_table(table_name) do
+                    locks_for_parent = locks_for_table(table_name, connection: alternate_connection)
+                    locks_for_child = locks_for_table("bogus_table_child", connection: alternate_connection)
+
+                    aggregate_failures do
+                      expect(locks_for_parent.size).to eq(1)
+                      expect(locks_for_child.size).to eq(1)
+                      expect(locks_for_parent.first.pid).to_not be_nil
+                      expect(locks_for_parent.first.pid).to eq(locks_for_child.first.pid)
+                    end
+                  end
+                end
+              ensure
+                thread.join
+              end
+            end
+
             it "fails lock acquisition quickly if Postgres doesn't grant an exclusive lock but then retries" do
               stub_const("PgHaMigrations::LOCK_TIMEOUT_SECONDS", 1)