RUBY-3712 Expose atClusterTime in snapshot sessions (#3035)

Author: comandeo-mongo

lib/mongo/operation/shared/executable.rb

@@ -59,7 +59,7 @@ def do_execute(connection, context, options = {})
                       end
                     end
 
-                    if session.snapshot? && !session.snapshot_timestamp
+                    if session.snapshot?
                       session.snapshot_timestamp = result.snapshot_timestamp
                     end
                   end

lib/mongo/session.rb

@@ -74,6 +74,8 @@ class Session
     #     and *:nearest*.
     # @option options [ true | false ] :snapshot Set up the session for
     #   snapshot reads.
+    # @option options [ BSON::Timestamp ] :snapshot_time The desired snapshot
+    #   time for snapshot reads. Only valid when :snapshot is true.
     #
     # @since 2.5.0
     # @api private
@@ -82,6 +84,14 @@ def initialize(server_session, client, options = {})
         raise ArgumentError, ':causal_consistency and :snapshot options cannot be both set on a session'
       end
 
+      if options[:snapshot_time] && !options[:snapshot]
+        raise ArgumentError, ':snapshot_time can only be set when :snapshot is true'
+      end
+
+      if options[:snapshot_time] && !options[:snapshot_time].is_a?(BSON::Timestamp)
+        raise ArgumentError, ':snapshot_time must be a BSON::Timestamp'
+      end
+
       if options[:implicit]
         unless server_session.nil?
           raise ArgumentError, 'Implicit session cannot reference server session during construction'
@@ -104,6 +114,7 @@ def initialize(server_session, client, options = {})
       @with_transaction_deadline = nil
       @with_transaction_timeout_ms = nil
       @inside_with_transaction = false
+      @snapshot_timestamp = options[:snapshot_time]
     end
 
     # @return [ Hash ] The options for this session.
@@ -1273,8 +1284,23 @@ def txn_num
       @server_session.txn_num
     end
 
+    # @return [ BSON::Timestamp | nil ] The snapshot time for this session.
+    #   nil if the session is not a snapshot session, or if it is a snapshot
+    #   session for which no :snapshot_time option was provided and no read
+    #   has yet captured atClusterTime from the server.
+    attr_reader :snapshot_timestamp
+
+    # Sets the snapshot time for the session. Once set, subsequent
+    # assignments are ignored: snapshotTime is established at most once per
+    # session, either from the :snapshot_time option at construction or from
+    # the atClusterTime returned by the first find/aggregate/distinct
+    # response. This keeps the property effectively read-only for callers,
+    # per the snapshot-sessions spec rationale.
+    #
     # @api private
-    attr_accessor :snapshot_timestamp
+    def snapshot_timestamp=(value)
+      @snapshot_timestamp ||= value
+    end
 
     # @return [ Integer | nil ] The deadline for the current transaction, if any.
     # @api private

spec/mongo/session_spec.rb

@@ -32,6 +32,53 @@
     it 'sets the cluster' do
       expect(session.cluster).to be(authorized_client.cluster)
     end
+
+    context 'when :snapshot_time is set without :snapshot' do
+      let(:options) do
+        { snapshot_time: BSON::Timestamp.new(0, 1) }
+      end
+
+      it 'raises ArgumentError' do
+        expect { session }.to raise_error(
+          ArgumentError, /:snapshot_time can only be set when :snapshot is true/
+        )
+      end
+    end
+
+    context 'when :snapshot_time is not a BSON::Timestamp' do
+      let(:options) do
+        { snapshot: true, snapshot_time: 12_345 }
+      end
+
+      it 'raises ArgumentError' do
+        expect { session }.to raise_error(
+          ArgumentError, /:snapshot_time must be a BSON::Timestamp/
+        )
+      end
+    end
+
+    context 'when :snapshot_time is set with :snapshot' do
+      let(:timestamp) { BSON::Timestamp.new(42, 7) }
+
+      let(:options) do
+        { snapshot: true, snapshot_time: timestamp }
+      end
+
+      it 'exposes the timestamp via snapshot_timestamp' do
+        expect(session.snapshot_timestamp).to eq(timestamp)
+      end
+    end
+  end
+
+  describe '#snapshot_timestamp=' do
+    let(:initial_timestamp) { BSON::Timestamp.new(1, 1) }
+    let(:later_timestamp) { BSON::Timestamp.new(2, 2) }
+    let(:options) { { snapshot: true, snapshot_time: initial_timestamp } }
+
+    it 'is a no-op once a snapshot timestamp is set' do
+      session.snapshot_timestamp = later_timestamp
+      expect(session.snapshot_timestamp).to eq(initial_timestamp)
+    end
   end
 
   describe '#inspect' do

spec/runners/unified/support_operations.rb

@@ -68,6 +68,11 @@ def end_session(op)
       session.end_session
     end
 
+    def get_snapshot_time(op)
+      session = entities.get(:session, op.use!('object'))
+      session.snapshot_timestamp
+    end
+
     def assert_session_dirty(op)
       consume_test_runner(op)
       use_arguments(op) do |args|

spec/runners/unified/test.rb

@@ -285,12 +285,7 @@ def generate_entities(es)
                    database.fs
                  when 'session'
                    client = entities.get(:client, spec.use!('client'))
-
-                   opts = if smc_opts = spec.use('sessionOptions')
-                            ::Utils.underscore_hash(smc_opts)
-                          else
-                            {}
-                          end
+                   opts = build_session_options(spec)
 
                    client.start_session(**opts).tap do |session|
                      session.advance_cluster_time(@cluster_time)
@@ -380,6 +375,20 @@ class << thread
       @entities_created = true
     end
 
+    # Builds the keyword options for Client#start_session from a session
+    # entity spec. When sessionOptions contains snapshotTime, the value is the
+    # name of a previously saved entity holding the actual BSON::Timestamp.
+    def build_session_options(spec)
+      smc_opts = spec.use('sessionOptions')
+      return {} unless smc_opts
+
+      opts = ::Utils.underscore_hash(smc_opts)
+      if opts[:snapshot_time].is_a?(String)
+        opts[:snapshot_time] = entities.get(:result, opts[:snapshot_time])
+      end
+      opts
+    end
+
     def set_initial_data
       @spec['initialData']&.each do |entity_spec|
         spec = UsingHash[entity_spec]