From 949da8bdcc77eb1e7fb43d18b6cc6bc83c96689f Mon Sep 17 00:00:00 2001
From: Stan Lo <stan.lo@shopify.com>
Date: Mon, 5 Aug 2024 15:09:01 +0100
Subject: [PATCH] Make indexing enhancement error-tolerant

This change rescues errors raised when indexing with enhancements
and logs them to stderr when the file's indexing is finished.
---
 .../lib/ruby_indexer/declaration_listener.rb  | 10 +++++-
 .../lib/ruby_indexer/enhancement.rb           |  2 ++
 lib/ruby_indexer/lib/ruby_indexer/index.rb    | 16 ++++++++-
 lib/ruby_indexer/test/enhancements_test.rb    | 34 +++++++++++++++++++
 4 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
index f01dab09a..327c96fee 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
@@ -8,6 +8,9 @@ class DeclarationListener
     OBJECT_NESTING = T.let(["Object"].freeze, T::Array[String])
     BASIC_OBJECT_NESTING = T.let(["BasicObject"].freeze, T::Array[String])
 
+    sig { returns(T::Array[String]) }
+    attr_reader :indexing_errors
+
     sig do
       params(
         index: Index,
@@ -36,6 +39,7 @@ def initialize(index, dispatcher, parse_result, file_path, enhancements: [])
 
       # A stack of namespace entries that represent where we currently are. Used to properly assign methods to an owner
       @owner_stack = T.let([], T::Array[Entry::Namespace])
+      @indexing_errors = T.let([], T::Array[String])
 
       dispatcher.register(
         self,
@@ -287,7 +291,11 @@ def on_call_node_enter(node)
         @visibility_stack.push(Entry::Visibility::PRIVATE)
       end
 
-      @enhancements.each { |aug| aug.on_call_node(@index, @owner_stack.last, node, @file_path) }
+      @enhancements.each do |enhancement|
+        enhancement.on_call_node(@index, @owner_stack.last, node, @file_path)
+      rescue StandardError
+        @indexing_errors << "Error occurred when indexing #{@file_path} with '#{enhancement.class.name}' enhancement"
+      end
     end
 
     sig { params(node: Prism::CallNode).void }
diff --git a/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb b/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb
index 59336c35a..fc06c8db1 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb
@@ -8,6 +8,8 @@ module Enhancement
 
     interface!
 
+    requires_ancestor { Object }
+
     # The `on_extend` indexing enhancement is invoked whenever an extend is encountered in the code. It can be used to
     # register for an included callback, similar to what `ActiveSupport::Concern` does in order to auto-extend the
     # `ClassMethods` modules
diff --git a/lib/ruby_indexer/lib/ruby_indexer/index.rb b/lib/ruby_indexer/lib/ruby_indexer/index.rb
index 803a3913e..c556e10a2 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/index.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/index.rb
@@ -317,11 +317,25 @@ def index_single(indexable_path, source = nil)
       dispatcher = Prism::Dispatcher.new
 
       result = Prism.parse(content)
-      DeclarationListener.new(self, dispatcher, result, indexable_path.full_path, enhancements: @enhancements)
+      listener = DeclarationListener.new(
+        self,
+        dispatcher,
+        result,
+        indexable_path.full_path,
+        enhancements: @enhancements,
+      )
       dispatcher.dispatch(result.value)
 
+      indexing_errors = listener.indexing_errors.uniq
+
       require_path = indexable_path.require_path
       @require_paths_tree.insert(require_path, indexable_path) if require_path
+
+      if indexing_errors.any?
+        indexing_errors.each do |error|
+          $stderr.puts error
+        end
+      end
     rescue Errno::EISDIR, Errno::ENOENT
       # If `path` is a directory, just ignore it and continue indexing. If the file doesn't exist, then we also ignore
       # it
diff --git a/lib/ruby_indexer/test/enhancements_test.rb b/lib/ruby_indexer/test/enhancements_test.rb
index 2e366d896..ba3101205 100644
--- a/lib/ruby_indexer/test/enhancements_test.rb
+++ b/lib/ruby_indexer/test/enhancements_test.rb
@@ -159,5 +159,39 @@ class User < ActiveRecord::Base
 
       assert_entry("posts", Entry::Method, "/fake/path/foo.rb:23-11:23-17")
     end
+
+    def test_error_handling_in_enhancement
+      enhancement_class = Class.new do
+        include Enhancement
+
+        def on_call_node(index, owner, node, file_path)
+          raise "Error"
+        end
+
+        class << self
+          def name
+            "TestEnhancement"
+          end
+        end
+      end
+
+      @index.register_enhancement(enhancement_class.new)
+
+      _stdout, stderr = capture_io do
+        index(<<~RUBY)
+          module ActiveSupport
+            module Concern
+              def self.extended(base)
+                base.class_eval("def new_method(a); end")
+              end
+            end
+          end
+        RUBY
+      end
+
+      assert_match(%r{Error occurred when indexing /fake/path/foo\.rb with 'TestEnhancement' enhancement}, stderr)
+      # The module should still be indexed
+      assert_entry("ActiveSupport::Concern", Entry::Module, "/fake/path/foo.rb:1-2:5-5")
+    end
   end
 end