feat: Add database name to annotations in multi-DB environments

lib/annotate_rb/model_annotator/annotation/annotation_builder.rb

@@ -5,7 +5,7 @@ module ModelAnnotator
     module Annotation
       class AnnotationBuilder
         class Annotation < Components::Base
-          attr_reader :version, :table_name, :table_comment, :max_size
+          attr_reader :version, :table_name, :table_comment, :max_size, :database_name
 
           def initialize(options, **input)
             @options = options
@@ -15,12 +15,13 @@ def initialize(options, **input)
             @table_comment = input[:table_comment]
             @max_size = input[:max_size]
             @model = input[:model]
+            @database_name = input[:database_name]
           end
 
           def body
             [
               MainHeader.new(version, @options[:include_version]),
-              SchemaHeader.new(table_name, table_comment, @options),
+              SchemaHeader.new(table_name, table_comment, database_name, @options),
               MarkdownHeader.new(max_size),
               *columns,
               IndexAnnotation::AnnotationBuilder.new(@model, @options).build,
@@ -63,10 +64,21 @@ def build
           table_name = @model.table_name
           table_comment = @model.connection.try(:table_comment, @model.table_name)
           max_size = @model.max_schema_info_width
+          database_name = @model.database_name if multi_db_environment?
 
           _annotation = Annotation.new(@options,
             version: version, table_name: table_name, table_comment: table_comment,
-            max_size: max_size, model: @model).build
+            max_size: max_size, model: @model, database_name: database_name).build
+        end
+
+        private
+
+        def multi_db_environment?
+          if defined?(::Rails) && ::Rails.env
+            ActiveRecord::Base.configurations.configs_for(env_name: ::Rails.env).size > 1
+          else
+            false
+          end
         end
       end
     end

lib/annotate_rb/model_annotator/annotation/schema_header.rb

@@ -20,20 +20,38 @@ def to_markdown
           end
         end
 
-        attr_reader :table_name, :table_comment
+        class DatabaseName < Components::Base
+          attr_reader :name
+
+          def initialize(name)
+            @name = name
+          end
+
+          def to_default
+            "# Database name: #{name}"
+          end
+
+          def to_markdown
+            "# Database name: `#{name}`"
+          end
+        end
+
+        attr_reader :table_name, :table_comment, :database_name
 
-        def initialize(table_name, table_comment, options)
+        def initialize(table_name, table_comment, database_name, options)
           @table_name = table_name
           @table_comment = table_comment
+          @database_name = database_name
           @options = options
         end
 
         def body
           [
             Components::BlankCommentLine.new,
             TableName.new(name),
+            (DatabaseName.new(database_name) if database_name),
             Components::BlankCommentLine.new
-          ]
+          ].compact
         end
 
         def to_default

lib/annotate_rb/model_annotator/model_wrapper.rb

@@ -38,6 +38,10 @@ def connection
         @klass.connection
       end
 
+      def database_name
+        connection.pool.db_config.name
+      end
+
       # Returns the unmodified model columns
       def raw_columns
         @raw_columns ||= @klass.columns

spec/integration/annotate_models_in_multi_db_spec.rb

@@ -24,4 +24,14 @@
     # Ensure "Model files unchanged." is included in the output
     expect(second_run_output).to include("Model files unchanged.")
   end
+
+  it "includes the database name in the annotation for secondary models" do
+    reset_database
+    run_migrations
+
+    run_command_and_stop("bundle exec annotaterb models", fail_on_error: true, exit_timeout: command_timeout_seconds)
+
+    content = read_file(dummyapp_model("secondary/test_default.rb"))
+    expect(content).to include("# Database name: secondary")
+  end
 end

spec/integration/annotate_single_file_spec.rb

@@ -29,6 +29,7 @@
 
     expect(last_command_started).to be_successfully_executed
     expect(annotated_test_default).to eq(expected_test_default)
+    expect(annotated_test_default).not_to include("# Database name:")
     expect(annotated_test_null_false).not_to eq(expected_test_null_false)
   end
 end

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

@@ -2,10 +2,11 @@
 
 RSpec.describe AnnotateRb::ModelAnnotator::Annotation::SchemaHeader do
   describe "#to_default" do
-    subject { described_class.new(table_name, table_comment, options).to_default }
+    subject { described_class.new(table_name, table_comment, database_name, options).to_default }
 
     let(:table_name) { "users" }
     let(:table_comment) {}
+    let(:database_name) {}
     let :options do
       AnnotateRb::Options.new({})
     end
@@ -20,6 +21,21 @@
 
     it { is_expected.to eq(expected_header) }
 
+    context "with database_name" do
+      let(:database_name) { "secondary" }
+
+      let(:expected_header) do
+        <<~HEADER.strip
+          #
+          # Table name: users
+          # Database name: secondary
+          #
+        HEADER
+      end
+
+      it { is_expected.to eq(expected_header) }
+    end
+
     context "with `with_comment: true`" do
       context "with `with_table_comments: true` and table has comments" do
         let :options do
@@ -126,10 +142,11 @@
   end
 
   describe "#to_markdown" do
-    subject { described_class.new(table_name, table_comment, options).to_markdown }
+    subject { described_class.new(table_name, table_comment, database_name, options).to_markdown }
 
     let(:table_name) { "users" }
     let(:table_comment) {}
+    let(:database_name) {}
     let :options do
       AnnotateRb::Options.new({})
     end
@@ -143,5 +160,20 @@
     end
 
     it { is_expected.to eq(expected_header) }
+
+    context "with database_name" do
+      let(:database_name) { "secondary" }
+
+      let(:expected_header) do
+        <<~HEADER.strip
+          #
+          # Table name: `users`
+          # Database name: `secondary`
+          #
+        HEADER
+      end
+
+      it { is_expected.to eq(expected_header) }
+    end
   end
 end