Add primary key indicator to model hover response (#316)

Closes #313

lib/ruby_lsp/ruby_lsp_rails/hover.rb

@@ -83,7 +83,8 @@ def generate_column_content(name)
 
         @response_builder.push(
           model[:columns].map do |name, type|
-            "**#{name}**: #{type}\n"
+            primary_key_suffix = " (PK)" if model[:primary_keys].include?(name)
+            "**#{name}**: #{type}#{primary_key_suffix}\n"
           end.join("\n"),
           category: :documentation,
         )

lib/ruby_lsp/ruby_lsp_rails/server.rb

@@ -91,6 +91,7 @@ def resolve_database_info_from_model(model_name)
         info = {
           result: {
             columns: const.columns.map { |column| [column.name, column.type] },
+            primary_keys: Array(const.primary_key),
           },
         }
 

test/dummy/app/models/composite_primary_key.rb

@@ -0,0 +1,5 @@
+# typed: true
+# frozen_string_literal: true
+
+class CompositePrimaryKey < ApplicationRecord
+end

test/dummy/db/migrate/20240403145625_create_composite_pk.rb

@@ -0,0 +1,11 @@
+class CreateCompositePk < ActiveRecord::Migration[7.1]
+  def change
+    create_table :composite_primary_keys, primary_key: [:order_id, :product_id] do |t|
+      t.integer :order_id
+      t.integer :product_id
+      t.text :note
+
+      t.timestamps
+    end
+  end
+end

test/dummy/db/schema.rb

@@ -10,7 +10,15 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[7.1].define(version: 2023_10_19_180159) do
+ActiveRecord::Schema[7.1].define(version: 2024_04_03_145625) do
+  create_table "composite_primary_keys", primary_key: ["order_id", "product_id"], force: :cascade do |t|
+    t.integer "order_id"
+    t.integer "product_id"
+    t.text "note"
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+  end
+
   create_table "posts", force: :cascade do |t|
     t.string "title"
     t.text "body"

test/ruby_lsp_rails/hover_test.rb

@@ -24,6 +24,7 @@ class HoverTest < ActiveSupport::TestCase
             ["created_at", "datetime"],
             ["updated_at", "datetime"],
           ],
+          primary_keys: ["id"],
         }
 
         RunnerClient.any_instance.stubs(model: expected_response)
@@ -44,7 +45,7 @@ class User < ApplicationRecord
           [Schema](file://#{dummy_root}/db/schema.rb)
 
 
-          **id**: integer
+          **id**: integer (PK)
 
           **first_name**: string
 
@@ -69,6 +70,7 @@ class User < ApplicationRecord
             ["created_at", "datetime"],
             ["updated_at", "datetime"],
           ],
+          primary_keys: ["id"],
         }
 
         RunnerClient.any_instance.stubs(model: expected_response)
@@ -85,7 +87,7 @@ class User < ApplicationRecord
         assert_equal(<<~CONTENT.chomp, response.contents.value)
           [Schema](file://#{dummy_root}/db/schema.rb)
 
-          **id**: integer
+          **id**: integer (PK)
 
           **first_name**: string
 
@@ -99,10 +101,54 @@ class User < ApplicationRecord
         CONTENT
       end
 
+      test "returns column information for models with composite primary keys" do
+        expected_response = {
+          schema_file: "#{dummy_root}/db/schema.rb",
+          columns: [
+            ["order_id", "integer"],
+            ["product_id", "integer"],
+            ["note", "string"],
+            ["created_at", "datetime"],
+            ["updated_at", "datetime"],
+          ],
+          primary_keys: ["order_id", "product_id"],
+        }
+
+        RunnerClient.any_instance.stubs(model: expected_response)
+
+        response = hover_on_source(<<~RUBY, { line: 3, character: 0 })
+          class CompositePrimaryKey < ApplicationRecord
+          end
+
+          CompositePrimaryKey
+        RUBY
+
+        assert_equal(<<~CONTENT.chomp, response.contents.value)
+          ```ruby
+          CompositePrimaryKey
+          ```
+
+          **Definitions**: [fake.rb](file:///fake.rb#L1,1-2,4)
+          [Schema](file://#{dummy_root}/db/schema.rb)
+
+
+          **order_id**: integer (PK)
+
+          **product_id**: integer (PK)
+
+          **note**: string
+
+          **created_at**: datetime
+
+          **updated_at**: datetime
+        CONTENT
+      end
+
       test "handles `db/structure.sql` instead of `db/schema.rb`" do
         expected_response = {
           schema_file: "#{dummy_root}/db/structure.sql",
           columns: [],
+          primary_keys: [],
         }
 
         RunnerClient.any_instance.stubs(model: expected_response)
@@ -124,6 +170,7 @@ class User < ApplicationRecord
         expected_response = {
           schema_file: nil,
           columns: [],
+          primary_keys: [],
         }
 
         RunnerClient.any_instance.stubs(model: expected_response)