From b26a4caa4e53503007461d79d8d55b6b1fdc861f Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Fri, 1 Nov 2024 14:19:14 -0400
Subject: [PATCH] Add support for checking for pending migrations and running
 them

---
 lib/ruby_lsp/ruby_lsp_rails/runner_client.rb | 23 ++++++++++++++++++
 lib/ruby_lsp/ruby_lsp_rails/server.rb        | 25 ++++++++++++++++++++
 test/ruby_lsp_rails/server_test.rb           | 22 +++++++++++++++++
 3 files changed, 70 insertions(+)

diff --git a/lib/ruby_lsp/ruby_lsp_rails/runner_client.rb b/lib/ruby_lsp/ruby_lsp_rails/runner_client.rb
index 4d96b510..a2ffaba9 100644
--- a/lib/ruby_lsp/ruby_lsp_rails/runner_client.rb
+++ b/lib/ruby_lsp/ruby_lsp_rails/runner_client.rb
@@ -189,6 +189,29 @@ def delegate_notification(server_addon_name:, request_name:, **params)
         )
       end
 
+      sig { returns(T.nilable(String)) }
+      def pending_migrations
+        response = make_request("pending_migrations")
+        response[:pending_migrations] if response
+      rescue IncompleteMessageError
+        log_message(
+          "Ruby LSP Rails failed to get pending migrations",
+          type: RubyLsp::Constant::MessageType::ERROR,
+        )
+        nil
+      end
+
+      sig { returns(T.nilable(T::Hash[Symbol, T.untyped])) }
+      def run_migrations
+        make_request("run_migrations")
+      rescue IncompleteMessageError
+        log_message(
+          "Ruby LSP Rails failed to run migrations",
+          type: RubyLsp::Constant::MessageType::ERROR,
+        )
+        nil
+      end
+
       # Delegates a request to a server add-on
       sig do
         params(
diff --git a/lib/ruby_lsp/ruby_lsp_rails/server.rb b/lib/ruby_lsp/ruby_lsp_rails/server.rb
index 24675580..4ea2276b 100644
--- a/lib/ruby_lsp/ruby_lsp_rails/server.rb
+++ b/lib/ruby_lsp/ruby_lsp_rails/server.rb
@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "open3"
 
 # NOTE: We should avoid printing to stderr since it causes problems. We never read the standard error pipe from the
 # client, so it will become full and eventually hang or crash. Instead, return a response with an `error` key.
@@ -109,6 +110,10 @@ def execute(request, params)
           send_message(resolve_database_info_from_model(params.fetch(:name)))
         when "association_target_location"
           send_message(resolve_association_target(params))
+        when "pending_migrations"
+          send_message({ result: { pending_migrations: pending_migrations_message } })
+        when "run_migrations"
+          send_message({ result: run_migrations })
         when "reload"
           ::Rails.application.reloader.reload!
         when "route_location"
@@ -252,6 +257,26 @@ def active_record_model?(const)
           !const.abstract_class?
         )
       end
+
+      def pending_migrations_message
+        return unless defined?(ActiveRecord)
+
+        ActiveRecord::Migration.check_all_pending!
+        nil
+      rescue ActiveRecord::PendingMigrationError => e
+        e.message
+      end
+
+      def run_migrations
+        # Running migrations invokes `load` which will repeatedly load the same files. It's not designed to be invoked
+        # multiple times within the same process. To avoid any memory bloat, we run migrations in a separate process
+        stdout, status = Open3.capture2(
+          { "VERBOSE" => "true" },
+          "bundle exec rake db:migrate",
+        )
+
+        { message: stdout, status: status.exitstatus }
+      end
     end
   end
 end
diff --git a/test/ruby_lsp_rails/server_test.rb b/test/ruby_lsp_rails/server_test.rb
index 4120d921..f750c1cf 100644
--- a/test/ruby_lsp_rails/server_test.rb
+++ b/test/ruby_lsp_rails/server_test.rb
@@ -176,6 +176,28 @@ def resolve_route_info(requirements)
     assert_equal("Hello\n", stderr)
   end
 
+  test "checking for pending migrations" do
+    capture_subprocess_io do
+      system("bundle exec rails g migration CreateStudents name:string")
+    end
+
+    @server.execute("pending_migrations", {})
+    message = response.dig(:result, :pending_migrations)
+    assert_match("You have 1 pending migration", message)
+    assert_match(%r{db/migrate/[\d]+_create_students\.rb}, message)
+  ensure
+    FileUtils.rm_rf("db") if File.directory?("db")
+  end
+
+  test "running migrations happens in a child process" do
+    Open3.expects(:capture2)
+      .with({ "VERBOSE" => "true" }, "bundle exec rake db:migrate")
+      .returns(["Running migrations...", mock(exitstatus: 0)])
+
+    @server.execute("run_migrations", {})
+    assert_equal({ message: "Running migrations...", status: 0 }, response[:result])
+  end
+
   private
 
   def response