Personal/rajitha/continuation token limit

src/Microsoft.Health.Fhir.CosmosDb/Constants.cs

@@ -9,5 +9,8 @@ internal static class Constants
     {
         public const string CollectionConfigurationName = "fhirCosmosDb";
         public const string CosmosDbResponseMessagesProperty = nameof(CosmosDbResponseMessagesProperty);
+        public const int ContinuationTokenMinLimit = 1;
+        public const int ContinuationTokenMaxLimit = 3;
+        public const int ContinuationTokenDefaultLimit = 3;
     }
 }

src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/CosmosDbHeaders.cs

@@ -14,5 +14,7 @@ public static class CosmosDbHeaders
         public const string RequestCharge = "x-ms-request-charge";
 
         public const string SubStatus = "x-ms-substatus";
+
+        public const string CosmosContinuationTokenSize = "x-ms-documentdb-responsecontinuationtokenlimitinkb";
     }
 }

src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/FhirCosmosResponseHandler.cs

@@ -60,7 +60,25 @@ public override async Task<ResponseMessage> SendAsync(RequestMessage request, Ca
 
         private void UpdateOptions(RequestMessage options)
         {
-            if (_cosmosDataStoreConfiguration.ContinuationTokenSizeLimitInKb != null)
+            IFhirRequestContext fhirRequestContext = _fhirRequestContextAccessor.RequestContext;
+            if (fhirRequestContext == null)
+            {
+                return;
+            }
+
+            if (fhirRequestContext.RequestHeaders.TryGetValue(CosmosDbHeaders.CosmosContinuationTokenSize, out var tokenSize))
+            {
+                var intTokenSize = int.TryParse(tokenSize, out var count) ? count : 0;
+                if (intTokenSize != 0 && intTokenSize >= Constants.ContinuationTokenMinLimit && intTokenSize <= Constants.ContinuationTokenMaxLimit)
+                {
+                    options.Headers[_continuationTokenLimitHeaderName] = intTokenSize.ToString();
+                }
+                else
+                {
+                    throw new BadRequestException(string.Format(Resources.InvalidCosmosContinuationTokenSize, tokenSize));
+                }
+            }
+            else if (_cosmosDataStoreConfiguration.ContinuationTokenSizeLimitInKb != null)
             {
                 options.Headers[_continuationTokenLimitHeaderName] = _cosmosDataStoreConfiguration.ContinuationTokenSizeLimitInKb.ToString();
             }

src/Microsoft.Health.Fhir.CosmosDb/Resources.resx

@@ -165,4 +165,7 @@
   <data name="ChainedExpressionSubqueryLimit" xml:space="preserve">
     <value>Sub-queries in a chained expression cannot return more than {0} results, please use a more selective criteria.</value>
   </data>
+  <data name="InvalidCosmosContinuationTokenSize" xml:space="preserve">
+    <value>ContinuationTokenSize '{0}'Kb specified in the request is invalid. Valid limits are between 1 - 3 Kb.</value>
+  </data>
 </root>

test/Microsoft.Health.Fhir.Shared.Tests.E2E/Rest/CustomHeadersTests.cs

@@ -8,6 +8,7 @@
 using System.Net.Http;
 using System.Threading.Tasks;
 using Microsoft.Health.Fhir.Core.Features;
+using Microsoft.Health.Fhir.CosmosDb.Features.Storage;
 using Microsoft.Health.Fhir.Tests.Common;
 using Microsoft.Health.Fhir.Tests.Common.FixtureParameters;
 using Microsoft.Health.Fhir.Tests.E2E.Common;
@@ -56,5 +57,38 @@ public async Task GivenNoRequestHeader_WhenSendRequest_TheServerShouldntReturnCo
             Assert.NotEmpty(response.Headers.GetValues(KnownHeaders.RequestId));
             Assert.False(response.Headers.Contains(KnownHeaders.CorrelationId));
         }
+
+        [Theory]
+        [InlineData("1")]
+        [InlineData("2")]
+        [InlineData("3")]
+        [Trait(Traits.Priority, Priority.One)]
+        [HttpIntegrationFixtureArgumentSets(DataStore.CosmosDb)]
+        public async Task GivenValidContinuationTokenLimitHeader_WhenSendRequest_TheServerShouldReturnSuccess(string tokenSize)
+        {
+            var message = new HttpRequestMessage(HttpMethod.Get, "Patient");
+            message.Headers.Add(CosmosDbHeaders.CosmosContinuationTokenSize, tokenSize);
+
+            using HttpResponseMessage response = await _client.HttpClient.SendAsync(message);
+
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+        }
+
+        [Theory]
+        [InlineData("invalid")]
+        [InlineData("5")]
+        [InlineData("-1")]
+        [InlineData("9999999999999999999999999999")]
+        [InlineData("1.0")]
+        [Trait(Traits.Priority, Priority.One)]
+        [HttpIntegrationFixtureArgumentSets(DataStore.CosmosDb)]
+        public async Task GivenInValidContinuationTokenLimitHeader_WhenSendRequest_ThenBadRequestShouldBeReturned(string tokenSize)
+        {
+            var message = new HttpRequestMessage(HttpMethod.Get, "Patient");
+            message.Headers.Add(CosmosDbHeaders.CosmosContinuationTokenSize, tokenSize);
+
+            using HttpResponseMessage response = await _client.HttpClient.SendAsync(message);
+            Assert.Equal(HttpStatusCode.BadRequest, response.StatusCode);
+        }
     }
 }