From 2d9a6660fcd0ecab2cf954cf8d586e7bccae86c4 Mon Sep 17 00:00:00 2001 From: mysticmind Date: Mon, 4 Dec 2023 11:09:25 +0530 Subject: [PATCH] Add funtionality to auto reload Npgsql types for added Postgres extensions - Add funtionality to auto reload Npgsql types for PG extensions added via schema objects - Add unit test This fixes Marten issue https://github.com/JasperFx/marten/issues/2515 --- src/Weasel.Core/Migrations/Database.cs | 2 +- .../creating_and_dropping_databases.cs | 6 +-- .../Migrations/migration_scenario_tests.cs | 51 ++++++++++++++++--- src/Weasel.Postgresql/PostgresqlDatabase.cs | 31 +++++++++++ 4 files changed, 78 insertions(+), 12 deletions(-) diff --git a/src/Weasel.Core/Migrations/Database.cs b/src/Weasel.Core/Migrations/Database.cs index a1c5738c..079a21d8 100644 --- a/src/Weasel.Core/Migrations/Database.cs +++ b/src/Weasel.Core/Migrations/Database.cs @@ -214,7 +214,7 @@ public async Task CreateMigrationAsync(CancellationToken ct = d return await SchemaMigration.DetermineAsync(conn, ct, objects).ConfigureAwait(false); } - public Task ApplyAllConfiguredChangesToDatabaseAsync( + public virtual Task ApplyAllConfiguredChangesToDatabaseAsync( AutoCreate? @override = null, ReconnectionOptions? reconnectionOptions = null, CancellationToken ct = default diff --git a/src/Weasel.Postgresql.Tests/Migrations/creating_and_dropping_databases.cs b/src/Weasel.Postgresql.Tests/Migrations/creating_and_dropping_databases.cs index 90f8f85e..79bed432 100644 --- a/src/Weasel.Postgresql.Tests/Migrations/creating_and_dropping_databases.cs +++ b/src/Weasel.Postgresql.Tests/Migrations/creating_and_dropping_databases.cs @@ -33,15 +33,15 @@ public async Task can_build_databases_once() (await theDatabases.FindOrCreateDatabase("three")).ShouldBeSameAs(three); } - public class Databases: SingleServerDatabaseCollection + public class Databases: SingleServerDatabaseCollection { public Databases() : base(ConnectionSource.ConnectionString) { } - protected override DatabaseWithTables buildDatabase(string databaseName, string connectionString) + protected override TestDatabase buildDatabase(string databaseName, string connectionString) { - return new DatabaseWithTables(databaseName, connectionString); + return new TestDatabase(databaseName, connectionString); } } } diff --git a/src/Weasel.Postgresql.Tests/Migrations/migration_scenario_tests.cs b/src/Weasel.Postgresql.Tests/Migrations/migration_scenario_tests.cs index 46e8b210..4aa2d814 100644 --- a/src/Weasel.Postgresql.Tests/Migrations/migration_scenario_tests.cs +++ b/src/Weasel.Postgresql.Tests/Migrations/migration_scenario_tests.cs @@ -1,6 +1,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Runtime.CompilerServices; using System.Threading.Tasks; using Npgsql; using Shouldly; @@ -16,11 +17,11 @@ namespace Weasel.Postgresql.Tests.Migrations; [Collection("migrations")] public class SchemaMigrationTests : IntegrationContext, IAsyncLifetime { - private readonly DatabaseWithTables theDatabase; + private readonly TestDatabase theDatabase; public SchemaMigrationTests() : base("migrations") { - theDatabase = new DatabaseWithTables(AutoCreate.None, "Migrations"); + theDatabase = new TestDatabase(AutoCreate.None, "Migrations"); } public override Task InitializeAsync() @@ -116,6 +117,14 @@ await theDatabase.ApplyAllConfiguredChangesToDatabaseAsync( connectionGlobalLock.Failed.ShouldBeTrue(); connectionGlobalLock.Retried.ShouldBeFalse(); } + + [Fact] + public async Task test_auto_reload_of_npgsql_types() + { + theDatabase.ExtensionFeatures["Extensions"].AddExtension("pgcrypto"); + await theDatabase.ApplyAllConfiguredChangesToDatabaseAsync(); + theDatabase.HasNpgsqlTypesReloaded.ShouldBeTrue(); + } } public class NamedTable: Table @@ -149,22 +158,42 @@ protected override IEnumerable schemaObjects() } } -public class DatabaseWithTables: PostgresqlDatabase +public class ExtensionFeature: FeatureSchemaBase { - public static DatabaseWithTables ForConnectionString(string connectionString) + public Dictionary ExtensionObjects { get; } = new(); + + public ExtensionFeature(string identifier, Migrator migrator) : base(identifier, migrator) + { + } + + protected override IEnumerable schemaObjects() + { + return ExtensionObjects.Values; + } + + public void AddExtension(string extensionName) + { + var extension = new Extension(extensionName); + ExtensionObjects[extensionName] = extension; + } +} + +public class TestDatabase: PostgresqlDatabase +{ + public static TestDatabase ForConnectionString(string connectionString) { var builder = new NpgsqlConnectionStringBuilder(connectionString); var identifier = builder.Database; - return new DatabaseWithTables(identifier, connectionString); + return new TestDatabase(identifier, connectionString); } - public DatabaseWithTables(AutoCreate autoCreate, string identifier) + public TestDatabase(AutoCreate autoCreate, string identifier) : base(new DefaultMigrationLogger(), autoCreate, new PostgresqlMigrator(), identifier, ConnectionSource.ConnectionString) { } - public DatabaseWithTables(string identifier, string connectionString) + public TestDatabase(string identifier, string connectionString) : base(new DefaultMigrationLogger(), AutoCreate.All, new PostgresqlMigrator(), identifier, connectionString) { } @@ -173,9 +202,15 @@ public DatabaseWithTables(string identifier, string connectionString) new LightweightCache(name => new NamedTableFeature(name, new PostgresqlMigrator())); + public LightweightCache ExtensionFeatures { get; } = + new LightweightCache(name => + new ExtensionFeature(name, new PostgresqlMigrator())); + public override IFeatureSchema[] BuildFeatureSchemas() { - return Features.OfType().ToArray(); + var featureSchemas = Features.OfType() + .Append(ExtensionFeatures.OfType()).ToArray(); + return featureSchemas; } } diff --git a/src/Weasel.Postgresql/PostgresqlDatabase.cs b/src/Weasel.Postgresql/PostgresqlDatabase.cs index ead17af2..b957edff 100644 --- a/src/Weasel.Postgresql/PostgresqlDatabase.cs +++ b/src/Weasel.Postgresql/PostgresqlDatabase.cs @@ -1,3 +1,4 @@ +using JasperFx.Core; using Npgsql; using Weasel.Core; using Weasel.Core.Migrations; @@ -7,6 +8,12 @@ namespace Weasel.Postgresql; public abstract class PostgresqlDatabase: DatabaseBase { + /// + /// Provide status of whether the Npgsql types were reloaded + /// Currently used for assertion in unit tests + /// + public bool HasNpgsqlTypesReloaded { get; private set; } + protected PostgresqlDatabase( IMigrationLogger logger, AutoCreate autoCreate, @@ -50,4 +57,28 @@ public async Task> SchemaTables(CancellationToken ct return await conn.ExistingTablesAsync(schemas: schemaNames, ct: ct).ConfigureAwait(false); } + + public override async Task ApplyAllConfiguredChangesToDatabaseAsync(AutoCreate? @override = null, + ReconnectionOptions? reconnectionOptions = null, CancellationToken ct = default) + { + HasNpgsqlTypesReloaded = false; + var schemaPatchDiff = await base.ApplyAllConfiguredChangesToDatabaseAsync(@override, reconnectionOptions, ct).ConfigureAwait(false); + + // If there were extensions installed, automatically reload Npgsql types cache. + var hasExtensions = AllObjects().OfType().Any(); + if (hasExtensions) + { + await ReloadNpgsqlTypes(ct).ConfigureAwait(false); + } + + return schemaPatchDiff; + } + + private async Task ReloadNpgsqlTypes(CancellationToken ct = default) + { + await using var conn = CreateConnection(); + await conn.OpenAsync(ct).ConfigureAwait(false); + await conn.ReloadTypesAsync().ConfigureAwait(false); + HasNpgsqlTypesReloaded = true; + } }