Added NpgsqlDataSource to Postgres SingleServerDatabaseCollection and…

… Migrator
JasperFx · Dec 7, 2023 · 036af5c · 036af5c
1 parent d1be272
commit 036af5c
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 21 deletions.
diff --git a/src/Weasel.Core/Migrations/Database.cs b/src/Weasel.Core/Migrations/Database.cs
@@ -52,37 +52,33 @@ public DatabaseBase(
         Migrator migrator,
         string identifier,
         string connectionString
-    )
+    ): this(logger, autoCreate, migrator, identifier, () => CreateConnection(connectionString))
     {
-        _logger = logger;
-        _connectionSource = () =>
-        {
-            var conn = new TConnection();
-            conn.ConnectionString = connectionString;
-
-            return conn;
-        };
-
-        AutoCreate = autoCreate;
-        Migrator = migrator;
-        Identifier = identifier;
     }
 
     public DatabaseBase(
         IMigrationLogger logger,
         AutoCreate autoCreate,
         Migrator migrator,
         string identifier,
-        Func<TConnection> connectionSource
+        Func<TConnection>? connectionSource
     )
     {
         _logger = logger;
-        _connectionSource = connectionSource;
+        _connectionSource = connectionSource ?? CreateConnection;
         AutoCreate = autoCreate;
         Migrator = migrator;
         Identifier = identifier;
     }
 
+    private static TConnection CreateConnection(string connectionString)
+    {
+        var conn = new TConnection();
+        conn.ConnectionString = connectionString;
+
+        return conn;
+    }
+
     public abstract IFeatureSchema[] BuildFeatureSchemas();
 
     public AutoCreate AutoCreate { get; set; }

diff --git a/src/Weasel.Core/Migrator.cs b/src/Weasel.Core/Migrator.cs
@@ -108,7 +108,7 @@ public async Task WriteTemplatedFile(string filename, Action<Migrator, TextWrite
     /// <param name="migration"></param>
     /// <param name="autoCreate"></param>
     /// <param name="logger"></param>
-    public async Task ApplyAllAsync(
+    public Task ApplyAllAsync(
         DbConnection conn,
         SchemaMigration migration,
         AutoCreate autoCreate,
@@ -118,23 +118,23 @@ public async Task ApplyAllAsync(
     {
         if (autoCreate == AutoCreate.None)
         {
-            return;
+            return Task.CompletedTask;
         }
 
         if (migration.Difference == SchemaPatchDifference.None)
         {
-            return;
+            return Task.CompletedTask;
         }
 
         if (!migration.Deltas.Any())
         {
-            return;
+            return Task.CompletedTask;
         }
 
         migration.AssertPatchingIsValid(autoCreate);
 
         logger ??= new DefaultMigrationLogger();
-        await executeDelta(migration, conn, autoCreate, logger, ct).ConfigureAwait(false);
+        return executeDelta(migration, conn, autoCreate, logger, ct);
     }
 
     protected abstract Task executeDelta(

diff --git a/src/Weasel.Postgresql.Tests/Migrations/migration_scenario_tests.cs b/src/Weasel.Postgresql.Tests/Migrations/migration_scenario_tests.cs
@@ -151,6 +151,14 @@ protected override IEnumerable<ISchemaObject> schemaObjects()
 
 public class DatabaseWithTables: PostgresqlDatabase
 {
+    public static DatabaseWithTables ForDataSource(NpgsqlDataSource dataSource)
+    {
+        var builder = new NpgsqlConnectionStringBuilder(dataSource.ConnectionString);
+        var identifier = builder.Database;
+
+        return new DatabaseWithTables(identifier, dataSource);
+    }
+
     public static DatabaseWithTables ForConnectionString(string connectionString)
     {
         var builder = new NpgsqlConnectionStringBuilder(connectionString);
@@ -169,6 +177,11 @@ public DatabaseWithTables(string identifier, string connectionString)
     {
     }
 
+    public DatabaseWithTables(string identifier, NpgsqlDataSource dataSource)
+        : base(new DefaultMigrationLogger(), AutoCreate.All, new PostgresqlMigrator(), identifier, dataSource)
+    {
+    }
+
     public LightweightCache<string, NamedTableFeature> Features { get; } =
         new LightweightCache<string, NamedTableFeature>(name =>
             new NamedTableFeature(name, new PostgresqlMigrator()));

diff --git a/src/Weasel.Postgresql/Migrations/SingleServerDatabaseCollection.cs b/src/Weasel.Postgresql/Migrations/SingleServerDatabaseCollection.cs
@@ -11,10 +11,17 @@ namespace Weasel.Postgresql.Migrations;
 /// <typeparam name="T"></typeparam>
 public abstract class SingleServerDatabaseCollection<T> where T : PostgresqlDatabase
 {
+    private readonly NpgsqlDataSource? npgsqlDataSource;
     private readonly TimedLock _lock = new();
     private readonly string _masterConnectionString;
     private ImHashMap<string, T> _databases = ImHashMap<string, T>.Empty;
 
+    protected SingleServerDatabaseCollection(NpgsqlDataSource npgsqlDataSource)
+    {
+        this.npgsqlDataSource = npgsqlDataSource;
+        _masterConnectionString = npgsqlDataSource.ConnectionString;
+    }
+
     protected SingleServerDatabaseCollection(string masterConnectionString)
     {
         _masterConnectionString = masterConnectionString;
@@ -48,7 +55,9 @@ public virtual async ValueTask<T> FindOrCreateDatabase(string databaseName, Canc
                 return database;
             }
 
-            await using var conn = new NpgsqlConnection(_masterConnectionString);
+            await using var conn =
+                npgsqlDataSource?.CreateConnection()
+                ?? new NpgsqlConnection(_masterConnectionString);
             await conn.OpenAsync(ct).ConfigureAwait(false);
 
             if (DropAndRecreate)

diff --git a/src/Weasel.Postgresql/PostgresqlDatabase.cs b/src/Weasel.Postgresql/PostgresqlDatabase.cs
@@ -17,6 +17,16 @@ string connectionString
     {
     }
 
+    protected PostgresqlDatabase(
+        IMigrationLogger logger,
+        AutoCreate autoCreate,
+        Migrator migrator,
+        string identifier,
+        NpgsqlDataSource dataSource
+    ): base(logger, autoCreate, migrator, identifier, dataSource.CreateConnection)
+    {
+    }
+
     protected PostgresqlDatabase(
         IMigrationLogger logger,
         AutoCreate autoCreate,