Add header support for caching (#2650)

## Why make this change?

Closes #2604

## What is this change?

Inside of the `SqlQueryStructure`, there is a private constructor, that
all of the public constructors will call. This private constructor
offers a single point where all of the queries will bottleneck, and
therefore we add the cache control information to the query structure at
this point. Then when we are in the `SqlQueryEngine`, we can check for
this cache control information and handle the query appropriately.

The cache control options can be found

here: #2253

and here:
https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Cache-Control#request_directives

But for reference, they are `no-cache`, `no-store`, `only-if-cached`,
which will mean we do not get from the cache, we do not store in the
cache, and if there is a cache miss we return gateway timeout,
respectively.

## How was this tested?

Run against test suite.

## Sample Request(s)

To test you need to include the relevant cache headers in the request,
"no-cache", "no-store", or "only-if-cached"

---------

Co-authored-by: Copilot <[email protected]>
Co-authored-by: Aniruddh Munde <[email protected]>
Co-authored-by: RubenCerna2079 <[email protected]>

Author: 4 people

src/Config/ObjectModel/RuntimeConfig.cs

@@ -22,7 +22,7 @@ public record RuntimeConfig
 
     public RuntimeOptions? Runtime { get; init; }
 
-    public RuntimeEntities Entities { get; init; }
+    public virtual RuntimeEntities Entities { get; init; }
 
     public DataSourceFiles? DataSourceFiles { get; init; }
 
@@ -325,7 +325,7 @@ public RuntimeConfig(string Schema, DataSource DataSource, RuntimeOptions Runtim
     /// <param name="dataSourceName">Name of datasource.</param>
     /// <returns>DataSource object.</returns>
     /// <exception cref="DataApiBuilderException">Not found exception if key is not found.</exception>
-    public DataSource GetDataSourceFromDataSourceName(string dataSourceName)
+    public virtual DataSource GetDataSourceFromDataSourceName(string dataSourceName)
     {
         CheckDataSourceNamePresent(dataSourceName);
         return _dataSourceNameToDataSource[dataSourceName];
@@ -430,7 +430,7 @@ Runtime is not null && Runtime.Host is not null
     /// <param name="entityName">Name of the entity to check cache configuration.</param>
     /// <returns>Number of seconds (ttl) that a cache entry should be valid before cache eviction.</returns>
     /// <exception cref="DataApiBuilderException">Raised when an invalid entity name is provided or if the entity has caching disabled.</exception>
-    public int GetEntityCacheEntryTtl(string entityName)
+    public virtual int GetEntityCacheEntryTtl(string entityName)
     {
         if (!Entities.TryGetValue(entityName, out Entity? entityConfig))
         {
@@ -464,7 +464,7 @@ public int GetEntityCacheEntryTtl(string entityName)
     /// - whether the datasource is SQL and session context is disabled.
     /// </summary>
     /// <returns>Whether cache operations should proceed.</returns>
-    public bool CanUseCache()
+    public virtual bool CanUseCache()
     {
         bool setSessionContextEnabled = DataSource.GetTypedOptions<MsSqlOptions>()?.SetSessionContext ?? true;
         return IsCachingEnabled && !setSessionContextEnabled;

src/Core/Configurations/RuntimeConfigProvider.cs

@@ -117,7 +117,7 @@ public string ConfigFilePath
     /// <returns>The RuntimeConfig instance.</returns>
     /// <remark>Dont use this method if environment variable references need to be retained.</remark>
     /// <exception cref="DataApiBuilderException">Thrown when the loader is unable to load an instance of the config from its known location.</exception>
-    public RuntimeConfig GetConfig()
+    public virtual RuntimeConfig GetConfig()
     {
         if (_configLoader.RuntimeConfig is not null)
         {

src/Core/Resolvers/BaseQueryStructure.cs

@@ -19,7 +19,7 @@ public class BaseQueryStructure
         /// <summary>
         /// The Entity name associated with this query as appears in the config file.
         /// </summary>
-        public string EntityName { get; set; }
+        public virtual string EntityName { get; set; }
 
         /// <summary>
         /// The alias of the entity as used in the generated query.
@@ -74,7 +74,7 @@ public class BaseQueryStructure
         /// DbPolicyPredicates is a string that represents the filter portion of our query
         /// in the WHERE Clause added by virtue of the database policy.
         /// </summary>
-        public Dictionary<EntityActionOperation, string?> DbPolicyPredicatesForOperations { get; set; } = new();
+        public virtual Dictionary<EntityActionOperation, string?> DbPolicyPredicatesForOperations { get; set; } = new();
 
         public const string PARAM_NAME_PREFIX = "@";
 
@@ -156,7 +156,7 @@ public string CreateTableAlias()
         /// <summary>
         /// Returns the SourceDefinitionDefinition for the entity(table/view) of this query.
         /// </summary>
-        public SourceDefinition GetUnderlyingSourceDefinition()
+        public virtual SourceDefinition GetUnderlyingSourceDefinition()
         {
             return MetadataProvider.GetSourceDefinition(EntityName);
         }

src/Core/Resolvers/Sql Query Structures/SqlQueryStructure.cs

@@ -92,6 +92,20 @@ public class SqlQueryStructure : BaseSqlQueryStructure
         /// </summary>
         public GroupByMetadata GroupByMetadata { get; private set; }
 
+        public virtual string? CacheControlOption { get; set; }
+
+        public const string CACHE_CONTROL = "Cache-Control";
+
+        public const string CACHE_CONTROL_NO_STORE = "no-store";
+
+        public const string CACHE_CONTROL_NO_CACHE = "no-cache";
+
+        public const string CACHE_CONTROL_ONLY_IF_CACHED = "only-if-cached";
+
+        public HashSet<string> cacheControlHeaderOptions = new(
+            new[] { CACHE_CONTROL_NO_STORE, CACHE_CONTROL_NO_CACHE, CACHE_CONTROL_ONLY_IF_CACHED },
+            StringComparer.OrdinalIgnoreCase);
+
         /// <summary>
         /// Generate the structure for a SQL query based on GraphQL query
         /// information.
@@ -150,6 +164,7 @@ public SqlQueryStructure(
             : this(sqlMetadataProvider,
                   authorizationResolver,
                   gQLFilterParser,
+                  gQLFilterParser.GetHttpContextFromMiddlewareContext(ctx).Request.Headers,
                   predicates: null,
                   entityName: entityName,
                   counter: counter)
@@ -217,6 +232,7 @@ public SqlQueryStructure(
             }
 
             HttpContext httpContext = GraphQLFilterParser.GetHttpContextFromMiddlewareContext(ctx);
+
             // Process Authorization Policy of the entity being processed.
             AuthorizationPolicyHelpers.ProcessAuthorizationPolicies(EntityActionOperation.Read, queryStructure: this, httpContext, authorizationResolver, sqlMetadataProvider);
 
@@ -255,6 +271,7 @@ public SqlQueryStructure(
             : this(sqlMetadataProvider,
                 authorizationResolver,
                 gQLFilterParser,
+                httpRequestHeaders: httpContext?.Request.Headers,
                 predicates: null,
                 entityName: context.EntityName,
                 counter: new IncrementingInteger(),
@@ -380,10 +397,10 @@ private SqlQueryStructure(
             : this(sqlMetadataProvider,
                   authorizationResolver,
                   gQLFilterParser,
+                  gQLFilterParser.GetHttpContextFromMiddlewareContext(ctx).Request.Headers,
                   predicates: null,
                   entityName: entityName,
-                  counter: counter
-                  )
+                  counter: counter)
         {
             _ctx = ctx;
             IOutputType outputType = schemaField.Type;
@@ -541,6 +558,7 @@ private SqlQueryStructure(
             ISqlMetadataProvider metadataProvider,
             IAuthorizationResolver authorizationResolver,
             GQLFilterParser gQLFilterParser,
+            IHeaderDictionary? httpRequestHeaders,
             List<Predicate>? predicates = null,
             string entityName = "",
             IncrementingInteger? counter = null,
@@ -559,6 +577,25 @@ private SqlQueryStructure(
             ColumnLabelToParam = new();
             FilterPredicates = string.Empty;
             OrderByColumns = new();
+            AddCacheControlOptions(httpRequestHeaders);
+        }
+
+        private void AddCacheControlOptions(IHeaderDictionary? httpRequestHeaders)
+        {
+            // Set the cache control based on the request header if it exists.
+            if (httpRequestHeaders is not null && httpRequestHeaders.TryGetValue(CACHE_CONTROL, out Microsoft.Extensions.Primitives.StringValues cacheControlOption))
+            {
+                CacheControlOption = cacheControlOption;
+            }
+
+            if (!string.IsNullOrEmpty(CacheControlOption) &&
+                !cacheControlHeaderOptions.Contains(CacheControlOption))
+            {
+                throw new DataApiBuilderException(
+                    message: "Request Header Cache-Control is invalid: " + CacheControlOption,
+                    statusCode: HttpStatusCode.BadRequest,
+                    subStatusCode: DataApiBuilderException.SubStatusCodes.BadRequest);
+            }
         }
 
         /// <summary>

src/Core/Resolvers/SqlQueryEngine.cs

@@ -12,12 +12,14 @@
 using Azure.DataApiBuilder.Core.Services;
 using Azure.DataApiBuilder.Core.Services.Cache;
 using Azure.DataApiBuilder.Core.Services.MetadataProviders;
+using Azure.DataApiBuilder.Service.Exceptions;
 using Azure.DataApiBuilder.Service.GraphQLBuilder;
 using Azure.DataApiBuilder.Service.GraphQLBuilder.Queries;
 using HotChocolate.Resolvers;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.Extensions.Logging;
+using ZiggyCreatures.Caching.Fusion;
 using static Azure.DataApiBuilder.Service.GraphQLBuilder.GraphQLStoredProcedureBuilder;
 
 namespace Azure.DataApiBuilder.Core.Resolvers
@@ -328,12 +330,13 @@ public object ResolveList(JsonElement array, IObjectField fieldSchema, ref IMeta
                 // We want to avoid caching token metadata because token metadata can change frequently and we want to avoid caching it.
                 if (!dbPolicyConfigured && entityCacheEnabled)
                 {
-                    DatabaseQueryMetadata queryMetadata = new(queryText: queryString, dataSource: dataSourceName, queryParameters: structure.Parameters);
-                    JsonElement result = await _cache.GetOrSetAsync<JsonElement>(queryExecutor, queryMetadata, cacheEntryTtl: runtimeConfig.GetEntityCacheEntryTtl(entityName: structure.EntityName));
-                    byte[] jsonBytes = JsonSerializer.SerializeToUtf8Bytes(result);
-                    JsonDocument cacheServiceResponse = JsonDocument.Parse(jsonBytes);
-
-                    return cacheServiceResponse;
+                    return await GetResultInCacheScenario(
+                    runtimeConfig,
+                    structure,
+                    queryString,
+                    dataSourceName,
+                    queryExecutor
+                    );
                 }
             }
 
@@ -353,6 +356,82 @@ public object ResolveList(JsonElement array, IObjectField fieldSchema, ref IMeta
             return response;
         }
 
+        private async Task<JsonDocument?> GetResultInCacheScenario(RuntimeConfig runtimeConfig, SqlQueryStructure structure, string queryString, string dataSourceName, IQueryExecutor queryExecutor)
+        {
+            DatabaseQueryMetadata queryMetadata = new(queryText: queryString, dataSource: dataSourceName, queryParameters: structure.Parameters);
+            JsonElement? result;
+            MaybeValue<JsonElement?>? maybeResult;
+            switch (structure.CacheControlOption?.ToLowerInvariant())
+            {
+                // Do not get result from cache even if it exists, still cache result.
+                case SqlQueryStructure.CACHE_CONTROL_NO_CACHE:
+                    result = await queryExecutor.ExecuteQueryAsync(
+                        sqltext: queryMetadata.QueryText,
+                        parameters: queryMetadata.QueryParameters,
+                        dataReaderHandler: queryExecutor.GetJsonResultAsync<JsonElement>,
+                        httpContext: _httpContextAccessor.HttpContext!,
+                        args: null,
+                        dataSourceName: queryMetadata.DataSource);
+                    _cache.Set<JsonElement?>(
+                        queryMetadata,
+                        cacheEntryTtl: runtimeConfig.GetEntityCacheEntryTtl(entityName: structure.EntityName),
+                        result);
+                    return ParseResultIntoJsonDocument(result);
+
+                // Do not store result even if valid, still get from cache if available.
+                case SqlQueryStructure.CACHE_CONTROL_NO_STORE:
+                    maybeResult = _cache.TryGet<JsonElement?>(queryMetadata);
+                    // maybeResult is a nullable wrapper so we must check hasValue at outer and inner layer.
+                    if (maybeResult.HasValue && maybeResult.Value.HasValue)
+                    {
+                        result = maybeResult.Value.Value;
+                    }
+                    else
+                    {
+                        result = await queryExecutor.ExecuteQueryAsync(
+                            sqltext: queryMetadata.QueryText,
+                            parameters: queryMetadata.QueryParameters,
+                            dataReaderHandler: queryExecutor.GetJsonResultAsync<JsonElement>,
+                            httpContext: _httpContextAccessor.HttpContext!,
+                            args: null,
+                            dataSourceName: queryMetadata.DataSource);
+                    }
+
+                    return ParseResultIntoJsonDocument(result);
+
+                // Only return query response if it exists in cache, return gateway timeout otherwise.
+                case SqlQueryStructure.CACHE_CONTROL_ONLY_IF_CACHED:
+                    maybeResult = _cache.TryGet<JsonElement?>(queryMetadata);
+                    // maybeResult is a nullable wrapper so we must check hasValue at outer and inner layer.
+                    if (maybeResult.HasValue && maybeResult.Value.HasValue)
+                    {
+                        result = maybeResult.Value.Value;
+                    }
+                    else
+                    {
+                        throw new DataApiBuilderException(
+                            message: "Header 'only-if-cached' was used but item was not found in cache.",
+                            statusCode: System.Net.HttpStatusCode.GatewayTimeout,
+                            subStatusCode: DataApiBuilderException.SubStatusCodes.ItemNotFound);
+                    }
+
+                    return ParseResultIntoJsonDocument(result);
+
+                default:
+                    result = await _cache.GetOrSetAsync<JsonElement>(
+                        queryExecutor,
+                        queryMetadata,
+                        cacheEntryTtl: runtimeConfig.GetEntityCacheEntryTtl(entityName: structure.EntityName));
+                    return ParseResultIntoJsonDocument(result);
+            }
+        }
+
+        private static JsonDocument? ParseResultIntoJsonDocument(JsonElement? result)
+        {
+            byte[] jsonBytes = JsonSerializer.SerializeToUtf8Bytes(result);
+            return JsonDocument.Parse(jsonBytes);
+        }
+
         // <summary>
         // Given the SqlExecuteStructure structure, obtains the query text and executes it against the backend.
         // Unlike a normal query, result from database may not be JSON. Instead we treat output as SqlMutationEngine does (extract by row).

src/Core/Services/Cache/DabCacheService.cs

@@ -81,6 +81,41 @@ public DabCacheService(IFusionCache cache, ILogger<DabCacheService>? logger, IHt
         return result;
     }
 
+    /// <summary>
+    /// Try to get cacheValue from the cache with the derived cache key.
+    /// </summary>
+    /// <typeparam name="JsonElement">The type of value in the cache</typeparam>
+    /// <param name="queryMetadata">Metadata used to create a cache key or fetch a response from the database.</param>
+    /// <returns>JSON Response</returns>
+    public MaybeValue<JsonElement>? TryGet<JsonElement>(DatabaseQueryMetadata queryMetadata)
+    {
+        string cacheKey = CreateCacheKey(queryMetadata);
+        return _cache.TryGet<JsonElement>(key: cacheKey);
+    }
+
+    /// <summary>
+    /// Store cacheValue into the cache with the derived cache key.
+    /// </summary>
+    /// <typeparam name="JsonElement">The type of value in the cache</typeparam>
+    /// <param name="queryMetadata">Metadata used to create a cache key or fetch a response from the database.</param>
+    /// <param name="cacheEntryTtl">Number of seconds the cache entry should be valid before eviction.</param>
+    /// <param name="cacheValue"">The value to store in the cache.</param>
+    public void Set<JsonElement>(
+        DatabaseQueryMetadata queryMetadata,
+        int cacheEntryTtl,
+        JsonElement? cacheValue)
+    {
+        string cacheKey = CreateCacheKey(queryMetadata);
+        _cache.Set(
+            key: cacheKey,
+            value: cacheValue,
+            (FusionCacheEntryOptions options) =>
+            {
+                options.SetSize(EstimateCacheEntrySize(cacheKey: cacheKey, cacheValue: cacheValue?.ToString()));
+                options.SetDuration(duration: TimeSpan.FromSeconds(cacheEntryTtl));
+            });
+    }
+
     /// <summary>
     /// Attempts to fetch response from cache. If there is a cache miss, invoke executeQueryAsync Func to get a response
     /// </summary>

src/Core/Services/MetadataProviders/SqlMetadataProvider.cs

@@ -81,7 +81,7 @@ public abstract class SqlMetadataProvider<ConnectionT, DataAdapterT, CommandT> :
         /// <summary>
         /// Maps an entity name to a DatabaseObject.
         /// </summary>
-        public Dictionary<string, DatabaseObject> EntityToDatabaseObject { get; set; } =
+        public virtual Dictionary<string, DatabaseObject> EntityToDatabaseObject { get; set; } =
             new(StringComparer.InvariantCulture);
 
         protected readonly ILogger<ISqlMetadataProvider> _logger;