Add client provider support (#794)

* Introduce mongo client provider

* Properly conditionally activate the installers

* Field annotation

---------

Co-authored-by: Daniel Marbach <[email protected]>

Author: danielmarbach

src/NServiceBus.Storage.MongoDB.AcceptanceTests/When_custom_provider_registered.cs

@@ -0,0 +1,87 @@
+#nullable enable
+
+namespace NServiceBus.AcceptanceTests;
+
+using System;
+using System.Diagnostics.CodeAnalysis;
+using System.Threading.Tasks;
+using AcceptanceTesting;
+using EndpointTemplates;
+using Microsoft.Extensions.DependencyInjection;
+using MongoDB.Driver;
+using NUnit.Framework;
+using Storage.MongoDB;
+
+public class When_custom_provider_registered : NServiceBusAcceptanceTest
+{
+    [Test]
+    public async Task Should_be_used()
+    {
+        Context context = await Scenario.Define<Context>()
+            .WithEndpoint<EndpointWithCustomProvider>(b => b.When(session => session.SendLocal(new StartSaga1 { DataId = Guid.NewGuid() })))
+            .Done(c => c.SagaReceivedMessage)
+            .Run();
+
+        Assert.That(context.ProviderWasCalled, Is.True);
+    }
+
+    public class Context : ScenarioContext
+    {
+        public bool SagaReceivedMessage { get; set; }
+        public bool ProviderWasCalled { get; set; }
+    }
+
+    public class EndpointWithCustomProvider : EndpointConfigurationBuilder
+    {
+        public EndpointWithCustomProvider() =>
+            EndpointSetup<DefaultServer>(config =>
+            {
+                config.RegisterComponents(c =>
+                    c.AddSingleton<IMongoClientProvider>(b => new CustomProvider(b.GetRequiredService<Context>())));
+            });
+
+        public class JustASaga(Context testContext) : Saga<JustASagaData>, IAmStartedByMessages<StartSaga1>
+        {
+            public Task Handle(StartSaga1 message, IMessageHandlerContext context)
+            {
+                Data.DataId = message.DataId;
+                testContext.SagaReceivedMessage = true;
+                MarkAsComplete();
+                return Task.CompletedTask;
+            }
+
+            protected override void ConfigureHowToFindSaga(SagaPropertyMapper<JustASagaData> mapper) => mapper.ConfigureMapping<StartSaga1>(m => m.DataId).ToSaga(s => s.DataId);
+        }
+
+        public class CustomProvider(Context testContext) : IMongoClientProvider
+        {
+            [field: AllowNull, MaybeNull]
+            public IMongoClient Client
+            {
+                get
+                {
+                    if (field is not null)
+                    {
+                        return field;
+                    }
+
+                    var containerConnectionString = Environment.GetEnvironmentVariable("NServiceBusStorageMongoDB_ConnectionString");
+
+                    field = string.IsNullOrWhiteSpace(containerConnectionString) ? new MongoClient() : new MongoClient(containerConnectionString);
+                    testContext.ProviderWasCalled = true;
+                    return field;
+                }
+            }
+        }
+
+        public class JustASagaData : ContainSagaData
+        {
+            public virtual Guid DataId { get; set; }
+        }
+    }
+
+    public class StartSaga1 : ICommand
+    {
+        public Guid DataId { get; set; }
+    }
+}

src/NServiceBus.Storage.MongoDB.PersistenceTests/PersistenceTestsConfiguration.cs

@@ -7,7 +7,6 @@
 using MongoDB.Bson;
 using MongoDB.Bson.Serialization;
 using MongoDB.Bson.Serialization.Serializers;
-using MongoDB.Driver;
 using NServiceBus.Outbox;
 using NServiceBus.Sagas;
 using Persistence;
@@ -40,15 +39,9 @@ public async Task Configure(CancellationToken cancellationToken = default)
         BsonSerializer.TryRegisterSerializer(new GuidSerializer(GuidRepresentation.Standard));
 
         var memberMapCache = new MemberMapCache();
-        var databaseSettings = new MongoDatabaseSettings
-        {
-            ReadConcern = ReadConcern.Majority,
-            ReadPreference = ReadPreference.Primary,
-            WriteConcern = WriteConcern.WMajority
-        };
 
         SagaStorageFeature.RegisterSagaEntityClassMappings(SagaMetadataCollection, memberMapCache);
-        await SagaInstaller.CreateInfrastructureForSagaDataTypes(ClientProvider.Client, databaseSettings, memberMapCache, databaseName,
+        await SagaInstaller.CreateInfrastructureForSagaDataTypes(ClientProvider.Client, MongoPersistence.DefaultDatabaseSettings, memberMapCache, databaseName,
             MongoPersistence.DefaultCollectionNamingConvention, MongoPersistence.DefaultCollectionSettings, SagaMetadataCollection, cancellationToken);
 
         SagaStorage = new SagaPersister(SagaPersister.DefaultVersionElementName, memberMapCache);
@@ -60,7 +53,7 @@ await SagaInstaller.CreateInfrastructureForSagaDataTypes(ClientProvider.Client,
         OutboxStorageFeature.RegisterOutboxClassMappings();
         await OutboxInstaller.CreateInfrastructureForOutboxTypes(ClientProvider.Client, databaseName, MongoPersistence.DefaultDatabaseSettings, MongoPersistence.DefaultCollectionNamingConvention, MongoPersistence.DefaultCollectionSettings, TimeSpan.FromDays(7), cancellationToken);
 
-        OutboxStorage = new OutboxPersister(ClientProvider.Client, databaseName, MongoPersistence.DefaultCollectionNamingConvention);
+        OutboxStorage = new OutboxPersister(ClientProvider.Client, databaseName, MongoPersistence.DefaultDatabaseSettings, MongoPersistence.DefaultCollectionNamingConvention, MongoPersistence.DefaultCollectionSettings);
     }
 
     public async Task Cleanup(CancellationToken cancellationToken = default) =>

src/NServiceBus.Storage.MongoDB.Tests/ApprovalFiles/APIApprovals.Approve.approved.txt

@@ -28,6 +28,10 @@ namespace NServiceBus
 }
 namespace NServiceBus.Storage.MongoDB
 {
+    public interface IMongoClientProvider
+    {
+        MongoDB.Driver.IMongoClient Client { get; }
+    }
     public interface IMongoSynchronizedStorageSession
     {
         MongoDB.Driver.IClientSessionHandle? MongoSession { get; }

src/NServiceBus.Storage.MongoDB.Tests/Outbox/OutboxTestsConfiguration.cs

@@ -41,7 +41,7 @@ public async Task Configure()
         await OutboxInstaller.CreateInfrastructureForOutboxTypes(ClientProvider.Client, DatabaseName, MongoPersistence.DefaultDatabaseSettings, CollectionNamingConvention,
             MongoPersistence.DefaultCollectionSettings, TimeSpan.FromHours(1));
 
-        OutboxStorage = new OutboxPersister(ClientProvider.Client, DatabaseName, CollectionNamingConvention);
+        OutboxStorage = new OutboxPersister(ClientProvider.Client, DatabaseName, MongoPersistence.DefaultDatabaseSettings, CollectionNamingConvention, MongoPersistence.DefaultCollectionSettings);
     }
 
     public async Task Cleanup() => await ClientProvider.Client.DropDatabaseAsync(DatabaseName);

src/NServiceBus.Storage.MongoDB.Tests/SubscriptionStorageTests.cs

@@ -18,12 +18,10 @@ public async Task OneTimeSetUp()
     {
         DatabaseName = "Test_" + DateTime.UtcNow.Ticks.ToString(CultureInfo.InvariantCulture);
 
-        var subscriptionCollection = ClientProvider.Client.GetDatabase(DatabaseName, MongoPersistence.DefaultDatabaseSettings)
-            .GetCollection<EventSubscription>("eventsubscriptions", MongoPersistence.DefaultCollectionSettings);
-
         await SubscriptionInstaller.CreateInfrastructureForSubscriptionTypes(ClientProvider.Client, MongoPersistence.DefaultDatabaseSettings, DatabaseName, MongoPersistence.DefaultCollectionSettings, _ => "eventsubscriptions");
 
-        var subscriptionPersister = new SubscriptionPersister(subscriptionCollection);
+        var subscriptionPersister = new SubscriptionPersister(ClientProvider.Client, DatabaseName, MongoPersistence.DefaultDatabaseSettings,
+            _ => "eventsubscriptions", MongoPersistence.DefaultCollectionSettings);
         storage = subscriptionPersister;
     }
 

src/NServiceBus.Storage.MongoDB/Configuration/DefaultMongoClientProvider.cs

@@ -0,0 +1,17 @@
+namespace NServiceBus.Storage.MongoDB;
+
+using System.Diagnostics.CodeAnalysis;
+using global::MongoDB.Driver;
+
+sealed class DefaultMongoClientProvider : IMongoClientProvider
+{
+    [field: AllowNull, MaybeNull]
+    public IMongoClient Client
+    {
+        get
+        {
+            field ??= new MongoClient();
+            return field;
+        }
+    }
+}

src/NServiceBus.Storage.MongoDB/Configuration/IMongoClientProvider.cs

@@ -0,0 +1,17 @@
+namespace NServiceBus.Storage.MongoDB;
+
+using global::MongoDB.Driver;
+
+/// <summary>
+/// Provides a mongo client via dependency injection. A custom implementation can be registered on the container and will be picked up by the persistence.
+/// <remarks>
+/// The client provided will not be disposed by the persistence. It is the responsibility of the provider to take care of proper resource disposal if necessary.
+/// </remarks>
+/// </summary>
+public interface IMongoClientProvider
+{
+    /// <summary>
+    /// The mongo client to use.
+    /// </summary>
+    IMongoClient Client { get; }
+}

src/NServiceBus.Storage.MongoDB/Configuration/MongoClientProvidedByConfiguration.cs

@@ -0,0 +1,8 @@
+namespace NServiceBus.Storage.MongoDB;
+
+using global::MongoDB.Driver;
+
+class MongoClientProvidedByConfiguration(IMongoClient client) : IMongoClientProvider
+{
+    public IMongoClient Client { get; } = client;
+}

src/NServiceBus.Storage.MongoDB/Configuration/MongoSettingsExtensions.cs

@@ -19,7 +19,7 @@ public static PersistenceExtensions<MongoPersistence> MongoClient(
         ArgumentNullException.ThrowIfNull(persistenceExtensions);
         ArgumentNullException.ThrowIfNull(mongoClient);
 
-        persistenceExtensions.GetSettings().Set(SettingsKeys.MongoClient, () => mongoClient);
+        persistenceExtensions.GetSettings().Set<IMongoClientProvider>(new MongoClientProvidedByConfiguration(mongoClient));
         return persistenceExtensions;
     }
 

src/NServiceBus.Storage.MongoDB/Configuration/SettingsKeys.cs

@@ -6,7 +6,6 @@ static class SettingsKeys
     public const string VersionElementName = baseName + nameof(VersionElementName);
     public const string CollectionNamingConvention = baseName + nameof(CollectionNamingConvention);
     public const string DatabaseName = baseName + nameof(DatabaseName);
-    public const string MongoClient = baseName + nameof(MongoClient);
     public const string UseTransactions = baseName + nameof(UseTransactions);
     public const string TimeToKeepOutboxDeduplicationData = baseName + nameof(TimeToKeepOutboxDeduplicationData);
 }