From 9a7b0d8abb2e610697129974b7e437d028215eb3 Mon Sep 17 00:00:00 2001
From: Danil Nurgaliev <danil.nurgaliev@toptal.com>
Date: Thu, 28 Dec 2023 17:26:38 +0300
Subject: [PATCH] Add before request filter

---
 CHANGELOG.md                      |  1 +
 README.md                         | 17 +++++++++++++++++
 lib/chewy.rb                      |  8 ++------
 lib/chewy/config.rb               |  4 +++-
 lib/chewy/elastic_client.rb       | 31 +++++++++++++++++++++++++++++++
 spec/chewy/elastic_client_spec.rb | 26 ++++++++++++++++++++++++++
 spec/chewy_spec.rb                | 11 +++++++----
 7 files changed, 87 insertions(+), 11 deletions(-)
 create mode 100644 lib/chewy/elastic_client.rb
 create mode 100644 spec/chewy/elastic_client_spec.rb

diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb2863c9..c4bc165c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### New Features
 
 * [#894](https://github.com/toptal/chewy/pull/894): Way of cleaning redis from artifacts left by `delayed_sidekiq` strategy which could potentially cause flaky tests. ([@Drowze](https://github.com/Drowze))
+* [#919](https://github.com/toptal/chewy/pull/919): Add pre-request filter ([@konalegi][https://github.com/konalegi])
 
 ### Changes
 
diff --git a/README.md b/README.md
index dbfe974c..8c0baec8 100644
--- a/README.md
+++ b/README.md
@@ -1281,6 +1281,23 @@ If you use `DatabaseCleaner` in your tests with [the `transaction` strategy](htt
 Chewy.use_after_commit_callbacks = !Rails.env.test?
 ```
 
+### Pre-request Filter
+
+Should you need to inspect the query prior to it being dispatched to ElasticSearch during any queries, you can use the `before_es_request_filter`. `before_es_request_filter` is a callable object, as demonstrated below:
+
+```ruby
+Chewy.before_es_request_filter = -> (method_name, args, kw_args) { ... }
+```
+
+While using the `before_es_request_filter`, please consider the following:
+
+* `before_es_request_filter` acts as a simple proxy before any request made via the    `ElasticSearch::Client`. The arguments passed to this filter include:
+  * `method_name` -  The name of the method being called. Examples are search, count, bulk and etc.
+  * `args` and `kw_args` - These are the positional arguments provided in the method call.
+* The operation is synchronous, so avoid executing any heavy or time-consuming operations within the filter to prevent performance degradation.
+* The return value of the proc is disregarded. This filter is intended for inspection or modification of the query rather than generating a response.
+* Any exception raised inside the callback will propagate upward and halt the execution of the query. It is essential to handle potential errors adequately to ensure the stability of your search functionality.
+
 ## Contributing
 
 1. Fork it (http://github.com/toptal/chewy/fork)
diff --git a/lib/chewy.rb b/lib/chewy.rb
index cac44772..919f9bac 100644
--- a/lib/chewy.rb
+++ b/lib/chewy.rb
@@ -49,6 +49,7 @@ def try_require(path)
 require 'chewy/fields/root'
 require 'chewy/journal'
 require 'chewy/railtie' if defined?(Rails::Railtie)
+require 'chewy/elastic_client'
 
 ActiveSupport.on_load(:active_record) do
   include Chewy::Index::Observe::ActiveRecordMethods
@@ -97,12 +98,7 @@ def derive_name(index_name)
     # Main elasticsearch-ruby client instance
     #
     def client
-      Chewy.current[:chewy_client] ||= begin
-        client_configuration = configuration.deep_dup
-        client_configuration.delete(:prefix) # used by Chewy, not relevant to Elasticsearch::Client
-        block = client_configuration[:transport_options].try(:delete, :proc)
-        ::Elasticsearch::Client.new(client_configuration, &block)
-      end
+      Chewy.current[:chewy_client] ||= Chewy::ElasticClient.new
     end
 
     # Sends wait_for_status request to ElasticSearch with status
diff --git a/lib/chewy/config.rb b/lib/chewy/config.rb
index 4d7a695f..65c60b36 100644
--- a/lib/chewy/config.rb
+++ b/lib/chewy/config.rb
@@ -38,7 +38,9 @@ class Config
                   # for type mappings like `_all`.
                   :default_root_options,
                   # Default field type for any field in any Chewy type. Defaults to 'text'.
-                  :default_field_type
+                  :default_field_type,
+                  # Callback called on each search request to be done into ES
+                  :before_es_request_filter
 
     attr_reader :transport_logger, :transport_tracer,
                 # Chewy search request DSL base class, used by every index.
diff --git a/lib/chewy/elastic_client.rb b/lib/chewy/elastic_client.rb
new file mode 100644
index 00000000..d46f9ef0
--- /dev/null
+++ b/lib/chewy/elastic_client.rb
@@ -0,0 +1,31 @@
+module Chewy
+  # Replacement for Chewy.client
+  class ElasticClient
+    def self.build_es_client(configuration = Chewy.configuration)
+      client_configuration = configuration.deep_dup
+      client_configuration.delete(:prefix) # used by Chewy, not relevant to Elasticsearch::Client
+      block = client_configuration[:transport_options].try(:delete, :proc)
+      ::Elasticsearch::Client.new(client_configuration, &block)
+    end
+
+    def initialize(elastic_client = self.class.build_es_client)
+      @elastic_client = elastic_client
+    end
+
+    private
+
+    def method_missing(name, *args, **kwargs, &block)
+      inspect_payload(name, args, kwargs)
+
+      @elastic_client.__send__(name, *args, **kwargs, &block)
+    end
+
+    def respond_to_missing?(name, _include_private = false)
+      @elastic_client.respond_to?(name) || super
+    end
+
+    def inspect_payload(name, args, kwargs)
+      Chewy.config.before_es_request_filter&.call(name, args, kwargs)
+    end
+  end
+end
diff --git a/spec/chewy/elastic_client_spec.rb b/spec/chewy/elastic_client_spec.rb
new file mode 100644
index 00000000..79d3946b
--- /dev/null
+++ b/spec/chewy/elastic_client_spec.rb
@@ -0,0 +1,26 @@
+require 'spec_helper'
+
+describe Chewy::ElasticClient do
+  describe 'payload inspection' do
+    let(:filter) { instance_double('Proc') }
+    let!(:filter_previous_value) { Chewy.before_es_request_filter }
+
+    before do
+      Chewy.massacre
+      stub_index(:products) do
+        field :id, type: :integer
+      end
+      ProductsIndex.create
+      Chewy.before_es_request_filter = filter
+    end
+
+    after do
+      Chewy.before_es_request_filter = filter_previous_value
+    end
+
+    it 'call filter with the request body' do
+      expect(filter).to receive(:call).with(:search, [{body: {size: 0}, index: ['products']}], {})
+      Chewy.client.search({index: ['products'], body: {size: 0}}).to_a
+    end
+  end
+end
diff --git a/spec/chewy_spec.rb b/spec/chewy_spec.rb
index 3023b71e..0c3e0d0e 100644
--- a/spec/chewy_spec.rb
+++ b/spec/chewy_spec.rb
@@ -57,18 +57,21 @@
 
     before do
       Chewy.current[:chewy_client] = nil
-      allow(Chewy).to receive_messages(configuration: {transport_options: {proc: faraday_block}})
+    end
+
+    specify do
+      expect(Chewy).to receive_messages(configuration: {transport_options: {proc: faraday_block}})
 
-      allow(Elasticsearch::Client).to receive(:new).with(expected_client_config) do |*_args, &passed_block|
+      expect(Elasticsearch::Client).to receive(:new).with(expected_client_config) do |*_args, &passed_block|
         # RSpec's `with(..., &block)` was used previously, but doesn't actually do
         # any verification of the passed block (even of its presence).
         expect(passed_block.source_location).to eq(faraday_block.source_location)
 
         mock_client
       end
-    end
 
-    its(:client) { is_expected.to eq(mock_client) }
+      expect(Chewy.client).to be_a(Chewy::ElasticClient)
+    end
 
     after { Chewy.current[:chewy_client] = initial_client }
   end