Implement method autocompletion in Active Record migrations

Author: janko

lib/ruby_lsp/ruby_lsp_rails/addon.rb

@@ -135,7 +135,9 @@ def create_definition_listener(response_builder, uri, node_context, dispatcher)
       # @override
       #: (ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem] response_builder, NodeContext node_context, Prism::Dispatcher dispatcher, URI::Generic uri) -> void
       def create_completion_listener(response_builder, node_context, dispatcher, uri)
-        Completion.new(@rails_runner_client, response_builder, node_context, dispatcher, uri)
+        return unless @global_state
+
+        Completion.new(@rails_runner_client, response_builder, node_context, @global_state.index, dispatcher, uri)
       end
 
       #: (Array[{uri: String, type: Integer}] changes) -> void

lib/ruby_lsp/ruby_lsp_rails/completion.rb

@@ -7,11 +7,13 @@ class Completion
       include Requests::Support::Common
 
       # @override
-      #: (RunnerClient client, ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem] response_builder, NodeContext node_context, Prism::Dispatcher dispatcher, URI::Generic uri) -> void
-      def initialize(client, response_builder, node_context, dispatcher, uri)
+      #: (RunnerClient client, ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem] response_builder, NodeContext node_context, RubyIndexer::Index index, Prism::Dispatcher dispatcher, URI::Generic uri) -> void
+      def initialize(client, response_builder, node_context, index, dispatcher, uri)
         @response_builder = response_builder
         @client = client
         @node_context = node_context
+        @index = index
+        @path = uri.to_standardized_path #: String?
         dispatcher.register(
           self,
           :on_call_node_enter,
@@ -21,11 +23,12 @@ def initialize(client, response_builder, node_context, dispatcher, uri)
       #: (Prism::CallNode node) -> void
       def on_call_node_enter(node)
         call_node = @node_context.call_node
-        return unless call_node
+        receiver = call_node&.receiver
 
-        receiver = call_node.receiver
-        if call_node.name == :where && receiver.is_a?(Prism::ConstantReadNode)
+        if call_node&.name == :where && receiver.is_a?(Prism::ConstantReadNode)
           handle_active_record_where_completions(node: node, receiver: receiver)
+        elsif active_record_migration?
+          handle_active_record_migration_completions(node: node)
         end
       end
 
@@ -62,6 +65,45 @@ def handle_active_record_where_completions(node:, receiver:)
         end
       end
 
+      #: (node: Prism::CallNode) -> void
+      def handle_active_record_migration_completions(node:)
+        return if @path.nil?
+
+        db_configs = @client.db_configs
+        return if db_configs.nil?
+
+        db_config = db_configs.values.find do |config|
+          config[:migrations_paths].any? do |path|
+            File.join(@client.rails_root, path) == File.dirname(@path)
+          end
+        end
+        return if db_config.nil?
+
+        range = range_from_location(node.location)
+
+        @index.method_completion_candidates(node.message, db_config[:adapter_class]).each do |entry|
+          next unless entry.public?
+
+          entry_name = entry.name
+          owner_name = entry.owner&.name
+
+          label_details = Interface::CompletionItemLabelDetails.new(
+            description: entry.file_name,
+            detail: entry.decorated_parameters,
+          )
+          @response_builder << Interface::CompletionItem.new(
+            label: entry_name,
+            filter_text: entry_name,
+            label_details: label_details,
+            text_edit: Interface::TextEdit.new(range: range, new_text: entry_name),
+            kind: Constant::CompletionItemKind::METHOD,
+            data: {
+              owner_name: owner_name,
+            },
+          )
+        end
+      end
+
       #: (arguments: Array[Prism::Node]) -> Hash[String, Prism::Node]
       def index_call_node_args(arguments:)
         indexed_call_node_args = {}
@@ -79,6 +121,28 @@ def index_call_node_args(arguments:)
         end
         indexed_call_node_args
       end
+
+      # Checks that we're on instance level of a `ActiveRecord::Migration` subclass.
+      #
+      #: -> bool
+      def active_record_migration?
+        nesting_nodes = @node_context.instance_variable_get(:@nesting_nodes).reverse
+        class_node = nesting_nodes.find { |node| node.is_a?(Prism::ClassNode) }
+        return false unless class_node
+
+        superclass = class_node.superclass
+        return false unless superclass.is_a?(Prism::CallNode)
+
+        receiver = superclass.receiver
+        return false unless receiver.is_a?(Prism::ConstantPathNode)
+        return false unless receiver.slice == "ActiveRecord::Migration"
+
+        def_node = nesting_nodes.find { |n| n.is_a?(Prism::DefNode) }
+        return false if def_node.receiver
+
+        true
+      end
+
     end
   end
 end

lib/ruby_lsp/ruby_lsp_rails/runner_client.rb

@@ -133,6 +133,17 @@ def register_server_addon(server_addon_path)
         nil
       end
 
+      #: -> Hash[Symbol, untyped]?
+      def db_configs
+        make_request("db_configs")
+      rescue MessageError
+        log_message(
+          "Ruby LSP Rails failed to get database configurations",
+          type: RubyLsp::Constant::MessageType::ERROR,
+        )
+        nil
+      end
+
       #: (String name) -> Hash[Symbol, untyped]?
       def model(name)
         make_request("model", name: name)

lib/ruby_lsp/ruby_lsp_rails/server.rb

@@ -307,6 +307,10 @@ def execute(request, params)
           with_request_error_handling(request) do
             send_result(resolve_database_info_from_model(params.fetch(:name)))
           end
+        when "db_configs"
+          with_request_error_handling(request) do
+            send_result(resolve_database_configurations)
+          end
         when "association_target"
           with_request_error_handling(request) do
             send_result(resolve_association_target(params))
@@ -423,6 +427,18 @@ def resolve_database_info_from_model(model_name)
         info
       end
 
+      #: -> Hash[Symbol | String, untyped]?
+      def resolve_database_configurations
+        return unless defined?(ActiveRecord)
+
+        ActiveRecord::Base.connection_handler.connection_pools.each_with_object({}) do |pool, hash|
+          hash[pool.db_config.name] = {
+            migrations_paths: Array(pool.migrations_paths),
+            adapter_class: pool.db_config.adapter_class.name,
+          }
+        end
+      end
+
       #: (Hash[Symbol | String, untyped]) -> Hash[Symbol | String, untyped]?
       def resolve_association_target(params)
         const = ActiveSupport::Inflector.safe_constantize(params[:model_name]) # rubocop:disable Sorbet/ConstantsFromStrings

test/ruby_lsp_rails/completion_test.rb

@@ -56,20 +56,61 @@ class CompletionTest < ActiveSupport::TestCase
         assert_equal(0, response.size)
       end
 
+      test "on_call_node_enter provides completion for migration files" do
+        source = <<~RUBY
+          # typed: false
+          class FooBar < ActiveRecord::Migration[8.0]
+            def change
+              create
+            end
+          end
+        RUBY
+        position = { line: 3, character: 10 }
+        uri = Kernel.URI("file://#{dummy_root}/db/migrate/123456789_foo_bar.rb")
+
+        response = with_ready_server(source, uri) do |server|
+          index_gem(server.global_state.index, "activerecord")
+          text_document_completion(server, uri, position)
+        end
+
+        assert_includes response.map(&:label), "create_table"
+      end
+
       private
 
-      def generate_completions_for_source(source, position)
-        with_server(source) do |server, uri|
+      def generate_completions_for_source(source, position, uri = Kernel.URI("file:///fake.rb"))
+        with_ready_server(source, uri) do |server, uri|
+          text_document_completion(server, uri, position)
+        end
+      end
+
+      def with_ready_server(source, uri)
+        with_server(source, uri) do |server|
           sleep(0.1) while RubyLsp::Addon.addons.first.instance_variable_get(:@rails_runner_client).is_a?(NullClient)
 
-          server.process_message(
-            id: 1,
-            method: "textDocument/completion",
-            params: { textDocument: { uri: uri }, position: position },
-          )
+          yield server
+        end
+      end
+
+      def text_document_completion(server, uri, position)
+        server.process_message(
+          id: 1,
+          method: "textDocument/completion",
+          params: { textDocument: { uri: uri }, position: position },
+        )
+
+        result = pop_result(server)
+        result.response
+      end
 
-          result = pop_result(server)
-          result.response
+      def index_gem(index, gem_name)
+        spec = Gem::Specification.find_by_name(gem_name)
+        spec.require_paths.each do |require_path|
+          load_path_entry = File.join(spec.full_gem_path, require_path)
+          Dir.glob(File.join(load_path_entry, "**", "*.rb")).map! do |path|
+            uri = URI::Generic.from_path(path: path, load_path_entry: load_path_entry)
+            index.index_file(uri)
+          end
         end
       end
     end

test/ruby_lsp_rails/runner_client_test.rb

@@ -132,6 +132,18 @@ class RunnerClientTest < ActiveSupport::TestCase
         end
       end
 
+      test "fetches database configurations" do
+        assert_equal(
+          {
+            primary: {
+              migrations_paths: ["db/migrate"],
+              adapter_class: "ActiveRecord::ConnectionAdapters::SQLite3Adapter",
+            },
+          },
+          @client.db_configs,
+        )
+      end
+
       test "delegate notification" do
         @client.expects(:send_notification).with(
           "server_addon/delegate",

test/ruby_lsp_rails/server_test.rb

@@ -124,6 +124,14 @@ def <(other)
     assert_match %r{test/dummy/app/models/country.rb:3$}, location
   end
 
+  test "resolve database configurations" do
+    @server.execute("db_configs", {})
+    migrations_paths = response[:result][:primary][:migrations_paths]
+    adapter_class = response[:result][:primary][:adapter_class]
+    assert_includes migrations_paths, "#{dummy_root}/db/migrate"
+    assert_equal "ActiveRecord::ConnectionAdapters::SQLite3Adapter", adapter_class
+  end
+
   test "route location returns the location for a valid route" do
     @server.execute("route_location", { name: "user_path" })
     location = response[:result][:location]