From 41ffca79dfe09cd6b15a71d652b13531383c5207 Mon Sep 17 00:00:00 2001
From: Keith McHugh <kmcq@github.com>
Date: Tue, 26 Sep 2023 05:42:33 -0700
Subject: [PATCH] Add transform_row and scalar_parser documentation and make
 them easier to use (#118)

* Make it easier to set a `scalar_parser` on a Query

Also fix an issue with `Query#transform_rows`.

Co-authored-by: fionaochs <fionaochs@github.com>

* Add examples of transform_row and scalar_parser to README

Co-authored-by: fionaochs <fionaochs@github.com>

* Allow changing scalar_parser to something new

---------

Co-authored-by: fionaochs <fionaochs@github.com>
---
 README.md                 | 13 ++++++++++++
 lib/trino/client/query.rb |  9 +++++---
 spec/client_spec.rb       | 43 ++++++++++++++++++++++++++++++++++-----
 3 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 27c95107..7a6dadbe 100644
--- a/README.md
+++ b/README.md
@@ -64,6 +64,19 @@ query = client.query("select * from sys.node")
 query_id = query.query_info.query_id
 query.each_row {|row| ... }  # when a thread is processing the query,
 client.kill(query_id)  # another thread / process can kill the query.
+
+# Use Query#transform_row to parse Trino ROW types into Ruby Hashes.
+# You can also set a scalar_parser to parse scalars how you'd like them.
+scalar_parser = -> (data, type) { (type === 'json') ? JSON.parse(data) : data }
+client.query("select * from sys.node") do |q|
+  q.scalar_parser = scalar_parser
+
+  # get query results. it feeds more rows until
+  # query execution finishes:
+  q.each_row {|row|
+    p q.transform_row(row)
+  }
+end
 ```
 
 ## Build models
diff --git a/lib/trino/client/query.rb b/lib/trino/client/query.rb
index 02d6ad39..8b94f325 100644
--- a/lib/trino/client/query.rb
+++ b/lib/trino/client/query.rb
@@ -58,6 +58,8 @@ def self.transform_row(column_value_parsers, row)
       row_object
     end
 
+    attr_accessor :scalar_parser
+
     def initialize(api)
       @api = api
     end
@@ -101,13 +103,14 @@ def columns
     end
 
     def column_value_parsers
-      @column_value_parsers ||= columns.map {|column|
-        ColumnValueParser.new(column)
+      @column_value_parsers ||= {}
+      @column_value_parsers[scalar_parser] ||= columns.map {|column|
+        ColumnValueParser.new(column, scalar_parser)
       }
     end
 
     def transform_rows
-      rows.map(&:transform_row)
+      rows.map { |row| transform_row(row) }
     end
 
     def transform_row(row)
diff --git a/spec/client_spec.rb b/spec/client_spec.rb
index e0bc56f5..41196caa 100644
--- a/spec/client_spec.rb
+++ b/spec/client_spec.rb
@@ -59,29 +59,62 @@
       ]
       client.stub(:run).and_return([columns, rows])
 
+      query = Trino::Client::Query.new(nil)
+      query.stub(:columns).and_return(columns)
+      query.stub(:rows).and_return(rows)
+
+      # For this test, we'll use scalar_parser to add 2 to every integer
+      query.scalar_parser = ->(data, type) { (type == 'integer') ? data + 2 : data }
+
       columns, rows = client.run('fake query')
-      column_value_parsers = columns.map { |column| Trino::Client::ColumnValueParser.new(column) }
-      transformed_rows = rows.map { |row| Trino::Client::Query.transform_row(column_value_parsers, row) }
+      transformed_rows = query.transform_rows
 
       expect(transformed_rows[0]).to eq({
         "animal" => "dog",
-        "score" => 1,
+        "score" => 3,
         "name" => "Lassie",
         "foods" => ["kibble", "peanut butter"],
         "traits" => {
           "breed" => "spaniel",
+          "num_spots" => 4,
+        },
+      })
+
+      expect(transformed_rows[1]).to eq({
+        "animal" => "horse",
+        "score" => 7,
+        "name" => "Mr. Ed",
+        "foods" => ["hay", "sugar cubes"],
+        "traits" => {
+          "breed" => "some horse",
           "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" => 5,
+        "score" => 6,
         "name" => "Mr. Ed",
         "foods" => ["hay", "sugar cubes"],
         "traits" => {
           "breed" => "some horse",
-          "num_spots" => 0,
+          "num_spots" => 1,
         },
       })
     end