Add definition support for :through associations

Author: joshuay03

lib/ruby_lsp/ruby_lsp_rails/definition.rb

@@ -80,7 +80,7 @@ def handle_possible_dsl(node)
         return unless arguments
 
         if Support::Associations::ALL.include?(message)
-          handle_association(call_node)
+          handle_association(node, call_node)
         elsif Support::Callbacks::ALL.include?(message)
           handle_callback(node, call_node, arguments)
           handle_if_unless_conditional(node, call_node, arguments)
@@ -125,18 +125,35 @@ def handle_validation(node, call_node, arguments)
         collect_definitions(name)
       end
 
-      #: (Prism::CallNode node) -> void
-      def handle_association(node)
-        first_argument = node.arguments&.arguments&.first
+      #: ((Prism::SymbolNode | Prism::StringNode) clicked_node, Prism::CallNode call_node) -> void
+      def handle_association(clicked_node, call_node)
+        arguments = call_node.arguments&.arguments
+        return unless arguments
+
+        first_argument = arguments.first
         return unless first_argument.is_a?(Prism::SymbolNode)
 
         association_name = first_argument.unescaped
 
+        through_element = arguments
+          .filter_map { |arg| arg.elements if arg.is_a?(Prism::KeywordHashNode) }
+          .flatten
+          .find { |elem| elem.key.value == "through" }
+
+        if clicked_node == first_argument
+          handle_association_name(association_name)
+        elsif through_element && clicked_node == through_element.value
+          through_association_name = through_element.value.unescaped
+          handle_association_name(through_association_name)
+        end
+      end
+
+      #: (String association_name) -> void
+      def handle_association_name(association_name)
         result = @client.association_target(
           model_name: @nesting.join("::"),
           association_name: association_name,
         )
-
         return unless result
 
         @response_builder << Support::LocationBuilder.line_location_from_s(result.fetch(:location))

test/dummy/app/models/country.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
 
 class Country < ApplicationRecord
+  has_one :flag, dependent: :destroy
 end

test/dummy/app/models/flag.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class Flag < ApplicationRecord
+  belongs_to :country
+end

test/dummy/app/models/user.rb

@@ -5,7 +5,8 @@ class User < ApplicationRecord
   validates :first_name, presence: true
   has_one :profile
   scope :adult, -> { where(age: 18..) }
-  has_one :location, class_name: "Country"
+  belongs_to :location, class_name: "Country"
+  has_one :country_flag, through: :location, source: :flag
 
   attr_readonly :last_name
 

test/dummy/db/migrate/20250703132109_create_flags.rb

@@ -0,0 +1,9 @@
+class CreateFlags < ActiveRecord::Migration[8.0]
+  def change
+    create_table :flags do |t|
+      t.references :country, null: false, foreign_key: true
+
+      t.timestamps
+    end
+  end
+end

test/dummy/db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[8.0].define(version: 2024_10_25_225348) do
+ActiveRecord::Schema[8.0].define(version: 2025_07_03_132109) do
   create_table "composite_primary_keys", primary_key: ["order_id", "product_id"], force: :cascade do |t|
     t.integer "order_id"
     t.integer "product_id"
@@ -25,6 +25,13 @@
     t.datetime "updated_at", null: false
   end
 
+  create_table "flags", force: :cascade do |t|
+    t.integer "country_id", null: false
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["country_id"], name: "index_flags_on_country_id"
+  end
+
   create_table "memberships", force: :cascade do |t|
     t.integer "user_id", null: false
     t.integer "organization_id", null: false
@@ -58,6 +65,7 @@
     t.index ["country_id"], name: "index_users_on_country_id"
   end
 
+  add_foreign_key "flags", "countries"
   add_foreign_key "memberships", "organizations"
   add_foreign_key "memberships", "users"
   add_foreign_key "users", "countries"

test/ruby_lsp_rails/definition_test.rb

@@ -53,6 +53,46 @@ class Organization < ActiveRecord::Base
         assert_equal(2, response[0].range.end.line)
       end
 
+      test "recognizes main association on has_many :through association" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 4, character: 12 })
+          # typed: false
+
+          class Organization < ActiveRecord::Base
+            has_many :memberships
+            has_many :users, through: :memberships
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "user.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
+      test "recognizes through association on has_many :through association" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 4, character: 29 })
+          # typed: false
+
+          class Organization < ActiveRecord::Base
+            has_many :memberships
+            has_many :users, through: :memberships
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "membership.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
       test "recognizes belongs_to model associations" do
         response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 14 })
           # typed: false
@@ -91,6 +131,46 @@ class User < ActiveRecord::Base
         assert_equal(2, response[0].range.end.line)
       end
 
+      test "recognizes main association on has_one :through association" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 4, character: 19 })
+          # typed: false
+
+          class User < ActiveRecord::Base
+            belongs_to :location, class_name: "Country"
+            has_one :country_flag, through: :location, source: :flag
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "flag.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
+      test "recognizes through association on has_one :through association" do
+        response = generate_definitions_for_source(<<~RUBY, { line: 4, character: 35 })
+          # typed: false
+
+          class User < ActiveRecord::Base
+            belongs_to :location, class_name: "Country"
+            has_one :country_flag, through: :location, source: :flag
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+
+        assert_equal(
+          URI::Generic.from_path(path: File.join(dummy_root, "app", "models", "country.rb")).to_s,
+          response[0].uri,
+        )
+        assert_equal(2, response[0].range.start.line)
+        assert_equal(2, response[0].range.end.line)
+      end
+
       test "recognizes has_and_belongs_to_many model associations" do
         response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 27 })
           # typed: false
@@ -111,11 +191,11 @@ class Profile < ActiveRecord::Base
       end
 
       test "handles class_name argument for associations" do
-        response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 11 })
+        response = generate_definitions_for_source(<<~RUBY, { line: 3, character: 14 })
           # typed: false
 
           class User < ActiveRecord::Base
-            has_one :location, class_name: "Country"
+            belongs_to :location, class_name: "Country"
           end
         RUBY