GODRIVER-1104 Add imperative sessions API (#364)

Author: Divjot Arora

mongo/client.go

@@ -766,7 +766,7 @@ func (c *Client) ListDatabaseNames(ctx context.Context, filter interface{}, opts
 //
 // Any error returned by the fn callback will be returned without any modifications.
 func WithSession(ctx context.Context, sess Session, fn func(SessionContext) error) error {
-	return fn(contextWithSession(ctx, sess))
+	return fn(NewSessionContext(ctx, sess))
 }
 
 // UseSession creates a new Session and uses it to create a new SessionContext, which is used to call the fn callback.
@@ -789,13 +789,7 @@ func (c *Client) UseSessionWithOptions(ctx context.Context, opts *options.Sessio
 	}
 
 	defer defaultSess.EndSession(ctx)
-
-	sessCtx := sessionContext{
-		Context: context.WithValue(ctx, sessionKey{}, defaultSess),
-		Session: defaultSess,
-	}
-
-	return fn(sessCtx)
+	return fn(NewSessionContext(ctx, defaultSess))
 }
 
 // Watch returns a change stream for all changes on the deployment. See

mongo/crud_examples_test.go

@@ -614,6 +614,41 @@ func ExampleClient_StartSession_withTransaction() {
 	fmt.Printf("result: %v\n", result)
 }
 
+func ExampleNewSessionContext() {
+	var client *mongo.Client
+
+	// Create a new Session and SessionContext.
+	sess, err := client.StartSession()
+	if err != nil {
+		panic(err)
+	}
+	defer sess.EndSession(context.TODO())
+	sessCtx := mongo.NewSessionContext(context.TODO(), sess)
+
+	// Start a transaction and sessCtx as the Context parameter to InsertOne and FindOne so both operations will be
+	// run in the transaction.
+	if err = sess.StartTransaction(); err != nil {
+		panic(err)
+	}
+
+	coll := client.Database("db").Collection("coll")
+	res, err := coll.InsertOne(sessCtx, bson.D{{"x", 1}})
+	if err != nil {
+		panic(err)
+	}
+
+	var result bson.M
+	if err = coll.FindOne(sessCtx, bson.D{{"_id", res.InsertedID}}).Decode(&result); err != nil {
+		panic(err)
+	}
+	fmt.Printf("result: %v\n", result)
+
+	// Commit the transaction so the inserted document will be stored.
+	if err = sess.CommitTransaction(sessCtx); err != nil {
+		panic(err)
+	}
+}
+
 // Cursor examples
 
 func ExampleCursor_All() {

mongo/integration/sessions_test.go

@@ -254,6 +254,68 @@ func TestSessions(t *testing.T) {
 			})
 		}
 	})
+
+	mt.Run("imperative API", func(mt *mtest.T) {
+		mt.Run("round trip Session object", func(mt *mtest.T) {
+			// Rountrip a Session object through NewSessionContext/ContextFromSession and assert that it is correctly
+			// stored/retrieved.
+
+			sess, err := mt.Client.StartSession()
+			assert.Nil(mt, err, "StartSession error: %v", err)
+			defer sess.EndSession(mtest.Background)
+
+			sessCtx := mongo.NewSessionContext(mtest.Background, sess)
+			assert.Equal(mt, sess.ID(), sessCtx.ID(), "expected Session ID %v, got %v", sess.ID(), sessCtx.ID())
+
+			gotSess := mongo.SessionFromContext(sessCtx)
+			assert.NotNil(mt, gotSess, "expected SessionFromContext to return non-nil value, got nil")
+			assert.Equal(mt, sess.ID(), gotSess.ID(), "expected Session ID %v, got %v", sess.ID(), gotSess.ID())
+		})
+
+		txnOpts := mtest.NewOptions().RunOn(
+			mtest.RunOnBlock{Topology: []mtest.TopologyKind{mtest.ReplicaSet}, MinServerVersion: "4.0"},
+			mtest.RunOnBlock{Topology: []mtest.TopologyKind{mtest.Sharded}, MinServerVersion: "4.2"},
+		)
+		mt.RunOpts("run transaction", txnOpts, func(mt *mtest.T) {
+			// Test that the imperative sessions API can be used to run a transaction.
+
+			createSessionContext := func(mt *mtest.T) mongo.SessionContext {
+				sess, err := mt.Client.StartSession()
+				assert.Nil(mt, err, "StartSession error: %v", err)
+
+				return mongo.NewSessionContext(mtest.Background, sess)
+			}
+
+			sessCtx := createSessionContext(mt)
+			sess := mongo.SessionFromContext(sessCtx)
+			assert.NotNil(mt, sess, "expected SessionFromContext to return non-nil value, got nil")
+			defer sess.EndSession(mtest.Background)
+
+			err := sess.StartTransaction()
+			assert.Nil(mt, err, "StartTransaction error: %v", err)
+
+			numDocs := 2
+			for i := 0; i < numDocs; i++ {
+				_, err = mt.Coll.InsertOne(sessCtx, bson.D{{"x", 1}})
+				assert.Nil(mt, err, "InsertOne error at index %d: %v", i, err)
+			}
+
+			// Assert that the collection count is 0 before committing and numDocs after. This tests that the InsertOne
+			// calls were actually executed in the transaction because the pre-commit count does not include them.
+			assertCollectionCount(mt, 0)
+			err = sess.CommitTransaction(sessCtx)
+			assert.Nil(mt, err, "CommitTransaction error: %v", err)
+			assertCollectionCount(mt, int64(numDocs))
+		})
+	})
+}
+
+func assertCollectionCount(mt *mtest.T, expectedCount int64) {
+	mt.Helper()
+
+	count, err := mt.Coll.CountDocuments(mtest.Background, bson.D{})
+	assert.Nil(mt, err, "CountDocuments error: %v", err)
+	assert.Equal(mt, expectedCount, count, "expected CountDocuments result %v, got %v", expectedCount, count)
 }
 
 type sessionFunction struct {

mongo/session.go

@@ -30,6 +30,10 @@ var withTransactionTimeout = 120 * time.Second
 // SessionContext combines the context.Context and mongo.Session interfaces. It should be used as the Context arguments
 // to operations that should be executed in a session. This type is not goroutine safe and must not be used concurrently
 // by multiple goroutines.
+//
+// There are two ways to create a SessionContext and use it in a session/transaction. The first is to use one of the
+// callback-based functions such as WithSession and UseSession. These functions create a SessionContext and pass it to
+// the provided callback. The other is to use NewSessionContext to explicitly create a SessionContext.
 type SessionContext interface {
 	context.Context
 	Session
@@ -43,6 +47,31 @@ type sessionContext struct {
 type sessionKey struct {
 }
 
+// NewSessionContext creates a new SessionContext associated with the given Context and Session parameters.
+func NewSessionContext(ctx context.Context, sess Session) SessionContext {
+	return &sessionContext{
+		Context: context.WithValue(ctx, sessionKey{}, sess),
+		Session: sess,
+	}
+}
+
+// SessionFromContext extracts the mongo.Session object stored in a Context. This can be used on a SessionContext that
+// was created implicitly through one of the callback-based session APIs or explicitly by calling NewSessionContext. If
+// there is no Session stored in the provided Context, nil is returned.
+func SessionFromContext(ctx context.Context) Session {
+	val := ctx.Value(sessionKey{})
+	if val == nil {
+		return nil
+	}
+
+	sess, ok := val.(Session)
+	if !ok {
+		return nil
+	}
+
+	return sess
+}
+
 // Session is an interface that represents a MongoDB logical session. Sessions can be used to enable causal consistency
 // for a group of operations or to execute operations in an ACID transaction. A new Session can be created from a Client
 // instance. A Session created from a Client must only be used to execute operations using that Client or a Database or
@@ -145,7 +174,7 @@ func (s *sessionImpl) WithTransaction(ctx context.Context, fn func(sessCtx Sessi
 			return nil, err
 		}
 
-		res, err := fn(contextWithSession(ctx, s))
+		res, err := fn(NewSessionContext(ctx, s))
 		if err != nil {
 			if s.clientSession.TransactionRunning() {
 				_ = s.AbortTransaction(ctx)
@@ -322,11 +351,3 @@ func sessionFromContext(ctx context.Context) *session.Client {
 
 	return nil
 }
-
-// contextWithSession creates a new SessionContext associated with the given Context and Session parameters.
-func contextWithSession(ctx context.Context, sess Session) SessionContext {
-	return &sessionContext{
-		Context: context.WithValue(ctx, sessionKey{}, sess),
-		Session: sess,
-	}
-}