Add support for annotating check constraints

Author: lovro-bikic

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

@@ -178,6 +178,10 @@ def get_schema_info(klass, header, options = {})
         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,31 @@ def get_foreign_key_info(klass, options = {})
       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|
+        cc_info << if options[:format_markdown]
+                     sprintf("# * `%s`: `(%s)`\n", check_constraint.name, check_constraint.expression)
+                   else
+                     sprintf("#  %-#{max_size}.#{max_size}s (%s)\n", check_constraint.name, check_constraint.expression)
+                   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 Constants
       :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

@@ -173,6 +173,12 @@ def add_options_to_parser(option_parser) # rubocop:disable Metrics/MethodLength
         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,24 @@ def mock_foreign_key(name, from_column, to_table, to_column = 'id', constraints
            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 = [])
+  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,
@@ -217,7 +225,7 @@ def mock_column(name, type, options = {})
     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
@@ -228,6 +236,10 @@ def mock_column(name, type, options = {})
       []
     end
 
+    let :check_constraints do
+      []
+    end
+
     context 'when option is not present' do
       let :options do
         {}
@@ -752,6 +764,73 @@ def mock_column(name, type, options = {})
               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')
+                    ]
+                  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)
+                      #  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
                 [
@@ -1488,6 +1567,44 @@ def mock_column(name, type, options = {})
                 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')
+                    ]
+                  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)`
+                      #
+                    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' }