[Gem] ES|QL Helper - Adds parser parameter to query

Author: picandocodigo

docs/helpers.asciidoc

@@ -211,11 +211,28 @@ require 'elasticsearch/helpers/esql_helper'
 response = Elasticsearch::Helpers::ESQLHelper.query(client, query)
 
 puts response
-{"duration_ms"=>3.5, "message"=>"Connected to 10.1.0.3", "event.duration"=>3450233, "client.ip"=>"172.21.2.162", "@timestamp"=>"2023-10-23T12:15:03.360Z"}
-{"duration_ms"=>2.8, "message"=>"Connected to 10.1.0.2", "event.duration"=>2764889, "client.ip"=>"172.21.2.113", "@timestamp"=>"2023-10-23T12:27:28.948Z"}
-{"duration_ms"=>1.2, "message"=>"Disconnected", "event.duration"=>1232382, "client.ip"=>"172.21.0.5", "@timestamp"=>"2023-10-23T13:33:34.937Z"}
-{"duration_ms"=>0.7, "message"=>"Connection error", "event.duration"=>725448, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:51:54.732Z"}
-{"duration_ms"=>8.3, "message"=>"Connection error", "event.duration"=>8268153, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:52:55.015Z"}
-{"duration_ms"=>5.0, "message"=>"Connection error", "event.duration"=>5033755, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:53:55.832Z"}
-{"duration_ms"=>1.8, "message"=>"Connected to 10.1.0.1", "event.duration"=>1756467, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:55:01.543Z"}
+[
+  {"duration_ms"=>3.5, "message"=>"Connected to 10.1.0.3", "event.duration"=>3450233, "client.ip"=>"172.21.2.162", "@timestamp"=>"2023-10-23T12:15:03.360Z"}
+  {"duration_ms"=>2.8, "message"=>"Connected to 10.1.0.2", "event.duration"=>2764889, "client.ip"=>"172.21.2.113", "@timestamp"=>"2023-10-23T12:27:28.948Z"}
+  {"duration_ms"=>1.2, "message"=>"Disconnected", "event.duration"=>1232382, "client.ip"=>"172.21.0.5", "@timestamp"=>"2023-10-23T13:33:34.937Z"}
+  {"duration_ms"=>0.7, "message"=>"Connection error", "event.duration"=>725448, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:51:54.732Z"}
+  {"duration_ms"=>8.3, "message"=>"Connection error", "event.duration"=>8268153, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:52:55.015Z"}
+  {"duration_ms"=>5.0, "message"=>"Connection error", "event.duration"=>5033755, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:53:55.832Z"}
+  {"duration_ms"=>1.8, "message"=>"Connected to 10.1.0.1", "event.duration"=>1756467, "client.ip"=>"172.21.3.15", "@timestamp"=>"2023-10-23T13:55:01.543Z"}
+]
+----
+
+Additionally, a block can be specified to work on the response data. Pass in a block to `query` and it will yield each item in the Array of responses.
+
+You could use this for example to convert '@timestamp' into a DateTime object:
+[source,ruby]
+----
+require 'elasticsearch/helpers/esql_helper'
+
+response = esql_helper.query(client, query).each do |r|
+             r['@timestamp'] = DateTime.parse(r['@timestamp'])
+           end
+
+response.first['@timestamp']
+# <DateTime: 2023-10-23T12:15:03+00:00 ((2460241j,44103s,360000000n),+0s,2299161j)>
 ----

elasticsearch/lib/elasticsearch/helpers/esql_helper.rb

@@ -22,19 +22,37 @@ module Helpers
     # @see https://www.elastic.co/guide/en/elasticsearch/reference/current/esql-query-api.html
     #
     module ESQLHelper
-      # By default, the `query` API returns a Hash response with a `columns` key whose value is an
-      # Array of { name: type } Hashes for each column and a `values` key whose value is an Array of
-      # Arrays with the values for each row.
+      # Query helper for ES|QL
+      #
+      # By default, the `esql.query` API returns a Hash response with the following keys:
+      # - `columns` with the value being an Array of { name: type } Hashes for each column
+      # - `values` with the value being an Array of Arrays with the values for each row.
       # This helper function returns an Array of hashes with the columns as keys and the respective
-      # values: { column['name'] => value }
+      # values: { column['name'] => value }.
+      #
+      # @param client [Elasticsearch::Client] an instance of the Client to use for the query.
+      # @param query [Hash, String] The query to be passed to the ES|QL query API.
+      # @param params [Hash] options to pass to the ES|QL query API.
+      # @param convert_datetime [Boolean] Converts datetime values from the response into DateTime
+      # objects in Ruby. Default: true.
+      #
+      # @example Using the ES|QL helper with a parser
+      #     response = Elasticsearch::Helpers::ESQLHelper.query(
+      #                  client,
+      #                  query,
+      #                  parser: { '@timestamp' => Proc.new { |t| DateTime.parse(t) } }
+      #                )
       #
-      def self.query(client, query, params = {})
+      # @see https://www.elastic.co/guide/en/elasticsearch/client/ruby-api/current/Helpers.html#_esql_helper
+      #
+      def self.query(client, query, params = {}, parser: {})
         response = client.esql.query({ body: { query: query }, format: 'json' }.merge(params))
-
         columns = response['columns']
         response['values'].map do |value|
           (value.length - 1).downto(0).map do |index|
-            { columns[index]['name'] => value[index] }
+            key = columns[index]['name']
+            value[index] = parser[key].call value[index] if parser[key]
+            { key => value[index] }
           end.reduce({}, :merge)
         end
       end

elasticsearch/spec/integration/helpers/esql_helper_spec.rb

@@ -21,6 +21,12 @@
   let(:index) { 'esql_helper_test' }
   let(:body) { { size: 12, query: { match_all: {} } } }
   let(:esql_helper) { Elasticsearch::Helpers::ESQLHelper }
+  let(:query) do
+    <<~ESQL
+        FROM #{index}
+        | EVAL duration_ms = ROUND(event.duration / 1000000.0, 1)
+    ESQL
+  end
 
   before do
     client.indices.create(
@@ -58,13 +64,31 @@
   end
 
   it 'returns an ESQL response as a relational key/value object' do
-    query = <<~ESQL
-        FROM #{index}
-        | EVAL duration_ms = ROUND(event.duration / 1000000.0, 1)
-    ESQL
     response = esql_helper.query(client, query)
-
     expect(response.count).to eq 7
     expect(response.first.keys).to eq ['duration_ms', 'message', 'event.duration', 'client.ip', '@timestamp']
+    response.each do |r|
+      expect(r['@timestamp']).to be_a String
+      expect(r['client.ip']).to be_a String
+      expect(r['message']).to be_a String
+      expect(r['event.duration']).to be_a Integer
+    end
+  end
+
+  it 'parses iterated objects when procs are passed in' do
+    require 'ipaddr'
+
+    parser = {
+      '@timestamp' => Proc.new { |t| DateTime.parse(t) },
+      'client.ip' => Proc.new { |i| IPAddr.new(i) },
+      'event.duration' => Proc.new { |d| d.to_s }
+    }
+    response = esql_helper.query(client, query, parser: parser)
+    response.each do |r|
+      expect(r['@timestamp']).to be_a DateTime
+      expect(r['client.ip']).to be_a IPAddr
+      expect(r['message']).to be_a String
+      expect(r['event.duration']).to be_a String
+    end
   end
 end