chore(ruby): Add support for dynamic mock server tests

Author: aakashanandg

spannerlib/wrappers/spannerlib-ruby/.rubocop.yml

@@ -26,7 +26,7 @@ RSpec/BeforeAfterAll:
   Enabled: false
 RSpec/DescribeClass:
   Exclude:
-    - 'spec/integration/**/*'
+    - 'spec/**/*'
 
 Style/StringLiterals:
   EnforcedStyle: double_quotes
@@ -35,20 +35,20 @@ Style/FormatStringToken:
   EnforcedStyle: template
 Metrics/ClassLength:
   Exclude:
-    - 'spec/mock_server/**/*'
+   - 'spec/**/*'
 
 Metrics/MethodLength:
   Exclude:
-    - 'spec/mock_server/**/*'
+  - 'spec/**/*'
 
 Metrics/AbcSize:
   Exclude:
-    - 'spec/mock_server/**/*'
+    - 'spec/**/*'
 
 Metrics/CyclomaticComplexity:
   Exclude:
-    - 'spec/mock_server/**/*'
+    - 'spec/**/*'
 
 Metrics/PerceivedComplexity:
   Exclude:
-    - 'spec/mock_server/**/*'
+    - 'spec/**/*'

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

@@ -38,13 +38,28 @@ def initialize
     @errors = {}
     @session_counter = 0
 
-    put_statement_result "SELECT 1", StatementResult.create_select1_result
-    put_statement_result "INSERT INTO test_table (id, name) VALUES ('1', 'Alice')", StatementResult.create_update_count_result(1)
+    @mock_file_path = ENV.fetch("MOCK_MSG_FILE", nil)
+    @req_file_path = ENV.fetch("MOCK_REQ_FILE", nil)
 
-    dialect_sql = "select option_value from information_schema.database_options where option_name='database_dialect'"
-    dialect_result = StatementResult.create_dialect_result
+    add_default_results
+  end
 
-    put_statement_result dialect_sql, dialect_result
+  def log_request(request)
+    @requests << request # Keep memory copy just in case
+    return unless @req_file_path
+
+    # We must manually serialize the Protobuf before Marshaling
+    data = {
+      class: request.class.name,
+      payload: request.class.encode(request)
+    }
+
+    # Append to file using atomic open
+    File.open(@req_file_path, "ab") do |f|
+      Marshal.dump(data, f)
+    end
+  rescue StandardError => e
+    warn "Failed to log request: #{e.message}"
   end
 
   def put_statement_result(sql, result)
@@ -57,12 +72,12 @@ def push_error(sql_or_method, error)
   end
 
   def create_session(request, _unused_call)
-    @requests << request
+    log_request(request)
     do_create_session(request.database, request.session)
   end
 
   def batch_create_sessions(request, _unused_call)
-    @requests << request
+    log_request(request)
     num_created = 0
     response = Google::Cloud::Spanner::V1::BatchCreateSessionsResponse.new
     while num_created < request.session_count
@@ -73,12 +88,12 @@ def batch_create_sessions(request, _unused_call)
   end
 
   def get_session(request, _unused_call)
-    @requests << request
+    log_request(request)
     @sessions[request.name]
   end
 
   def list_sessions(request, _unused_call)
-    @requests << request
+    log_request(request)
     response = Google::Cloud::Spanner::V1::ListSessionsResponse.new
     @sessions.each_value do |s|
       response.sessions << s
@@ -87,16 +102,18 @@ def list_sessions(request, _unused_call)
   end
 
   def delete_session(request, _unused_call)
-    @requests << request
+    log_request(request)
     @sessions.delete request.name
     Google::Protobuf::Empty.new
   end
 
   def execute_sql(request, _unused_call)
+    log_request(request)
     do_execute_sql request, false
   end
 
   def execute_streaming_sql(request, _unused_call)
+    log_request(request)
     do_execute_sql request, true
   end
 
@@ -111,7 +128,13 @@ 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
-    raise result.result if result.result_type == StatementResult::EXCEPTION
+
+    if result.result_type == StatementResult::EXCEPTION
+      raise GRPC::BadStatus.new(result.result.code, result.result.message) if result.result.is_a?(Google::Rpc::Status)
+
+      raise result.result
+
+    end
 
     if streaming
       result.each created_transaction
@@ -121,6 +144,7 @@ def do_execute_sql(request, streaming)
   end
 
   def execute_batch_dml(request, _unused_call)
+    log_request(request)
     @requests << request
     validate_session request.session
     created_transaction = do_create_transaction request.session if request.transaction&.begin
@@ -133,8 +157,14 @@ def execute_batch_dml(request, _unused_call)
     request.statements.each do |stmt|
       result = get_statement_result(stmt.sql).clone
       if result.result_type == StatementResult::EXCEPTION
-        status.code = result.result.code
-        status.message = result.result.message
+        err_proto = result.result
+        if err_proto.is_a?(Google::Rpc::Status)
+          status.code = err_proto.code
+          status.message = err_proto.message
+        else
+          status.code = GRPC::Core::StatusCodes::UNKNOWN
+          status.message = err_proto.to_s
+        end
         break
       end
       if first
@@ -149,16 +179,19 @@ def execute_batch_dml(request, _unused_call)
   end
 
   def read(request, _unused_call)
+    log_request(request)
     @requests << request
     raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
   end
 
   def streaming_read(request, _unused_call)
+    log_request(request)
     @requests << request
     raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
   end
 
   def begin_transaction(request, _unused_call)
+    log_request(request)
     @requests << request
     raise @errors[__method__.to_s].pop if @errors[__method__.to_s] && !@errors[__method__.to_s].empty?
 
@@ -167,13 +200,15 @@ def begin_transaction(request, _unused_call)
   end
 
   def commit(request, _unused_call)
+    log_request(request)
     @requests << request
     validate_session request.session
     validate_transaction request.session, request.transaction_id
     Google::Cloud::Spanner::V1::CommitResponse.new commit_timestamp: Google::Protobuf::Timestamp.new(seconds: Time.now.to_i)
   end
 
   def rollback(request, _unused_call)
+    log_request(request)
     @requests << request
     validate_session request.session
     name = "#{request.session}/transactions/#{request.transaction_id}"
@@ -182,16 +217,19 @@ def rollback(request, _unused_call)
   end
 
   def partition_query(request, _unused_call)
+    log_request(request)
     @requests << request
     raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
   end
 
   def partition_read(request, _unused_call)
+    log_request(request)
     @requests << request
     raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
   end
 
   def get_database(request, _unused_call)
+    log_request(request)
     @requests << request
     raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented"
   end
@@ -208,18 +246,50 @@ def abort_next_transaction
   end
 
   def get_statement_result(sql)
+    load_dynamic_mocks!
+
     unless @statement_results.key? sql
       @statement_results.each do |key, value|
         return value if key.end_with?("%") && sql.start_with?(key.chop)
       end
+      available_keys = @statement_results.keys.join(", ")
       raise GRPC::BadStatus.new(
         GRPC::Core::StatusCodes::INVALID_ARGUMENT,
-        "There's no result registered for #{sql}"
+        "No result registered for '#{sql}'. Available: [#{available_keys}]"
       )
     end
     @statement_results[sql]
   end
 
+  def load_dynamic_mocks!
+    return unless @mock_file_path && File.exist?(@mock_file_path)
+
+    begin
+      # Read binary data from file
+      content = File.binread(@mock_file_path)
+      return if content.empty?
+
+      # Deserialize the Ruby Hash containing new mocks
+      # rubocop:disable Security/MarshalLoad
+      new_mocks = Marshal.load(content)
+      # rubocop:enable Security/MarshalLoad
+
+      # Merge into existing results
+      @statement_results.merge!(new_mocks)
+    rescue StandardError => e
+      # Ignore read errors (race conditions during writing)
+      warn "Failed to load mocks: #{e.message}"
+    end
+  end
+
+  def add_default_results
+    put_statement_result "SELECT 1", StatementResult.create_select1_result
+    put_statement_result "INSERT INTO test_table (id, name) VALUES ('1', 'Alice')", StatementResult.create_update_count_result(1)
+
+    dialect_sql = "select option_value from information_schema.database_options where option_name='database_dialect'"
+    put_statement_result dialect_sql, StatementResult.create_dialect_result
+  end
+
   def delete_all_sessions
     @sessions.clear
   end

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

@@ -22,6 +22,72 @@ class StatementResult
   UPDATE_COUNT = 2
   EXCEPTION = 3
 
+  # -------------------------------------------------------------------------
+  # CUSTOM SERIALIZATION LOGIC
+  # -------------------------------------------------------------------------
+  # This method is called by Marshal.dump
+  def marshal_dump
+    payload = @result
+    class_name = nil
+
+    # If the result is a Protobuf, encode it to a binary string
+    if defined?(Google::Protobuf::MessageExts) && @result.is_a?(Google::Protobuf::MessageExts)
+      payload = @result.class.encode(@result)
+      class_name = @result.class.name
+    end
+
+    # We return an array of data we want to save
+    [@result_type, payload, class_name]
+  end
+
+  # This method is called by Marshal.load
+  def marshal_load(array)
+    type, payload, class_name = array
+    @result_type = type
+
+    if class_name
+      # It was a Protobuf! Rehydrate it from the string.
+      klass = Object.const_get(class_name)
+      @result = klass.decode(payload)
+    else
+      # It was a primitive (Integer, Exception, etc.)
+      @result = payload
+    end
+  end
+
+  def self.create_single_value_result(col_name, col_type_code, value, value_type_field)
+    # 1. Define Metadata (Column Name and Type)
+    metadata = Google::Cloud::Spanner::V1::ResultSetMetadata.new(
+      row_type: Google::Cloud::Spanner::V1::StructType.new(
+        fields: [
+          Google::Cloud::Spanner::V1::StructType::Field.new(
+            name: col_name,
+            type: Google::Cloud::Spanner::V1::Type.new(code: col_type_code)
+          )
+        ]
+      )
+    )
+
+    # 2. Define Row Data
+    # value_type_field is usually :string_value, :bool_value, etc.
+    row = Google::Protobuf::ListValue.new(
+      values: [Google::Protobuf::Value.new(value_type_field => value)]
+    )
+
+    # 3. Wrap in ResultSet
+    result_set = Google::Cloud::Spanner::V1::ResultSet.new(
+      metadata: metadata,
+      rows: [row]
+    )
+
+    new(result_set, VALID)
+  end
+
+  def self.create_exception_result(error_proto)
+    # error_proto should be a Google::Rpc::Status object
+    new(error_proto, EXCEPTION)
+  end
+
   def self.create_select1_result
     create_single_int_result_set "Col1", 1
   end
@@ -132,7 +198,12 @@ def self.random_value_or_null(value, null_fraction_divisor)
 
   attr_reader :result_type
 
-  def initialize(result)
+  def initialize(result, explicit_type = nil)
+    if explicit_type
+      @result = result
+      @result_type = explicit_type
+      return
+    end
     case result
     when Google::Cloud::Spanner::V1::ResultSet
       @result_type = QUERY