chore(ruby): support multi-statement SQL strings in Ruby wrapper (googleapis#655)

* chore(ruby): support multi-statement SQL strings in Ruby wrapper

Author: aakashanandg

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

@@ -116,6 +116,7 @@ class Message < FFI::Struct
   attach_function :ResultSetStats, %i[int64 int64 int64], Message.by_value
   attach_function :CloseRows, %i[int64 int64 int64], Message.by_value
   attach_function :Release, [:int64], :void
+  attach_function :NextResultSet, %i[int64 int64 int64], Message.by_value
 
   # --- Ruby-friendly Wrappers ---
 
@@ -272,6 +273,12 @@ def self.close_rows(pool_id, conn_id, rows_id)
     message = CloseRows(pool_id, conn_id, rows_id)
     handle_status_response(message, "CloseRows")
   end
+
+  def self.next_result_set(pool_id, conn_id, rows_id)
+    message = NextResultSet(pool_id, conn_id, rows_id)
+    # This returns Metadata for the next result set, or nil if no more sets exist.
+    handle_data_response(message, "NextResultSet")
+  end
 end
 
 # rubocop:enable Metrics/ModuleLength

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

@@ -14,6 +14,8 @@
 
 # frozen_string_literal: true
 
+require "google/spanner/v1/result_set_pb"
+
 module SpannerLib
   class Rows
     include Enumerable
@@ -24,6 +26,7 @@ def initialize(connection, rows_id)
       @connection = connection
       @id = rows_id
       @closed = false
+      @metadata = nil
     end
 
     def each
@@ -41,26 +44,46 @@ def next
 
       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?)
+      return nil if row_data.nil? || row_data.empty?
+
+      row_data
+    end
+
+    def next_result_set
+      return nil if @closed
+
+      @metadata = nil
+
+      res = SpannerLib.next_result_set(connection.pool_id, connection.conn_id, id)
+
+      if res.nil? || res.empty?
         close
         return nil
       end
-
-      row_data
+      @metadata = Google::Cloud::Spanner::V1::ResultSetMetadata.decode(res)
     end
 
     def metadata
-      SpannerLib.metadata(connection.pool_id, connection.conn_id, id)
+      return @metadata if @metadata
+
+      raw = SpannerLib.metadata(connection.pool_id, connection.conn_id, id)
+      return nil if raw.nil? || raw.empty?
+
+      @metadata = Google::Cloud::Spanner::V1::ResultSetMetadata.decode(raw)
     end
 
     def result_set_stats
-      SpannerLib.result_set_stats(connection.pool_id, connection.conn_id, id)
+      raw = SpannerLib.result_set_stats(connection.pool_id, connection.conn_id, id)
+      return nil if raw.nil? || raw.empty?
+
+      Google::Cloud::Spanner::V1::ResultSetStats.decode(raw)
     end
 
     def close
       return if @closed
 
       SpannerLib.close_rows(connection.pool_id, connection.conn_id, id)
+      @metadata = nil
       @closed = true
     end
   end

spannerlib/wrappers/spannerlib-ruby/spec/mock_server/spanner_mock_server.rb

@@ -48,15 +48,20 @@ def log_request(request)
     @requests << request
     return unless @req_file_path
 
-    data = {
-      class: request.class.name,
-      payload: request.class.encode(request)
-    }
-    File.open(@req_file_path, "ab") do |f|
-      Marshal.dump(data, f)
+    # Guard clause: Ensure we only attempt to encode Protobuf objects.
+    return unless request.class.respond_to?(:encode)
+
+    begin
+      data = {
+        class: request.class.name,
+        payload: request.class.encode(request)
+      }
+      File.open(@req_file_path, "ab") do |f|
+        Marshal.dump(data, f)
+      end
+    rescue StandardError => e
+      warn "Failed to log request: #{e.message}"
     end
-  rescue StandardError => e
-    warn "Failed to log request: #{e.message}"
   end
 
   def put_statement_result(sql, result)
@@ -123,20 +128,28 @@ def do_execute_sql(request, streaming)
 
     raise @errors[request.sql].pop if @errors[request.sql] && !@errors[request.sql].empty?
 
-    result = get_statement_result(request.sql).clone
-
-    if result.result_type == StatementResult::EXCEPTION
-      err_proto = result.result
-      raise GRPC::BadStatus.new(err_proto.code, err_proto.message) if err_proto.is_a?(Google::Rpc::Status)
-
-      raise err_proto
-
-    end
+    results_or_single = get_statement_result(request.sql)
+    results = results_or_single.is_a?(Array) ? results_or_single : [results_or_single]
 
+    # Handle streaming vs non-streaming response
+    # Streaming: return an enumerator that yields all result sets
     if streaming
-      result.each created_transaction
+      Enumerator.new do |yielder|
+        results.each do |entry|
+          result = entry.clone
+          raise_if_error!(result)
+
+          result.each(created_transaction) do |partial|
+            yielder << partial
+          end
+        end
+      end
     else
-      result.result created_transaction
+      # Non-streaming: only return the first result set
+      entry = results.first
+      result = entry.clone
+      raise_if_error!(result)
+      result.result(created_transaction)
     end
   end
 
@@ -331,4 +344,14 @@ def do_create_transaction(session)
     end
     transaction
   end
+
+  def raise_if_error!(result)
+    return unless result.result_type == StatementResult::EXCEPTION
+
+    err_proto = result.result
+
+    raise GRPC::BadStatus.new(err_proto.code, err_proto.message) if err_proto.is_a?(Google::Rpc::Status)
+
+    raise err_proto
+  end
 end

spannerlib/wrappers/spannerlib-ruby/spec/spannerlib_ruby_spec.rb

@@ -352,6 +352,36 @@ def find_transaction_options(requests)
     err = _ { @conn.execute_batch(batch_req) }.must_raise StandardError
     _(err.message).must_match(/Permission denied/)
   end
+
+  it "can execute a multi-statement query with multiple result sets" do
+    sql = "SELECT 1; SELECT 2"
+    set_mock_result("SELECT 1", StatementResult.create_select1_result)
+
+    set_mock_result(" SELECT 2", StatementResult.create_single_int_result_set("Col1", 2))
+
+    req = Google::Cloud::Spanner::V1::ExecuteSqlRequest.new(sql: sql)
+    rows = @conn.execute(req)
+
+    row1 = rows.next
+    _(row1).wont_be_nil
+    decoded1 = Google::Protobuf::ListValue.decode(row1)
+    _(decoded1.values[0].string_value).must_equal "1"
+
+    _(rows.next).must_be_nil
+
+    metadata = rows.next_result_set
+    _(metadata).wont_be_nil
+    _(metadata.row_type.fields[0].name).must_equal "Col1"
+
+    row2 = rows.next
+    _(row2).wont_be_nil
+    decoded2 = Google::Protobuf::ListValue.decode(row2)
+    _(decoded2.values[0].string_value).must_equal "2"
+
+    _(rows.next).must_be_nil
+
+    _(rows.next_result_set).must_be_nil
+  end
 end
 # rubocop:enable RSpec/NoExpectationExample
 # rubocop:enable Style/GlobalVars