pg_ha_migrations 1.7.0 → 1.8.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9111677a3084d7b769a43f1c7078822c9ae84420443893f54dc1e228bfa037d1
-  data.tar.gz: 7f1db28ace4c6416980c1f5c3461669f43bd7cb116996c2af7759d48515a8899
+  metadata.gz: 6858f02b9a874bbaf79865c789a490c9aa1140240537d0eded8f48486c6348f3
+  data.tar.gz: 733394b7f83821f71821816777765d982d1580761522e0751534d3e3c62f598b
 SHA512:
-  metadata.gz: 30c567438be90db49faf206696a51176ea397d448913194ac7b2de02349294cced3d3c94fc08e3b5225eb7b1678d997af716de4f3da9cf15491745ae92af2a22
-  data.tar.gz: dba6fa0e40b690a838a2ae26bd48a0cf18b0acda0922a74a14ae8dec295f1c3a48181a11f0ed4bc6cb46f24e9dadd03dcad86ec805c83c00ecd0cf998176cead
+  metadata.gz: 55f3f8e3730fc11183e71cbf6ba9d32dabf1c8b0ec532da6861b47d045cffb348184350831f564999ecee22b601931ab577ca20f8c4e831e0d4b776babe21a58
+  data.tar.gz: e260ebe93cafba7f3119c41bd2594a797293713f95d072378dbc5733f176d86ce8a93c8a24b53b6fbb9ec77756f4db9ab5fdb92f6080de6ac0fb56e158e5b5d3

data/.github/workflows/ci.yml

@@ -10,13 +10,15 @@ jobs:
           - 13
           - 14
           - 15
+          - 16
         ruby:
-          - 3.0
-          - 3.1
-          - 3.2
+          - "3.0"
+          - "3.1"
+          - "3.2"
         gemfile:
           - rails_6.1
           - rails_7.0
+          - rails_7.1
     name: PostgreSQL ${{ matrix.pg }} - Ruby ${{ matrix.ruby }} - ${{ matrix.gemfile }}
     runs-on: ubuntu-latest
     env: # $BUNDLE_GEMFILE must be set at the job level, so it is set for all steps

data/Appraisals

@@ -1,7 +1,11 @@
 appraise "rails-6.1" do
-  gem "rails", "6.1.0"
+  gem "rails", "6.1.7.6"
 end
 
 appraise "rails-7.0" do
-  gem "rails", "7.0.1"
+  gem "rails", "7.0.8"
+end
+
+appraise "rails-7.1" do
+  gem "rails", "7.1.0"
 end

data/Dockerfile

@@ -6,6 +6,6 @@ RUN apt-get update && apt-get install -y curl ca-certificates gnupg lsb-release
 
 RUN curl https://www.postgresql.org/media/keys/ACCC4CF8.asc | gpg --dearmor | tee /etc/apt/trusted.gpg.d/apt.postgresql.org.gpg >/dev/null
 
-RUN echo "deb http://apt.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg main" > /etc/apt/sources.list.d/pgdg.list
+RUN echo "deb https://apt-archive.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg-archive main" > /etc/apt/sources.list.d/pgdg.list
 
-RUN apt update && apt-get install -y postgresql-$PG_MAJOR-partman
+RUN apt update && apt-get install -y postgresql-$PG_MAJOR-partman=4.7.4-2.pgdg110+1

data/README.md

@@ -68,7 +68,7 @@ When `unsafe_*` migration methods support checks of this type you can bypass the
 
 Similarly we believe the `force: true` option to ActiveRecord's `create_table` method is always unsafe, and therefore we disallow it even when calling `unsafe_create_table`. This option won't be enabled by default until 2.0, but you can opt-in by setting `config.allow_force_create_table = false` [in your configuration initializer](#configuration).
 
-[Running multiple DDL statements inside a transaction acquires exclusive locks on all of the modified objects](https://medium.com/paypal-tech/postgresql-at-scale-database-schema-changes-without-downtime-20d3749ed680#cc22). For that reason, this gem [disables DDL transactions](./lib/pg_ha_migrations.rb:8) by default. You can change this by resetting `ActiveRecord::Migration.disable_ddl_transaction` in your application.
+[Running multiple DDL statements inside a transaction acquires exclusive locks on all of the modified objects](https://medium.com/paypal-tech/postgresql-at-scale-database-schema-changes-without-downtime-20d3749ed680#cc22). For that reason, this gem [disables DDL transactions](./lib/pg_ha_migrations/hacks/disable_ddl_transaction.rb) by default. You can change this by resetting `ActiveRecord::Migration.disable_ddl_transaction` in your application.
 
 The following functionality is currently unsupported:
 
@@ -76,6 +76,11 @@ The following functionality is currently unsupported:
 - Generators
 - schema.rb
 
+Compatibility notes:
+
+- While some features may work with other versions, this gem is currently tested against PostgreSQL 11+ and Partman 4.x
+- There is a [bug](https://github.com/rails/rails/pull/41490) in early versions of Rails 6.1 when using `algorithm: :concurrently`. To add / remove indexes concurrently, please upgrade to at least Rails 6.1.4.
+
 #### safe\_create\_table
 
 Safely creates a new table.
@@ -166,6 +171,14 @@ Unsafely make a column not nullable.
 unsafe_make_column_not_nullable :table, :column
 ```
 
+#### safe\_add\_index\_on\_empty\_table
+
+Safely add an index on a table with zero rows. This will raise an error if the table contains data.
+
+```ruby
+safe_add_index_on_empty_table :table, :column
+```
+
 #### safe\_add\_concurrent\_index
 
 Add an index concurrently.
@@ -188,6 +201,41 @@ Safely remove an index. Migrations that contain this statement must also include
 safe_remove_concurrent_index :table, :name => :index_name
 ```
 
+#### safe\_add\_concurrent\_partitioned\_index
+
+Add an index to a natively partitioned table concurrently, as described in the [table partitioning docs](https://www.postgresql.org/docs/current/ddl-partitioning.html):
+
+> To avoid long lock times, it is possible to use `CREATE INDEX ON ONLY` the partitioned table; such an index is marked invalid, and the partitions do not get the index applied automatically.
+> The indexes on partitions can be created individually using `CONCURRENTLY`, and then attached to the index on the parent using `ALTER INDEX .. ATTACH PARTITION`.
+> Once indexes for all partitions are attached to the parent index, the parent index is marked valid automatically.
+
+```ruby
+# Assuming this table has partitions child1 and child2, the following indexes will be created:
+#   - index_partitioned_table_on_column
+#   - index_child1_on_column (attached to index_partitioned_table_on_column)
+#   - index_child2_on_column (attached to index_partitioned_table_on_column)
+safe_add_concurrent_partitioned_index :partitioned_table, :column
+```
+
+Add a composite index using the `hash` index type with custom name for the parent index when the parent table contains sub-partitions.
+
+```ruby
+# Assuming this table has partitions child1 and child2, and child1 has sub-partitions sub1 and sub2,
+# the following indexes will be created:
+#   - custom_name_idx
+#   - index_child1_on_column1_column2 (attached to custom_name_idx)
+#   - index_sub1_on_column1_column2 (attached to index_child1_on_column1_column2)
+#   - index_sub2_on_column1_column2 (attached to index_child1_on_column1_column2)
+#   - index_child2_on_column1_column2 (attached to custom_name_idx)
+safe_add_concurrent_partitioned_index :partitioned_table, [:column1, :column2], name: "custom_name_idx", using: :hash
+```
+
+Note:
+
+This method runs multiple DDL statements non-transactionally.
+Creating or attaching an index on a child table could fail.
+In such cases an exception will be raised, and an `INVALID` index will be left on the parent table.
+
 #### safe\_add\_unvalidated\_check\_constraint
 
 Safely add a `CHECK` constraint. The constraint will not be immediately validated on existing rows to avoid a full table scan while holding an exclusive lock. After adding the constraint, you'll need to use `safe_validate_check_constraint` to validate existing rows.
@@ -387,7 +435,7 @@ safe_partman_reapply_privileges :table
 
 #### safely\_acquire\_lock\_for\_table
 
-Safely acquire a lock for a table.
+Safely acquire an access exclusive lock for a table.
 
 ```ruby
 safely_acquire_lock_for_table(:table) do
@@ -395,6 +443,19 @@ safely_acquire_lock_for_table(:table) do
 end
 ```
 
+Safely acquire a lock for a table in a different mode.
+
+```ruby
+safely_acquire_lock_for_table(:table, mode: :share) do
+  ...
+end
+```
+
+Note:
+
+We enforce that only one table (or a table and its partitions) can be locked at a time.
+Attempting to acquire a nested lock on a different table will result in an error.
+
 #### adjust\_lock\_timeout
 
 Adjust lock timeout.
@@ -423,6 +484,22 @@ Set maintenance work mem.
 safe_set_maintenance_work_mem_gb 1
 ```
 
+#### ensure\_small\_table!
+
+Ensure a table on disk is below the default threshold (10 megabytes).
+This will raise an error if the table is too large.
+
+```ruby
+ensure_small_table! :table
+```
+
+Ensure a table on disk is below a custom threshold and is empty.
+This will raise an error if the table is too large and/or contains data.
+
+```ruby
+ensure_small_table! :table, empty: true, threshold: 100.megabytes
+```
+
 ### Configuration
 
 The gem can be configured in an initializer.

data/docker-compose.yml

@@ -4,7 +4,7 @@ services:
     build:
       context: .
       args:
-        - PGVERSION=${PGVERSION:-15}
+        - PGVERSION=${PGVERSION:-16}
     ports:
       - "5432:5432"
     environment:

data/gemfiles/rails_6.1.gemfile

@@ -2,6 +2,6 @@
 
 source "https://rubygems.org"
 
-gem "rails", "6.1.0"
+gem "rails", "6.1.7.6"
 
 gemspec path: "../"

data/gemfiles/rails_7.0.gemfile

@@ -2,6 +2,6 @@
 
 source "https://rubygems.org"
 
-gem "rails", "7.0.1"
+gem "rails", "7.0.8"
 
 gemspec path: "../"

data/gemfiles/rails_7.1.gemfile

@@ -0,0 +1,7 @@
+# This file was generated by Appraisal
+
+source "https://rubygems.org"
+
+gem "rails", "7.1.0"
+
+gemspec path: "../"

data/lib/pg_ha_migrations/allowed_versions.rb

@@ -1,7 +1,7 @@
 require "active_record/migration/compatibility"
 
 module PgHaMigrations::AllowedVersions
-  ALLOWED_VERSIONS = [4.2, 5.0, 5.1, 5.2, 6.0, 6.1, 7.0].map do |v|
+  ALLOWED_VERSIONS = [4.2, 5.0, 5.1, 5.2, 6.0, 6.1, 7.0, 7.1].map do |v|
     begin
       ActiveRecord::Migration[v]
     rescue ArgumentError

data/lib/pg_ha_migrations/blocking_database_transactions.rb

@@ -1,13 +1,18 @@
 module PgHaMigrations
   class BlockingDatabaseTransactions
     LongRunningTransaction = Struct.new(:database, :current_query, :state, :transaction_age, :tables_with_locks) do
+      def initialize(*args)
+        super
+
+        self.tables_with_locks = tables_with_locks.map { |args| Table.new(*args) }.select(&:present?)
+      end
+
       def description
-        locked_tables = tables_with_locks.compact
         [
           database,
-          locked_tables.size > 0 ? "tables (#{locked_tables.join(', ')})" : nil,
+          tables_with_locks.size > 0 ? "tables (#{tables_with_locks.map(&:fully_qualified_name).join(', ')})" : nil,
           "#{idle? ? "currently idle " : ""}transaction open for #{transaction_age}",
-          "#{idle? ? "last " : ""}query: #{current_query}"
+          "#{idle? ? "last " : ""}query: #{current_query}",
         ].compact.join(" | ")
       end
 
@@ -43,7 +48,7 @@ module PgHaMigrations
           psa.#{query_column} as current_query,
           psa.state,
           clock_timestamp() - psa.xact_start AS transaction_age,
-          array_agg(distinct c.relname) AS tables_with_locks
+          array_agg(distinct array[c.relname, ns.nspname, l.mode]) AS tables_with_locks
         FROM pg_stat_activity psa -- Cluster wide
           LEFT JOIN pg_locks l ON (psa.#{pid_column} = l.pid)  -- Cluster wide
           LEFT JOIN pg_class c ON ( -- Database wide
@@ -56,7 +61,7 @@ module PgHaMigrations
             l.locktype != 'relation'
             OR (
                ns.nspname != 'pg_catalog'
-               AND c.relkind = 'r'
+               AND c.relkind IN ('r', 'p') -- 'r' is a standard table; 'p' is a partition parent
             )
           )
           AND psa.xact_start < clock_timestamp() - ?::interval

data/lib/pg_ha_migrations/hacks/add_index_on_only.rb

@@ -0,0 +1,30 @@
+require "active_record/connection_adapters/postgresql_adapter"
+require "active_record/connection_adapters/postgresql/schema_creation"
+
+module PgHaMigrations
+  module ActiveRecordHacks
+    module IndexAlgorithms
+      def index_algorithms
+        super.merge(only: "ONLY")
+      end
+    end
+
+    module CreateIndexDefinition
+      def visit_CreateIndexDefinition(o)
+        if o.algorithm == "ONLY"
+          o.algorithm = nil
+
+          quoted_index = quote_column_name(o.index.name)
+          quoted_table = quote_table_name(o.index.table)
+
+          super.sub("#{quoted_index} ON #{quoted_table}", "#{quoted_index} ON ONLY #{quoted_table}")
+        else
+          super
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(PgHaMigrations::ActiveRecordHacks::IndexAlgorithms)
+ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaCreation.prepend(PgHaMigrations::ActiveRecordHacks::CreateIndexDefinition)

data/lib/pg_ha_migrations/hacks/disable_ddl_transaction.rb

@@ -12,4 +12,3 @@ module PgHaMigrations
 end
 
 ActiveRecord::Migration.singleton_class.prepend(PgHaMigrations::ActiveRecordHacks::DisableDdlTransaction)
-

data/lib/pg_ha_migrations/lock_mode.rb

@@ -0,0 +1,100 @@
+module PgHaMigrations
+  class LockMode
+    include Comparable
+
+    MODE_CONFLICTS = ActiveSupport::OrderedHash.new
+
+    MODE_CONFLICTS[:access_share] = %i[
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:row_share] = %i[
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:row_exclusive] = %i[
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:share_update_exclusive] = %i[
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:share] = %i[
+      row_exclusive
+      share_update_exclusive
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:share_row_exclusive] = %i[
+      row_exclusive
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:exclusive] = %i[
+      row_share
+      row_exclusive
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:access_exclusive] = %i[
+      access_share
+      row_share
+      row_exclusive
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    attr_reader :mode
+
+    delegate :to_s, to: :mode
+
+    def initialize(mode)
+      @mode = mode
+        .to_s
+        .underscore
+        .delete_suffix("_lock")
+        .to_sym
+
+      if !MODE_CONFLICTS.keys.include?(@mode)
+        raise ArgumentError, "Unrecognized lock mode #{@mode.inspect}. Valid modes: #{MODE_CONFLICTS.keys}"
+      end
+    end
+
+    def to_sql
+      mode
+        .to_s
+        .upcase
+        .gsub("_", " ")
+    end
+
+    def <=>(other)
+      MODE_CONFLICTS.keys.index(mode) <=> MODE_CONFLICTS.keys.index(other.mode)
+    end
+
+    def conflicts_with?(other)
+      MODE_CONFLICTS[mode].include?(other.mode)
+    end
+  end
+end

data/lib/pg_ha_migrations/relation.rb

@@ -0,0 +1,155 @@
+module PgHaMigrations
+  Relation = Struct.new(:name, :schema, :mode) do
+    def self.connection
+      ActiveRecord::Base.connection
+    end
+
+    delegate :inspect, to: :name
+    delegate :connection, to: :class
+
+    def initialize(name, schema, mode=nil)
+      super(name, schema)
+
+      self.mode = LockMode.new(mode) if mode.present?
+    end
+
+    def conflicts_with?(other)
+      self == other && (
+        mode.nil? || other.mode.nil? || mode.conflicts_with?(other.mode)
+      )
+    end
+
+    def fully_qualified_name
+      @fully_qualified_name ||= [
+        PG::Connection.quote_ident(schema),
+        PG::Connection.quote_ident(name),
+      ].join(".")
+    end
+
+    def present?
+      name.present? && schema.present?
+    end
+
+    def ==(other)
+      other.is_a?(Relation) && name == other.name && schema == other.schema
+    end
+  end
+
+  class Table < Relation
+    def self.from_table_name(table, mode=nil)
+      pg_name = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(table.to_s)
+
+      schema_conditional = if pg_name.schema
+        "#{connection.quote(pg_name.schema)}"
+      else
+        "ANY (current_schemas(false))"
+      end
+
+      schema = connection.select_value(<<~SQL)
+        SELECT schemaname
+        FROM pg_tables
+        WHERE tablename = #{connection.quote(pg_name.identifier)} AND schemaname = #{schema_conditional}
+        ORDER BY array_position(current_schemas(false), schemaname)
+        LIMIT 1
+      SQL
+
+      raise UndefinedTableError, "Table #{pg_name.quoted} does not exist#{" in search path" unless pg_name.schema}" unless schema.present?
+
+      new(pg_name.identifier, schema, mode)
+    end
+
+    def natively_partitioned?
+      return @natively_partitioned if defined?(@natively_partitioned)
+
+      @natively_partitioned = !!connection.select_value(<<~SQL)
+        SELECT true
+        FROM pg_partitioned_table, pg_class, pg_namespace
+        WHERE pg_class.oid = pg_partitioned_table.partrelid
+          AND pg_class.relnamespace = pg_namespace.oid
+          AND pg_class.relname = #{connection.quote(name)}
+          AND pg_namespace.nspname = #{connection.quote(schema)}
+      SQL
+    end
+
+    def partitions(include_sub_partitions: false, include_self: false)
+      tables = connection.structs_from_sql(self.class, <<~SQL)
+        SELECT child.relname AS name, child_ns.nspname AS schema, NULLIF('#{mode}', '') AS mode
+        FROM pg_inherits
+          JOIN pg_class parent        ON pg_inherits.inhparent = parent.oid
+          JOIN pg_class child         ON pg_inherits.inhrelid  = child.oid
+          JOIN pg_namespace parent_ns ON parent.relnamespace = parent_ns.oid
+          JOIN pg_namespace child_ns  ON child.relnamespace = child_ns.oid
+        WHERE parent.relname = #{connection.quote(name)}
+          AND parent_ns.nspname = #{connection.quote(schema)}
+        ORDER BY child.oid -- Ensure consistent ordering for tests
+      SQL
+
+      if include_sub_partitions
+        sub_partitions = tables.each_with_object([]) do |table, arr|
+          arr.concat(table.partitions(include_sub_partitions: true))
+        end
+
+        tables.concat(sub_partitions)
+      end
+
+      tables.prepend(self) if include_self
+
+      tables
+    end
+
+    def has_rows?
+      connection.select_value("SELECT EXISTS (SELECT 1 FROM #{fully_qualified_name} LIMIT 1)")
+    end
+
+    def total_bytes
+      connection.select_value(<<~SQL)
+        SELECT pg_total_relation_size(pg_class.oid)
+        FROM pg_class, pg_namespace
+        WHERE pg_class.relname = #{connection.quote(name)}
+          AND pg_namespace.nspname = #{connection.quote(schema)}
+      SQL
+    end
+  end
+
+  class Index < Relation
+    MAX_NAME_SIZE = 63 # bytes
+
+    def self.from_table_and_columns(table, columns)
+      name = connection.index_name(table.name, columns)
+
+      # modified from https://github.com/rails/rails/pull/47753
+      if name.bytesize > MAX_NAME_SIZE
+        hashed_identifier = "_#{OpenSSL::Digest::SHA256.hexdigest(name).first(10)}"
+        description = name.sub("index_#{table.name}_on", "idx_on")
+
+        short_limit = MAX_NAME_SIZE - hashed_identifier.bytesize
+        short_description = description.mb_chars.limit(short_limit).to_s
+
+        name = "#{short_description}#{hashed_identifier}"
+      end
+
+      new(name, table)
+    end
+
+    attr_accessor :table
+
+    def initialize(name, table)
+      super(name, table.schema)
+
+      self.table = table
+
+      connection.send(:validate_index_length!, table.name, name)
+    end
+
+    def valid?
+      !!connection.select_value(<<~SQL)
+        SELECT pg_index.indisvalid
+        FROM pg_index, pg_class, pg_namespace
+        WHERE pg_class.oid = pg_index.indexrelid
+          AND pg_class.relnamespace = pg_namespace.oid
+          AND pg_namespace.nspname = #{connection.quote(schema)}
+          AND pg_class.relname = #{connection.quote(name)}
+      SQL
+    end
+  end
+end

data/lib/pg_ha_migrations/safe_statements.rb

@@ -1,14 +1,4 @@
 module PgHaMigrations::SafeStatements
-  PARTITION_TYPES = %i[range list hash]
-
-  PARTMAN_UPDATE_CONFIG_OPTIONS = %i[
-    infinite_time_partitions
-    inherit_privileges
-    premake
-    retention
-    retention_keep_table
-  ]
-
   def safe_added_columns_without_default_value
     @safe_added_columns_without_default_value ||= []
   end
@@ -154,6 +144,22 @@ module PgHaMigrations::SafeStatements
     end
   end
 
+  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
+
+    # Avoids taking out an unnecessary SHARE lock if the table does have data
+    ensure_small_table!(table, empty: true)
+
+    safely_acquire_lock_for_table(table, mode: :share) do
+      # Ensure data wasn't written in the split second after the first check
+      ensure_small_table!(table, empty: true)
+
+      unsafe_add_index(table, columns, **options)
+    end
+  end
+
   def safe_add_concurrent_index(table, columns, options={})
     unsafe_add_index(table, columns, **options.merge(:algorithm => :concurrently))
   end
@@ -170,6 +176,90 @@ module PgHaMigrations::SafeStatements
     unsafe_remove_index(table, **options.merge(:algorithm => :concurrently))
   end
 
+  def safe_add_concurrent_partitioned_index(
+    table,
+    columns,
+    name: nil,
+    if_not_exists: nil,
+    using: nil,
+    unique: nil,
+    where: nil,
+    comment: nil
+  )
+
+    if ActiveRecord::Base.connection.postgresql_version < 11_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Concurrent partitioned index creation not supported on Postgres databases before version 11"
+    end
+
+    parent_table = PgHaMigrations::Table.from_table_name(table)
+
+    raise PgHaMigrations::InvalidMigrationError, "Table #{parent_table.inspect} is not a partitioned table" unless parent_table.natively_partitioned?
+
+    parent_index = if name.present?
+      PgHaMigrations::Index.new(name, parent_table)
+    else
+      PgHaMigrations::Index.from_table_and_columns(parent_table, columns)
+    end
+
+    # Short-circuit when if_not_exists: true and index already valid
+    return if if_not_exists && parent_index.valid?
+
+    child_indexes = parent_table.partitions.map do |child_table|
+      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
+
+    child_indexes.each do |child_index|
+      add_index_method = if child_index.table.natively_partitioned?
+        :safe_add_concurrent_partitioned_index
+      else
+        :safe_add_concurrent_index
+      end
+
+      send(
+        add_index_method,
+        child_index.table.fully_qualified_name,
+        columns,
+        name: child_index.name,
+        if_not_exists: if_not_exists,
+        using: using,
+        unique: unique,
+        where: where,
+      )
+    end
+
+    # Avoid taking out an unnecessary lock if there are no child tables to attach
+    if child_indexes.present?
+      safely_acquire_lock_for_table(parent_table.fully_qualified_name) do
+        child_indexes.each do |child_index|
+          say_with_time "Attaching index #{child_index.inspect} to #{parent_index.inspect}" do
+            connection.execute(<<~SQL)
+              ALTER INDEX #{parent_index.fully_qualified_name}
+              ATTACH PARTITION #{child_index.fully_qualified_name}
+            SQL
+          end
+        end
+      end
+    end
+
+    raise PgHaMigrations::InvalidMigrationError, "Unexpected state. Parent index #{parent_index.inspect} is invalid" unless parent_index.valid?
+  end
+
   def safe_set_maintenance_work_mem_gb(gigabytes)
     unsafe_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
   end
@@ -223,8 +313,8 @@ module PgHaMigrations::SafeStatements
   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?
 
-    unless PARTITION_TYPES.include?(type)
-      raise ArgumentError, "Expected <type> to be symbol in #{PARTITION_TYPES} but received #{type.inspect}"
+    unless PgHaMigrations::PARTITION_TYPES.include?(type)
+      raise ArgumentError, "Expected <type> to be symbol in #{PgHaMigrations::PARTITION_TYPES} but received #{type.inspect}"
     end
 
     if ActiveRecord::Base.connection.postgresql_version < 10_00_00
@@ -351,7 +441,7 @@ module PgHaMigrations::SafeStatements
       retention_keep_table: retention_keep_table,
     }.compact
 
-    unsafe_partman_update_config(create_parent_options[:parent_table], **update_config_options)
+    unsafe_partman_update_config(table, **update_config_options)
   end
 
   def safe_partman_update_config(table, **options)
@@ -363,7 +453,7 @@ module PgHaMigrations::SafeStatements
   end
 
   def unsafe_partman_update_config(table, **options)
-    invalid_options = options.keys - PARTMAN_UPDATE_CONFIG_OPTIONS
+    invalid_options = options.keys - PgHaMigrations::PARTMAN_UPDATE_CONFIG_OPTIONS
 
     raise ArgumentError, "Unrecognized argument(s): #{invalid_options}" unless invalid_options.empty?
 
@@ -402,38 +492,13 @@ module PgHaMigrations::SafeStatements
   end
 
   def _fully_qualified_table_name_for_partman(table)
-    identifiers = table.to_s.split(".")
-
-    raise PgHaMigrations::InvalidMigrationError, "Expected table to be in the format <table> or <schema>.<table> but received #{table}" if identifiers.size > 2
-
-    identifiers.each { |identifier| _validate_partman_identifier(identifier) }
-
-    schema_conditional = if identifiers.size > 1
-      "'#{identifiers.first}'"
-    else
-      "ANY (current_schemas(false))"
-    end
+    table = PgHaMigrations::Table.from_table_name(table)
 
-    schema = connection.select_value(<<~SQL)
-      SELECT schemaname
-      FROM pg_tables
-      WHERE tablename = '#{identifiers.last}' AND schemaname = #{schema_conditional}
-      ORDER BY array_position(current_schemas(false), schemaname)
-      LIMIT 1
-    SQL
-
-    raise PgHaMigrations::InvalidMigrationError, "Could not find table #{table}" unless schema.present?
-
-    _validate_partman_identifier(schema)
-
-    # Quoting is unneeded since _validate_partman_identifier ensures the schema / table use standard naming conventions
-    "#{schema}.#{identifiers.last}"
-  end
-
-  def _validate_partman_identifier(identifier)
-    if identifier.to_s !~ /^[a-z_][a-z_\d]*$/
-      raise PgHaMigrations::InvalidMigrationError, "Partman requires schema / table names to be lowercase with underscores"
-    end
+    [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
@@ -463,17 +528,39 @@ module PgHaMigrations::SafeStatements
     super(conn, direction)
   end
 
-  def safely_acquire_lock_for_table(table, &block)
+  def safely_acquire_lock_for_table(table, mode: :access_exclusive, &block)
+    nested_target_table = Thread.current[__method__]
+
     _check_postgres_adapter!
-    table = table.to_s
-    quoted_table_name = connection.quote_table_name(table)
+
+    target_table = PgHaMigrations::Table.from_table_name(table, 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}."
+      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)
 
     successfully_acquired_lock = false
 
     until successfully_acquired_lock
       while (
         blocking_transactions = PgHaMigrations::BlockingDatabaseTransactions.find_blocking_transactions("#{PgHaMigrations::LOCK_TIMEOUT_SECONDS} seconds")
-        blocking_transactions.any? { |query| query.tables_with_locks.include?(table) }
+        blocking_transactions.any? do |query|
+          query.tables_with_locks.any? do |locked_table|
+            target_tables.any? do |target_table|
+              target_table.conflicts_with?(locked_table)
+            end
+          end
+        end
       )
         say "Waiting on blocking transactions:"
         blocking_transactions.each do |blocking_transaction|
@@ -486,13 +573,13 @@ module PgHaMigrations::SafeStatements
         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 #{quoted_table_name};")
+            connection.execute("LOCK #{target_table.fully_qualified_name} IN #{target_table.mode.to_sql} MODE;")
           end
           successfully_acquired_lock = true
         rescue ActiveRecord::StatementInvalid => e
           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 an exclusive lock on the #{quoted_table_name} table."
+            say "Timed out trying to acquire #{target_table.mode.to_sql} lock on the #{target_table.fully_qualified_name} table."
             say "Sleeping for #{sleep_seconds}s to allow potentially queued up queries to finish before continuing."
             sleep(sleep_seconds)
 
@@ -507,6 +594,8 @@ module PgHaMigrations::SafeStatements
         end
       end
     end
+  ensure
+    Thread.current[__method__] = nil unless nested_target_table
   end
 
   def adjust_lock_timeout(timeout_seconds = PgHaMigrations::LOCK_TIMEOUT_SECONDS, &block)
@@ -548,4 +637,16 @@ module PgHaMigrations::SafeStatements
       end
     end
   end
+
+  def ensure_small_table!(table, empty: false, threshold: PgHaMigrations::SMALL_TABLE_THRESHOLD_BYTES)
+    table = PgHaMigrations::Table.from_table_name(table)
+
+    if empty && table.has_rows?
+      raise PgHaMigrations::InvalidMigrationError, "Table #{table.inspect} has rows"
+    end
+
+    if table.total_bytes > threshold
+      raise PgHaMigrations::InvalidMigrationError, "Table #{table.inspect} is larger than #{threshold} bytes"
+    end
+  end
 end

data/lib/pg_ha_migrations/unsafe_statements.rb

@@ -80,6 +80,16 @@ module PgHaMigrations::UnsafeStatements
       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"
     end
 
+    validated_table = PgHaMigrations::Table.from_table_name(table)
+
+    validated_index = if options[:name]
+      PgHaMigrations::Index.new(options[:name], validated_table)
+    else
+      PgHaMigrations::Index.from_table_and_columns(validated_table, column_names)
+    end
+
+    options[:name] = validated_index.name
+
     execute_ancestor_statement(:add_index, table, column_names, **options)
   end
 

data/lib/pg_ha_migrations/version.rb

@@ -1,3 +1,3 @@
 module PgHaMigrations
-  VERSION = "1.7.0"
+  VERSION = "1.8.0"
 end

data/lib/pg_ha_migrations.rb

@@ -2,6 +2,8 @@ require "pg_ha_migrations/version"
 require "rails"
 require "active_record"
 require "active_record/migration"
+require "active_record/connection_adapters/postgresql/utils"
+require "active_support/core_ext/numeric/bytes"
 require "relation_to_struct"
 require "ruby2_keywords"
 
@@ -30,6 +32,17 @@ module PgHaMigrations
 
   LOCK_TIMEOUT_SECONDS = 5
   LOCK_FAILURE_RETRY_DELAY_MULTLIPLIER = 5
+  SMALL_TABLE_THRESHOLD_BYTES = 10.megabytes
+
+  PARTITION_TYPES = %i[range list hash]
+
+  PARTMAN_UPDATE_CONFIG_OPTIONS = %i[
+    infinite_time_partitions
+    inherit_privileges
+    premake
+    retention
+    retention_keep_table
+  ]
 
   # Safe versus unsafe in this context specifically means the following:
   # - Safe operations will not block for long periods of time.
@@ -52,11 +65,17 @@ module PgHaMigrations
 
   # This gem only supports the PostgreSQL adapter at this time.
   UnsupportedAdapter = Class.new(StandardError)
+
+  # Some methods need to inspect the attributes of a table. In such cases,
+  # this error will be raised if the table does not exist
+  UndefinedTableError = Class.new(StandardError)
 end
 
+require "pg_ha_migrations/relation"
 require "pg_ha_migrations/blocking_database_transactions"
 require "pg_ha_migrations/blocking_database_transactions_reporter"
 require "pg_ha_migrations/partman_config"
+require "pg_ha_migrations/lock_mode"
 require "pg_ha_migrations/unsafe_statements"
 require "pg_ha_migrations/safe_statements"
 require "pg_ha_migrations/dependent_objects_checks"
@@ -64,6 +83,7 @@ require "pg_ha_migrations/allowed_versions"
 require "pg_ha_migrations/railtie"
 require "pg_ha_migrations/hacks/disable_ddl_transaction"
 require "pg_ha_migrations/hacks/cleanup_unnecessary_output"
+require "pg_ha_migrations/hacks/add_index_on_only"
 
 module PgHaMigrations::AutoIncluder
   def inherited(klass)

data/pg_ha_migrations.gemspec

@@ -32,12 +32,12 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rake", ">= 12.3.3"
   spec.add_development_dependency "rspec", "~> 3.0"
   spec.add_development_dependency "pg"
-  spec.add_development_dependency "db-query-matchers", "~> 0.11.0"
+  spec.add_development_dependency "db-query-matchers", "~> 0.12.0"
   spec.add_development_dependency "pry"
   spec.add_development_dependency "pry-byebug"
   spec.add_development_dependency "appraisal", "~> 2.2.0"
 
-  spec.add_dependency "rails", ">= 6.1", "< 7.1"
+  spec.add_dependency "rails", ">= 6.1", "< 7.2"
   spec.add_dependency "relation_to_struct", ">= 1.5.1"
   spec.add_dependency "ruby2_keywords"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pg_ha_migrations
 version: !ruby/object:Gem::Version
-  version: 1.7.0
+  version: 1.8.0
 platform: ruby
 authors:
 - celeen
@@ -14,7 +14,7 @@ authors:
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2023-08-10 00:00:00.000000000 Z
+date: 2024-01-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rake
@@ -64,14 +64,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 0.11.0
+        version: 0.12.0
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 0.11.0
+        version: 0.12.0
 - !ruby/object:Gem::Dependency
   name: pry
   requirement: !ruby/object:Gem::Requirement
@@ -123,7 +123,7 @@ dependencies:
         version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.1'
+        version: '7.2'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
@@ -133,7 +133,7 @@ dependencies:
         version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7.1'
+        version: '7.2'
 - !ruby/object:Gem::Dependency
   name: relation_to_struct
   requirement: !ruby/object:Gem::Requirement
@@ -188,15 +188,19 @@ files:
 - gemfiles/.bundle/config
 - gemfiles/rails_6.1.gemfile
 - gemfiles/rails_7.0.gemfile
+- gemfiles/rails_7.1.gemfile
 - lib/pg_ha_migrations.rb
 - lib/pg_ha_migrations/allowed_versions.rb
 - lib/pg_ha_migrations/blocking_database_transactions.rb
 - lib/pg_ha_migrations/blocking_database_transactions_reporter.rb
 - lib/pg_ha_migrations/dependent_objects_checks.rb
+- lib/pg_ha_migrations/hacks/add_index_on_only.rb
 - lib/pg_ha_migrations/hacks/cleanup_unnecessary_output.rb
 - lib/pg_ha_migrations/hacks/disable_ddl_transaction.rb
+- lib/pg_ha_migrations/lock_mode.rb
 - lib/pg_ha_migrations/partman_config.rb
 - lib/pg_ha_migrations/railtie.rb
+- lib/pg_ha_migrations/relation.rb
 - lib/pg_ha_migrations/safe_statements.rb
 - lib/pg_ha_migrations/unsafe_statements.rb
 - lib/pg_ha_migrations/version.rb