mermaid_rails_erd 1.0.1

data/lib/mermaid_rails_erd/model_data_collector.rb

@@ -0,0 +1,157 @@
+# frozen_string_literal: true
+
+require_relative "column_info"
+
+module MermaidRailsErd
+  class ModelDataCollector
+    attr_reader :models_data, :tables, :models_no_tables, :models, :invalid_associations, :polymorphic_associations, :regular_associations
+
+    def initialize(model_loader)
+      @models_data = {}
+      @polymorphic_associations = []
+      @regular_associations = []
+      @invalid_associations = []
+      @polymorphic_targets = Hash.new { |h, k| h[k] = [] }
+      @tables = {}
+      @models_no_tables = []
+      @models = model_loader.load
+    end
+
+    def collect
+      (base_models & models_with_tables).each do |model|
+        model_data = { model: model, associations: [] }
+
+        # Register polymorphic interfaces that this model implements
+        model.reflect_on_all_associations.each do |assoc|
+          if (interface_name = assoc.options[:as])
+            register_polymorphic_target(interface_name, model)
+          end
+
+          if assoc.options[:polymorphic]
+            @polymorphic_associations << { model: model, association: assoc }
+          else
+            @regular_associations << { model: model, association: assoc }
+          end
+
+          model_data[:associations] << assoc
+        end
+
+        @models_data[model.name] = model_data
+
+        collect_table_for_model(model)
+      end
+
+      (base_models & models_without_tables).each do |model|
+        @models_no_tables << model
+      end
+
+      self
+    end
+
+    def get_model_data(model_name)
+      @models_data[model_name]
+    end
+
+    # Register a model as implementing a polymorphic interface
+    # @param interface_name [String, Symbol] Name of the polymorphic interface
+    # @param model [Class] Model that implements the interface
+    def register_polymorphic_target(interface_name, model)
+      @polymorphic_targets[interface_name.to_s] << model
+    end
+
+    # Returns all models that implement the given polymorphic interface
+    # @param polymorphic_name [String, Symbol] Name of the polymorphic interface
+    # @return [Array<Class>] Array of models that implement the interface
+    def polymorphic_targets_for(polymorphic_name)
+      @polymorphic_targets[polymorphic_name.to_s] || []
+    end
+
+    # Reset all registered polymorphic targets (useful for testing)
+    def reset_polymorphic_targets
+      @polymorphic_targets.clear
+    end
+
+    # Collect table information from a model
+    # @param model [Class] ActiveRecord model to collect table information from
+    def collect_table_for_model(model)
+      # Skip if we've already collected this table
+      return if @tables.key?(model.table_name)
+
+      # Collect basic table structure with columns
+      @tables[model.table_name] = model.columns.map do |col|
+        annotations = []
+        annotations << "PK" if col.name == model.primary_key
+
+        ColumnInfo.new(col.name, annotations, col.sql_type, col.type, col.null)
+      end
+    end
+
+    # Update table definitions with foreign key annotations
+    # @param relationships [Array<Relationship>] List of relationships to extract FKs from
+    # @return [Hash] Updated tables hash
+    def update_foreign_keys(relationships)
+      fk_mapping = extract_foreign_keys(relationships)
+
+      @tables.each do |table_name, columns|
+        next unless fk_mapping[table_name]
+
+        begin
+          columns.each do |col|
+            col.annotations << "FK" if fk_mapping[table_name].include?(col.name) && !col.annotations.include?("FK")
+          end
+        rescue StandardError => e
+          puts "Error updating FK annotations for #{table_name}: #{e.message}"
+        end
+      end
+
+      @tables
+    end
+
+    # Register an invalid association that couldn't be resolved
+    # @param model [Class] Model containing the invalid association
+    # @param assoc [ActiveRecord::Reflection::Association] Invalid association
+    # @param reason [String] Reason why the association is invalid
+    def register_invalid_association(model, assoc, reason = nil)
+      @invalid_associations << {
+        model: model,
+        association: assoc,
+        reason: reason,
+        label: "#{model.name}##{assoc.name}",
+      }
+    end
+
+    private
+
+    def base_models
+      models.reject { |model| (model < model.base_class) }
+    end
+
+    def models_with_tables
+      models.select(&:table_exists?)
+    end
+
+    def models_without_tables
+      models.reject(&:table_exists?)
+    end
+
+    # Extract foreign key information from relationships
+    # @param relationships [Array<Relationship>] List of relationships to extract FKs from
+    # @return [Hash] Mapping of table names to arrays of FK column names
+    def extract_foreign_keys(relationships)
+      fk_mapping = {}
+
+      relationships.each do |rel|
+        # Use the relationship attributes directly instead of parsing the label
+        table_name = rel.fk_table
+        fk_column = rel.fk_column
+
+        fk_mapping[table_name] ||= []
+        fk_mapping[table_name] << fk_column unless fk_mapping[table_name].include?(fk_column)
+      rescue StandardError => e
+        puts "Error extracting FK from relationship: #{e.message}"
+      end
+
+      fk_mapping
+    end
+  end
+end

data/lib/mermaid_rails_erd/model_loader.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module MermaidRailsErd
+  class ModelLoader
+    def load
+      raise MermaidRailsErd::Error, "Rails is not loaded" unless defined?(Rails)
+
+      Rails.application.eager_load! unless Rails.configuration.cache_classes
+      ActiveRecord::Base.descendants.reject(&:abstract_class?)
+    end
+  end
+end

data/lib/mermaid_rails_erd/parsed_data.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module MermaidRailsErd
+  # Class to hold parsed ERD data with delegation to model_data_collector
+  # Provides structured access to relationships, tables, and model data
+  class ParsedData
+    attr_reader :relationships, :tables, :model_data_collector
+
+    # Initialize with collected ERD data
+    # @param relationships [Array] Array of relationship objects
+    # @param tables [Hash] Hash of table definitions keyed by table name
+    # @param model_data_collector [ModelDataCollector] Collector instance with model data
+    def initialize(relationships, tables, model_data_collector)
+      @relationships = relationships
+      @tables = tables
+      @model_data_collector = model_data_collector
+    end
+
+    # Delegated methods from model_data_collector for better IDE support
+
+    # @return [Hash] Hash of model data keyed by model name
+    def models_data
+      model_data_collector.models_data
+    end
+
+    # @return [Array] Array of models without tables
+    def models_no_tables
+      model_data_collector.models_no_tables
+    end
+
+    # @return [Array] Array of all loaded models
+    def models
+      model_data_collector.models
+    end
+
+    # @return [Array<Hash>] Array of invalid associations with details
+    def invalid_associations
+      model_data_collector.invalid_associations
+    end
+
+    # @return [Array<Hash>] Array of polymorphic associations
+    def polymorphic_associations
+      model_data_collector.polymorphic_associations
+    end
+
+    # @return [Array<Hash>] Array of regular (non-polymorphic) associations
+    def regular_associations
+      model_data_collector.regular_associations
+    end
+  end
+end

data/lib/mermaid_rails_erd/polymorphic_targets_resolver.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require_relative "relationship"
+
+module MermaidRailsErd
+  class PolymorphicTargetsResolver
+    attr_reader :model_data_collector
+
+    def initialize(model_data_collector)
+      @model_data_collector = model_data_collector
+    end
+
+    def resolve(name, from_table, rel_type)
+      # Get all models that implement the polymorphic interface
+      target_models = model_data_collector.polymorphic_targets_for(name)
+
+      # Create relationships for each target model
+      target_models.map do |target|
+        fk_column = "#{name}_id"
+        Relationship.new(
+          from_table,
+          target.table_name,
+          fk_column,
+          rel_type,
+          nil, # Let the Relationship generate the label
+          from_table, # fk_table
+          fk_column, # fk_column
+          target.table_name, # pk_table
+          "id", # pk_column
+          # Add (polymorphic) to the label
+          true, # is_polymorphic
+          "polymorphic", # extra_label
+        )
+      end
+    end
+  end
+end

data/lib/mermaid_rails_erd/railtie.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module MermaidRailsErd
+  class Railtie < Rails::Railtie
+    rake_tasks do
+      load File.expand_path("../tasks/mermaid_rails_erd.rake", __dir__)
+    end
+  end
+end

data/lib/mermaid_rails_erd/relationship.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module MermaidRailsErd
+  class Relationship
+    attr_reader :from_table, :to_table, :foreign_key, :relationship_type, :label, :fk_table, :fk_column, :pk_table, :pk_column, :is_polymorphic, :extra_label
+
+    # Provide aliases for backward compatibility
+    alias from from_table
+    alias to to_table
+    alias symbol relationship_type
+
+    def initialize(from, to, foreign_key, relationship_type, label = nil,
+                   fk_table = nil, fk_column = nil, pk_table = nil, pk_column = nil,
+                   is_polymorphic = false, extra_label = nil)
+      @from_table = from
+      @to_table = to
+      @foreign_key = foreign_key
+      @relationship_type = relationship_type
+      @is_polymorphic = is_polymorphic
+      @extra_label = extra_label
+
+      # Store FK/PK information
+      @fk_table = fk_table || from
+      @fk_column = fk_column || foreign_key
+      @pk_table = pk_table || to
+      @pk_column = pk_column || "id"
+
+      # Generate label if not provided
+      @label = label || generate_label
+    end
+
+    def key
+      [from_table, to_table, foreign_key].sort.join("::")
+    end
+
+    def generate_label
+      base_label = "#{@fk_table}.#{@fk_column} FK → #{@pk_table}.#{@pk_column} PK"
+      @is_polymorphic ? "#{base_label} (#{@extra_label})" : base_label
+    end
+  end
+end

data/lib/mermaid_rails_erd/relationship_builders/base_relationship_builder.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+require "set"
+require_relative "../relationship"
+
+module MermaidRailsErd
+  module RelationshipBuilders
+    class BaseRelationshipBuilder
+      attr_reader :symbol_mapper, :association_resolver, :model_data_collector
+
+      # Class variable to track one-to-one relationships across builder instances
+      @@one_to_one_keys = Set.new
+
+      def initialize(symbol_mapper:, association_resolver:, model_data_collector: nil)
+        @symbol_mapper = symbol_mapper
+        @association_resolver = association_resolver
+        @model_data_collector = model_data_collector
+      end
+
+      def build(model, assoc)
+        raise NotImplementedError, "Subclasses must implement #build"
+      end
+
+      protected
+
+      def resolve_association_model(model, assoc)
+        # Let the association_resolver handle all the resolution logic
+        # including dynamic model creation when needed
+        association_resolver.resolve(assoc)
+      rescue StandardError => e
+        puts "  ERROR resolving association class for #{model.name}##{assoc.name}: #{e.class} - #{e.message}"
+        puts e.backtrace.join("\n")
+        # Register invalid association if model_data_collector is available
+        register_invalid_association(model, assoc, "#{e.class} - #{e.message}")
+        nil
+      end
+
+      def safe_foreign_key(model, assoc)
+        # Skip through associations entirely
+        return nil if assoc.options[:through]
+
+        # Try to get the foreign key
+        begin
+          assoc.foreign_key
+        rescue StandardError => e
+          puts "  WARNING: Cannot determine foreign key for #{model.name}##{assoc.name} - skipping"
+          register_invalid_association(model, assoc, "Cannot determine foreign key: #{e.class} - #{e.message}")
+          nil
+        end
+      end
+
+      def skip_duplicate_one_to_one?(model, assoc, to_table_info)
+        return false unless %i[has_one belongs_to].include?(assoc.macro)
+
+        # Skip the check for polymorphic associations
+        return false if assoc.options[:polymorphic]
+
+        return false unless to_table_info
+
+        # Create a unique key for this one-to-one relationship
+        rel_key = [model.table_name, to_table_info[:table_name], "1:1"].sort.join("::")
+
+        # Check if we've already processed this relationship
+        return true if @@one_to_one_keys.include?(rel_key)
+
+        # Mark this relationship as processed
+        @@one_to_one_keys << rel_key
+        false
+      end
+
+      def log_missing_table_warning(model, assoc, reason = "table does not exist")
+        if assoc.options[:class_name]
+          assoc.options[:class_name].to_s.tableize
+        else
+          assoc.name.to_s.tableize
+        end
+        puts "  WARNING: Could not create relationship for #{model.name}##{assoc.name} - #{reason}"
+        register_invalid_association(model, assoc, reason)
+        []
+      end
+
+      private
+
+      def register_invalid_association(model, assoc, reason)
+        return unless @model_data_collector
+
+        @model_data_collector.register_invalid_association(model, assoc, reason)
+      end
+    end
+  end
+end

data/lib/mermaid_rails_erd/relationship_builders/belongs_to_relationship_builder.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require_relative "base_relationship_builder"
+
+module MermaidRailsErd
+  module RelationshipBuilders
+    class BelongsToRelationshipBuilder < BaseRelationshipBuilder
+      def build(model, assoc)
+        from_table = model.table_name
+        fk = safe_foreign_key(model, assoc)
+
+        # Skip if we couldn't determine the foreign key
+        return [] unless fk
+
+        to_table_info = resolve_association_model(model, assoc)
+
+        # Skip if this is a duplicate one-to-one relationship
+        return [] if skip_duplicate_one_to_one?(model, assoc, to_table_info)
+
+        if to_table_info
+          # FK is on source table for belongs_to
+          [Relationship.new(
+            from_table,
+            to_table_info[:table_name],
+            fk,
+            "}o--||",
+            nil, # Let the Relationship generate the label
+            from_table, # fk_table
+            fk, # fk_column
+            to_table_info[:table_name], # pk_table
+            to_table_info[:primary_key], # pk_column
+          )]
+        else
+          log_missing_table_warning(model, assoc)
+        end
+      end
+    end
+  end
+end

data/lib/mermaid_rails_erd/relationship_builders/habtm_relationship_builder.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require_relative "base_relationship_builder"
+require_relative "../relationship"
+
+module MermaidRailsErd
+  module RelationshipBuilders
+    class HasAndBelongsToManyRelationshipBuilder < BaseRelationshipBuilder
+      def initialize(symbol_mapper:, association_resolver:, printed_tables: Set.new, model_data_collector: nil)
+        super(symbol_mapper: symbol_mapper, association_resolver: association_resolver, model_data_collector: model_data_collector)
+        @printed_tables = printed_tables
+      end
+
+      def build(model, assoc)
+        from_table = model.table_name
+        to_table_info = resolve_association_model(model, assoc)
+
+        return log_missing_table_warning(model, assoc, "target model does not exist") unless to_table_info
+
+        join_table = assoc.join_table.to_s
+
+        # Check if we need to verify the table existence
+        unless @printed_tables.include?(join_table)
+          begin
+            ActiveRecord::Base.connection.columns(join_table)
+            @printed_tables << join_table
+          rescue StandardError => e
+            return log_missing_table_warning(model, assoc, "join table #{join_table} is missing: #{e.message}")
+          end
+        end
+
+        # Try to get the foreign keys
+        source_fk = nil
+        target_fk = nil
+
+        begin
+          source_fk = assoc.foreign_key
+        rescue StandardError => e
+          puts "  WARNING: Could not determine foreign key for #{model.name} in HABTM: #{e.message}"
+          register_invalid_association(model, assoc, "Could not determine foreign key: #{e.message}")
+        end
+
+        begin
+          target_fk = assoc.association_foreign_key
+        rescue StandardError => e
+          puts "  WARNING: Could not determine association foreign key for #{model.name}##{assoc.name}: #{e.message}"
+          register_invalid_association(model, assoc, "Could not determine association foreign key: #{e.message}")
+        end
+
+        # Skip if we couldn't determine both foreign keys
+        return log_missing_table_warning(model, assoc, "could not determine foreign keys") if !source_fk || !target_fk
+
+        # If we reach here, the join table exists, so create relationships
+        [
+          Relationship.new(
+            join_table,
+            from_table,
+            source_fk,
+            "}o--||",
+            nil, # Let the Relationship generate the label
+            join_table, # fk_table
+            source_fk, # fk_column
+            from_table, # pk_table
+            model.primary_key, # pk_column
+          ),
+          Relationship.new(
+            join_table,
+            to_table_info[:table_name],
+            target_fk,
+            "}o--||",
+            nil, # Let the Relationship generate the label
+            join_table, # fk_table
+            target_fk, # fk_column
+            to_table_info[:table_name], # pk_table
+            to_table_info[:primary_key], # pk_column
+          ),
+        ]
+      end
+    end
+  end
+end

data/lib/mermaid_rails_erd/relationship_builders/has_many_relationship_builder.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require_relative "base_relationship_builder"
+
+module MermaidRailsErd
+  module RelationshipBuilders
+    class HasManyRelationshipBuilder < BaseRelationshipBuilder
+      def build(model, assoc)
+        from_table = model.table_name
+        fk = safe_foreign_key(model, assoc)
+
+        # Skip if we couldn't determine the foreign key
+        return [] unless fk
+
+        to_table_info = resolve_association_model(model, assoc)
+
+        if to_table_info
+          # FK is on target table for has_many
+          [Relationship.new(
+            to_table_info[:table_name],
+            from_table,
+            fk,
+            "}o--||",
+            nil, # Let the Relationship generate the label
+            to_table_info[:table_name], # fk_table
+            fk, # fk_column
+            from_table, # pk_table
+            model.primary_key, # pk_column
+          )]
+        else
+          log_missing_table_warning(model, assoc)
+        end
+      end
+    end
+  end
+end

data/lib/mermaid_rails_erd/relationship_builders/has_one_relationship_builder.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require_relative "base_relationship_builder"
+
+module MermaidRailsErd
+  module RelationshipBuilders
+    class HasOneRelationshipBuilder < BaseRelationshipBuilder
+      def build(model, assoc)
+        from_table = model.table_name
+        rel_type = symbol_mapper.map(assoc.macro)
+        fk = safe_foreign_key(model, assoc)
+
+        # Skip if we couldn't determine the foreign key
+        return [] unless fk
+
+        to_table_info = resolve_association_model(model, assoc)
+
+        # Skip if this is a duplicate one-to-one relationship
+        return [] if skip_duplicate_one_to_one?(model, assoc, to_table_info)
+
+        if to_table_info
+          [Relationship.new(
+            from_table,
+            to_table_info[:table_name],
+            fk,
+            rel_type,
+            nil, # Let the Relationship generate the label
+            to_table_info[:table_name], # fk_table
+            fk, # fk_column
+            from_table, # pk_table
+            model.primary_key, # pk_column
+          )]
+        else
+          log_missing_table_warning(model, assoc)
+        end
+      end
+    end
+  end
+end

data/lib/mermaid_rails_erd/relationship_registry.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require_relative "relationship_builders/belongs_to_relationship_builder"
+require_relative "relationship_builders/has_many_relationship_builder"
+require_relative "relationship_builders/has_one_relationship_builder"
+require_relative "relationship_builders/habtm_relationship_builder"
+require_relative "model_data_collector"
+
+module MermaidRailsErd
+  class RelationshipRegistry
+    attr_reader :builders, :polymorphic_resolver, :model_data_collector
+
+    def initialize(
+      symbol_mapper:,
+      association_resolver:,
+      polymorphic_resolver:,
+      model_data_collector:,
+      printed_tables: Set.new
+    )
+      @polymorphic_resolver = polymorphic_resolver
+      @model_data_collector = model_data_collector
+      @printed_tables = printed_tables
+
+      @builders = {
+        belongs_to: RelationshipBuilders::BelongsToRelationshipBuilder.new(
+          symbol_mapper: symbol_mapper,
+          association_resolver: association_resolver,
+          model_data_collector: model_data_collector,
+        ),
+        has_many: RelationshipBuilders::HasManyRelationshipBuilder.new(
+          symbol_mapper: symbol_mapper,
+          association_resolver: association_resolver,
+          model_data_collector: model_data_collector,
+        ),
+        has_one: RelationshipBuilders::HasOneRelationshipBuilder.new(
+          symbol_mapper: symbol_mapper,
+          association_resolver: association_resolver,
+          model_data_collector: model_data_collector,
+        ),
+        has_and_belongs_to_many: RelationshipBuilders::HasAndBelongsToManyRelationshipBuilder.new(
+          symbol_mapper: symbol_mapper,
+          association_resolver: association_resolver,
+          printed_tables: printed_tables,
+          model_data_collector: model_data_collector,
+        ),
+      }
+    end
+
+    def build_relationships(model, assoc)
+      # Check for polymorphic association first
+      if assoc.options[:polymorphic]
+        from_table = model.table_name
+        rel_type = builders[assoc.macro].symbol_mapper.map(assoc.macro)
+        return polymorphic_resolver.resolve(assoc.name.to_s, from_table, rel_type)
+      end
+
+      # Delegate to the appropriate builder
+      builder = builders[assoc.macro]
+      return builder.build(model, assoc) if builder
+
+      # If no builder exists for this macro type, return an empty array
+      []
+    end
+
+    def build_all_relationships
+      relationships = []
+
+      # Process polymorphic associations first
+      @model_data_collector.polymorphic_associations.each do |data|
+        model = data[:model]
+        assoc = data[:association]
+        relationships.concat(build_relationships(model, assoc))
+      end
+
+      # Then process regular associations
+      @model_data_collector.regular_associations.each do |data|
+        model = data[:model]
+        assoc = data[:association]
+        relationships.concat(build_relationships(model, assoc))
+      end
+
+      relationships
+    end
+  end
+end