From 1b625fbc3d1a9361f4b73f1d56a1289315e2240f Mon Sep 17 00:00:00 2001
From: Olle Jonsson <olle.jonsson@gmail.com>
Date: Thu, 8 Feb 2024 16:32:26 +0100
Subject: [PATCH] feat: add COMMENT support to indexes

This reads out any COMMENT added to an INDEX object and adds its after
the last output on its "details" line.
---
 lib/annotate/annotate_models.rb           | 23 ++++++++++---
 spec/lib/annotate/annotate_models_spec.rb | 40 ++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/lib/annotate/annotate_models.rb b/lib/annotate/annotate_models.rb
index dc2901a3..49963f3e 100644
--- a/lib/annotate/annotate_models.rb
+++ b/lib/annotate/annotate_models.rb
@@ -36,6 +36,10 @@ module AnnotateModels
     using: {
       default: 'USING',
       markdown: '_using_'
+    },
+    comment: {
+      default: 'COMMENT',
+      markdown: '_comment_'
     }
   }.freeze
 
@@ -295,12 +299,22 @@ def index_using_info(index, format = :default)
       end
     end
 
+    def index_comment_info(index, format = :default)
+      value = index.try(:comment).try(:to_s)
+      if value.blank?
+        ''
+      else
+        " #{INDEX_CLAUSES[:comment][format]} #{value}"
+      end
+    end
+
     def final_index_string_in_markdown(index)
       details = sprintf(
-        "%s%s%s",
+        "%s%s%s%s",
         index_unique_info(index, :markdown),
         index_where_info(index, :markdown),
-        index_using_info(index, :markdown)
+        index_using_info(index, :markdown),
+        index_comment_info(index, :markdown)
       ).strip
       details = " (#{details})" unless details.blank?
 
@@ -314,12 +328,13 @@ def final_index_string_in_markdown(index)
 
     def final_index_string(index, max_size)
       sprintf(
-        "#  %-#{max_size}.#{max_size}s %s%s%s%s",
+        "#  %-#{max_size}.#{max_size}s %s%s%s%s%s",
         index.name,
         "(#{index_columns_info(index).join(',')})",
         index_unique_info(index),
         index_where_info(index),
-        index_using_info(index)
+        index_using_info(index),
+        index_comment_info(index)
       ).rstrip + "\n"
     end
 
diff --git a/spec/lib/annotate/annotate_models_spec.rb b/spec/lib/annotate/annotate_models_spec.rb
index 09647461..1b21bcb2 100644
--- a/spec/lib/annotate/annotate_models_spec.rb
+++ b/spec/lib/annotate/annotate_models_spec.rb
@@ -28,7 +28,8 @@ def mock_index(name, params = {})
            unique:        params[:unique] || false,
            orders:        params[:orders] || {},
            where:         params[:where],
-           using:         params[:using])
+           using:         params[:using],
+           comment:       params[:comment])
   end
 
   def mock_foreign_key(name, from_column, to_table, to_column = 'id', constraints = {})
@@ -539,6 +540,43 @@ def mock_column(name, type, options = {})
                   end
                 end
 
+                context 'when an index has a comment' do
+                  let :columns do
+                    [
+                      mock_column(:id, :integer),
+                      mock_column(:foreign_thing_id, :integer)
+                    ]
+                  end
+
+                  let :indexes do
+                    [
+                      mock_index('index_rails_02e851e3b9', columns: ['id']),
+                      mock_index('index_rails_02e851e3ba', columns: ['foreign_thing_id'], comment: 'This is a comment')
+                    ]
+                  end
+
+                  let :expected_result do
+                    <<~EOS
+                      # Schema Info
+                      #
+                      # Table name: users
+                      #
+                      #  id               :integer          not null, primary key
+                      #  foreign_thing_id :integer          not null
+                      #
+                      # Indexes
+                      #
+                      #  index_rails_02e851e3b9  (id)
+                      #  index_rails_02e851e3ba  (foreign_thing_id) COMMENT This is a comment
+                      #
+                    EOS
+                  end
+
+                  it 'returns schema info with index information' do
+                    is_expected.to eq expected_result
+                  end
+                end
+
                 context 'when one of indexes includes ordered index key' do
                   let :columns do
                     [