From 033e1232c93e073d0a9a827807a962f65d2dac11 Mon Sep 17 00:00:00 2001
From: Ryo Takahashi <suzunedev@gmail.com>
Date: Sat, 27 May 2023 18:32:15 +0900
Subject: [PATCH] Add annotation for table comment

---
 lib/annotate/annotate_models.rb           | 13 +++++--
 spec/lib/annotate/annotate_models_spec.rb | 44 ++++++++++++++++++++---
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/lib/annotate/annotate_models.rb b/lib/annotate/annotate_models.rb
index ff01a3d9..5e3d53d4 100644
--- a/lib/annotate/annotate_models.rb
+++ b/lib/annotate/annotate_models.rb
@@ -190,13 +190,16 @@ def get_schema_info(klass, header, options = {}) # rubocop:disable Metrics/Metho
     end
 
     def get_schema_header_text(klass, options = {})
+      table_comment = ""
+      table_comment = "(#{klass.connection.table_comment(klass.table_name)})" if with_table_comment?(klass, options)
+
       info = "#\n"
       if options[:format_markdown]
-        info << "# Table name: `#{klass.table_name}`\n"
+        info << "# Table name: `#{klass.table_name}#{table_comment}`\n"
         info << "#\n"
         info << "# ### Columns\n"
       else
-        info << "# Table name: #{klass.table_name}\n"
+        info << "# Table name: #{klass.table_name}#{table_comment}\n"
       end
       info << "#\n"
     end
@@ -798,6 +801,12 @@ def with_comments?(klass, options)
         klass.columns.any? { |col| !col.comment.nil? }
     end
 
+    def with_table_comment?(klass, options)
+      options[:with_comment] &&
+        klass.connection.respond_to?(:table_comment) &&
+        klass.connection.table_comment(klass.table_name).present?
+    end
+
     def max_schema_info_width(klass, options)
       cols = columns(klass, options)
 
diff --git a/spec/lib/annotate/annotate_models_spec.rb b/spec/lib/annotate/annotate_models_spec.rb
index 33b22017..2e2be78c 100644
--- a/spec/lib/annotate/annotate_models_spec.rb
+++ b/spec/lib/annotate/annotate_models_spec.rb
@@ -47,20 +47,21 @@ def mock_check_constraint(name, expression)
            expression: expression)
   end
 
-  def mock_connection(indexes = [], foreign_keys = [], check_constraints = [])
+  def mock_connection(table_comment, indexes = [], foreign_keys = [], check_constraints = [])
     double('Conn',
            indexes:      indexes,
            foreign_keys: foreign_keys,
            check_constraints: check_constraints,
            supports_foreign_keys?: true,
            supports_check_constraints?: true,
-           table_exists?: true)
+           table_exists?: true,
+           table_comment: table_comment)
   end
 
   # rubocop:disable Metrics/ParameterLists
-  def mock_class(table_name, primary_key, columns, indexes = [], foreign_keys = [], check_constraints = [])
+  def mock_class(table_name, primary_key, columns, indexes = [], foreign_keys = [], check_constraints = [], table_comments: {})
     options = {
-      connection:       mock_connection(indexes, foreign_keys, check_constraints),
+      connection:       mock_connection(table_comments[table_name], indexes, foreign_keys, check_constraints),
       table_exists?:    true,
       table_name:       table_name,
       primary_key:      primary_key,
@@ -1205,6 +1206,41 @@ def mock_column(name, type, options = {})
                   end
                 end
 
+                context 'when table have comment' do
+                  let :klass do
+                    mock_class(:users, primary_key, columns, indexes, foreign_keys, check_constraints, table_comments: { users: 'Users' })
+                  end
+
+                  let :columns do
+                    [
+                      mock_column(:id,         :integer, limit: 8,  comment: 'ID'),
+                      mock_column(:active,     :boolean, limit: 1,  comment: 'Active'),
+                      mock_column(:name,       :string,  limit: 50, comment: 'Name'),
+                      mock_column(:notes,      :text,    limit: 55, comment: 'Notes'),
+                      mock_column(:no_comment, :text,    limit: 20, comment: nil)
+                    ]
+                  end
+
+                  let :expected_result do
+                    <<~EOS
+                      # Schema Info
+                      #
+                      # Table name: users(Users)
+                      #
+                      #  id(ID)         :integer          not null, primary key
+                      #  active(Active) :boolean          not null
+                      #  name(Name)     :string(50)       not null
+                      #  notes(Notes)   :text(55)         not null
+                      #  no_comment     :text(20)         not null
+                      #
+                    EOS
+                  end
+
+                  it 'returns schema info in Markdown format' do
+                    is_expected.to eq expected_result
+                  end
+                end
+
                 context 'when columns have multibyte comments' do
                   let :columns do
                     [