CSHARP-2596: AsQueryable with session (usable during transaction)

Author: origindotnet

src/MongoDB.Driver/IMongoCollectionExtensions.cs

@@ -73,7 +73,24 @@ public static IMongoQueryable<TDocument> AsQueryable<TDocument>(this IMongoColle
             Ensure.IsNotNull(collection, nameof(collection));
 
             aggregateOptions = aggregateOptions ?? new AggregateOptions();
-            var provider = new MongoQueryProviderImpl<TDocument>(collection, aggregateOptions);
+            var provider = new MongoQueryProviderImpl<TDocument>(null, collection, aggregateOptions);
+            return new MongoQueryableImpl<TDocument, TDocument>(provider);
+        }
+
+        /// <summary>
+        /// Creates a queryable source of documents.
+        /// </summary>
+        /// <typeparam name="TDocument">The type of the document.</typeparam>
+        /// <param name="collection">The collection.</param>
+        /// <param name="session">The session.</param>
+        /// <param name="aggregateOptions">The aggregate options</param>
+        /// <returns>A queryable source of documents.</returns>
+        public static IMongoQueryable<TDocument> AsQueryable<TDocument>(this IMongoCollection<TDocument> collection, IClientSessionHandle session, AggregateOptions aggregateOptions = null)
+        {
+            Ensure.IsNotNull(collection, nameof(collection));
+
+            aggregateOptions = aggregateOptions ?? new AggregateOptions();
+            var provider = new MongoQueryProviderImpl<TDocument>(session, collection, aggregateOptions);
             return new MongoQueryableImpl<TDocument, TDocument>(provider);
         }
 

src/MongoDB.Driver/Linq/AggregateQueryableExecutionModel.cs

@@ -81,18 +81,32 @@ public override string ToString()
             return sb.ToString();
         }
 
-        internal override object Execute<TInput>(IMongoCollection<TInput> collection, AggregateOptions options)
+        internal override object Execute<TInput>(IClientSessionHandle session, IMongoCollection<TInput> collection, AggregateOptions options)
         {
             var pipeline = CreatePipeline<TInput>();
 
-            return collection.Aggregate(pipeline, options, CancellationToken.None);
+            if (session == null)
+            {
+                return collection.Aggregate(pipeline, options, CancellationToken.None);
+            }
+            else
+            {
+                return collection.Aggregate(session, pipeline, options, CancellationToken.None);
+            }
         }
 
-        internal override Task ExecuteAsync<TInput>(IMongoCollection<TInput> collection, AggregateOptions options, CancellationToken cancellationToken)
+        internal override Task ExecuteAsync<TInput>(IClientSessionHandle session, IMongoCollection<TInput> collection, AggregateOptions options, CancellationToken cancellationToken)
         {
             var pipeline = CreatePipeline<TInput>();
 
-            return collection.AggregateAsync(pipeline, options, cancellationToken);
+            if (session == null)
+            {
+                return collection.AggregateAsync(pipeline, options, cancellationToken);
+            }
+            else
+            {
+                return collection.AggregateAsync(session, pipeline, options, cancellationToken);
+            }
         }
 
         private BsonDocumentStagePipelineDefinition<TInput, TOutput> CreatePipeline<TInput>()

src/MongoDB.Driver/Linq/MongoQueryProviderImpl.cs

@@ -31,10 +31,12 @@ namespace MongoDB.Driver.Linq
     internal sealed class MongoQueryProviderImpl<TDocument> : IMongoQueryProvider
     {
         private readonly IMongoCollection<TDocument> _collection;
+        private readonly IClientSessionHandle _session;
         private readonly AggregateOptions _options;
 
-        public MongoQueryProviderImpl(IMongoCollection<TDocument> collection, AggregateOptions options)
+        public MongoQueryProviderImpl(IClientSessionHandle session, IMongoCollection<TDocument> collection, AggregateOptions options)
         {
+            _session = session; // can be null
             _collection = Ensure.IsNotNull(collection, nameof(collection));
             _options = Ensure.IsNotNull(options, nameof(options));
         }
@@ -107,12 +109,12 @@ public QueryableExecutionModel GetExecutionModel(Expression expression)
 
         internal object ExecuteModel(QueryableExecutionModel model)
         {
-            return model.Execute(_collection, _options);
+            return model.Execute(_session, _collection, _options);
         }
 
         private Task ExecuteModelAsync(QueryableExecutionModel model, CancellationToken cancellationToken)
         {
-            return model.ExecuteAsync(_collection, _options, cancellationToken);
+            return model.ExecuteAsync(_session, _collection, _options, cancellationToken);
         }
 
         private Expression Prepare(Expression expression)

src/MongoDB.Driver/Linq/QueryableExecutionModel.cs

@@ -35,8 +35,8 @@ internal QueryableExecutionModel()
         {
         }
 
-        internal abstract Task ExecuteAsync<TInput>(IMongoCollection<TInput> collection, AggregateOptions options, CancellationToken cancellationToken);
+        internal abstract Task ExecuteAsync<TInput>(IClientSessionHandle session, IMongoCollection<TInput> collection, AggregateOptions options, CancellationToken cancellationToken);
 
-        internal abstract object Execute<TInput>(IMongoCollection<TInput> collection, AggregateOptions options);
+        internal abstract object Execute<TInput>(IClientSessionHandle session, IMongoCollection<TInput> collection, AggregateOptions options);
     }
 }

tests/MongoDB.Driver.Tests/Linq/MongoQueryableTests.cs

@@ -22,6 +22,7 @@
 using MongoDB.Driver;
 using MongoDB.Driver.Core.TestHelpers.XunitExtensions;
 using MongoDB.Driver.Linq;
+using MongoDB.Driver.Tests;
 using MongoDB.Driver.Tests.Linq;
 using Xunit;
 
@@ -1716,6 +1717,31 @@ public void Where_method_with_predicated_any()
                 "{ $match : { 'G' : { '$elemMatch' : { 'D' : \"Don't\" } } } }");
         }
 
+        [Fact]
+        public void AsQueryable_in_transaction()
+        {
+            using (var session = DriverTestConfiguration.Client.StartSession())
+            {
+                session.StartTransaction();
+                try
+                {
+                    __collection.InsertOne(session, new Root());
+
+                    var result_not_in_transaction = CreateQuery(null).Count(); // checks AsQueryable with null session (outside transaction)
+
+                    result_not_in_transaction.Should().Be(2);
+
+                    var result_in_transaction = CreateQuery(session).Count(); // checks AsQueryable with current session (inside transaction)
+
+                    result_in_transaction.Should().Be(3);
+                }
+                finally
+                {
+                    session.AbortTransaction();
+                }
+            }
+        }
+
         private List<T> Assert<T>(IMongoQueryable<T> queryable, int resultCount, params string[] expectedStages)
         {
             var executionModel = (AggregateQueryableExecutionModel<T>)queryable.GetExecutionModel();
@@ -1738,6 +1764,11 @@ private IMongoQueryable<Root> CreateQuery()
             return __collection.AsQueryable();
         }
 
+        private IMongoQueryable<Root> CreateQuery(IClientSessionHandle session)
+        {
+            return __collection.AsQueryable(session);
+        }
+
         private IMongoQueryable<Other> CreateOtherQuery()
         {
             return __otherCollection.AsQueryable();