Move ESQL executor retriable method to private. Unit tests are organized in a better way.

Author: mashhurs

lib/logstash/inputs/elasticsearch/esql.rb

@@ -19,20 +19,6 @@ def initialize(client, plugin)
           @query = @plugin_params["query"]
         end
 
-        # Execute a retryable operation with proper error handling
-        # @param job_name [String] Name of the job for logging purposes
-        # @yield The block to execute
-        # @return [Boolean] true if successful, false otherwise
-        def retryable(job_name, &block)
-          stud_try = ::LogStash::Helpers::LoggableTry.new(logger, job_name)
-          stud_try.try((@retries + 1).times) { yield }
-        rescue => e
-          error_details = {:message => e.message, :cause => e.cause}
-          error_details[:backtrace] = e.backtrace if logger.debug?
-          logger.error("#{job_name} failed with ", error_details)
-          false
-        end
-
         # Execute the ESQL query and process results
         # @param output_queue [Queue] The queue to push processed events to
         def do_run(output_queue)
@@ -43,16 +29,28 @@ def do_run(output_queue)
           # retriable already printed error details
           return if response == false
 
-          puts "response class: #{response.class}"
-          puts "response: #{response.inspect}"
-          unless response&.headers&.dig("warning")
+          if response&.headers&.dig("warning")
             logger.warn("ES|QL executor received warning", {:message => response.headers["warning"]})
           end
           if response['values'] && response['columns']
             process_response(response['values'], response['columns'], output_queue)
           end
         end
 
+        # Execute a retryable operation with proper error handling
+        # @param job_name [String] Name of the job for logging purposes
+        # @yield The block to execute
+        # @return [Boolean] true if successful, false otherwise
+        def retryable(job_name, &block)
+          stud_try = ::LogStash::Helpers::LoggableTry.new(logger, job_name)
+          stud_try.try((@retries + 1).times) { yield }
+        rescue => e
+          error_details = {:message => e.message, :cause => e.cause}
+          error_details[:backtrace] = e.backtrace if logger.debug?
+          logger.error("#{job_name} failed with ", error_details)
+          false
+        end
+
         private
 
         # Process the ESQL response and push events to the output queue

spec/inputs/elasticsearch_esql_spec.rb

@@ -2,114 +2,95 @@
 require "logstash/devutils/rspec/spec_helper"
 require "logstash/inputs/elasticsearch"
 require "elasticsearch"
-require 'logstash/plugin_mixins/ecs_compatibility_support/spec_helper'
 
-describe LogStash::Inputs::Elasticsearch, :ecs_compatibility_support do
-  let(:plugin) { described_class.new(config) }
+describe LogStash::Inputs::Elasticsearch::Esql do
   let(:client) { instance_double(Elasticsearch::Client) }
-  let(:queue) { Queue.new }
-  let(:cluster_info) { { "version" => { "number" => "8.11.0", "build_flavor" => "default" } } }
-  
-  let(:config) do
+  let(:esql_client) { double("esql-client") }
+  let(:plugin) { instance_double(LogStash::Inputs::Elasticsearch, params: plugin_config) }
+  let(:plugin_config) do
     {
       "query" => "FROM test-index | STATS count() BY field",
-      "response_type" => "esql",
       "retries" => 3
     }
   end
+  let(:esql_executor) { described_class.new(client, plugin) }
 
-  describe "#initialize" do
+  describe "when initializes" do
     it "sets up the ESQL client with correct parameters" do
-      expect(plugin.instance_variable_get(:@query)).to eq(config["query"])
-      expect(plugin.instance_variable_get(:@response_type)).to eq(config["response_type"])
-      expect(plugin.instance_variable_get(:@retries)).to eq(config["retries"])
+      expect(esql_executor.instance_variable_get(:@query)).to eq(plugin_config["query"])
+      expect(esql_executor.instance_variable_get(:@retries)).to eq(plugin_config["retries"])
     end
   end
 
-  describe "#register" do
-    before(:each) do
-      Elasticsearch::Client.send(:define_method, :ping) { }
-      allow_any_instance_of(Elasticsearch::Client).to receive(:info).and_return(cluster_info)
+  describe "when faces error while retrying" do
+    it "retries the given block the specified number of times" do
+      attempts = 0
+      result = esql_executor.retryable("Test Job") do
+        attempts += 1
+        raise StandardError if attempts < 3
+        "success"
+      end
+      expect(attempts).to eq(3)
+      expect(result).to eq("success")
     end
-    it "creates ES|QL executor" do
-      plugin.register
-      expect(plugin.instance_variable_get(:@query_executor)).to be_an_instance_of(LogStash::Inputs::Elasticsearch::Esql)
+
+    it "returns false if the block fails all attempts" do
+      result = esql_executor.retryable("Test Job") do
+        raise StandardError
+      end
+      expect(result).to eq(false)
     end
   end
 
-  describe "#validation" do
+  describe "when executing chain of processes" do
+    let(:output_queue) { Queue.new }
+    let(:response) { { 'values' => [%w[foo bar]], 'columns' => [{ 'name' => 'id'}, { 'name' => 'val'}] } }
 
-    describe "LS version" do
-      context "when compatible" do
-        before(:each) do
-          stub_const("LogStash::VERSION", "8.11.0")
-        end
-
-        it "does not raise an error" do
-          expect { plugin.send(:validate_ls_version_for_esql_support!) }.not_to raise_error
-        end
-      end
-
-      context "when incompatible" do
-        before(:each) do
-          stub_const("LOGSTASH_VERSION", "8.10.0")
-        end
+    before do
+      allow(esql_executor).to receive(:retryable).and_yield
+      allow(client).to receive_message_chain(:esql, :query).and_return(response)
+      allow(plugin).to receive(:decorate_and_push_to_queue)
+    end
 
-        it "raises a runtime error" do
-          expect { plugin.send(:validate_ls_version_for_esql_support!) }
-            .to raise_error(RuntimeError, /Current version of Logstash does not include Elasticsearch client which supports ES|QL. Please upgrade Logstash to at least 8.17.4/)
-        end
-      end
+    it "executes the ESQL query and processes the results" do
+      allow(response).to receive(:headers).and_return({})
+      esql_executor.do_run(output_queue)
+      expect(plugin).to have_received(:decorate_and_push_to_queue).with(output_queue, {'id' => 'foo', 'val' => 'bar'})
     end
 
-    describe "ES version" do
-      before(:each) do
-        allow(plugin).to receive(:es_version).and_return("8.10.5")
-      end
+    it "logs a warning if the response contains a warning header" do
+      allow(response).to receive(:headers).and_return({"warning" => "some warning"})
+      expect(esql_executor.logger).to receive(:warn).with("ES|QL executor received warning", {:message => "some warning"})
+      esql_executor.do_run(output_queue)
+    end
 
-      context "when incompatible" do
-        it "raises a runtime error" do
-          expect { plugin.send(:validate_es_for_esql_support!) }
-            .to raise_error(RuntimeError, /Connected Elasticsearch 8.10.5 version does not supports ES|QL. Please upgrade it./)
-        end
-      end
+    it "does not log a warning if the response does not contain a warning header" do
+      allow(response).to receive(:headers).and_return({})
+      expect(esql_executor.logger).not_to receive(:warn)
+      esql_executor.do_run(output_queue)
     end
+  end
 
-    describe "ES|QL query"  do
-      context "when query is valid" do
-        it "does not raise an error" do
-          expect { plugin.send(:validate_esql_query!) }.not_to raise_error
-        end
-      end
 
-      context "when query is empty" do
-        let(:config) do
-          {
-            "query" => " "
-          }
-        end
+  describe "when starts processing the response" do
+    let(:output_queue) { Queue.new }
+    let(:values) { [%w[foo bar]] }
+    let(:columns) { [{'name' => 'id'}, {'name' => 'val'}] }
 
-        # TODO: make shared spec
-        it "raises a configuration error" do
-          expect { plugin.send(:validate_esql_query!) }
-            .to raise_error(LogStash::ConfigurationError, /`query` cannot be empty/)
-        end
-      end
+    it "processes the ESQL response and pushes events to the output queue" do
+      allow(plugin).to receive(:decorate_and_push_to_queue)
+      esql_executor.send(:process_response, values, columns, output_queue)
+      expect(plugin).to have_received(:decorate_and_push_to_queue).with(output_queue, {'id' => 'foo', 'val' => 'bar'})
+    end
+  end
 
-      context "when query doesn't align with ES syntax" do
-        let(:config) do
-          {
-            "query" => "RANDOM query"
-          }
-        end
+  describe "when maps column and values" do
+    let(:columns) { [{'name' => 'id'}, {'name' => 'val'}] }
+    let(:values) { %w[foo bar] }
 
-        it "raises a configuration error" do
-          source_commands = %w[FROM ROW SHOW]
-          expect { plugin.send(:validate_esql_query!) }
-            .to raise_error(LogStash::ConfigurationError, "`query` needs to start with any of #{source_commands}")
-        end
-      end
+    it "maps column names to their corresponding values" do
+      result = esql_executor.send(:map_column_and_values, columns, values)
+      expect(result).to eq({'id' => 'foo', 'val' => 'bar'})
     end
   end
-
-end 
+end

spec/inputs/elasticsearch_spec.rb

@@ -1370,4 +1370,110 @@ def extract_transport(client) # on 7.x client.transport is a ES::Transport::Clie
     client.transport.respond_to?(:transport) ? client.transport.transport : client.transport
   end
 
+  context "#ESQL" do
+    let(:config) do
+      {
+        "query" => "FROM test-index | STATS count() BY field",
+        "response_type" => "esql",
+        "retries" => 3
+      }
+    end
+    let(:es_version) { "8.11.0" }
+
+    before(:each) do
+      # ES|QL supported |elasticsearch-ruby v8 client is available from 8.17.4
+      # this is a safeguard to let tests succeed in <8.17.4 versions, see validation test cases for unsupported behavior
+      stub_const("LOGSTASH_VERSION", "8.17.4")
+    end
+
+    describe "#initialize" do
+      it "sets up the ESQL client with correct parameters" do
+        expect(plugin.instance_variable_get(:@query)).to eq(config["query"])
+        expect(plugin.instance_variable_get(:@response_type)).to eq(config["response_type"])
+        expect(plugin.instance_variable_get(:@retries)).to eq(config["retries"])
+      end
+    end
+
+    describe "#register" do
+      before(:each) do
+        Elasticsearch::Client.send(:define_method, :ping) { }
+        allow_any_instance_of(Elasticsearch::Client).to receive(:info).and_return(cluster_info)
+      end
+      it "creates ES|QL executor" do
+        plugin.register
+        expect(plugin.instance_variable_get(:@query_executor)).to be_an_instance_of(LogStash::Inputs::Elasticsearch::Esql)
+      end
+    end
+
+    describe "#validation" do
+
+      describe "LS version" do
+        context "when compatible" do
+
+          it "does not raise an error" do
+            expect { plugin.send(:validate_ls_version_for_esql_support!) }.not_to raise_error
+          end
+        end
+
+        context "when incompatible" do
+          before(:each) do
+            stub_const("LOGSTASH_VERSION", "8.10.0")
+          end
+
+          it "raises a runtime error" do
+            expect { plugin.send(:validate_ls_version_for_esql_support!) }
+              .to raise_error(RuntimeError, /Current version of Logstash does not include Elasticsearch client which supports ES|QL. Please upgrade Logstash to at least 8.17.4/)
+          end
+        end
+      end
+
+      describe "ES version" do
+        before(:each) do
+          allow(plugin).to receive(:es_version).and_return("8.10.5")
+        end
+
+        context "when incompatible" do
+          it "raises a runtime error" do
+            expect { plugin.send(:validate_es_for_esql_support!) }
+              .to raise_error(RuntimeError, /Connected Elasticsearch 8.10.5 version does not supports ES|QL. Please upgrade it./)
+          end
+        end
+      end
+
+      describe "ES|QL query"  do
+        context "when query is valid" do
+          it "does not raise an error" do
+            expect { plugin.send(:validate_esql_query!) }.not_to raise_error
+          end
+        end
+
+        context "when query is empty" do
+          let(:config) do
+            {
+              "query" => " "
+            }
+          end
+
+          it "raises a configuration error" do
+            expect { plugin.send(:validate_esql_query!) }
+              .to raise_error(LogStash::ConfigurationError, /`query` cannot be empty/)
+          end
+        end
+
+        context "when query doesn't align with ES syntax" do
+          let(:config) do
+            {
+              "query" => "RANDOM query"
+            }
+          end
+
+          it "raises a configuration error" do
+            source_commands = %w[FROM ROW SHOW]
+            expect { plugin.send(:validate_esql_query!) }
+              .to raise_error(LogStash::ConfigurationError, "`query` needs to start with any of #{source_commands}")
+          end
+        end
+      end
+    end
+  end
 end