Add support for annotating check constraints (#868)

README.md

@@ -224,6 +224,7 @@ you can do so with a simple environment variable, instead of editing the
         -a, --active-admin               Annotate active_admin models
         -v, --version                    Show the current version of this gem
         -m, --show-migration             Include the migration version number in the annotation
+        -c, --show-check-constraints     List the table's check constraints in the annotation
         -k, --show-foreign-keys          List the table's foreign key constraints in the annotation
             --ck, --complete-foreign-keys
                                          Complete foreign key names in the annotation

lib/annotate/annotate_models.rb

@@ -131,7 +131,7 @@ module AnnotateModels
     # to create a comment block containing a line for
     # each column. The line contains the column name,
     # the type (and length), and any optional attributes
-    def get_schema_info(klass, header, options = {})
+    def get_schema_info(klass, header, options = {}) # rubocop:disable Metrics/MethodLength
       info = "# #{header}\n"
       info << get_schema_header_text(klass, options)
 
@@ -178,6 +178,10 @@ module AnnotateModels
         info << get_foreign_key_info(klass, options)
       end
 
+      if options[:show_check_constraints] && klass.table_exists?
+        info << get_check_constraint_info(klass, options)
+      end
+
       info << get_schema_footer_text(klass, options)
     end
 
@@ -352,6 +356,35 @@ module AnnotateModels
       fk_info
     end
 
+    def get_check_constraint_info(klass, options = {})
+      cc_info = if options[:format_markdown]
+                  "#\n# ### Check Constraints\n#\n"
+                else
+                  "#\n# Check Constraints\n#\n"
+                end
+
+      return '' unless klass.connection.respond_to?(:supports_check_constraints?) &&
+        klass.connection.supports_check_constraints? && klass.connection.respond_to?(:check_constraints)
+
+      check_constraints = klass.connection.check_constraints(klass.table_name)
+      return '' if check_constraints.empty?
+
+      max_size = check_constraints.map { |check_constraint| check_constraint.name.size }.max + 1
+      check_constraints.sort_by(&:name).each do |check_constraint|
+        expression = check_constraint.expression ? "(#{check_constraint.expression.squish})" : nil
+
+        cc_info << if options[:format_markdown]
+                     cc_info_markdown = sprintf("# * `%s`", check_constraint.name)
+                     cc_info_markdown << sprintf(": `%s`", expression) if expression
+                     cc_info_markdown << "\n"
+                   else
+                     sprintf("#  %-#{max_size}.#{max_size}s %s", check_constraint.name, expression).rstrip + "\n"
+                   end
+      end
+
+      cc_info
+    end
+
     # Add a schema block to a file. If the file already contains
     # a schema info block (a comment starting with "== Schema Information"),
     # check if it matches the block that is already there. If so, leave it be.

lib/annotate/constants.rb

@@ -18,7 +18,8 @@ module Annotate
       :trace, :timestamp, :exclude_serializers, :classified_sort,
       :show_foreign_keys, :show_complete_foreign_keys,
       :exclude_scaffolds, :exclude_controllers, :exclude_helpers,
-      :exclude_sti_subclasses, :ignore_unknown_models, :with_comment
+      :exclude_sti_subclasses, :ignore_unknown_models, :with_comment,
+      :show_check_constraints
     ].freeze
 
     OTHER_OPTIONS = [

lib/annotate/parser.rb

@@ -48,7 +48,7 @@ module Annotate
       end
     end
 
-    def add_options_to_parser(option_parser) # rubocop:disable Metrics/MethodLength
+    def add_options_to_parser(option_parser) # rubocop:disable Metrics/MethodLength, Metrics/AbcSize
       has_set_position = {}
 
       option_parser.banner = 'Usage: annotate [options] [model_file]*'
@@ -173,6 +173,12 @@ module Annotate
         env['include_version'] = 'yes'
       end
 
+      option_parser.on('-c',
+                       '--show-check-constraints',
+                       "List the table's check constraints in the annotation") do
+        env['show_check_constraints'] = 'yes'
+      end
+
       option_parser.on('-k',
                        '--show-foreign-keys',
                        "List the table's foreign key constraints in the annotation") do

lib/generators/annotate/templates/auto_annotate_models.rake

@@ -17,6 +17,7 @@ if Rails.env.development?
       'position_in_fixture'         => 'before',
       'position_in_factory'         => 'before',
       'position_in_serializer'      => 'before',
+      'show_check_constraints'      => 'false',
       'show_foreign_keys'           => 'true',
       'show_complete_foreign_keys'  => 'false',
       'show_indexes'                => 'true',

lib/tasks/annotate_models.rake

@@ -18,6 +18,7 @@ task annotate_models: :environment do
   options[:position_in_factory] = Annotate::Helpers.fallback(ENV['position_in_factory'], ENV['position'])
   options[:position_in_test] = Annotate::Helpers.fallback(ENV['position_in_test'], ENV['position'])
   options[:position_in_serializer] = Annotate::Helpers.fallback(ENV['position_in_serializer'], ENV['position'])
+  options[:show_check_constraints] = Annotate::Helpers.true?(ENV['show_check_constraints'])
   options[:show_foreign_keys] = Annotate::Helpers.true?(ENV['show_foreign_keys'])
   options[:show_complete_foreign_keys] = Annotate::Helpers.true?(ENV['show_complete_foreign_keys'])
   options[:show_indexes] = Annotate::Helpers.true?(ENV['show_indexes'])

spec/lib/annotate/annotate_models_spec.rb

@@ -41,16 +41,25 @@ describe AnnotateModels do
            on_update:    constraints[:on_update])
   end
 
-  def mock_connection(indexes = [], foreign_keys = [])
+  def mock_check_constraint(name, expression)
+    double('CheckConstraintDefinition',
+           name:       name,
+           expression: expression)
+  end
+
+  def mock_connection(indexes = [], foreign_keys = [], check_constraints = [])
     double('Conn',
            indexes:      indexes,
            foreign_keys: foreign_keys,
-           supports_foreign_keys?: true)
+           check_constraints: check_constraints,
+           supports_foreign_keys?: true,
+           supports_check_constraints?: true)
   end
 
-  def mock_class(table_name, primary_key, columns, indexes = [], foreign_keys = [])
+  # rubocop:disable Metrics/ParameterLists
+  def mock_class(table_name, primary_key, columns, indexes = [], foreign_keys = [], check_constraints = [])
     options = {
-      connection:       mock_connection(indexes, foreign_keys),
+      connection:       mock_connection(indexes, foreign_keys, check_constraints),
       table_exists?:    true,
       table_name:       table_name,
       primary_key:      primary_key,
@@ -62,6 +71,7 @@ describe AnnotateModels do
 
     double('An ActiveRecord class', options)
   end
+  # rubocop:enable Metrics/ParameterLists
 
   def mock_column(name, type, options = {})
     default_options = {
@@ -221,7 +231,7 @@ describe AnnotateModels do
     end
 
     let :klass do
-      mock_class(:users, primary_key, columns, indexes, foreign_keys)
+      mock_class(:users, primary_key, columns, indexes, foreign_keys, check_constraints)
     end
 
     let :indexes do
@@ -232,6 +242,10 @@ describe AnnotateModels do
       []
     end
 
+    let :check_constraints do
+      []
+    end
+
     context 'when option is not present' do
       let :options do
         {}
@@ -391,7 +405,7 @@ describe AnnotateModels do
               end
             end
 
-            context 'with Globalize gem' do
+            context 'with Globalize gem' do # rubocop:disable RSpec/MultipleMemoizedHelpers
               let :translation_klass do
                 double('Folder::Post::Translation',
                        to_s: 'Folder::Post::Translation',
@@ -756,6 +770,82 @@ describe AnnotateModels do
               end
             end
 
+            context 'when check constraints exist' do
+              let :columns do
+                [
+                  mock_column(:id, :integer),
+                  mock_column(:age, :integer)
+                ]
+              end
+
+              context 'when option "show_check_constraints" is true' do
+                let :options do
+                  { show_check_constraints: true }
+                end
+
+                context 'when check constraints are defined' do
+                  let :check_constraints do
+                    [
+                      mock_check_constraint('alive', 'age < 150'),
+                      mock_check_constraint('must_be_adult', 'age >= 18'),
+                      mock_check_constraint('missing_expression', nil),
+                      mock_check_constraint('multiline_test', <<~SQL)
+                        CASE
+                          WHEN (age >= 18) THEN (age <= 21)
+                          ELSE true
+                        END
+                      SQL
+                    ]
+                  end
+
+                  let :expected_result do
+                    <<~EOS
+                      # Schema Info
+                      #
+                      # Table name: users
+                      #
+                      #  id  :integer          not null, primary key
+                      #  age :integer          not null
+                      #
+                      # Check Constraints
+                      #
+                      #  alive               (age < 150)
+                      #  missing_expression
+                      #  multiline_test      (CASE WHEN (age >= 18) THEN (age <= 21) ELSE true END)
+                      #  must_be_adult       (age >= 18)
+                      #
+                    EOS
+                  end
+
+                  it 'returns schema info with check constraint information' do
+                    is_expected.to eq expected_result
+                  end
+                end
+
+                context 'when check constraint is not defined' do
+                  let :check_constraints do
+                    []
+                  end
+
+                  let :expected_result do
+                    <<~EOS
+                      # Schema Info
+                      #
+                      # Table name: users
+                      #
+                      #  id  :integer          not null, primary key
+                      #  age :integer          not null
+                      #
+                    EOS
+                  end
+
+                  it 'returns schema info without check constraint information' do
+                    is_expected.to eq expected_result
+                  end
+                end
+              end
+            end
+
             context 'when foreign keys exist' do
               let :columns do
                 [
@@ -1492,6 +1582,53 @@ describe AnnotateModels do
                 end
               end
 
+              context 'when option "show_check_constraints" is true' do
+                let :options do
+                  { format_markdown: true, show_check_constraints: true }
+                end
+
+                context 'when check constraints are defined' do
+                  let :check_constraints do
+                    [
+                      mock_check_constraint('min_name_length', 'LENGTH(name) > 2'),
+                      mock_check_constraint('missing_expression', nil),
+                      mock_check_constraint('multiline_test', <<~SQL)
+                        CASE
+                          WHEN (age >= 18) THEN (age <= 21)
+                          ELSE true
+                        END
+                      SQL
+                    ]
+                  end
+
+                  let :expected_result do
+                    <<~EOS
+                      # == Schema Information
+                      #
+                      # Table name: `users`
+                      #
+                      # ### Columns
+                      #
+                      # Name        | Type               | Attributes
+                      # ----------- | ------------------ | ---------------------------
+                      # **`id`**    | `integer`          | `not null, primary key`
+                      # **`name`**  | `string(50)`       | `not null`
+                      #
+                      # ### Check Constraints
+                      #
+                      # * `min_name_length`: `(LENGTH(name) > 2)`
+                      # * `missing_expression`
+                      # * `multiline_test`: `(CASE WHEN (age >= 18) THEN (age <= 21) ELSE true END)`
+                      #
+                    EOS
+                  end
+
+                  it 'returns schema info with check constraint information in Markdown format' do
+                    is_expected.to eq expected_result
+                  end
+                end
+              end
+
               context 'when option "show_foreign_keys" is true' do
                 let :options do
                   { format_markdown: true, show_foreign_keys: true }

spec/lib/annotate/parser_spec.rb

@@ -260,6 +260,17 @@ module Annotate # rubocop:disable Metrics/ModuleLength
       end
     end
 
+    %w[-c --show-check-constraints].each do |option|
+      describe option do
+        let(:env_key) { 'show_check_constraints' }
+        let(:set_value) { 'yes' }
+        it 'sets the ENV variable' do
+          expect(ENV).to receive(:[]=).with(env_key, set_value)
+          Parser.parse([option])
+        end
+      end
+    end
+
     %w[-k --show-foreign-keys].each do |option|
       describe option do
         let(:env_key) { 'show_foreign_keys' }