Merge pull request #37585 from sebastian-palma/upsert-all-partitioned-indexes

Add support for partitioned indexes in PostgreSQL 11+

activerecord/CHANGELOG.md

@@ -1,3 +1,7 @@
+*   Add support for PostgreSQL 11+ partitioned indexes when using `upsert_all`.
+
+    *Sebastián Palma*
+
 *   Adds support for `if_not_exists` to `add_column` and `if_exists` to `remove_column`.
 
     Applications can set their migrations to ignore exceptions raised when adding a column that already exists or when removing a column that does not exist.

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -283,6 +283,10 @@ module ActiveRecord
         false
       end
 
+      def supports_partitioned_indexes?
+        false
+      end
+
       # Does this adapter support index sort order?
       def supports_index_sort_order?
         false

activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb

@@ -75,7 +75,7 @@ module ActiveRecord
             INNER JOIN pg_index d ON t.oid = d.indrelid
             INNER JOIN pg_class i ON d.indexrelid = i.oid
             LEFT JOIN pg_namespace n ON n.oid = i.relnamespace
-            WHERE i.relkind = 'i'
+            WHERE i.relkind IN ('i', 'I')
               AND i.relname = #{index[:name]}
               AND t.relname = #{table[:name]}
               AND n.nspname = #{index[:schema]}
@@ -93,7 +93,7 @@ module ActiveRecord
             INNER JOIN pg_index d ON t.oid = d.indrelid
             INNER JOIN pg_class i ON d.indexrelid = i.oid
             LEFT JOIN pg_namespace n ON n.oid = i.relnamespace
-            WHERE i.relkind = 'i'
+            WHERE i.relkind IN ('i', 'I')
               AND d.indisprimary = 'f'
               AND t.relname = #{scope[:name]}
               AND n.nspname = #{scope[:schema]}

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -157,6 +157,10 @@ module ActiveRecord
         true
       end
 
+      def supports_partitioned_indexes?
+        database_version >= 110_000
+      end
+
       def supports_partial_index?
         true
       end

activerecord/test/cases/adapters/postgresql/schema_test.rb

@@ -45,6 +45,8 @@ class SchemaTest < ActiveRecord::PostgreSQLTestCase
   PK_TABLE_NAME = "table_with_pk"
   UNMATCHED_SEQUENCE_NAME = "unmatched_primary_key_default_value_seq"
   UNMATCHED_PK_TABLE_NAME = "table_with_unmatched_sequence_for_pk"
+  PARTITIONED_TABLE = "measurements"
+  PARTITIONED_TABLE_INDEX = "index_measurements_on_logdate_and_city_id"
 
   class Thing1 < ActiveRecord::Base
     self.table_name = "test_schema.things"
@@ -311,6 +313,12 @@ class SchemaTest < ActiveRecord::PostgreSQLTestCase
       assert @connection.index_name_exists?(TABLE_NAME, INDEX_E_NAME)
       assert @connection.index_name_exists?(TABLE_NAME, INDEX_E_NAME)
       assert_not @connection.index_name_exists?(TABLE_NAME, "missing_index")
+
+      if supports_partitioned_indexes?
+        create_partitioned_table
+        create_partitioned_table_index
+        assert @connection.index_name_exists?(PARTITIONED_TABLE, PARTITIONED_TABLE_INDEX)
+      end
     end
   end
 
@@ -329,6 +337,13 @@ class SchemaTest < ActiveRecord::PostgreSQLTestCase
   def test_dump_indexes_for_table_with_scheme_specified_in_name
     indexes = @connection.indexes("#{SCHEMA_NAME}.#{TABLE_NAME}")
     assert_equal 5, indexes.size
+
+    if supports_partitioned_indexes?
+      create_partitioned_table
+      create_partitioned_table_index
+      indexes = @connection.indexes("#{SCHEMA_NAME}.#{PARTITIONED_TABLE}")
+      assert_equal 1, indexes.size
+    end
   end
 
   def test_with_uppercase_index_name
@@ -337,6 +352,15 @@ class SchemaTest < ActiveRecord::PostgreSQLTestCase
     with_schema_search_path SCHEMA_NAME do
       assert_nothing_raised { @connection.remove_index "things", name: "things_Index" }
     end
+
+    if supports_partitioned_indexes?
+      create_partitioned_table
+      @connection.execute "CREATE INDEX \"#{PARTITIONED_TABLE}_Index\" ON #{SCHEMA_NAME}.#{PARTITIONED_TABLE} (logdate, city_id)"
+
+      with_schema_search_path SCHEMA_NAME do
+        assert_nothing_raised { @connection.remove_index PARTITIONED_TABLE, name: "#{PARTITIONED_TABLE}_Index" }
+      end
+    end
   end
 
   def test_remove_index_when_schema_specified
@@ -351,6 +375,22 @@ class SchemaTest < ActiveRecord::PostgreSQLTestCase
 
     @connection.execute "CREATE INDEX \"things_Index\" ON #{SCHEMA_NAME}.things (name)"
     assert_raises(ArgumentError) { @connection.remove_index "#{SCHEMA2_NAME}.things", name: "#{SCHEMA_NAME}.things_Index" }
+
+    if supports_partitioned_indexes?
+      create_partitioned_table
+
+      @connection.execute "CREATE INDEX \"#{PARTITIONED_TABLE}_Index\" ON #{SCHEMA_NAME}.#{PARTITIONED_TABLE} (logdate, city_id)"
+      assert_nothing_raised { @connection.remove_index PARTITIONED_TABLE, name: "#{SCHEMA_NAME}.#{PARTITIONED_TABLE}_Index" }
+
+      @connection.execute "CREATE INDEX \"#{PARTITIONED_TABLE}_Index\" ON #{SCHEMA_NAME}.#{PARTITIONED_TABLE} (logdate, city_id)"
+      assert_nothing_raised { @connection.remove_index "#{SCHEMA_NAME}.#{PARTITIONED_TABLE}", name: "#{PARTITIONED_TABLE}_Index" }
+
+      @connection.execute "CREATE INDEX \"#{PARTITIONED_TABLE}_Index\" ON #{SCHEMA_NAME}.#{PARTITIONED_TABLE} (logdate, city_id)"
+      assert_nothing_raised { @connection.remove_index "#{SCHEMA_NAME}.#{PARTITIONED_TABLE}", name: "#{SCHEMA_NAME}.#{PARTITIONED_TABLE}_Index" }
+
+      @connection.execute "CREATE INDEX \"#{PARTITIONED_TABLE}_Index\" ON #{SCHEMA_NAME}.#{PARTITIONED_TABLE} (logdate, city_id)"
+      assert_raises(ArgumentError) { @connection.remove_index "#{SCHEMA2_NAME}.#{PARTITIONED_TABLE}", name: "#{SCHEMA_NAME}.#{PARTITIONED_TABLE}_Index" }
+    end
   end
 
   def test_primary_key_with_schema_specified
@@ -473,6 +513,14 @@ class SchemaTest < ActiveRecord::PostgreSQLTestCase
     def bind_param(value)
       ActiveRecord::Relation::QueryAttribute.new(nil, value, ActiveRecord::Type::Value.new)
     end
+
+    def create_partitioned_table
+      @connection.execute "CREATE TABLE #{SCHEMA_NAME}.\"#{PARTITIONED_TABLE}\" (city_id integer not null, logdate date not null) PARTITION BY LIST (city_id)"
+    end
+
+    def create_partitioned_table_index
+      @connection.execute "CREATE INDEX #{PARTITIONED_TABLE_INDEX} ON #{SCHEMA_NAME}.#{PARTITIONED_TABLE} (logdate, city_id)"
+    end
 end
 
 class SchemaForeignKeyTest < ActiveRecord::PostgreSQLTestCase

activerecord/test/cases/helper.rb

@@ -60,6 +60,7 @@ end
 %w[
   supports_savepoints?
   supports_partial_index?
+  supports_partitioned_indexes?
   supports_insert_returning?
   supports_insert_on_duplicate_skip?
   supports_insert_on_duplicate_update?

activerecord/test/cases/insert_all_test.rb

@@ -280,6 +280,21 @@ class InsertAllTest < ActiveRecord::TestCase
     end
   end
 
+  def test_upsert_all_works_with_partitioned_indexes
+    skip unless supports_insert_on_duplicate_update? && supports_insert_conflict_target? && supports_partitioned_indexes?
+
+    require "models/measurement"
+
+    Measurement.upsert_all([{ city_id: "1", logdate: 1.days.ago, peaktemp: 1, unitsales: 1 },
+                            { city_id: "2", logdate: 2.days.ago, peaktemp: 2, unitsales: 2 },
+                            { city_id: "2", logdate: 3.days.ago, peaktemp: 0, unitsales: 0 }],
+                            unique_by: %i[logdate city_id])
+    assert_equal [[1.day.ago.to_date, 1, 1]],
+                 Measurement.where(city_id: 1).pluck(:logdate, :peaktemp, :unitsales)
+    assert_equal [[2.days.ago.to_date, 2, 2], [3.days.ago.to_date, 0, 0]],
+                 Measurement.where(city_id: 2).pluck(:logdate, :peaktemp, :unitsales)
+  end
+
   private
     def capture_log_output
       output = StringIO.new

activerecord/test/models/measurement.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+class Measurement < ActiveRecord::Base
+end

activerecord/test/schema/postgresql_specific_schema.rb

@@ -108,4 +108,18 @@ _SQL
     t.uuid :uuid, primary_key: true, **uuid_default
     t.string :title
   end
+
+  if supports_partitioned_indexes?
+    create_table(:measurements, id: false, force: true, options: "PARTITION BY LIST (city_id)") do |t|
+      t.string :city_id, null: false
+      t.date :logdate, null: false
+      t.integer :peaktemp
+      t.integer :unitsales
+      t.index [:logdate, :city_id], unique: true
+    end
+    create_table(:measurements_toronto, id: false, force: true,
+                                        options: "PARTITION OF measurements FOR VALUES IN (1)")
+    create_table(:measurements_concepcion, id: false, force: true,
+                                           options: "PARTITION OF measurements FOR VALUES IN (2)")
+  end
 end