Merge pull request rails#44803 from steve-abrams/activerecord-index-include-option

Add support for `:include` index option

Author: yahonda

activerecord/CHANGELOG.md

@@ -81,6 +81,24 @@
 
     *Leonardo Luarte*
 
+*   Add `:include` index option
+
+    Add support for including non-key columns in indexes for PostgreSQL
+    with the `INCLUDE` parameter.
+
+    ```ruby
+    add_index(:users, :email, include: [:id, :created_at])
+    ```
+
+    will result in:
+
+    ```sql
+    CREATE INDEX index_users_on_email USING btree (email) INCLUDE (id,
+    created_at)
+    ```
+
+    *Steve Abrams*
+
 *   `ActiveRecord::Relation`’s `#any?`, `#none?`, and `#one?` methods take an optional pattern
     argument, more closely matching their `Enumerable` equivalents.
 

activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb

@@ -6,7 +6,7 @@ module ConnectionAdapters # :nodoc:
     # this type are typically created and returned by methods in database
     # adapters. e.g. ActiveRecord::ConnectionAdapters::MySQL::SchemaStatements#indexes
     class IndexDefinition # :nodoc:
-      attr_reader :table, :name, :unique, :columns, :lengths, :orders, :opclasses, :where, :type, :using, :comment, :valid
+      attr_reader :table, :name, :unique, :columns, :lengths, :orders, :opclasses, :where, :type, :using, :include, :comment, :valid
 
       def initialize(
         table, name,
@@ -18,6 +18,7 @@ def initialize(
         where: nil,
         type: nil,
         using: nil,
+        include: nil,
         comment: nil,
         valid: true
       )
@@ -31,6 +32,7 @@ def initialize(
         @where = where
         @type = type
         @using = using
+        @include = include
         @comment = comment
         @valid = valid
       end
@@ -47,12 +49,13 @@ def column_options
         }
       end
 
-      def defined_for?(columns = nil, name: nil, unique: nil, valid: nil, **options)
+      def defined_for?(columns = nil, name: nil, unique: nil, valid: nil, include: nil, **options)
         columns = options[:column] if columns.blank?
         (columns.nil? || Array(self.columns) == Array(columns).map(&:to_s)) &&
           (name.nil? || self.name == name.to_s) &&
           (unique.nil? || self.unique == unique) &&
-          (valid.nil? || self.valid == valid)
+          (valid.nil? || self.valid == valid) &&
+          (include.nil? || Array(self.include) == Array(include))
       end
 
       private

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

@@ -828,6 +828,16 @@ def rename_column(table_name, column_name, new_column_name)
       #
       # Note: Partial indexes are only supported for PostgreSQL and SQLite.
       #
+      # ====== Creating an index that includes additional columns
+      #
+      #   add_index(:accounts, :branch_id,  include: :party_id)
+      #
+      # generates:
+      #
+      #   CREATE INDEX index_accounts_on_branch_id ON accounts USING btree(branch_id) INCLUDE (party_id)
+      #
+      # Note: only supported by PostgreSQL.
+      #
       # ====== Creating an index with a specific method
       #
       #   add_index(:developers, :name, using: 'btree')
@@ -1385,7 +1395,7 @@ def update_table_definition(table_name, base) # :nodoc:
       end
 
       def add_index_options(table_name, column_name, name: nil, if_not_exists: false, internal: false, **options) # :nodoc:
-        options.assert_valid_keys(:unique, :length, :order, :opclass, :where, :type, :using, :comment, :algorithm)
+        options.assert_valid_keys(:unique, :length, :order, :opclass, :where, :type, :using, :comment, :algorithm, :include)
 
         column_names = index_column_names(column_name)
 
@@ -1404,6 +1414,7 @@ def add_index_options(table_name, column_name, name: nil, if_not_exists: false,
           where: options[:where],
           type: options[:type],
           using: options[:using],
+          include: options[:include],
           comment: options[:comment]
         )
 

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -444,6 +444,11 @@ def supports_partial_index?
         false
       end
 
+      # Does this adapter support including non-key columns?
+      def supports_index_include?
+        false
+      end
+
       # Does this adapter support expression indices?
       def supports_expression_index?
         false

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

@@ -5,6 +5,9 @@ module ConnectionAdapters
     module PostgreSQL
       class SchemaCreation < SchemaCreation # :nodoc:
         private
+          delegate :quoted_include_columns_for_index, :supports_index_include?,
+          to: :@conn
+
           def visit_AlterTable(o)
             sql = super
             sql << o.constraint_validations.map { |fk| visit_ValidateConstraint fk }.join(" ")
@@ -27,6 +30,12 @@ def visit_CheckConstraintDefinition(o)
             super.dup.tap { |sql| sql << " NOT VALID" unless o.validate? }
           end
 
+          def visit_CreateIndexDefinition(o)
+            super.dup.tap do |sql|
+              sql << " INCLUDE (#{quoted_include_columns(o.index.include)})" if supports_index_include? && o.index.include
+            end
+          end
+
           def visit_ValidateConstraint(name)
             "VALIDATE CONSTRAINT #{quote_column_name(name)}"
           end
@@ -115,6 +124,10 @@ def add_column_options!(sql, options)
             super
           end
 
+          def quoted_include_columns(o)
+            String === o ? o : quoted_include_columns_for_index(o)
+          end
+
           # Returns any SQL string to go between CREATE and TABLE. May be nil.
           def table_modifier_in_create(o)
             # A table cannot be both TEMPORARY and UNLOGGED, since all TEMPORARY

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

@@ -108,10 +108,11 @@ def indexes(table_name) # :nodoc:
             oid = row[4]
             comment = row[5]
             valid = row[6]
-            using, expressions, where = inddef.scan(/ USING (\w+?) \((.+?)\)(?: NULLS(?: NOT)? DISTINCT)?(?: WHERE (.+))?\z/m).flatten
+            using, expressions, where, include = inddef.scan(/ USING (\w+?) \((.+?)\)(?: NULLS(?: NOT)? DISTINCT)?(?: WHERE (.+?))?(?: INCLUDE \((.+)\))?\z/m).flatten
 
             orders = {}
             opclasses = {}
+            include_columns = include ? include.split(",").map(&:strip) : []
 
             if indkey.include?(0)
               columns = expressions
@@ -123,6 +124,9 @@ def indexes(table_name) # :nodoc:
                 AND a.attnum IN (#{indkey.join(",")})
               SQL
 
+              # prevent INCLUDE columns from being matched
+              columns.reject! { |c| include_columns.include?(c) }
+
               # add info on sort order (only desc order is explicitly specified, asc is the default)
               # and non-default opclasses
               expressions.scan(/(?<column>\w+)"?\s?(?<opclass>\w+_ops(_\w+)?)?\s?(?<desc>DESC)?\s?(?<nulls>NULLS (?:FIRST|LAST))?/).each do |column, opclass, desc, nulls|
@@ -144,6 +148,7 @@ def indexes(table_name) # :nodoc:
               opclasses: opclasses,
               where: where,
               using: using.to_sym,
+              include: include_columns.map(&:to_sym).presence,
               comment: comment.presence,
               valid: valid
             )
@@ -759,6 +764,15 @@ def foreign_key_column_for(table_name) # :nodoc:
           super
         end
 
+        def quoted_include_columns_for_index(column_names) # :nodoc:
+          return quote_column_name(column_names) if column_names.is_a?(Symbol)
+
+          quoted_columns = column_names.each_with_object({}) do |name, result|
+            result[name.to_sym] = quote_column_name(name).dup
+          end
+          add_options_for_index_columns(quoted_columns).values.join(", ")
+        end
+
         def schema_creation  # :nodoc:
           PostgreSQL::SchemaCreation.new(self)
         end

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -200,6 +200,10 @@ def supports_partial_index?
         true
       end
 
+      def supports_index_include?
+        database_version >= 11_00_00 # >= 11.0
+      end
+
       def supports_expression_index?
         true
       end

activerecord/lib/active_record/schema_dumper.rb

@@ -230,6 +230,7 @@ def index_parts(index)
         index_parts << "opclass: #{format_index_parts(index.opclasses)}" if index.opclasses.present?
         index_parts << "where: #{index.where.inspect}" if index.where
         index_parts << "using: #{index.using.inspect}" if !@connection.default_index_type?(index)
+        index_parts << "include: #{index.include.inspect}" if index.include
         index_parts << "type: #{index.type.inspect}" if index.type
         index_parts << "comment: #{index.comment.inspect}" if index.comment
         index_parts

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

@@ -315,6 +315,22 @@ def test_partial_index
         end
       end
 
+      def test_include_index
+        with_example_table do
+          @connection.add_index "ex", %w{ id }, name: "include", include: :number
+          index = @connection.indexes("ex").find { |idx| idx.name == "include" }
+          assert_equal [:number], index.include
+        end
+      end
+
+      def test_include_multiple_columns_index
+        with_example_table do
+          @connection.add_index "ex", %w{ id }, name: "include", include: [:number, :data]
+          index = @connection.indexes("ex").find { |idx| idx.name == "include" }
+          assert_equal [:number, :data], index.include
+        end
+      end
+
       def test_expression_index
         with_example_table do
           expr = "mod(id, 10), abs(number)"

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

@@ -768,3 +768,16 @@ def test_create_join_table
     assert_not @connection.table_exists?("test_schema.comments_posts")
   end
 end
+
+class SchemaIndexIncludeColumnsTest < ActiveRecord::PostgreSQLTestCase
+  include SchemaDumpingHelper
+
+  def test_schema_dumps_index_included_columns
+    index_definition = dump_table_schema("companies").split(/\n/).grep(/t\.index.*company_include_index/).first.strip
+    if ActiveRecord::Base.connection.supports_index_include?
+      assert_equal 't.index ["firm_id", "type"], name: "company_include_index", include: [:name, :account_id]', index_definition
+    else
+      assert_equal 't.index ["firm_ids", "type"], name: "company_include_index"', index_definition
+    end
+  end
+end