refactor(ruby): Introduce MessageHandler and Rows classes (#553)

* refactor: add new message handler in ruby wrapper

Author: aakashanandg

spannerlib/wrappers/spannerlib-ruby/Rakefile

@@ -42,4 +42,3 @@ task :compile do
 end
 
 task default: %i[compile spec rubocop]
-

spannerlib/wrappers/spannerlib-ruby/lib/spannerlib/connection.rb

@@ -15,6 +15,7 @@
 # frozen_string_literal: true
 
 require_relative "ffi"
+require_relative "rows"
 
 class Connection
   attr_reader :pool_id, :conn_id
@@ -35,8 +36,7 @@ def write_mutations(mutation_group)
                 else
                   mutation_group.to_s
                 end
-
-    SpannerLib.write_mutations(@pool_id, @conn_id, req_bytes)
+    SpannerLib.write_mutations(@pool_id, @conn_id, req_bytes, proto_klass: Google::Cloud::Spanner::V1::CommitResponse)
   end
 
   # Begin a read/write transaction on this connection. Accepts TransactionOptions proto or bytes.
@@ -52,7 +52,7 @@ def begin_transaction(transaction_options = nil)
 
   # Commit the current transaction. Returns CommitResponse bytes or nil.
   def commit
-    SpannerLib.commit(@pool_id, @conn_id)
+    SpannerLib.commit(@pool_id, @conn_id, proto_klass: Google::Cloud::Spanner::V1::CommitResponse)
   end
 
   # Rollback the current transaction.
@@ -68,7 +68,8 @@ def execute(request)
             else
               request.is_a?(String) ? request : request.to_s
             end
-    SpannerLib.execute(@pool_id, @conn_id, bytes)
+    rows_id = SpannerLib.execute(@pool_id, @conn_id, bytes)
+    SpannerLib::Rows.new(self, rows_id)
   end
 
   # Execute batch DML/DDL request. Returns ExecuteBatchDmlResponse bytes (or nil).
@@ -78,24 +79,8 @@ def execute_batch(request)
             else
               request.is_a?(String) ? request : request.to_s
             end
-    SpannerLib.execute_batch(@pool_id, @conn_id, bytes)
-  end
-
-  # Rows helpers — return raw message bytes (caller should parse them).
-  def metadata(rows_id)
-    SpannerLib.metadata(@pool_id, @conn_id, rows_id)
-  end
-
-  def next_rows(rows_id, num_rows, encoding = 0)
-    SpannerLib.next(@pool_id, @conn_id, rows_id, num_rows, encoding)
-  end
-
-  def result_set_stats(rows_id)
-    SpannerLib.result_set_stats(@pool_id, @conn_id, rows_id)
-  end
 
-  def close_rows(rows_id)
-    SpannerLib.close_rows(@pool_id, @conn_id, rows_id)
+    SpannerLib.execute_batch(@pool_id, @conn_id, bytes, proto_klass: Google::Cloud::Spanner::V1::ExecuteBatchDmlResponse)
   end
 
   # Closes this connection. Any active transaction on the connection is rolled back.

spannerlib/wrappers/spannerlib-ruby/lib/spannerlib/ffi.rb

@@ -23,6 +23,7 @@
 require "google/rpc/status_pb"
 
 require "ffi"
+require_relative "message_handler"
 
 module SpannerLib
   extend FFI::Library
@@ -128,45 +129,25 @@ def self.ensure_release(message)
     end
   end
 
-  def self.handle_object_id_response(message, func_name)
+  def self.handle_object_id_response(message, _func_name)
     ensure_release(message) do
-      if message[:code] != 0
-        error_msg = read_error_message(message)
-        raise "#{func_name} failed with code #{message[:code]}: #{error_msg}"
-      end
-      message[:objectId]
+      MessageHandler.new(message).object_id
     end
   end
 
-  def self.handle_status_response(message, func_name)
+  def self.handle_status_response(message, _func_name)
     ensure_release(message) do
-      if message[:code] != 0
-        error_msg = read_error_message(message)
-        raise "#{func_name} failed with code #{message[:code]}: #{error_msg}"
-      end
+      MessageHandler.new(message).throw_if_error!
     end
     nil
   end
 
-  # rubocop:disable Metrics/MethodLength
-  def self.handle_data_response(message, func_name)
+  def self.handle_data_response(message, _func_name, options = {})
+    proto_klass = options[:proto_klass]
     ensure_release(message) do
-      if message[:code] != 0
-        error_msg = read_error_message(message)
-        raise "#{func_name} failed with code #{message[:code]}: #{error_msg}"
-      end
-
-      len = message[:length]
-      ptr = message[:pointer]
-
-      if len.positive? && !ptr.null?
-        ptr.read_bytes(len)
-      else
-        ""
-      end
+      MessageHandler.new(message).data(proto_klass: proto_klass)
     end
   end
-  # rubocop:enable Metrics/MethodLength
 
   # rubocop:disable Metrics/MethodLength
   def self.read_error_message(message)
@@ -187,10 +168,11 @@ def self.read_error_message(message)
   end
   # rubocop:enable Metrics/MethodLength
 
-  def self.write_mutations(pool_id, conn_id, proto_bytes)
+  def self.write_mutations(pool_id, conn_id, proto_bytes, options = {})
+    proto_klass = options[:proto_klass]
     with_gobytes(proto_bytes) do |gobytes|
       message = WriteMutations(pool_id, conn_id, gobytes)
-      handle_data_response(message, "WriteMutations")
+      handle_data_response(message, "WriteMutations", proto_klass: proto_klass)
     end
   end
 
@@ -201,9 +183,10 @@ def self.begin_transaction(pool_id, conn_id, proto_bytes)
     end
   end
 
-  def self.commit(pool_id, conn_id)
+  def self.commit(pool_id, conn_id, options = {})
+    proto_klass = options[:proto_klass]
     message = Commit(pool_id, conn_id)
-    handle_data_response(message, "Commit")
+    handle_data_response(message, "Commit", proto_klass: proto_klass)
   end
 
   def self.rollback(pool_id, conn_id)
@@ -218,10 +201,11 @@ def self.execute(pool_id, conn_id, proto_bytes)
     end
   end
 
-  def self.execute_batch(pool_id, conn_id, proto_bytes)
+  def self.execute_batch(pool_id, conn_id, proto_bytes, options = {})
+    proto_klass = options[:proto_klass]
     with_gobytes(proto_bytes) do |gobytes|
       message = ExecuteBatch(pool_id, conn_id, gobytes)
-      handle_data_response(message, "ExecuteBatch")
+      handle_data_response(message, "ExecuteBatch", proto_klass: proto_klass)
     end
   end
 

spannerlib/wrappers/spannerlib-ruby/lib/spannerlib/message_handler.rb

@@ -0,0 +1,56 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# frozen_string_literal: true
+
+# lib/spannerlib/message_handler.rb
+
+require "spannerlib/exceptions"
+
+module SpannerLib
+  class MessageHandler
+    def initialize(message)
+      @message = message
+    end
+
+    def object_id
+      throw_if_error!
+      @message[:objectId]
+    end
+
+    # Returns the data payload from the message.
+    # If a proto_klass is provided, it decodes the bytes into a Protobuf object.
+    # Otherwise, it returns the raw bytes as a string.
+    def data(proto_klass: nil)
+      throw_if_error!
+
+      len = @message[:length]
+      ptr = @message[:pointer]
+
+      return (proto_klass ? proto_klass.new : "") unless len.positive? && !ptr.null?
+
+      bytes = ptr.read_string(len)
+
+      proto_klass ? proto_klass.decode(bytes) : bytes
+    end
+
+    def throw_if_error!
+      code = @message[:code]
+      return if code.zero?
+
+      error_msg = SpannerLib.read_error_message(@message)
+      raise SpannerLibException, "Call failed with code #{code}: #{error_msg}"
+    end
+  end
+end

spannerlib/wrappers/spannerlib-ruby/lib/spannerlib/rows.rb

@@ -0,0 +1,67 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# frozen_string_literal: true
+
+module SpannerLib
+  class Rows
+    include Enumerable
+
+    attr_reader :id, :connection
+
+    def initialize(connection, rows_id)
+      @connection = connection
+      @id = rows_id
+      @closed = false
+    end
+
+    def each
+      return enum_for(:each) unless block_given?
+
+      while (row = self.next)
+        yield row
+      end
+    ensure
+      close
+    end
+
+    def next
+      return nil if @closed
+
+      row_data = SpannerLib.next(connection.pool_id, connection.conn_id, id, 1, 0)
+
+      if row_data.nil? || row_data.empty? || (row_data.respond_to?(:values) && row_data.values.empty?)
+        close
+        return nil
+      end
+
+      row_data
+    end
+
+    def metadata
+      SpannerLib.metadata(connection.pool_id, connection.conn_id, id)
+    end
+
+    def result_set_stats
+      SpannerLib.result_set_stats(connection.pool_id, connection.conn_id, id)
+    end
+
+    def close
+      return if @closed
+
+      SpannerLib.close_rows(connection.pool_id, connection.conn_id, id)
+      @closed = true
+    end
+  end
+end