From 548a539c66885867e50832f7e08550e2145c5352 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.
---
 README.md                                |  1 +
 lib/ruby_lsp/ruby_lsp_rails/code_lens.rb | 46 +++++++++++++++++++++---
 test/ruby_lsp_rails/code_lens_test.rb    | 23 ++++++++++--
 3 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 3d41a4ab..0ed16887 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,7 @@ Ruby LSP Rails is a [Ruby LSP](https://github.com/Shopify/ruby-lsp) addon for ex
 * Navigate to associations, validations, callbacks and test cases using your editor's "Go to Symbol" feature, or outline view.
 * Jump to the definition of callbacks using your editor's "Go to Definition" feature.
 * Jump to the declaration of a route.
+* Code Lens allowing fast-forwarding or rewinding of migrations.
 
 ## Installation
 
diff --git a/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb b/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb
index 0f1f5807..e58da2e0 100644
--- a/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb
+++ b/lib/ruby_lsp/ruby_lsp_rails/code_lens.rb
@@ -10,6 +10,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)
@@ -34,6 +35,14 @@ module Rails
     # end
     # ````
     #
+    # # Example:
+    # ```ruby
+    # Run
+    # class AddFirstNameToUsers < ActiveRecord::Migration[7.1]
+    #   # ...
+    # end
+    # ````
+    #
     # The code lenses will be displayed above the class and above each test method.
     #
     # Note: When using the Test Explorer view, if your code contains a statement to pause execution (e.g. `debugger`) it
@@ -74,6 +83,7 @@ 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 = "#{test_command} #{@path}:#{line_number}"
@@ -84,12 +94,19 @@ 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 = "#{test_command} #{@path}"
           add_test_code_lens(node, name: class_name, command: command, kind: :group)
           @group_id_stack.push(@group_id)
           @group_id += 1
         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
       end
 
       sig { params(node: Prism::ClassNode).void }
@@ -104,11 +121,30 @@ def on_class_node_leave(node)
 
       sig { returns(String) }
       def test_command
-        if Gem.win_platform?
-          "ruby bin/rails test"
-        else
-          "bin/rails test"
-        end
+        "#{RbConfig.ruby} bin/rails test"
+      end
+
+      sig { returns(String) }
+      def migrate_command
+        "#{RbConfig.ruby} bin/rails db:migrate"
+      end
+
+      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
+
+        @response_builder << create_code_lens(
+          node,
+          title: "Run",
+          command_name: "rubyLsp.runTask",
+          arguments: [command],
+          data: { type: "migrate" },
+        )
       end
 
       sig { params(node: Prism::Node, name: String, command: String, kind: Symbol).void }
diff --git a/test/ruby_lsp_rails/code_lens_test.rb b/test/ruby_lsp_rails/code_lens_test.rb
index abf94f49..5b6d652d 100644
--- a/test/ruby_lsp_rails/code_lens_test.rb
+++ b/test/ruby_lsp_rails/code_lens_test.rb
@@ -8,6 +8,7 @@ module Rails
     class CodeLensTest < ActiveSupport::TestCase
       setup do
         GlobalState.any_instance.stubs(:test_library).returns("rails")
+        @ruby = Gem.win_platform? ? "ruby.exe" : "ruby"
       end
 
       test "does not create code lenses if rails is not the test library" do
@@ -270,13 +271,29 @@ class Test < ActiveSupport::TestCase
           end
         RUBY
 
-        assert_equal("ruby bin/rails test /fake.rb", response[0].command.arguments[2])
+        assert_match("#{ruby} bin/rails test /fake.rb", response[0].command.arguments[2])
+      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", response[0].command.title)
+        assert_match("#{ruby} bin/rails db:migrate VERSION=123456", response[0].command.arguments[0])
       end
 
       private
 
-      def generate_code_lens_for_source(source)
-        with_server(source) do |server, uri|
+      attr_reader :ruby
+
+      def generate_code_lens_for_source(source, file: "/fake.rb")
+        with_server(source, URI(file)) do |server, uri|
           server.process_message(
             id: 1,
             method: "textDocument/codeLens",