From 035679d85e5cde78b6a9b44f1a73fec38dd61ada Mon Sep 17 00:00:00 2001
From: Keith McHugh <kmcq@github.com>
Date: Thu, 21 Sep 2023 08:48:32 -0700
Subject: [PATCH] Allow changing scalar_parser to something new

---
 lib/trino/client/query.rb |  3 ++-
 spec/client_spec.rb       | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/lib/trino/client/query.rb b/lib/trino/client/query.rb
index f450db78..8b94f325 100644
--- a/lib/trino/client/query.rb
+++ b/lib/trino/client/query.rb
@@ -103,7 +103,8 @@ def columns
     end
 
     def column_value_parsers
-      @column_value_parsers ||= columns.map {|column|
+      @column_value_parsers ||= {}
+      @column_value_parsers[scalar_parser] ||= columns.map {|column|
         ColumnValueParser.new(column, scalar_parser)
       }
     end
diff --git a/spec/client_spec.rb b/spec/client_spec.rb
index 89fa0e7d..41196caa 100644
--- a/spec/client_spec.rb
+++ b/spec/client_spec.rb
@@ -90,6 +90,33 @@
           "num_spots" => 2,
         },
       })
+
+      # And to show that you can change the scalar_parser, now we only add 1 to each integer.
+      query.scalar_parser = ->(data, type) { (type == 'integer') ? data + 1 : data }
+
+      transformed_rows = query.transform_rows
+
+      expect(transformed_rows[0]).to eq({
+        "animal" => "dog",
+        "score" => 2,
+        "name" => "Lassie",
+        "foods" => ["kibble", "peanut butter"],
+        "traits" => {
+          "breed" => "spaniel",
+          "num_spots" => 3,
+        },
+      })
+
+      expect(transformed_rows[1]).to eq({
+        "animal" => "horse",
+        "score" => 6,
+        "name" => "Mr. Ed",
+        "foods" => ["hay", "sugar cubes"],
+        "traits" => {
+          "breed" => "some horse",
+          "num_spots" => 1,
+        },
+      })
     end
 
     it 'empty results' do