src/Connected.Data/Storage/ConnectionProvider.cs

@@ -35,7 +35,7 @@ internal sealed class ConnectionProvider : IConnectionProvider, IAsyncDisposable
 		if (TransactionService.State == MiddlewareTransactionState.Completed)
 			Mode = StorageConnectionMode.Isolated;
 
-		return args is ISchemaSynchronizationContext context ? ResolveSingle(context) : await ResolveMultiple<TEntity>(args);
+		return args is ISchemaSynchronizationContext context ? await ResolveSingle(context) : await ResolveMultiple<TEntity>(args);
 	}
 	/// <summary>
 	/// This method is called if the supplied arguments already provided connection type on which they will perform operations.
@@ -48,9 +48,9 @@ internal sealed class ConnectionProvider : IConnectionProvider, IAsyncDisposable
 	/// <param name="behavior"></param>
 	/// <returns></returns>
 	/// <exception cref="NullReferenceException"></exception>
-	private ImmutableList<IStorageConnection> ResolveSingle(ISchemaSynchronizationContext args)
+	private async Task<ImmutableList<IStorageConnection>> ResolveSingle(ISchemaSynchronizationContext args)
 	{
-		return new List<IStorageConnection> { EnsureConnection(args.ConnectionType, args.ConnectionString) }.ToImmutableList();
+		return new List<IStorageConnection> { await EnsureConnection(args.ConnectionType, args.ConnectionString) }.ToImmutableList();
 	}
 
 	private async Task<ImmutableList<IStorageConnection>> ResolveMultiple<TEntity>(StorageContextArgs args)
@@ -76,7 +76,7 @@ internal sealed class ConnectionProvider : IConnectionProvider, IAsyncDisposable
 			/*
 			 * Default connection is always used regardless of sharding support
 			 */
-			EnsureConnection(connectionMiddleware.ConnectionType, connectionMiddleware.DefaultConnectionString)
+			await EnsureConnection(connectionMiddleware.ConnectionType, connectionMiddleware.DefaultConnectionString)
 		};
 
 		if (shardingMiddleware is not null)
@@ -92,14 +92,14 @@ internal sealed class ConnectionProvider : IConnectionProvider, IAsyncDisposable
 				if (Type.GetType(node.ConnectionType) is not Type connectionType)
 					throw new NullReferenceException(node.ConnectionType);
 
-				result.Add(EnsureConnection(connectionType, node.ConnectionString));
+				result.Add(await EnsureConnection(connectionType, node.ConnectionString));
 			}
 		}
 
 		return result.ToImmutableList();
 	}
 
-	private IStorageConnection EnsureConnection(Type connectionType, string connectionString)
+	private async Task<IStorageConnection> EnsureConnection(Type connectionType, string connectionString)
 	{
 		if (Mode == StorageConnectionMode.Shared
 			&& Connections.FirstOrDefault(f => f.GetType() == connectionType
@@ -108,15 +108,15 @@ internal sealed class ConnectionProvider : IConnectionProvider, IAsyncDisposable
 			return existing;
 		}
 		else
-			return CreateConnection(connectionType, connectionString, Mode);
+			return await CreateConnection(connectionType, connectionString, Mode);
 	}
 
-	private IStorageConnection CreateConnection(Type connectionType, string connectionString, StorageConnectionMode behavior)
+	private async Task<IStorageConnection> CreateConnection(Type connectionType, string connectionString, StorageConnectionMode behavior)
 	{
 		if (Context.GetService(connectionType) is not IStorageConnection newConnection)
 			throw new NullReferenceException(connectionType.Name);
 
-		newConnection.Initialize(new StorageConnectionArgs(connectionString, behavior));
+		await newConnection.Initialize(new StorageConnectionArgs(connectionString, behavior));
 
 		if (behavior == StorageConnectionMode.Shared)
 			Connections.Add(newConnection);