Add completion for AR .where queries using AR model's column names

Author: ChallaHalla

lib/ruby_lsp/ruby_lsp_rails/addon.rb

@@ -13,6 +13,7 @@
 require_relative "code_lens"
 require_relative "document_symbol"
 require_relative "definition"
+require_relative "completion"
 require_relative "indexing_enhancement"
 
 module RubyLsp
@@ -119,6 +120,18 @@ def create_definition_listener(response_builder, uri, node_context, dispatcher)
         Definition.new(@rails_runner_client, response_builder, node_context, index, dispatcher)
       end
 
+      sig do
+        override.params(
+          response_builder: ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem],
+          node_context: NodeContext,
+          dispatcher: Prism::Dispatcher,
+          uri: URI::Generic,
+        ).void
+      end
+      def create_completion_listener(response_builder, node_context, dispatcher, uri)
+        Completion.new(@rails_runner_client, response_builder, node_context, dispatcher, uri)
+      end
+
       sig { params(changes: T::Array[{ uri: String, type: Integer }]).void }
       def workspace_did_change_watched_files(changes)
         if changes.any? { |c| c[:uri].end_with?("db/schema.rb") || c[:uri].end_with?("structure.sql") }

lib/ruby_lsp/ruby_lsp_rails/completion.rb

@@ -0,0 +1,91 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  module Rails
+    class Completion
+      extend T::Sig
+      include Requests::Support::Common
+
+      sig do
+        override.params(
+          client: RunnerClient,
+          response_builder: ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem],
+          node_context: NodeContext,
+          dispatcher: Prism::Dispatcher,
+          uri: URI::Generic,
+        ).void
+      end
+      def initialize(client, response_builder, node_context, dispatcher, uri)
+        @response_builder = response_builder
+        @client = client
+        @node_context = node_context
+        dispatcher.register(
+          self,
+          :on_call_node_enter,
+        )
+      end
+
+      sig { params(node: Prism::CallNode).void }
+      def on_call_node_enter(node)
+        if @node_context.call_node && @node_context.call_node&.name == :where
+          handle_active_record_where_completions(node)
+        end
+      end
+
+      private
+
+      sig { params(node: Prism::CallNode).void }
+      def handle_active_record_where_completions(node)
+        receiver = T.must(@node_context.call_node).receiver
+        return if receiver.nil?
+        return unless receiver.is_a?(Prism::ConstantReadNode)
+
+        resolved_class = @client.model(receiver.name.to_s)
+        return if resolved_class.nil?
+
+        arguments = T.must(@node_context.call_node).arguments&.arguments
+        indexed_call_node_args = T.let({}, T::Hash[String, Prism::Node])
+
+        if arguments&.is_a?(Array)
+          indexed_call_node_args = index_call_node_args(arguments: arguments)
+          binding.irb
+          return if indexed_call_node_args.values.any? { |v| v == node }
+        end
+
+        resolved_class[:columns].each do |column|
+          next unless column[0].start_with?(node.name.to_s)
+          next if indexed_call_node_args.key?(column[0])
+
+          @response_builder << Interface::CompletionItem.new(
+            label: column[0],
+            filter_text: column[0],
+            label_details: Interface::CompletionItemLabelDetails.new(
+              description: "Filter #{receiver.name} records by #{column[0]}",
+            ),
+            text_edit: Interface::TextEdit.new(range: range_from_location(node.location), new_text: "#{column[0]}: "),
+            kind: Constant::CompletionItemKind::FIELD,
+          )
+        end
+      end
+
+      sig { params(arguments: T::Array[Prism::Node]).returns(T::Hash[String, Prism::Node]) }
+      def index_call_node_args(arguments:)
+        indexed_call_node_args = {}
+        arguments.each do |argument|
+          next unless argument.is_a?(Prism::KeywordHashNode)
+
+          argument.elements.each do |e|
+            next unless e.is_a?(Prism::AssocNode)
+
+            key = e.key
+            if key.is_a?(Prism::SymbolNode)
+              indexed_call_node_args[key.value] = e.value
+            end
+          end
+        end
+        indexed_call_node_args
+      end
+    end
+  end
+end

test/ruby_lsp_rails/completion_test.rb

@@ -0,0 +1,59 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module RubyLsp
+  module Rails
+    class CompletionTest < ActiveSupport::TestCase
+      test "Does not suggest column if it already exists within .where as an arg and parantheses are not closed" do
+        response = generate_completions_for_source(<<~RUBY, { line: 1, character: 37 })
+          # typed: false
+          User.where(id:, first_name:, first_na
+        RUBY
+
+        assert_equal(0, response.size)
+      end
+
+      test "Provides suggestions when typing column name partially" do
+        response = generate_completions_for_source(<<~RUBY, { line: 1, character: 17 })
+          # typed: false
+          User.where(first_
+        RUBY
+
+        assert_equal(1, response.size)
+        assert_equal("first_name", response[0].label)
+        assert_equal("first_name", response[0].filter_text)
+        assert_equal(11, response[0].text_edit.range.start.character)
+        assert_equal(1, response[0].text_edit.range.start.line)
+        assert_equal(17, response[0].text_edit.range.end.character)
+        assert_equal(1, response[0].text_edit.range.end.line)
+      end
+
+      test "Does not provide suggestion when typing argument value" do
+        response = generate_completions_for_source(<<~RUBY, { line: 1, character: 20 })
+          # typed: false
+          User.where(id: creat
+        RUBY
+        assert_equal(0, response.size)
+      end
+
+      private
+
+      def generate_completions_for_source(source, position)
+        with_server(source) do |server, uri|
+          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 },
+          )
+
+          result = pop_result(server)
+          result.response
+        end
+      end
+    end
+  end
+end