Add RECOMPILE query hint to package dependents SQL (#8096)

Progress on #8078

Author: joelverhagen

src/NuGetGallery.Core/Entities/EntitiesContext.cs

@@ -56,7 +56,18 @@ public EntitiesContext(DbConnection connection, bool readOnly)
            ReadOnly = readOnly;
         }
 
+        public IDisposable WithQueryHint(string queryHint)
+        {
+            if (QueryHint != null)
+            {
+                throw new InvalidOperationException("A query hint is already applied.");
+            }
+
+            return new QueryHintScope(this, queryHint);
+        }
+
         public bool ReadOnly { get; private set; }
+        public string QueryHint { get; private set; }
         public DbSet<Package> Packages { get; set; }
         public DbSet<PackageDeprecation> Deprecations { get; set; }
         public DbSet<PackageRegistration> PackageRegistrations { get; set; }
@@ -503,5 +514,21 @@ protected override void OnModelCreating(DbModelBuilder modelBuilder)
                 .WillCascadeOnDelete(false);
         }
 #pragma warning restore 618
+
+        private class QueryHintScope : IDisposable
+        {
+            private readonly EntitiesContext _entitiesContext;
+
+            public QueryHintScope(EntitiesContext entitiesContext, string queryHint)
+            {
+                _entitiesContext = entitiesContext ?? throw new ArgumentNullException(nameof(entitiesContext));
+                _entitiesContext.QueryHint = queryHint;
+            }
+
+            public void Dispose()
+            {
+                _entitiesContext.QueryHint = null;
+            }
+        }
     }
 }

src/NuGetGallery.Core/Entities/IReadOnlyEntitiesContext.cs

@@ -14,5 +14,19 @@ public interface IReadOnlyEntitiesContext : IDisposable
         DbSet<T> Set<T>() where T : class;
 
         void SetCommandTimeout(int? seconds);
+
+        /// <summary>
+        /// Sets the <see cref="QueryHint"/> for queries to the provided value until the returned disposable is disposed.
+        /// Note that this method MUST NOT be called with user input.
+        /// </summary>
+        /// <param name="queryHint">The query hint.</param>
+        /// <example>Provide "RECOMPILE" to disable query plan caching.</example>
+        /// <returns>A disposable that determines the lifetime of the provided query hint.</returns>
+        IDisposable WithQueryHint(string queryHint);
+
+        /// <summary>
+        /// The current query hint to use for queries. Can be set using <see cref="WithQueryHint(string)"/>.
+        /// </summary>
+        string QueryHint { get; }
     }
 }

src/NuGetGallery.Core/Entities/QueryHintInterceptor.cs

@@ -0,0 +1,37 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Data.Common;
+using System.Data.Entity.Infrastructure.Interception;
+using System.Diagnostics.CodeAnalysis;
+
+namespace NuGetGallery
+{
+    /// <summary>
+    /// A global (static) interceptor for Entity Framework queries. This is used for
+    /// <see cref="IReadOnlyEntitiesContext.WithQueryHint(string)"/> to set a query hint for a short window of time on
+    /// a specific entity context. Inspired by: https://stackoverflow.com/a/45170243
+    /// </summary>
+    public class QueryHintInterceptor : DbCommandInterceptor
+    {
+        [SuppressMessage("Microsoft.Security", "CA2100:Review SQL queries for security vulnerabilities", Justification = "The query hint is not customer input.")]
+        public override void ReaderExecuting(DbCommand command, DbCommandInterceptionContext<DbDataReader> interceptionContext)
+        {
+            foreach (var dbContext in interceptionContext.DbContexts)
+            {
+                if (dbContext is IReadOnlyEntitiesContext entitiesContext)
+                {
+                    var queryHint = entitiesContext.QueryHint;
+                    if (queryHint != null)
+                    {
+                        command.CommandText += $" OPTION ( {queryHint} )";
+                    }
+
+                    break;
+                }
+            }
+
+            base.ReaderExecuting(command, interceptionContext);
+        }
+    }
+}

src/NuGetGallery.Core/Entities/ReadOnlyEntitiesContext.cs

@@ -30,6 +30,8 @@ public DbSet<Package> Packages
             }
         }
 
+        public string QueryHint => _entitiesContext.QueryHint;
+
         DbSet<T> IReadOnlyEntitiesContext.Set<T>()
         {
             return _entitiesContext.Set<T>();
@@ -44,5 +46,7 @@ public void Dispose()
         {
             _entitiesContext.Dispose();
         }
+
+        public IDisposable WithQueryHint(string queryHint) => _entitiesContext.WithQueryHint(queryHint);
     }
 }

src/NuGetGallery.Core/NuGetGallery.Core.csproj

@@ -121,6 +121,7 @@
     <Compile Include="DisposableAction.cs" />
     <Compile Include="Entities\DatabaseWrapper.cs" />
     <Compile Include="Entities\DbContextTransactionWrapper.cs" />
+    <Compile Include="Entities\QueryHintInterceptor.cs" />
     <Compile Include="Entities\ReadOnlyEntitiesContext.cs" />
     <Compile Include="Entities\IReadOnlyEntitiesContext.cs" />
     <Compile Include="Entities\ReadOnlyEntityRepository.cs" />

src/NuGetGallery.Services/PackageManagement/PackageService.cs

@@ -155,8 +155,13 @@ public PackageDependents GetPackageDependents(string id)
             }
 
             PackageDependents result = new PackageDependents();
-            result.TopPackages = GetListOfDependents(id);
-            result.TotalPackageCount = GetDependentCount(id);
+            
+            using (_entitiesContext.WithQueryHint("RECOMPILE"))
+            {
+                result.TopPackages = GetListOfDependents(id);
+                result.TotalPackageCount = GetDependentCount(id);
+            }
+
             return result;
         }
 

src/NuGetGallery/App_Data/Files/Content/AB-Test-Configuration.json

@@ -1,5 +1,5 @@
 {
   "PreviewSearchPercentage": 0,
   "PreviewHijackPercentage": 0,
-  "DependentsPercentage": 0
+  "DependentsPercentage": 100
 }

src/NuGetGallery/App_Data/Files/Content/flags.json

@@ -48,8 +48,8 @@
       "Domains": []
     },
     "NuGetGallery.PackageDependents": {
-      "All": false,
-      "SiteAdmins": true,
+      "All": true,
+      "SiteAdmins": false,
       "Accounts": [],
       "Domains": []
     }

src/NuGetGallery/App_Start/DefaultDependenciesModule.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Data.Common;
 using System.Data.Entity;
+using System.Data.Entity.Infrastructure.Interception;
 using System.IO;
 using System.Linq;
 using System.Net;
@@ -172,6 +173,8 @@ protected override void Load(ContainerBuilder builder)
                 .As<ICacheService>()
                 .InstancePerLifetimeScope();
 
+            DbInterception.Add(new QueryHintInterceptor());
+
             var galleryDbConnectionFactory = CreateDbConnectionFactory(
                 loggerFactory,
                 nameof(EntitiesContext),

src/VerifyMicrosoftPackage/Fakes/FakeEntitiesContext.cs

@@ -27,6 +27,8 @@ class FakeEntitiesContext : IEntitiesContext
         public DbSet<Package> Packages { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
         public DbSet<PackageRename> PackageRenames { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
 
+        public string QueryHint => throw new NotImplementedException();
+
         public void DeleteOnCommit<T>(T entity) where T : class
         {
             throw new NotImplementedException();
@@ -56,5 +58,10 @@ public void SetCommandTimeout(int? seconds)
         {
             throw new NotImplementedException();
         }
+
+        public IDisposable WithQueryHint(string queryHint)
+        {
+            throw new NotImplementedException();
+        }
     }
 }