From d39502674e95d3a0474517c491b3ebc269bc4e72 Mon Sep 17 00:00:00 2001
From: Bjarke Berg <mail@bergmania.dk>
Date: Fri, 22 Sep 2023 10:38:55 +0200
Subject: [PATCH] Added our own db context pool, that basically bypass the bool
 until umbraco is in Run mode (#14852)

---
 ...mbracoEFCoreServiceCollectionExtensions.cs | 18 ++++++-
 .../UmbracoPooledDbContextFactory.cs          | 47 +++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 src/Umbraco.Cms.Persistence.EFCore/Factories/UmbracoPooledDbContextFactory.cs

diff --git a/src/Umbraco.Cms.Persistence.EFCore/Extensions/UmbracoEFCoreServiceCollectionExtensions.cs b/src/Umbraco.Cms.Persistence.EFCore/Extensions/UmbracoEFCoreServiceCollectionExtensions.cs
index da9c2e59ef..ded5be40fd 100644
--- a/src/Umbraco.Cms.Persistence.EFCore/Extensions/UmbracoEFCoreServiceCollectionExtensions.cs
+++ b/src/Umbraco.Cms.Persistence.EFCore/Extensions/UmbracoEFCoreServiceCollectionExtensions.cs
@@ -1,11 +1,13 @@
 using Microsoft.EntityFrameworkCore;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.Extensions.Options;
 using Umbraco.Cms.Core;
 using Umbraco.Cms.Core.Configuration.Models;
 using Umbraco.Cms.Core.DistributedLocking;
+using Umbraco.Cms.Core.Services;
+using Umbraco.Cms.Persistence.EFCore.Factories;
 using Umbraco.Cms.Persistence.EFCore.Locking;
-using Umbraco.Cms.Persistence.EFCore.Migrations;
 using Umbraco.Cms.Persistence.EFCore.Scoping;
 
 namespace Umbraco.Extensions;
@@ -17,6 +19,13 @@ public static class UmbracoEFCoreServiceCollectionExtensions
     public static IServiceCollection AddUmbracoEFCoreContext<T>(this IServiceCollection services, DefaultEFCoreOptionsAction? defaultEFCoreOptionsAction = null)
         where T : DbContext
     {
+        var optionsBuilder = new DbContextOptionsBuilder<T>();
+        services.TryAddSingleton<IDbContextFactory<T>>(
+            sp =>
+            {
+                SetupDbContext(defaultEFCoreOptionsAction, sp, optionsBuilder);
+                return new UmbracoPooledDbContextFactory<T>(sp.GetRequiredService<IRuntimeState>(),optionsBuilder.Options);
+            });
         services.AddPooledDbContextFactory<T>((provider, builder) => SetupDbContext(defaultEFCoreOptionsAction, provider, builder));
         services.AddTransient(services => services.GetRequiredService<IDbContextFactory<T>>().CreateDbContext());
 
@@ -39,6 +48,13 @@ public static class UmbracoEFCoreServiceCollectionExtensions
             connectionString = connectionString.Replace(Constants.System.DataDirectoryPlaceholder, dataDirectory);
         }
 
+        var optionsBuilder = new DbContextOptionsBuilder<T>();
+        services.TryAddSingleton<IDbContextFactory<T>>(
+            sp =>
+            {
+                SetupDbContext(defaultEFCoreOptionsAction, sp, optionsBuilder);
+                return new UmbracoPooledDbContextFactory<T>(sp.GetRequiredService<IRuntimeState>(),optionsBuilder.Options);
+            });
         services.AddPooledDbContextFactory<T>(options => defaultEFCoreOptionsAction?.Invoke(options, providerName, connectionString));
         services.AddTransient(services => services.GetRequiredService<IDbContextFactory<T>>().CreateDbContext());
 
diff --git a/src/Umbraco.Cms.Persistence.EFCore/Factories/UmbracoPooledDbContextFactory.cs b/src/Umbraco.Cms.Persistence.EFCore/Factories/UmbracoPooledDbContextFactory.cs
new file mode 100644
index 0000000000..de0f7db200
--- /dev/null
+++ b/src/Umbraco.Cms.Persistence.EFCore/Factories/UmbracoPooledDbContextFactory.cs
@@ -0,0 +1,47 @@
+using Microsoft.EntityFrameworkCore;
+using Microsoft.EntityFrameworkCore.Infrastructure;
+using Umbraco.Cms.Core;
+using Umbraco.Cms.Core.Services;
+
+namespace Umbraco.Cms.Persistence.EFCore.Factories;
+
+/// <inheritdoc/>
+internal class UmbracoPooledDbContextFactory<TContext> : PooledDbContextFactory<TContext>
+    where TContext : DbContext
+{
+    private readonly IRuntimeState _runtimeState;
+    private readonly DbContextOptions<TContext> _options;
+
+    /// <inheritdoc/>
+    public UmbracoPooledDbContextFactory(IRuntimeState runtimeState, DbContextOptions<TContext> options, int poolSize = 1024 /*DbContextPool<DbContext>.DefaultPoolSize*/) : base(options, poolSize)
+    {
+        _runtimeState = runtimeState;
+        _options = options;
+    }
+
+    /// <inheritdoc/>
+    public override TContext CreateDbContext()
+    {
+        if (_runtimeState.Level == RuntimeLevel.Run)
+        {
+            return base.CreateDbContext();
+        }
+        else
+        {
+            return (TContext?)Activator.CreateInstance(typeof(TContext), _options) ?? throw new InvalidOperationException("Unable to create DbContext");
+        }
+    }
+
+    /// <inheritdoc/>
+    public override async Task<TContext> CreateDbContextAsync(CancellationToken cancellationToken = default)
+    {
+        if (_runtimeState.Level == RuntimeLevel.Run)
+        {
+            return await base.CreateDbContextAsync(cancellationToken);
+        }
+        else
+        {
+            return (TContext?)Activator.CreateInstance(typeof(TContext), _options) ?? throw new InvalidOperationException("Unable to create DbContext");
+        }
+    }
+}