Annotate model check constraints

lib/annotate_rb/model_annotator.rb

@@ -26,5 +26,6 @@ module ModelAnnotator
     autoload :AnnotatedFile, "annotate_rb/model_annotator/annotated_file"
     autoload :FileParser, "annotate_rb/model_annotator/file_parser"
     autoload :ZeitwerkClassGetter, "annotate_rb/model_annotator/zeitwerk_class_getter"
+    autoload :CheckConstraintAnnotation, "annotate_rb/model_annotator/check_constraint_annotation"
   end
 end

lib/annotate_rb/model_annotator/annotation_builder.rb

@@ -44,6 +44,10 @@ def build
           @info += ForeignKeyAnnotation::AnnotationBuilder.new(@model, @options).build
         end
 
+        if @options[:show_check_constraints] && @model.table_exists?
+          @info += CheckConstraintAnnotation::AnnotationBuilder.new(@model, @options).build
+        end
+
         @info += schema_footer_text
 
         @info

lib/annotate_rb/model_annotator/check_constraint_annotation.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    module CheckConstraintAnnotation
+      autoload :AnnotationBuilder, "annotate_rb/model_annotator/check_constraint_annotation/annotation_builder"
+    end
+  end
+end

lib/annotate_rb/model_annotator/check_constraint_annotation/annotation_builder.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    module CheckConstraintAnnotation
+      class AnnotationBuilder
+        def initialize(model, options)
+          @model = model
+          @options = options
+        end
+
+        def build
+          constraint_info = if @options[:format_markdown]
+            "#\n# ### Check Constraints\n#\n"
+          else
+            "#\n# Check Constraints\n#\n"
+          end
+
+          return "" unless @model.connection.respond_to?(:supports_check_constraints?) &&
+            @model.connection.supports_check_constraints? && @model.connection.respond_to?(:check_constraints)
+
+          check_constraints = @model.connection.check_constraints(@model.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
+
+            constraint_info += if @options[:format_markdown]
+              cc_info_in_markdown(check_constraint.name, expression)
+            else
+              cc_info_string(check_constraint.name, expression, max_size)
+            end
+          end
+
+          constraint_info
+        end
+
+        private
+
+        def cc_info_in_markdown(name, expression)
+          cc_info_markdown = sprintf("# * `%s`", name)
+          cc_info_markdown += sprintf(": `%s`", expression) if expression
+          cc_info_markdown += "\n"
+
+          cc_info_markdown
+        end
+
+        def cc_info_string(name, expression, max_size)
+          # standard:disable Lint/FormatParameterMismatch
+          sprintf("#  %-#{max_size}.#{max_size}s %s", name, expression).rstrip + "\n"
+          # standard:enable Lint/FormatParameterMismatch
+        end
+      end
+    end
+  end
+end

lib/annotate_rb/options.rb

@@ -45,6 +45,7 @@ def from(options = {}, state = {})
       ignore_unknown_models: false, # ModelAnnotator
       include_version: false, # ModelAnnotator
       show_complete_foreign_keys: false, # ModelAnnotator
+      show_check_constraints: false, # ModelAnnotator
       show_foreign_keys: true, # ModelAnnotator
       show_indexes: true, # ModelAnnotator
       simple_indexes: false, # ModelAnnotator
@@ -109,6 +110,7 @@ def from(options = {}, state = {})
       :ignore_model_sub_dir,
       :ignore_unknown_models,
       :include_version,
+      :show_check_constraints,
       :show_complete_foreign_keys,
       :show_foreign_keys,
       :show_indexes,

lib/annotate_rb/parser.rb

@@ -198,6 +198,12 @@ def add_model_options_to_parser(option_parser)
         @options[:simple_indexes] = true
       end
 
+      option_parser.on("-c",
+        "--show-check-constraints",
+        "List the table's check constraints in the annotation") do
+        @options[:show_check_constraints] = true
+      end
+
       option_parser.on("--hide-limit-column-types VALUES",
         "don't show limit for given column types, separated by commas (i.e., `integer,boolean,text`)") do |values|
         @options[:hide_limit_column_types] = values.to_s

spec/lib/annotate_rb/model_annotator/annotation_builder_spec.rb

@@ -1967,6 +1967,77 @@
         end
       end
     end
+
+    context 'when "show_check_constraints" is true' do
+      let :klass do
+        mock_class_with_custom_connection(:users, primary_key, columns, custom_connection)
+      end
+
+      let :custom_connection do
+        check_constraints = [
+          mock_check_constraint("must_be_us_adult", "age >= 18")
+        ]
+
+        mock_connection([], [], check_constraints)
+      end
+
+      let :primary_key do
+        :id
+      end
+
+      let :options do
+        {show_check_constraints: true}
+      end
+
+      let :columns do
+        [
+          mock_column("id", :integer),
+          mock_column("age", :integer)
+        ]
+      end
+
+      let :expected_result do
+        <<~EOS
+          # == Schema Information
+          #
+          # Table name: users
+          #
+          #  id  :integer          not null, primary key
+          #  age :integer          not null
+          #
+          # Check Constraints
+          #
+          #  must_be_us_adult  (age >= 18)
+          #
+        EOS
+      end
+
+      it "returns schema info with check constraints" do
+        is_expected.to eq expected_result
+      end
+
+      context "when option is set to false" do
+        let(:options) do
+          {show_check_constraints: false}
+        end
+
+        let :expected_result do
+          <<~EOS
+            # == Schema Information
+            #
+            # Table name: users
+            #
+            #  id  :integer          not null, primary key
+            #  age :integer          not null
+            #
+          EOS
+        end
+
+        it "returns schema info without check constraints" do
+          is_expected.to eq expected_result
+        end
+      end
+    end
   end
 
   describe "#schema_header_text" do

spec/lib/annotate_rb/model_annotator/check_constraint_annotation/annotation_builder_spec.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+RSpec.describe AnnotateRb::ModelAnnotator::CheckConstraintAnnotation::AnnotationBuilder do
+  include AnnotateTestHelpers
+
+  describe "#build" do
+    subject { described_class.new(model, options).build }
+
+    let(:model) do
+      instance_double(
+        AnnotateRb::ModelAnnotator::ModelWrapper,
+        connection: connection,
+        table_name: "Foo"
+      )
+    end
+    let(:connection) do
+      mock_connection([], [], check_constraints)
+    end
+    let(:options) { AnnotateRb::Options.new }
+    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
+      <<~RESULT
+        #
+        # Check Constraints
+        #
+        #  alive               (age < 150)
+        #  missing_expression
+        #  multiline_test      (CASE WHEN (age >= 18) THEN (age <= 21) ELSE true END)
+        #  must_be_adult       (age >= 18)
+      RESULT
+    end
+
+    it "annotates the check constraints" do
+      is_expected.to eq(expected_result)
+    end
+
+    context "when model connection does not support check constraints" do
+      let(:connection) do
+        conn_options = {supports_check_constraints?: false}
+
+        mock_connection([], [], [], conn_options)
+      end
+
+      it { is_expected.to be_empty }
+    end
+
+    context "when check constraints is empty" do
+      let(:connection) do
+        conn_options = {supports_check_constraints?: true}
+
+        mock_connection([], [], [], conn_options)
+      end
+
+      it { is_expected.to be_empty }
+    end
+
+    context "when there are check constraints using markdown" do
+      let(:options) { AnnotateRb::Options.new({format_markdown: true}) }
+      let(:expected_result) do
+        <<~RESULT
+          #
+          # ### Check Constraints
+          #
+          # * `alive`: `(age < 150)`
+          # * `missing_expression`
+          # * `multiline_test`: `(CASE WHEN (age >= 18) THEN (age <= 21) ELSE true END)`
+          # * `must_be_adult`: `(age >= 18)`
+        RESULT
+      end
+
+      it { is_expected.to eq(expected_result) }
+    end
+
+    context "when it is just the header using markdown" do
+      let(:options) { AnnotateRb::Options.new({format_markdown: true}) }
+      let(:connection) do
+        mock_connection([], [], [])
+      end
+
+      it { is_expected.to be_empty }
+    end
+  end
+end

spec/lib/annotate_rb/parser_spec.rb

@@ -295,6 +295,16 @@ module AnnotateRb # rubocop:disable Metrics/ModuleLength
       end
     end
 
+    %w[-c --show-check-constraints].each do |option|
+      describe option do
+        let(:args) { [option] }
+
+        it "sets show_check_constraints to true" do
+          expect(result).to include(show_check_constraints: true)
+        end
+      end
+    end
+
     describe "--model-dir" do
       let(:option) { "--model-dir" }
       let(:set_value) { "some_dir/" }

spec/support/annotate_test_helpers.rb

@@ -36,11 +36,16 @@ 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 = [])
-    double("Conn",
+  def mock_connection(indexes = [], foreign_keys = [], check_constraints = [], options = {})
+    double_options = {
       indexes: indexes,
+      check_constraints: check_constraints,
       foreign_keys: foreign_keys,
-      supports_foreign_keys?: true)
+      supports_foreign_keys?: true,
+      supports_check_constraints?: true
+    }.merge(options)
+
+    double("Conn", double_options)
   end
 
   def mock_connection_with_table_fields(indexes, foreign_keys, table_exists, table_comment)
@@ -97,4 +102,10 @@ def mock_column(name, type, options = {})
 
     double("Column", stubs)
   end
+
+  def mock_check_constraint(name, expression)
+    double("CheckConstraintDefinition",
+      name: name,
+      expression: expression)
+  end
 end