Support opt-in display of PostgreSQL enum types in annotations

lib/annotate_rb/model_annotator.rb

@@ -27,6 +27,7 @@ module ModelAnnotator
     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"
+    autoload :EnumAnnotation, "annotate_rb/model_annotator/enum_annotation"
     autoload :FileToParserMapper, "annotate_rb/model_annotator/file_to_parser_mapper"
     autoload :Components, "annotate_rb/model_annotator/components"
     autoload :Annotation, "annotate_rb/model_annotator/annotation"

lib/annotate_rb/model_annotator/annotation/annotation_builder.rb

@@ -27,6 +27,7 @@ def body
               IndexAnnotation::AnnotationBuilder.new(@model, @options).build,
               ForeignKeyAnnotation::AnnotationBuilder.new(@model, @options).build,
               CheckConstraintAnnotation::AnnotationBuilder.new(@model, @options).build,
+              EnumAnnotation::AnnotationBuilder.new(@model, @options).build,
               SchemaFooter.new
             ]
           end

lib/annotate_rb/model_annotator/column_annotation/attributes_builder.rb

@@ -84,6 +84,13 @@ def build
             end
           end
 
+          if column_type == "enum" && @options[:show_enums]
+            enum_type_name = @column.sql_type
+            if enum_type_name.present? && enum_type_name != "enum"
+              attrs << "enum_type: #{enum_type_name}"
+            end
+          end
+
           # Check if the column is a virtual column and print the function
           if @options[:show_virtual_columns] && @column.virtual?
             # Any whitespace in the function gets reduced to a single space

lib/annotate_rb/model_annotator/column_annotation/column_wrapper.rb

@@ -26,6 +26,10 @@ def type
           @column.type
         end
 
+        def sql_type
+          @column.sql_type.to_s
+        end
+
         def column_type_string
           if (@column.respond_to?(:bigint?) && @column.bigint?) || /\Abigint\b/ =~ @column.sql_type
             "bigint"

lib/annotate_rb/model_annotator/enum_annotation.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    module EnumAnnotation
+      autoload :AnnotationBuilder, "annotate_rb/model_annotator/enum_annotation/annotation_builder"
+      autoload :Annotation, "annotate_rb/model_annotator/enum_annotation/annotation"
+      autoload :EnumComponent, "annotate_rb/model_annotator/enum_annotation/enum_component"
+    end
+  end
+end

lib/annotate_rb/model_annotator/enum_annotation/annotation.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    module EnumAnnotation
+      class Annotation
+        HEADER_TEXT = "Enums"
+
+        def initialize(enums)
+          @enums = enums
+        end
+
+        def body
+          [
+            Components::BlankCommentLine.new,
+            Components::Header.new(HEADER_TEXT),
+            Components::BlankCommentLine.new,
+            *@enums
+          ]
+        end
+
+        def to_markdown
+          body.map(&:to_markdown).join("\n")
+        end
+
+        def to_rdoc
+          body.map(&:to_rdoc).join("\n")
+        end
+
+        def to_yard
+          body.map(&:to_yard).join("\n")
+        end
+
+        def to_default
+          body.map(&:to_default).join("\n")
+        end
+      end
+    end
+  end
+end

lib/annotate_rb/model_annotator/enum_annotation/annotation_builder.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    module EnumAnnotation
+      class AnnotationBuilder
+        def initialize(model, options)
+          @model = model
+          @options = options
+        end
+
+        def build
+          return Components::NilComponent.new unless @options[:show_enums]
+
+          enum_types = @model.enum_types
+          return Components::NilComponent.new if enum_types.empty?
+
+          # Filter to only enum types used by this table's columns
+          table_enum_types = @model.columns.select { |col| col.type == :enum }
+            .map { |col| col.sql_type.to_s }
+            .uniq
+
+          relevant_enums = enum_types
+            .filter_map { |name, values| [name.to_s, values] if table_enum_types.include?(name.to_s) }
+          return Components::NilComponent.new if relevant_enums.empty?
+
+          max_size = relevant_enums.map { |name, _values| name.size }.max + 1
+
+          components = relevant_enums.sort_by { |name, _values| name }.map do |name, values|
+            EnumComponent.new(name, values, max_size)
+          end
+
+          Annotation.new(components)
+        end
+      end
+    end
+  end
+end

lib/annotate_rb/model_annotator/enum_annotation/enum_component.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module AnnotateRb
+  module ModelAnnotator
+    module EnumAnnotation
+      class EnumComponent < Components::Base
+        attr_reader :name, :values, :max_size
+
+        def initialize(name, values, max_size)
+          @name = name
+          @values = values
+          @max_size = max_size
+        end
+
+        def to_default
+          # standard:disable Lint/FormatParameterMismatch
+          sprintf("#  %-#{max_size}.#{max_size}s %s", name, values.join(", ")).rstrip
+          # standard:enable Lint/FormatParameterMismatch
+        end
+
+        def to_markdown
+          sprintf("# * `%s`: `%s`", name, values.join(", "))
+        end
+      end
+    end
+  end
+end

lib/annotate_rb/model_annotator/model_wrapper.rb

@@ -219,6 +219,23 @@ def ignored_translation_table_columns
         ]
       end
 
+      def enum_types
+        @enum_types ||=
+          if connection.respond_to?(:enum_types)
+            begin
+              # enum values may be a String or an Array depending on the Rails version.
+              # See: https://github.com/rails/rails/pull/54141
+              connection.enum_types.map do |name, values|
+                [name, values.is_a?(Array) ? values : values.to_s.split(",")]
+              end
+            rescue ActiveRecord::StatementInvalid
+              []
+            end
+          else
+            []
+          end
+      end
+
       def migration_version
         return 0 unless @options[:include_version]
 

lib/annotate_rb/options.rb

@@ -47,6 +47,7 @@ def from(options = {}, state = {})
       include_version: false, # ModelAnnotator
       show_complete_foreign_keys: false, # ModelAnnotator
       show_check_constraints: false, # ModelAnnotator
+      show_enums: false, # ModelAnnotator
       show_foreign_keys: true, # ModelAnnotator
       show_indexes: true, # ModelAnnotator
       show_indexes_comments: false, # ModelAnnotator
@@ -121,6 +122,7 @@ def from(options = {}, state = {})
       :ignore_unknown_models,
       :include_version,
       :show_check_constraints,
+      :show_enums,
       :show_complete_foreign_keys,
       :show_foreign_keys,
       :show_indexes,

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

@@ -97,6 +97,108 @@
       end
     end
 
+    context "integration test with enums" do
+      let(:klass) do
+        primary_key = :id
+        columns = [
+          mock_column("id", :integer),
+          mock_column("status", :enum, sql_type: "order_status"),
+          mock_column("name", :string)
+        ]
+        indexes = [
+          mock_index("index_rails_02e851e3b7", columns: ["id"])
+        ]
+        foreign_keys = []
+        check_constraints = []
+        enum_types = [
+          ["order_status", ["pending", "shipped", "delivered"]],
+          ["unused_enum", ["a", "b"]]
+        ]
+
+        custom_connection = mock_connection(indexes, foreign_keys, check_constraints, enum_types: enum_types)
+        mock_class_with_custom_connection(:orders, primary_key, columns, custom_connection)
+      end
+
+      let(:options) do
+        AnnotateRb::Options.new({
+          show_indexes: true,
+          show_enums: true
+        })
+      end
+
+      let(:expected_result) do
+        <<~EOS
+          # == Schema Information
+          #
+          # Table name: orders
+          #
+          #  id     :integer          not null, primary key
+          #  status :enum             not null, enum_type: order_status
+          #  name   :string           not null
+          #
+          # Indexes
+          #
+          #  index_rails_02e851e3b7  (id)
+          #
+          # Enums
+          #
+          #  order_status  pending, shipped, delivered
+          #
+        EOS
+      end
+
+      it "includes enum annotation in the output" do
+        is_expected.to eq expected_result
+      end
+    end
+
+    context "integration test with enums in markdown format" do
+      let(:klass) do
+        primary_key = :id
+        columns = [
+          mock_column("id", :integer),
+          mock_column("status", :enum, sql_type: "order_status")
+        ]
+        enum_types = [
+          ["order_status", ["pending", "shipped", "delivered"]]
+        ]
+
+        custom_connection = mock_connection([], [], [], enum_types: enum_types)
+        mock_class_with_custom_connection(:orders, primary_key, columns, custom_connection)
+      end
+
+      let(:options) do
+        AnnotateRb::Options.new({
+          format_markdown: true,
+          show_enums: true
+        })
+      end
+
+      let(:expected_result) do
+        <<~EOS
+          # ## Schema Information
+          #
+          # Table name: `orders`
+          #
+          # ### Columns
+          #
+          # Name          | Type               | Attributes
+          # ------------- | ------------------ | ---------------------------
+          # **`id`**      | `integer`          | `not null, primary key`
+          # **`status`**  | `enum`             | `not null, enum_type: order_status`
+          #
+          # ### Enums
+          #
+          # * `order_status`: `pending, shipped, delivered`
+          #
+        EOS
+      end
+
+      it "includes enum annotation in markdown format" do
+        is_expected.to eq expected_result
+      end
+    end
+
     context "with primary key and using globalize gem" do
       let :options do
         AnnotateRb::Options.new({})

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

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+RSpec.describe AnnotateRb::ModelAnnotator::EnumAnnotation::AnnotationBuilder do
+  include AnnotateTestHelpers
+
+  describe "#build" do
+    subject { described_class.new(model, options).build }
+    let(:default_format) { subject.to_default }
+    let(:markdown_format) { subject.to_markdown }
+
+    let(:model) do
+      instance_double(
+        AnnotateRb::ModelAnnotator::ModelWrapper,
+        enum_types: enum_types,
+        table_name: "Foo",
+        columns: columns
+      )
+    end
+    let(:columns) do
+      [
+        mock_column("billing_method", :enum, sql_type: "billing_method"),
+        mock_column("status", :enum, sql_type: "order_status"),
+        mock_column("name", :string, sql_type: "character varying")
+      ]
+    end
+    let(:enum_types) do
+      [
+        ["billing_method", ["agency_bill", "direct_bill_to_insured"]],
+        ["order_status", ["pending", "shipped", "delivered"]],
+        ["unused_enum", ["a", "b"]]
+      ]
+    end
+    let(:options) { AnnotateRb::Options.new({show_enums: true}) }
+
+    context "when show_enums option is false" do
+      let(:options) { AnnotateRb::Options.new({show_enums: false}) }
+      it { is_expected.to be_a(AnnotateRb::ModelAnnotator::Components::NilComponent) }
+    end
+
+    context "when enum_types is empty" do
+      let(:enum_types) { [] }
+      it { is_expected.to be_a(AnnotateRb::ModelAnnotator::Components::NilComponent) }
+    end
+
+    context "when table has no enum columns" do
+      let(:columns) do
+        [
+          mock_column("name", :string, sql_type: "character varying"),
+          mock_column("age", :integer, sql_type: "integer")
+        ]
+      end
+
+      it { is_expected.to be_a(AnnotateRb::ModelAnnotator::Components::NilComponent) }
+    end
+
+    context "using default format" do
+      let(:expected_result) do
+        <<~RESULT.strip
+          #
+          # Enums
+          #
+          #  billing_method  agency_bill, direct_bill_to_insured
+          #  order_status    pending, shipped, delivered
+        RESULT
+      end
+
+      it "annotates the enum types" do
+        expect(default_format).to eq(expected_result)
+      end
+    end
+
+    context "using markdown format" do
+      let(:expected_result) do
+        <<~RESULT.strip
+          #
+          # ### Enums
+          #
+          # * `billing_method`: `agency_bill, direct_bill_to_insured`
+          # * `order_status`: `pending, shipped, delivered`
+        RESULT
+      end
+
+      it "annotates the enum types" do
+        expect(markdown_format).to eq(expected_result)
+      end
+    end
+  end
+end