DI changes to support mulitple buckets

couchbaselabs · Jan 23, 2025 · 6cfdfc3 · 6cfdfc3
1 parent 3ea6603
commit 6cfdfc3
Show file tree

Hide file tree

Showing 25 changed files with 504 additions and 189 deletions.
diff --git a/samples/ContosoUniversity/Program.cs b/samples/ContosoUniversity/Program.cs
@@ -12,9 +12,9 @@
 builder.Services.AddControllersWithViews();
 
 builder.Services.AddDbContext<SchoolContext>(options=> options
-    .UseCouchbase<INamedBucketProvider>(new ClusterOptions()
-        .WithCredentials("Ajax", "GE9jk9i28L2Psg@")
-        .WithConnectionString("couchbases://cb.umolxgoqkdzpvdvo.cloud.couchbase.com"),
+    .UseCouchbase(new ClusterOptions()
+        .WithCredentials("Administrator", "password")
+        .WithConnectionString("couchbases://localhost"),
         couchbaseDbContextOptions =>
         {
             couchbaseDbContextOptions.Bucket = "Universities";

diff --git a/src/Couchbase.EntityFrameworkCore/CouchbaseDbContextOptionsExtensions.cs b/src/Couchbase.EntityFrameworkCore/CouchbaseDbContextOptionsExtensions.cs
@@ -11,26 +11,42 @@ namespace Couchbase.EntityFrameworkCore;
 
 public static class CouchbaseDbContextOptionsExtensions
 {
-  public static DbContextOptionsBuilder UseCouchbase<TNamedBucketProvider>(
+  public static DbContextOptionsBuilder UseCouchbase(
       this DbContextOptionsBuilder optionsBuilder, 
       ClusterOptions clusterOptions,
-      Action<CouchbaseDbContextOptionsBuilder>? couchbaseActionOptions = null) where TNamedBucketProvider : class, INamedBucketProvider
+      Action<CouchbaseDbContextOptionsBuilder>? couchbaseActionOptions = null)
   {
     var couchbaseDbContextOptionsBuilder = new CouchbaseDbContextOptionsBuilder(optionsBuilder, clusterOptions);
     couchbaseActionOptions?.Invoke(couchbaseDbContextOptionsBuilder);
 
-    var extension = CouchbaseDbContextOptionsBuilderExtensions.GetOrCreateExtension<TNamedBucketProvider>(optionsBuilder, clusterOptions, couchbaseDbContextOptionsBuilder);
+    var extension = CouchbaseDbContextOptionsBuilderExtensions.GetOrCreateExtension(optionsBuilder, clusterOptions, couchbaseDbContextOptionsBuilder);
+    ((IDbContextOptionsBuilderInfrastructure) optionsBuilder).AddOrUpdateExtension(extension);
+    CouchbaseDbContextOptionsBuilderExtensions.ConfigureWarnings(optionsBuilder);
+
+    return optionsBuilder;
+  }
+
+  public static DbContextOptionsBuilder UseCouchbase(
+    this DbContextOptionsBuilder optionsBuilder,
+    string connectionString,
+    Action<CouchbaseDbContextOptionsBuilder>? couchbaseActionOptions = null)
+  {
+    var clusterOptions = new ClusterOptions().WithConnectionString(connectionString);
+    var couchbaseDbContextOptionsBuilder = new CouchbaseDbContextOptionsBuilder(optionsBuilder, clusterOptions);
+    couchbaseActionOptions?.Invoke(couchbaseDbContextOptionsBuilder);
+
+    var extension = CouchbaseDbContextOptionsBuilderExtensions.GetOrCreateExtension(optionsBuilder, clusterOptions, couchbaseDbContextOptionsBuilder);
     ((IDbContextOptionsBuilderInfrastructure) optionsBuilder).AddOrUpdateExtension(extension);
     CouchbaseDbContextOptionsBuilderExtensions.ConfigureWarnings(optionsBuilder);
 
     return optionsBuilder;
   }
 
   #nullable disable
-  private static CouchbaseOptionsExtension<TNamedBucketProvider> GetOrCreateExtension<TNamedBucketProvider>(DbContextOptionsBuilder optionsBuilder, ClusterOptions clusterOptions, CouchbaseDbContextOptionsBuilder couchbaseDbContextOptionsBuilder) where TNamedBucketProvider : class, INamedBucketProvider
-    => optionsBuilder.Options.FindExtension<CouchbaseOptionsExtension<TNamedBucketProvider>>() is CouchbaseOptionsExtension<TNamedBucketProvider> existing
-      ? new CouchbaseOptionsExtension<TNamedBucketProvider>(existing)
-      : new CouchbaseOptionsExtension<TNamedBucketProvider>(clusterOptions, couchbaseDbContextOptionsBuilder);
+  private static CouchbaseOptionsExtension GetOrCreateExtension(DbContextOptionsBuilder optionsBuilder, ClusterOptions clusterOptions, CouchbaseDbContextOptionsBuilder couchbaseDbContextOptionsBuilder)
+    => optionsBuilder.Options.FindExtension<CouchbaseOptionsExtension>() is CouchbaseOptionsExtension existing
+      ? new CouchbaseOptionsExtension(existing)
+      : new CouchbaseOptionsExtension(couchbaseDbContextOptionsBuilder);
 
   private static void ConfigureWarnings(DbContextOptionsBuilder optionsBuilder)
   {

diff --git a/src/Couchbase.EntityFrameworkCore/Extensions/CouchbaseDbContextOptionsBuilderExtensions.cs b/src/Couchbase.EntityFrameworkCore/Extensions/CouchbaseDbContextOptionsBuilderExtensions.cs
@@ -19,19 +19,19 @@ public static DbContextOptionsBuilder UseCouchbaseProvider<TNamedBucketProvider>
         var couchbaseDbContextOptionsBuilder = new CouchbaseDbContextOptionsBuilder(optionsBuilder, clusterOptions);
         couchbaseOptionsAction?.Invoke(couchbaseDbContextOptionsBuilder);
 
-        var extension = GetOrCreateExtension<TNamedBucketProvider>(optionsBuilder, clusterOptions, couchbaseDbContextOptionsBuilder);
+        var extension = GetOrCreateExtension(optionsBuilder, clusterOptions, couchbaseDbContextOptionsBuilder);
         ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension);
         ConfigureWarnings(optionsBuilder);
 
         return optionsBuilder;
     }
 
-    internal static CouchbaseOptionsExtension<TNamedBucketProvider> GetOrCreateExtension<TNamedBucketProvider>(DbContextOptionsBuilder optionsBuilder,
-        ClusterOptions clusterOptions, CouchbaseDbContextOptionsBuilder couchbaseDbContextOptionsBuilder) where TNamedBucketProvider : class, INamedBucketProvider
-        => optionsBuilder.Options.FindExtension<CouchbaseOptionsExtension<TNamedBucketProvider>>()
-            is CouchbaseOptionsExtension<TNamedBucketProvider> existing
-            ? new CouchbaseOptionsExtension<TNamedBucketProvider>(existing)
-            : new CouchbaseOptionsExtension<TNamedBucketProvider>(clusterOptions, couchbaseDbContextOptionsBuilder);
+    internal static CouchbaseOptionsExtension GetOrCreateExtension(DbContextOptionsBuilder optionsBuilder,
+        ClusterOptions clusterOptions, CouchbaseDbContextOptionsBuilder couchbaseDbContextOptionsBuilder)
+        => optionsBuilder.Options.FindExtension<CouchbaseOptionsExtension>()
+            is CouchbaseOptionsExtension existing
+            ? new CouchbaseOptionsExtension(existing)
+            : new CouchbaseOptionsExtension(couchbaseDbContextOptionsBuilder);
 
     internal static void ConfigureWarnings(DbContextOptionsBuilder optionsBuilder)
     {

diff --git a/src/Couchbase.EntityFrameworkCore/Extensions/CouchbaseServiceCollectionExtensions.cs b/src/Couchbase.EntityFrameworkCore/Extensions/CouchbaseServiceCollectionExtensions.cs
@@ -26,42 +26,27 @@ namespace Couchbase.EntityFrameworkCore.Extensions;
 
 public static class CouchbaseServiceCollectionExtensions
 {
-    public static IServiceCollection AddCouchbase<TContext, TNamedBucketProvider>(
+    public static IServiceCollection AddCouchbase<TContext>(
         this IServiceCollection serviceCollection,
         ClusterOptions clusterOptions,
         Action<ICouchbaseDbContextOptionsBuilder>? couchbaseOptionsAction = null,
-        Action<DbContextOptionsBuilder>? optionsAction = null) where TNamedBucketProvider : class, INamedBucketProvider
+        Action<DbContextOptionsBuilder>? optionsAction = null)
         where TContext : DbContext
         => serviceCollection.AddDbContext<TContext>((_, options) =>
         {
             optionsAction?.Invoke(options);
-            options.UseCouchbase<TNamedBucketProvider>(clusterOptions, couchbaseOptionsAction);
+            options.UseCouchbase(clusterOptions, couchbaseOptionsAction);
         });
 
-    public static IServiceCollection AddEntityFrameworkCouchbaseProvider<TNamedBucketProvider>(this IServiceCollection serviceCollection,
-        CouchbaseOptionsExtension<TNamedBucketProvider> optionsExtension) where TNamedBucketProvider : class, INamedBucketProvider
+    public static IServiceCollection AddEntityFrameworkCouchbaseProvider(this IServiceCollection serviceCollection,
+        CouchbaseOptionsExtension optionsExtension)
     {
-       serviceCollection.AddCouchbase(options =>
-        {
-            options.WithConnectionString(optionsExtension.ClusterOptions.ConnectionString);
-            options.WithCredentials(optionsExtension.ClusterOptions.UserName, optionsExtension.ClusterOptions.Password);
-        });
-
-        serviceCollection.AddCouchbaseBucket<TNamedBucketProvider>(optionsExtension.DbContextOptionsBuilder.Bucket,
-            builder =>
-            {
-                //The default scope is replaced with the mapped entity name
-                builder
-                    .AddScope(optionsExtension.DbContextOptionsBuilder.Scope)
-                    .AddCollection<INamedCollectionProvider>("_default");
-            });
-
         serviceCollection.AddLogging(); //this should be injectable from the app side
 
         var builder = new EntityFrameworkRelationalServicesBuilder(serviceCollection)
             .TryAdd<IRelationalTypeMappingSource, CouchbaseTypeMappingSource>()
             .TryAdd<IDatabase, CouchbaseDatabaseWrapper>()
-            .TryAdd<IDatabaseProvider, DatabaseProvider<CouchbaseOptionsExtension<TNamedBucketProvider>>>()
+            .TryAdd<IDatabaseProvider, DatabaseProvider<CouchbaseOptionsExtension>>()
             .TryAdd<LoggingDefinitions, CouchbaseLoggingDefinitions>()
             .TryAdd<IModificationCommandBatchFactory, CouchbaseModificationCommandBatchFactory>()
             .TryAdd<IUpdateSqlGenerator, CouchbaseUpdateSqlGenerator>()
@@ -85,6 +70,8 @@ public static IServiceCollection AddEntityFrameworkCouchbaseProvider<TNamedBucke
                 //.TryAddScoped<QueryContext, RelationalQueryContext>()
                 .TryAddScoped<ICouchbaseClientWrapper, CouchbaseClientWrapper>()
                 .TryAddScoped<IRelationalCommandBuilder, RelationalCommandBuilder>()
+                .TryAddScoped<ICouchbaseDbContextOptionsBuilder,
+                    CouchbaseDbContextOptionsBuilder>(b=>optionsExtension.DbContextOptionsBuilder)
             );
 
         builder.TryAddCoreServices();

diff --git a/src/Couchbase.EntityFrameworkCore/Infrastructure/CouchbaseDbContextOptionsBuilder.cs b/src/Couchbase.EntityFrameworkCore/Infrastructure/CouchbaseDbContextOptionsBuilder.cs
@@ -4,12 +4,24 @@ namespace Couchbase.EntityFrameworkCore.Infrastructure;
 
 public class CouchbaseDbContextOptionsBuilder : ICouchbaseDbContextOptionsBuilder
 {
+    private readonly string _connectionString;
+
+    public CouchbaseDbContextOptionsBuilder(DbContextOptionsBuilder dbContextOptionsBuilder, string connectionString)
+    {
+        _connectionString = connectionString;
+        OptionsBuilder = dbContextOptionsBuilder;
+        ClusterOptions = new ClusterOptions().WithConnectionString(connectionString);
+    }
+
     public CouchbaseDbContextOptionsBuilder(DbContextOptionsBuilder dbContextOptionsBuilder, ClusterOptions clusterOptions)
     {
         OptionsBuilder = dbContextOptionsBuilder;
         ClusterOptions = clusterOptions;
     }
 
+    //TODO temp
+    public string ConnectionString => ClusterOptions.ConnectionString! + $"?bucket={Bucket}";
+
     private DbContextOptionsBuilder OptionsBuilder { get; }
 
     public ClusterOptions ClusterOptions { get; }
@@ -27,7 +39,9 @@ public interface ICouchbaseDbContextOptionsBuilder
 
     ClusterOptions ClusterOptions { get; }
 
-    public string Bucket { get; }
+    public string ConnectionString { get; }
+
+    public string Bucket { get; set; }
 
-    public string Scope { get; }
+    public string Scope { get; set; }
 }
diff --git a/src/Couchbase.EntityFrameworkCore/Infrastructure/Internal/CouchbaseOptionsExtension.cs b/src/Couchbase.EntityFrameworkCore/Infrastructure/Internal/CouchbaseOptionsExtension.cs
@@ -8,51 +8,58 @@
 
 namespace Couchbase.EntityFrameworkCore.Infrastructure.Internal;
 
-//TODO potentially implement IDbContextOptionsExtension instead of deriving from RelationalOptionsExtension
-public class CouchbaseOptionsExtension<TNamedBucketProvider> : RelationalOptionsExtension where TNamedBucketProvider : class, INamedBucketProvider
+public class CouchbaseOptionsExtension: RelationalOptionsExtension
 {
-    private readonly ClusterOptions _clusterOptions;
     private readonly CouchbaseDbContextOptionsBuilder _couchbaseDbContextOptionsBuilder;
     private CouchbaseOptionsExtensionInfo? _info;
 
-    public CouchbaseOptionsExtension(ClusterOptions clusterOptions, CouchbaseDbContextOptionsBuilder couchbaseDbContextOptionsBuilder)
+    public CouchbaseOptionsExtension(CouchbaseDbContextOptionsBuilder couchbaseDbContextOptionsBuilder)
     {
-        _clusterOptions = clusterOptions;
         _couchbaseDbContextOptionsBuilder = couchbaseDbContextOptionsBuilder;
     }
-    protected internal CouchbaseOptionsExtension(CouchbaseOptionsExtension<TNamedBucketProvider> copyFrom)
+    protected internal CouchbaseOptionsExtension(CouchbaseOptionsExtension copyFrom)
         : base(copyFrom)
     {
+        _couchbaseDbContextOptionsBuilder = copyFrom.CouchbaseDbContextOptionsBuilder;
     }
 
-    public ClusterOptions ClusterOptions => _clusterOptions;
+    public CouchbaseDbContextOptionsBuilder? CouchbaseDbContextOptionsBuilder => _couchbaseDbContextOptionsBuilder;
 
-    public override string? ConnectionString => _clusterOptions.ConnectionString;
+    public override string? ConnectionString => _couchbaseDbContextOptionsBuilder.ConnectionString;
 
     public override DbContextOptionsExtensionInfo Info => _info ??= new CouchbaseOptionsExtensionInfo(this);
 
     public CouchbaseDbContextOptionsBuilder DbContextOptionsBuilder => _couchbaseDbContextOptionsBuilder;
 
     public override void ApplyServices(IServiceCollection services)
     {
+        services.AddCouchbase(options =>
+        {
+            options.WithConnectionString(_couchbaseDbContextOptionsBuilder.ClusterOptions.ConnectionString);
+            options.WithCredentials(_couchbaseDbContextOptionsBuilder.ClusterOptions.UserName, _couchbaseDbContextOptionsBuilder.ClusterOptions.Password);
+        });
+
+        services.AddKeyedCouchbase(_couchbaseDbContextOptionsBuilder.ConnectionString, options =>
+        {
+            options.WithLogging(_couchbaseDbContextOptionsBuilder.ClusterOptions.Logging);
+            options.WithConnectionString(_couchbaseDbContextOptionsBuilder.ClusterOptions.ConnectionString);
+            options.WithCredentials(_couchbaseDbContextOptionsBuilder.ClusterOptions.UserName, _couchbaseDbContextOptionsBuilder.ClusterOptions.Password);
+        });
+
         services.AddEntityFrameworkCouchbaseProvider(this);
     }
-
-    private TNamedBucketProvider BucketProvider;
 
     public override void Validate(IDbContextOptions options)
     {
         // You can add any validation logic here, if necessary.
     }
 
-    protected override RelationalOptionsExtension Clone() => new CouchbaseOptionsExtension<TNamedBucketProvider>(this);
+    protected override RelationalOptionsExtension Clone() => new CouchbaseOptionsExtension(this);
 
 
     public class CouchbaseOptionsExtensionInfo : DbContextOptionsExtensionInfo
     {
-        private readonly ClusterOptions _clusterOptions;
-
-        public CouchbaseOptionsExtensionInfo(CouchbaseOptionsExtension<TNamedBucketProvider> extension)
+        public CouchbaseOptionsExtensionInfo(CouchbaseOptionsExtension extension)
             : base(extension)
         {
         }
@@ -70,7 +77,7 @@ public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
             debugInfo["Couchbase:ConnectionString"] = ConnectionString;
         }
 
-        public override CouchbaseOptionsExtension<TNamedBucketProvider> Extension => (CouchbaseOptionsExtension<TNamedBucketProvider>)base.Extension;
+        public override CouchbaseOptionsExtension Extension => (CouchbaseOptionsExtension)base.Extension;
         private string? ConnectionString => Extension.Connection == null ?
             Extension.ConnectionString :
             Extension.Connection.ConnectionString;

diff --git a/src/Couchbase.EntityFrameworkCore/Query/Internal/CouchbaseQueryEnumerable.cs b/src/Couchbase.EntityFrameworkCore/Query/Internal/CouchbaseQueryEnumerable.cs
@@ -1,4 +1,5 @@
 using System.Collections;
+using Couchbase.EntityFrameworkCore.Infrastructure;
 using Couchbase.Extensions.DependencyInjection;
 using Couchbase.Query;
 using Microsoft.EntityFrameworkCore;
@@ -7,6 +8,7 @@
 using Microsoft.EntityFrameworkCore.Query.Internal;
 using Microsoft.EntityFrameworkCore.Storage;
 using Microsoft.EntityFrameworkCore.Storage.Internal;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Couchbase.EntityFrameworkCore.Query.Internal;
 
@@ -15,23 +17,26 @@ public class CouchbaseQueryEnumerable<T> : IEnumerable<T>, IAsyncEnumerable<T>
     private readonly RelationalQueryContext _relationalQueryContext;
     private readonly RelationalCommandCache _relationalCommandCache;
     private readonly DbContext _dbContext;
-    private readonly IClusterProvider _clusterProvider;
     private readonly bool _standAloneStateManager;
-
+    private readonly IServiceProvider _serviceProvider;
+    private readonly ICouchbaseDbContextOptionsBuilder _couchbaseDbContextOptionsBuilder;
+
     public CouchbaseQueryEnumerable(
         RelationalQueryContext relationalQueryContext, 
         RelationalCommandCache relationalCommandCache,
         bool standAloneStateManager,
-        IClusterProvider clusterProvider)
+        IServiceProvider serviceProvider,
+        ICouchbaseDbContextOptionsBuilder couchbaseDbContextOptionsBuilder)
     {
 
         _dbContext = relationalQueryContext.Context;
-        _clusterProvider = clusterProvider;
         _standAloneStateManager = standAloneStateManager;
+        _serviceProvider = serviceProvider;
+        _couchbaseDbContextOptionsBuilder = couchbaseDbContextOptionsBuilder;
         _relationalQueryContext = relationalQueryContext;
         _relationalCommandCache = relationalCommandCache;
     }
-    
+
     public IEnumerator<T> GetEnumerator()
     {
         var queryOptions = new QueryOptions();
@@ -40,8 +45,9 @@ public IEnumerator<T> GetEnumerator()
             queryOptions.Parameter(parameter.Key, parameter.Value);
         }
         var command = _relationalCommandCache.RentAndPopulateRelationalCommand(_relationalQueryContext);
-        var cluster = _clusterProvider.GetClusterAsync().GetAwaiter().GetResult();
-        var result = cluster.QueryAsync<T>(command.CommandText, queryOptions).GetAwaiter().GetResult();
+        var bucketProvider = _serviceProvider.GetKeyedService<IBucketProvider>(_couchbaseDbContextOptionsBuilder.ConnectionString);
+        var bucket = bucketProvider.GetBucketAsync(_couchbaseDbContextOptionsBuilder.Bucket).GetAwaiter().GetResult();
+        var result = bucket.Cluster.QueryAsync<T>(command.CommandText, queryOptions).GetAwaiter().GetResult();
 
         _relationalQueryContext.InitializeStateManager(_standAloneStateManager);
 
@@ -77,9 +83,10 @@ IEnumerator IEnumerable.GetEnumerator()
     {
         var command = _relationalCommandCache.RentAndPopulateRelationalCommand(_relationalQueryContext);
         var queryOptions = GetParameters(command);
-        var cluster = await _clusterProvider.GetClusterAsync().ConfigureAwait(false);
-        var result = await cluster.QueryAsync<T>(command.CommandText, queryOptions).ConfigureAwait(false);
-
+        var bucketProvider = _serviceProvider.GetKeyedService<IBucketProvider>(_couchbaseDbContextOptionsBuilder.ConnectionString);
+        var bucket = await bucketProvider.GetBucketAsync(_couchbaseDbContextOptionsBuilder.Bucket);
+        var result = await bucket.Cluster.QueryAsync<T>(command.CommandText, queryOptions).ConfigureAwait(false);
+
         _relationalQueryContext.InitializeStateManager(_standAloneStateManager);
 
         var model = _dbContext.Model;
@@ -112,7 +119,7 @@ private QueryOptions GetParameters(IRelationalCommand command)
         {
             var key = parameter.Key;
             var value = parameter.Value;
-            
+
             foreach (var compositeParameter in command.Parameters)
             {
                 if (compositeParameter is CompositeRelationalParameter actualParameter)