From 096cb2bd7ecb5c6d8616ea4e0a1f5b2b8534e8b1 Mon Sep 17 00:00:00 2001
From: Gannon McGibbon <gannon.mcgibbon@gmail.com>
Date: Wed, 31 Jan 2024 20:52:59 -0600
Subject: [PATCH] Add code lens for running migrations

Adds code lens to run migrations to specific versions in the terminal.
This is convenient because it allows developers to quickly rollback or
fast-forward to specific schema versions with a click.
---
 lib/ruby_lsp/ruby_lsp_rails/code_lens.rb | 57 ++++++++++++++++++++++--
 test/ruby_lsp_rails/code_lens_test.rb    | 18 +++++++-
 2 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb b/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb
index 51d41096..e49beda5 100644
--- a/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb
+++ b/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb
@@ -9,6 +9,7 @@ module Rails
     # - Run tests in the VS Terminal
     # - Run tests in the VS Code Test Explorer
     # - Debug tests
+    # - Run migrations in the VS Terminal
     #
     # The
     # [code lens](https://microsoft.github.io/language-server-protocol/specification#textDocument_codeLens)
@@ -31,13 +32,23 @@ module Rails
     # end
     # ````
     #
+    # # Example:
+    # ```ruby
+    # Run in terminal
+    # class AddFirstNameToUsers < ActiveRecord::Migration[7.1]
+    #   # ...
+    # end
+    # ````
+    #
     # The code lenses will be displayed above the class and above each test method.
     class CodeLens < ::RubyLsp::Listener
       extend T::Sig
       extend T::Generic
 
       ResponseType = type_member { { fixed: T::Array[::RubyLsp::Interface::CodeLens] } }
-      BASE_COMMAND = "bin/rails test"
+      MIGRATE_COMMAND = "bin/rails db:migrate"
+      TEST_COMMAND = "bin/rails test"
+      BASE_COMMAND = TEST_COMMAND # TODO: Deprecate?
 
       sig { override.returns(ResponseType) }
       attr_reader :_response
@@ -78,7 +89,7 @@ def on_call_node_enter(node)
         return unless content && !content.empty?
 
         line_number = node.location.start_line
-        command = "#{BASE_COMMAND} #{@path}:#{line_number}"
+        command = "#{TEST_COMMAND} #{@path}:#{line_number}"
         add_test_code_lens(node, name: content, command: command, kind: :example)
       end
 
@@ -86,9 +97,10 @@ def on_call_node_enter(node)
       sig { params(node: Prism::DefNode).void }
       def on_def_node_enter(node)
         method_name = node.name.to_s
+
         if method_name.start_with?("test_")
           line_number = node.location.start_line
-          command = "#{BASE_COMMAND} #{@path}:#{line_number}"
+          command = "#{TEST_COMMAND} #{@path}:#{line_number}"
           add_test_code_lens(node, name: method_name, command: command, kind: :example)
         end
       end
@@ -96,11 +108,18 @@ def on_def_node_enter(node)
       sig { params(node: Prism::ClassNode).void }
       def on_class_node_enter(node)
         class_name = node.constant_path.slice
+        superclass_name = node.superclass&.slice
+
         if class_name.end_with?("Test")
-          command = "#{BASE_COMMAND} #{@path}"
+          command = "#{TEST_COMMAND} #{@path}"
           add_test_code_lens(node, name: class_name, command: command, kind: :group)
         end
 
+        if superclass_name&.start_with?("ActiveRecord::Migration")
+          command = "#{MIGRATE_COMMAND} VERSION=#{migration_version}"
+          add_migrate_code_lens(node, name: class_name, command: command)
+        end
+
         @group_id_stack.push(@group_id)
         @group_id += 1
       end
@@ -112,6 +131,36 @@ def on_class_node_leave(node)
 
       private
 
+      sig { returns(T.nilable(String)) }
+      def migration_version
+        File.basename(T.must(@path)).split("_").first
+      end
+
+      sig { params(node: Prism::Node, name: String, command: String).void }
+      def add_migrate_code_lens(node, name:, command:)
+        return unless @path
+
+        arguments = [
+          @path,
+          name,
+          command,
+          {
+            start_line: node.location.start_line - 1,
+            start_column: node.location.start_column,
+            end_line: node.location.end_line - 1,
+            end_column: node.location.end_column,
+          },
+        ]
+
+        @_response << create_code_lens(
+          node,
+          title: "Run in terminal",
+          command_name: "rubyLsp.runMigrationInTerminal",
+          arguments: arguments,
+          data: { type: "migrate" },
+        )
+      end
+
       sig { params(node: Prism::Node, name: String, command: String, kind: Symbol).void }
       def add_test_code_lens(node, name:, command:, kind:)
         return unless @path
diff --git a/test/ruby_lsp_rails/code_lens_test.rb b/test/ruby_lsp_rails/code_lens_test.rb
index ed4bff2a..0ea0ec4f 100644
--- a/test/ruby_lsp_rails/code_lens_test.rb
+++ b/test/ruby_lsp_rails/code_lens_test.rb
@@ -182,10 +182,24 @@ class NestedTest < ActiveSupport::TestCase
         assert_empty(data)
       end
 
+      test "recognizes migrations" do
+        response = generate_code_lens_for_source(<<~RUBY, file: "file://db/migrate/123456_add_first_name_to_users.rb")
+          class AddFirstNameToUsers < ActiveRecord::Migration[7.1]
+            def change
+              add_column(:users, :first_name, :string)
+            end
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+        assert_match("Run in terminal", response[0].command.title)
+        assert_equal("bin/rails db:migrate VERSION=123456", response[0].command.arguments[2])
+      end
+
       private
 
-      def generate_code_lens_for_source(source)
-        uri = URI("file:///fake.rb")
+      def generate_code_lens_for_source(source, file: "file:///fake.rb")
+        uri = URI(file)
         store = RubyLsp::Store.new
         store.set(uri: uri, source: source, version: 1)