diff --git a/src/Trax.Effect.Data.Postgres/Extensions/ModelBuilderExtensions.cs b/src/Trax.Effect.Data.Postgres/Extensions/ModelBuilderExtensions.cs
index 6f1eb77..904ed27 100644
--- a/src/Trax.Effect.Data.Postgres/Extensions/ModelBuilderExtensions.cs
+++ b/src/Trax.Effect.Data.Postgres/Extensions/ModelBuilderExtensions.cs
@@ -67,7 +67,20 @@ public static ModelBuilder AddPostgresEnums(this ModelBuilder modelBuilder)
/// This method is typically called when setting up the DbContextFactory in the
/// AddPostgresEffect extension method.
///
- public static NpgsqlDataSource BuildDataSource(string connectionString)
+ public static NpgsqlDataSource BuildDataSource(string connectionString) =>
+ BuildDataSource(connectionString, configure: null);
+
+ ///
+ /// Builds a PostgreSQL data source with enum mappings and optional custom configuration.
+ ///
+ /// The connection string to use
+ /// Optional action to configure the data source builder after enum mappings
+ /// (e.g., set ConnectionStringBuilder.MaxPoolSize or enable multiplexing)
+ /// A configured NpgsqlDataSource
+ public static NpgsqlDataSource BuildDataSource(
+ string connectionString,
+ Action? configure
+ )
{
var npgsqlDataSourceBuilder = new NpgsqlDataSourceBuilder(connectionString);
@@ -78,6 +91,8 @@ public static NpgsqlDataSource BuildDataSource(string connectionString)
npgsqlDataSourceBuilder.MapEnum("trax.work_queue_status");
npgsqlDataSourceBuilder.MapEnum("trax.misfire_policy");
+ configure?.Invoke(npgsqlDataSourceBuilder);
+
return npgsqlDataSourceBuilder.Build();
}
diff --git a/src/Trax.Effect.Data.Postgres/Extensions/ServiceExtensions.cs b/src/Trax.Effect.Data.Postgres/Extensions/ServiceExtensions.cs
index f029904..ab6d1a7 100644
--- a/src/Trax.Effect.Data.Postgres/Extensions/ServiceExtensions.cs
+++ b/src/Trax.Effect.Data.Postgres/Extensions/ServiceExtensions.cs
@@ -87,6 +87,25 @@ public static TraxEffectBuilder SkipMigrations(this TraxEffectBuilder builder)
public static TraxEffectBuilderWithData UsePostgres(
this TraxEffectBuilder configurationBuilder,
string connectionString
+ ) => configurationBuilder.UsePostgres(connectionString, configureDataSource: null);
+
+ ///
+ /// Adds PostgreSQL database support to the Trax.Effect system with custom data source configuration.
+ ///
+ /// The Trax.Core effect configuration builder
+ /// The connection string to the PostgreSQL database
+ /// Optional action to configure the
+ /// after enum mappings. Use this to tune connection pool settings, enable multiplexing, or add custom type mappings.
+ ///
+ /// effects.UsePostgres(connectionString, dataSource =>
+ /// dataSource.ConnectionStringBuilder.MaxPoolSize = 50);
+ ///
+ ///
+ /// The configuration builder for method chaining
+ public static TraxEffectBuilderWithData UsePostgres(
+ this TraxEffectBuilder configurationBuilder,
+ string connectionString,
+ Action? configureDataSource
)
{
// Migrate the database schema to the latest version (unless explicitly skipped)
@@ -94,7 +113,10 @@ string connectionString
DatabaseMigrator.Migrate(connectionString).Wait();
// Create a data source with enum mappings and register for disposal on shutdown
- var dataSource = ModelBuilderExtensions.BuildDataSource(connectionString);
+ var dataSource = ModelBuilderExtensions.BuildDataSource(
+ connectionString,
+ configureDataSource
+ );
configurationBuilder.ServiceCollection.AddSingleton(dataSource);
// Register the DbContextFactory
diff --git a/tests/Trax.Effect.Tests.Integration/UnitTests/Configuration/SkipMigrationsTests.cs b/tests/Trax.Effect.Tests.Integration/UnitTests/Configuration/SkipMigrationsTests.cs
index bb6e3df..ec4001d 100644
--- a/tests/Trax.Effect.Tests.Integration/UnitTests/Configuration/SkipMigrationsTests.cs
+++ b/tests/Trax.Effect.Tests.Integration/UnitTests/Configuration/SkipMigrationsTests.cs
@@ -137,4 +137,66 @@ public void UsePostgres_WithSkipMigrations_StillRegistersDataContext()
}
#endregion
+
+ #region UsePostgres with ConfigureDataSource
+
+ [Test]
+ public void UsePostgres_WithConfigureAction_DoesNotThrow()
+ {
+ var act = () =>
+ new ServiceCollection().AddTrax(trax =>
+ trax.AddEffects(effects =>
+ effects
+ .SkipMigrations()
+ .UsePostgres(
+ UnreachableConnectionString,
+ dataSource => dataSource.ConnectionStringBuilder.MaxPoolSize = 50
+ )
+ )
+ );
+
+ act.Should().NotThrow();
+ }
+
+ [Test]
+ public void UsePostgres_WithConfigureAction_StillRegistersDbContextFactory()
+ {
+ var services = new ServiceCollection();
+
+ services.AddTrax(trax =>
+ trax.AddEffects(effects =>
+ effects
+ .SkipMigrations()
+ .UsePostgres(
+ UnreachableConnectionString,
+ dataSource => dataSource.ConnectionStringBuilder.MaxPoolSize = 50
+ )
+ )
+ );
+
+ services
+ .Should()
+ .Contain(sd =>
+ sd.ServiceType.IsGenericType
+ && sd.ServiceType.GetGenericTypeDefinition()
+ == typeof(Microsoft.EntityFrameworkCore.IDbContextFactory<>)
+ );
+ }
+
+ [Test]
+ public void UsePostgres_WithNullConfigureAction_BehavesLikeDefault()
+ {
+ var act = () =>
+ new ServiceCollection().AddTrax(trax =>
+ trax.AddEffects(effects =>
+ effects
+ .SkipMigrations()
+ .UsePostgres(UnreachableConnectionString, configureDataSource: null)
+ )
+ );
+
+ act.Should().NotThrow();
+ }
+
+ #endregion
}