Support to recursively create partitioned index

lib/pg_ha_migrations/safe_statements.rb

@@ -183,7 +183,8 @@ def safe_add_concurrent_partitioned_index(
     using: nil,
     unique: nil,
     where: nil,
-    comment: nil
+    comment: nil,
+    recurse: nil
   )
     raise ArgumentError, "Expected <name_suffix> to be present" unless name_suffix.present?
 
@@ -200,16 +201,47 @@ def safe_add_concurrent_partitioned_index(
 
     _validate_index_name!(parent_index)
 
+    # Validate index names upfront to avoid unnecessary recursion
     child_tables_with_metadata = _partitions_for_table(parent_schema, parent_table).map do |child_schema, child_table|
-      raise PgHaMigrations::InvalidMigrationError, "Partitioned table #{parent_table.inspect} contains sub-partitions" if _partitioned_table?(child_schema, child_table)
-
       child_index = "#{child_table}_#{name_suffix}"
 
       _validate_index_name!(child_index)
 
       [child_schema, child_table, child_index]
     end
 
+    # When sub-partitioning detected:
+    #   - recursively create indexes when recurse is true
+    #   - do nothing when if_not_exists is true and index already valid
+    #   - otherwise raise helpful error
+    child_tables_with_metadata.each do |child_schema, child_table, child_index|
+      next unless _partitioned_table?(child_schema, child_table)
+
+      if recurse
+        safe_add_concurrent_partitioned_index(
+          "#{connection.quote_schema_name(child_schema)}.#{child_table}",
+          columns,
+          name_suffix: name_suffix,
+          if_not_exists: if_not_exists,
+          recurse: recurse,
+          using: using,
+          unique: unique,
+          where: where,
+        )
+      elsif !if_not_exists
+        message = "Partitioned table #{parent_table.inspect} contains sub-partitions. " \
+          "If sub-partition indexes have already been created, execute this method with 'if_not_exists: true'. " \
+          "If sub-partition indexes have not been created, consider executing this method with 'recurse: true'."
+
+        raise PgHaMigrations::InvalidMigrationError, message
+      elsif !_index_valid?(child_schema, child_index)
+        message = "Sub-partition #{child_table.inspect} does not have valid index #{child_index.inspect}. " \
+          "Consider executing this method with 'recurse: true'."
+
+        raise PgHaMigrations::InvalidMigrationError, message
+      end
+    end
+
     # TODO: take out ShareLock after issue #39 is implemented
     safely_acquire_lock_for_table(fully_qualified_parent_table) do
       # CREATE INDEX ON ONLY parent_table
@@ -227,6 +259,11 @@ def safe_add_concurrent_partitioned_index(
     end
 
     child_tables_with_metadata.each do |child_schema, child_table, child_index|
+      # If we get past all of the validation above, we can safely assume that
+      # indexes on sub-partitions have already been created recursviely or in a
+      # separate, explicit call to safe_add_concurrent_partitioned_index
+      next if _partitioned_table?(child_schema, child_table)
+
       # CREATE INDEX CONCURRENTLY ON child_table
       safe_add_concurrent_index(
         "#{connection.quote_schema_name(child_schema)}.#{child_table}",

spec/safe_statements_spec.rb

@@ -1692,24 +1692,116 @@ def up
             end
           end
 
-          it "raises error when name_suffix is not present" do
+          it "creates valid index when sub-partition indexes have already been created and if_not_exists is true" do
+            create_range_partitioned_table(:foos3, migration_klass, with_partman: true)
+            create_range_partitioned_table(:foos3_sub, migration_klass, with_partman: true)
+
+            setup_migration = Class.new(migration_klass) do
+              def up
+                unsafe_execute(<<~SQL)
+                  ALTER TABLE foos3
+                  ATTACH PARTITION foos3_sub
+                  FOR VALUES FROM ('2020-01-01') TO ('2020-02-01')
+                SQL
+
+                safe_add_concurrent_partitioned_index :foos3_sub, :updated_at, name_suffix: "updated_at_idx"
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
             test_migration = Class.new(migration_klass) do
               def up
-                safe_add_concurrent_partitioned_index :foos3, :updated_at, name_suffix: nil
+                safe_add_concurrent_partitioned_index :foos3, :updated_at, name_suffix: "updated_at_idx", if_not_exists: true
               end
             end
 
-            expect do
-              test_migration.suppress_messages { test_migration.migrate(:up) }
-            end.to raise_error(ArgumentError, "Expected <name_suffix> to be present")
+            allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+
+            test_migration.suppress_messages { test_migration.migrate(:up) }
+
+            aggregate_failures do
+              expect(ActiveRecord::Base.connection).to have_received(:execute).with(/CREATE INDEX IF NOT EXISTS "foos3_updated_at_idx" ON ONLY/).once
+              expect(ActiveRecord::Base.connection).to have_received(:execute).with(/CREATE INDEX CONCURRENTLY IF NOT EXISTS/).exactly(10).times
+              expect(ActiveRecord::Base.connection).to have_received(:execute).with(/ALTER INDEX/).exactly(11).times
+            end
+
+            aggregate_failures do
+              # look up indexes on child tables and parent table
+              partitions_for_table(:foos3).append("foos3").each do |table|
+                indexes = ActiveRecord::Base.connection.indexes(table)
+
+                expect(indexes.size).to eq(1)
+                expect(indexes.first).to have_attributes(
+                  table: table,
+                  name: "#{table}_updated_at_idx",
+                  columns: ["updated_at"],
+                  using: :btree,
+                )
+              end
+            end
           end
 
-          it "raises error when table is not partitioned" do
+          it "creates valid index when sub-partitioning detected and recurse is true" do
+            create_range_partitioned_table(:foos3, migration_klass, with_partman: true)
+            create_range_partitioned_table(:foos3_sub, migration_klass, with_partman: true)
+
             setup_migration = Class.new(migration_klass) do
               def up
-                safe_create_table :foos3 do |t|
-                  t.timestamps null: false
-                end
+                unsafe_execute(<<~SQL)
+                  ALTER TABLE foos3
+                  ATTACH PARTITION foos3_sub
+                  FOR VALUES FROM ('2020-01-01') TO ('2020-02-01')
+                SQL
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                safe_add_concurrent_partitioned_index :foos3, :updated_at, name_suffix: "updated_at_idx", recurse: true
+              end
+            end
+
+            allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+
+            test_migration.suppress_messages { test_migration.migrate(:up) }
+
+            aggregate_failures do
+              expect(ActiveRecord::Base.connection).to have_received(:execute).with(/CREATE INDEX "foos3_updated_at_idx" ON ONLY/).once
+              expect(ActiveRecord::Base.connection).to have_received(:execute).with(/CREATE INDEX "foos3_sub_updated_at_idx" ON ONLY/).once
+              expect(ActiveRecord::Base.connection).to have_received(:execute).with(/CREATE INDEX CONCURRENTLY/).exactly(20).times
+              expect(ActiveRecord::Base.connection).to have_received(:execute).with(/ALTER INDEX/).exactly(21).times
+            end
+
+            aggregate_failures do
+              # look up indexes on child tables, parent table, and sub-partitions
+              partitions_for_table(:foos3).append("foos3").concat(partitions_for_table(:foos3_sub)).each do |table|
+                indexes = ActiveRecord::Base.connection.indexes(table)
+
+                expect(indexes.size).to eq(1)
+                expect(indexes.first).to have_attributes(
+                  table: table,
+                  name: "#{table}_updated_at_idx",
+                  columns: ["updated_at"],
+                  using: :btree,
+                )
+              end
+            end
+          end
+
+          it "raises error when sub-partitioning detected and if_not_exists is false" do
+            create_range_partitioned_table(:foos3, migration_klass)
+            create_range_partitioned_table(:foos3_sub, migration_klass)
+
+            setup_migration = Class.new(migration_klass) do
+              def up
+                unsafe_execute(<<~SQL)
+                  ALTER TABLE foos3
+                  ATTACH PARTITION foos3_sub
+                  FOR VALUES FROM ('2020-01-01') TO ('2020-02-01')
+                SQL
               end
             end
 
@@ -1723,10 +1815,10 @@ def up
 
             expect do
               test_migration.suppress_messages { test_migration.migrate(:up) }
-            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos3\" is not a partitioned table")
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, /Partitioned table "foos3" contains sub-partitions/)
           end
 
-          it "raises error when sub-partitioning detected" do
+          it "raises error when sub-partition indexes have not been created and if_not_exists is true" do
             create_range_partitioned_table(:foos3, migration_klass)
             create_range_partitioned_table(:foos3_sub, migration_klass)
 
@@ -1742,6 +1834,40 @@ def up
 
             setup_migration.suppress_messages { setup_migration.migrate(:up) }
 
+            test_migration = Class.new(migration_klass) do
+              def up
+                safe_add_concurrent_partitioned_index :foos3, :updated_at, name_suffix: "updated_at_idx", if_not_exists: true
+              end
+            end
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, /Sub-partition "foos3_sub" does not have valid index "foos3_sub_updated_at_idx"/)
+          end
+
+          it "raises error when name_suffix is not present" do
+            test_migration = Class.new(migration_klass) do
+              def up
+                safe_add_concurrent_partitioned_index :foos3, :updated_at, name_suffix: nil
+              end
+            end
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to raise_error(ArgumentError, "Expected <name_suffix> to be present")
+          end
+
+          it "raises error when table is not partitioned" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                safe_create_table :foos3 do |t|
+                  t.timestamps null: false
+                end
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
             test_migration = Class.new(migration_klass) do
               def up
                 safe_add_concurrent_partitioned_index :foos3, :updated_at, name_suffix: "updated_at_idx"
@@ -1750,7 +1876,7 @@ def up
 
             expect do
               test_migration.suppress_messages { test_migration.migrate(:up) }
-            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Partitioned table \"foos3\" contains sub-partitions")
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos3\" is not a partitioned table")
           end
 
           it "raises error when index invalid" do