Add annotation for table comment

ctran · May 27, 2023 · 033e123 · 033e123
1 parent 3a78787
commit 033e123
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 6 deletions.
diff --git 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
@@ -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
                     [