Add support for table comments

dummyapp/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    annotaterb (4.3.0)
+    annotaterb (4.3.1)
 
 GEM
   remote: https://rubygems.org/

lib/annotate_rb/model_annotator/annotation_builder.rb

@@ -70,11 +70,11 @@ def schema_header_text
         info << "#"
 
         if @options[:format_markdown]
-          info << "# Table name: `#{@model.table_name}`"
+          info << "# Table name: `#{table_name}`"
           info << "#"
           info << "# ### Columns"
         else
-          info << "# Table name: #{@model.table_name}"
+          info << "# Table name: #{table_name}"
         end
         info << "#\n" # We want the last line break
 
@@ -94,6 +94,20 @@ def schema_footer_text
 
         info.join("\n")
       end
+
+      private
+
+      def table_name
+        table_name = @model.table_name
+        display_table_comments = @options[:with_comment] && @options[:with_table_comments]
+
+        if display_table_comments && @model.has_table_comments?
+          table_comment = "(#{@model.table_comments.gsub(/\n/, "\\n")})"
+          table_name = "#{table_name}#{table_comment}"
+        end
+
+        table_name
+      end
     end
   end
 end

lib/annotate_rb/model_annotator/column_annotation/annotation_builder.rb

@@ -26,7 +26,8 @@ def build
           column_attributes = AttributesBuilder.new(@column, @options, is_primary_key, column_indices, column_defaults).build
           formatted_column_type = TypeBuilder.new(@column, @options, column_defaults).build
 
-          col_name = if @model.with_comments? && @column.comment
+          display_column_comments = @options[:with_comment] && @options[:with_column_comments]
+          col_name = if display_column_comments && @model.with_comments? && @column.comment
             "#{@column.name}(#{@column.comment.gsub(/\n/, '\\n')})"
           else
             @column.name

lib/annotate_rb/model_annotator/model_wrapper.rb

@@ -49,6 +49,15 @@ def table_exists?
         @klass.table_exists?
       end
 
+      def table_comments
+        @klass.connection.table_comment(@klass.table_name)
+      end
+
+      def has_table_comments?
+        @klass.connection.respond_to?(:table_comment) &&
+          @klass.connection.table_comment(@klass.table_name).present?
+      end
+
       def column_defaults
         @klass.column_defaults
       end
@@ -110,8 +119,7 @@ def retrieve_indexes_from_table
       end
 
       def with_comments?
-        @with_comments ||= @options[:with_comment] &&
-          raw_columns.first.respond_to?(:comment) &&
+        @with_comments ||= raw_columns.first.respond_to?(:comment) &&
           raw_columns.map(&:comment).any? { |comment| !comment.nil? }
       end
 

lib/annotate_rb/options.rb

@@ -51,7 +51,9 @@ def from(options = {}, state = {})
       sort: false, # ModelAnnotator
       timestamp: false, # RouteAnnotator
       trace: false, # ModelAnnotator, but is part of Core
-      with_comment: true # ModelAnnotator
+      with_comment: true, # ModelAnnotator
+      with_column_comments: nil, # ModelAnnotator
+      with_table_comments: nil # ModelAnnotator
     }.freeze
 
     OTHER_OPTIONS = {
@@ -113,7 +115,9 @@ def from(options = {}, state = {})
       :sort,
       :timestamp,
       :trace,
-      :with_comment
+      :with_comment,
+      :with_column_comments,
+      :with_table_comments
     ].freeze
 
     OTHER_OPTION_KEYS = [
@@ -187,6 +191,10 @@ def load_defaults
       @options[:wrapper_open] ||= @options[:wrapper]
       @options[:wrapper_close] ||= @options[:wrapper]
 
+      # Set column and table comments to default to :with_comment, if not set
+      @options[:with_column_comments] = @options[:with_comment] if @options[:with_column_comments].nil?
+      @options[:with_table_comments] = @options[:with_comment] if @options[:with_table_comments].nil?
+
       self
     end
 

lib/annotate_rb/parser.rb

@@ -217,6 +217,31 @@ def add_model_options_to_parser(option_parser)
         "include database comments in model annotations") do
         @options[:with_comment] = true
       end
+
+      option_parser.on("--without-comment",
+        "include database comments in model annotations") do
+        @options[:with_comment] = false
+      end
+
+      option_parser.on("--with-column-comments",
+        "include column comments in model annotations") do
+        @options[:with_column_comments] = true
+      end
+
+      option_parser.on("--without-column-comments",
+        "exclude column comments in model annotations") do
+        @options[:with_column_comments] = false
+      end
+
+      option_parser.on("--with-table-comments",
+        "include table comments in model annotations") do
+        @options[:with_table_comments] = true
+      end
+
+      option_parser.on("--without-table-comments",
+        "exclude table comments in model annotations") do
+        @options[:with_table_comments] = false
+      end
     end
 
     def add_route_options_to_parser(option_parser)

spec/lib/annotate_rb/model_annotator/annotation_builder_spec.rb

@@ -1319,7 +1319,7 @@
       end
 
       let :options do
-        AnnotateRb::Options.new({classified_sort: false, with_comment: true})
+        AnnotateRb::Options.new({classified_sort: false, with_comment: true, with_column_comments: true})
       end
 
       let :columns do
@@ -1358,7 +1358,7 @@
       end
 
       let :options do
-        AnnotateRb::Options.new({classified_sort: false, with_comment: true})
+        AnnotateRb::Options.new({classified_sort: false, with_comment: true, with_column_comments: true})
       end
 
       let :columns do
@@ -1405,7 +1405,7 @@
       end
 
       let :options do
-        AnnotateRb::Options.new({classified_sort: false, with_comment: true})
+        AnnotateRb::Options.new({classified_sort: false, with_comment: true, with_column_comments: true})
       end
 
       let :columns do
@@ -1851,7 +1851,7 @@
       end
 
       let :options do
-        {format_rdoc: true, with_comment: true}
+        {format_rdoc: true, with_comment: true, with_column_comments: true}
       end
 
       let :columns do
@@ -1899,7 +1899,7 @@
       end
 
       let :options do
-        {format_markdown: true, with_comment: true}
+        {format_markdown: true, with_comment: true, with_column_comments: true}
       end
 
       let :columns do
@@ -1968,4 +1968,172 @@
       end
     end
   end
+
+  describe "#schema_header_text" do
+    subject do
+      described_class.new(klass, options).schema_header_text
+    end
+
+    let(:table_exists) { true }
+    let(:table_comment) { "" }
+
+    let(:connection) do
+      indexes = []
+      foreign_keys = []
+
+      mock_connection_with_table_fields(
+        indexes,
+        foreign_keys,
+        table_exists,
+        table_comment
+      )
+    end
+
+    let :klass do
+      primary_key = nil
+      columns = []
+
+      mock_class_with_custom_connection(
+        :users,
+        primary_key,
+        columns,
+        connection
+      )
+    end
+
+    context "with no options set" do
+      let :options do
+        AnnotateRb::Options.new({})
+      end
+
+      let(:expected_header) do
+        <<~HEADER
+          #
+          # Table name: users
+          #
+        HEADER
+      end
+
+      it "returns the schema header" do
+        is_expected.to eq(expected_header)
+      end
+    end
+
+    context "with `with_comment: true`" do
+      context "with `with_table_comments: true` and table has comments" do
+        let :options do
+          AnnotateRb::Options.new({with_comment: true, with_table_comments: true})
+        end
+
+        let(:table_comment) { "table_comments" }
+
+        let(:expected_header) do
+          <<~HEADER
+            #
+            # Table name: users(table_comments)
+            #
+          HEADER
+        end
+
+        it "returns the header with the table comment" do
+          is_expected.to eq(expected_header)
+        end
+      end
+
+      context "with `with_table_comments: true` and table does not have comments" do
+        let :options do
+          AnnotateRb::Options.new({with_comment: true, with_table_comments: true})
+        end
+
+        let :klass do
+          primary_key = nil
+          columns = []
+          indexes = []
+          foreign_keys = []
+
+          mock_class(
+            :users,
+            primary_key,
+            columns,
+            indexes,
+            foreign_keys
+          )
+        end
+
+        let(:expected_header) do
+          <<~HEADER
+            #
+            # Table name: users
+            #
+          HEADER
+        end
+
+        it "returns the header without table comments" do
+          is_expected.to eq(expected_header)
+        end
+      end
+
+      context "with `with_table_comments: false` and table has comments" do
+        let :options do
+          AnnotateRb::Options.new({with_comment: true, with_table_comments: false})
+        end
+
+        let(:table_comment) { "table_comments" }
+
+        let(:expected_header) do
+          <<~HEADER
+            #
+            # Table name: users
+            #
+          HEADER
+        end
+
+        it "returns the header without the table comment" do
+          is_expected.to eq(expected_header)
+        end
+      end
+    end
+
+    context "with `with_comment: false`" do
+      context "with `with_table_comments: true` and table has comments" do
+        let :options do
+          AnnotateRb::Options.new({with_comment: false, with_table_comments: true})
+        end
+
+        let(:table_comment) { "table_comments" }
+
+        let(:expected_header) do
+          <<~HEADER
+            #
+            # Table name: users
+            #
+          HEADER
+        end
+
+        it "returns the header without the table comment" do
+          is_expected.to eq(expected_header)
+        end
+      end
+
+      context "with `with_table_comments: false` and table has comments" do
+        let :options do
+          AnnotateRb::Options.new({with_comment: false, with_table_comments: false})
+        end
+
+        let(:table_comment) { "table_comments" }
+
+        let(:expected_header) do
+          <<~HEADER
+            #
+            # Table name: users
+            #
+          HEADER
+        end
+
+        it "returns the header without the table comment" do
+          is_expected.to eq(expected_header)
+        end
+      end
+    end
+  end
 end