From 9ee5fe53abe51f567edc6078021fd0655fffdfa8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?K=C3=A9vin=20Chalet?= <kevinchalet@gmail.com>
Date: Sat, 3 Feb 2024 17:46:17 +0100
Subject: [PATCH] Update the EF Core stores to resolve the navigation
 properties via the change tracker when possible

---
 Directory.Build.targets                       |  1 +
 ...ctEntityFrameworkCoreAuthorizationStore.cs | 23 ++++++++--
 ...OpenIddictEntityFrameworkCoreTokenStore.cs | 46 +++++++++++++++----
 3 files changed, 55 insertions(+), 15 deletions(-)

diff --git a/Directory.Build.targets b/Directory.Build.targets
index 51c5ac51..91bb05f5 100644
--- a/Directory.Build.targets
+++ b/Directory.Build.targets
@@ -80,6 +80,7 @@
     Condition=" ('$(TargetFrameworkIdentifier)' == '.NETCoreApp'  And $([MSBuild]::VersionGreaterThanOrEquals($(TargetFrameworkVersion), '3.0'))) Or
                 ('$(TargetFrameworkIdentifier)' == '.NETStandard' And $([MSBuild]::VersionGreaterThanOrEquals($(TargetFrameworkVersion), '2.1'))) ">
     <DefineConstants>$(DefineConstants);SUPPORTS_BCL_ASYNC_ENUMERABLE</DefineConstants>
+    <DefineConstants>$(DefineConstants);SUPPORTS_DBSET_VALUETASK_FINDASYNC</DefineConstants>
   </PropertyGroup>
 
   <PropertyGroup
diff --git a/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreAuthorizationStore.cs b/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreAuthorizationStore.cs
index 69edb1b0..bf09aa67 100644
--- a/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreAuthorizationStore.cs
+++ b/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreAuthorizationStore.cs
@@ -802,14 +802,27 @@ public class OpenIddictEntityFrameworkCoreAuthorizationStore<TAuthorization, TAp
 
         if (!string.IsNullOrEmpty(identifier))
         {
+#if SUPPORTS_DBSET_VALUETASK_FINDASYNC
+            authorization.Application = await Applications.FindAsync([ConvertIdentifierFromString(identifier)], cancellationToken);
+#else
+            // Warning: when targeting older TFMs, FindAsync() is deliberately not used to work around a breaking
+            // change introduced in Entity Framework Core 3.x (where a ValueTask instead of a Task is now returned).
+
             var key = ConvertIdentifierFromString(identifier);
 
-            // Warning: FindAsync() is deliberately not used to work around a breaking change introduced
-            // in Entity Framework Core 3.x (where a ValueTask instead of a Task is now returned).
-            authorization.Application = await Applications.AsQueryable()
-                .AsTracking()
-                .FirstOrDefaultAsync(application => application.Id!.Equals(key), cancellationToken) ??
+            authorization.Application = GetTrackedEntity() ?? await QueryAsync() ??
                 throw new InvalidOperationException(SR.GetResourceString(SR.ID0244));
+
+            TApplication? GetTrackedEntity() =>
+                (from entry in Context.ChangeTracker.Entries<TApplication>()
+                 where entry.Entity.Id is TKey identifier && identifier.Equals(key)
+                 select entry.Entity).FirstOrDefault();
+
+            Task<TApplication?> QueryAsync() =>
+                (from application in Applications.AsTracking()
+                 where application.Id!.Equals(key)
+                 select application).FirstOrDefaultAsync(cancellationToken);
+#endif
         }
 
         else
diff --git a/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreTokenStore.cs b/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreTokenStore.cs
index cf5920e8..4700ba07 100644
--- a/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreTokenStore.cs
+++ b/src/OpenIddict.EntityFrameworkCore/Stores/OpenIddictEntityFrameworkCoreTokenStore.cs
@@ -802,14 +802,27 @@ public class OpenIddictEntityFrameworkCoreTokenStore<TToken, TApplication, TAuth
 
         if (!string.IsNullOrEmpty(identifier))
         {
+#if SUPPORTS_DBSET_VALUETASK_FINDASYNC
+            token.Application = await Applications.FindAsync([ConvertIdentifierFromString(identifier)], cancellationToken);
+#else
+            // Warning: when targeting older TFMs, FindAsync() is deliberately not used to work around a breaking
+            // change introduced in Entity Framework Core 3.x (where a ValueTask instead of a Task is now returned).
+
             var key = ConvertIdentifierFromString(identifier);
 
-            // Warning: FindAsync() is deliberately not used to work around a breaking change introduced
-            // in Entity Framework Core 3.x (where a ValueTask instead of a Task is now returned).
-            token.Application = await Applications.AsQueryable()
-                .AsTracking()
-                .FirstOrDefaultAsync(application => application.Id!.Equals(key), cancellationToken) ??
+            token.Application = GetTrackedEntity() ?? await QueryAsync() ??
                 throw new InvalidOperationException(SR.GetResourceString(SR.ID0250));
+
+            TApplication? GetTrackedEntity() =>
+                (from entry in Context.ChangeTracker.Entries<TApplication>()
+                 where entry.Entity.Id is TKey identifier && identifier.Equals(key)
+                 select entry.Entity).FirstOrDefault();
+
+            Task<TApplication?> QueryAsync() =>
+                (from application in Applications.AsTracking()
+                 where application.Id!.Equals(key)
+                 select application).FirstOrDefaultAsync(cancellationToken);
+#endif
         }
 
         else
@@ -840,14 +853,27 @@ public class OpenIddictEntityFrameworkCoreTokenStore<TToken, TApplication, TAuth
 
         if (!string.IsNullOrEmpty(identifier))
         {
+#if SUPPORTS_DBSET_VALUETASK_FINDASYNC
+            token.Authorization = await Authorizations.FindAsync([ConvertIdentifierFromString(identifier)], cancellationToken);
+#else
+            // Warning: when targeting older TFMs, FindAsync() is deliberately not used to work around a breaking
+            // change introduced in Entity Framework Core 3.x (where a ValueTask instead of a Task is now returned).
+
             var key = ConvertIdentifierFromString(identifier);
 
-            // Warning: FindAsync() is deliberately not used to work around a breaking change introduced
-            // in Entity Framework Core 3.x (where a ValueTask instead of a Task is now returned).
-            token.Authorization = await Authorizations.AsQueryable()
-                .AsTracking()
-                .FirstOrDefaultAsync(authorization => authorization.Id!.Equals(key), cancellationToken) ??
+            token.Authorization = GetTrackedEntity() ?? await QueryAsync() ??
                 throw new InvalidOperationException(SR.GetResourceString(SR.ID0251));
+
+            TAuthorization? GetTrackedEntity() =>
+                (from entry in Context.ChangeTracker.Entries<TAuthorization>()
+                 where entry.Entity.Id is TKey identifier && identifier.Equals(key)
+                 select entry.Entity).FirstOrDefault();
+
+            Task<TAuthorization?> QueryAsync() =>
+                (from authorization in Authorizations.AsTracking()
+                 where authorization.Id!.Equals(key)
+                 select authorization).FirstOrDefaultAsync(cancellationToken);
+#endif
         }
 
         else