Streaming results in query flow to protect application memory

src/Config/ObjectModel/HostOptions.cs

@@ -18,12 +18,12 @@ public record HostOptions
     /// .net6: https://github.com/dotnet/runtime/blob/v8.0.0/src/libraries/System.Text.Json/src/System/Text/Json/Writer/JsonWriterHelper.cs#75
     /// ref: Json constant: https://github.com/dotnet/runtime/blob/main/src/libraries/System.Text.Json/src/System/Text/Json/JsonConstants.cs#L80
     /// </summary>
-    private const int MAX_RESPONSE_LENGTH_DAB_ENGINE_MB = 158;
+    public const int MAX_RESPONSE_LENGTH_DAB_ENGINE_MB = 158;
 
     /// <summary>
     /// Dab engine default response length. As of now this is same as max response length.
     /// </summary>
-    private const int DEFAULT_RESPONSE_LENGTH_DAB_ENGINE_MB = 158;
+    public const int DEFAULT_RESPONSE_LENGTH_DAB_ENGINE_MB = 158;
 
     [JsonPropertyName("cors")]
     public CorsOptions? Cors { get; init; }

src/Config/ObjectModel/RuntimeConfig.cs

@@ -495,9 +495,9 @@ public uint MaxPageSize()
         return (uint?)Runtime?.Pagination?.MaxPageSize ?? PaginationOptions.MAX_PAGE_SIZE;
     }
 
-    public int? MaxResponseSizeMB()
+    public int MaxResponseSizeMB()
     {
-        return Runtime?.Host?.MaxResponseSizeMB;
+        return Runtime?.Host?.MaxResponseSizeMB ?? HostOptions.MAX_RESPONSE_LENGTH_DAB_ENGINE_MB;
     }
 
     public bool MaxResponseSizeLogicEnabled()

src/Core/Resolvers/QueryExecutor.cs

@@ -36,6 +36,9 @@ public class QueryExecutor<TConnection> : IQueryExecutor
 
         private RetryPolicy _retryPolicy;
 
+        private int _maxResponseSizeMB;
+        private long _maxResponseSizeBytes;
+
         /// <summary>
         /// Dictionary that stores dataSourceName to its corresponding connection string builder.
         /// </summary>
@@ -51,6 +54,9 @@ public QueryExecutor(DbExceptionParser dbExceptionParser,
             ConnectionStringBuilders = new Dictionary<string, DbConnectionStringBuilder>();
             ConfigProvider = configProvider;
             HttpContextAccessor = httpContextAccessor;
+            _maxResponseSizeMB = configProvider.GetConfig().MaxResponseSizeMB();
+            _maxResponseSizeBytes = _maxResponseSizeMB * 1024 * 1024;
+
             _retryPolicyAsync = Policy
                 .Handle<DbException>(DbExceptionParser.IsTransientException)
                 .WaitAndRetryAsync(
@@ -242,7 +248,8 @@ public QueryExecutor(DbExceptionParser dbExceptionParser,
 
             try
             {
-                using DbDataReader dbDataReader = await cmd.ExecuteReaderAsync(CommandBehavior.CloseConnection);
+                using DbDataReader dbDataReader = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ?
+                    await cmd.ExecuteReaderAsync(CommandBehavior.SequentialAccess) : await cmd.ExecuteReaderAsync(CommandBehavior.CloseConnection);
                 if (dataReaderHandler is not null && dbDataReader is not null)
                 {
                     return await dataReaderHandler(dbDataReader, args);
@@ -318,7 +325,8 @@ public virtual DbCommand PrepareDbCommand(
 
             try
             {
-                using DbDataReader dbDataReader = cmd.ExecuteReader(CommandBehavior.CloseConnection);
+                using DbDataReader dbDataReader = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ?
+                    cmd.ExecuteReader(CommandBehavior.SequentialAccess) : cmd.ExecuteReader(CommandBehavior.CloseConnection);
                 if (dataReaderHandler is not null && dbDataReader is not null)
                 {
                     return dataReaderHandler(dbDataReader, args);
@@ -399,7 +407,6 @@ public bool Read(DbDataReader reader)
                 throw DbExceptionParser.Parse(e);
             }
         }
-
         /// <inheritdoc />
         public async Task<DbResultSet>
             ExtractResultSetFromDbDataReaderAsync(DbDataReader dbDataReader, List<string>? args = null)
@@ -636,6 +643,38 @@ public Dictionary<string, object> GetResultProperties(
             return resultProperties;
         }
 
+        /// <summary>
+        /// Reads data into jsonString.
+        /// </summary>
+        /// <param name="dbDataReader">DbDataReader.</param>
+        /// <param name="availableSize">Available buffer.</param>
+        /// <param name="resultJsonString">jsonString to read into.</param>
+        /// <returns>size of data read in bytes.</returns>
+        internal int StreamData(DbDataReader dbDataReader, long availableSize, StringBuilder resultJsonString)
+        {
+            long resultFieldSize = dbDataReader.GetChars(ordinal: 0, dataOffset: 0, buffer: null, bufferOffset: 0, length: 0);
+
+            // if the size of the field is less than available size, then we can read the entire field.
+            // else we throw exception.
+            ValidateSize(availableSize, resultFieldSize);
+
+            char[] buffer = new char[resultFieldSize];
+
+            // read entire field into buffer and reduce available size.
+            dbDataReader.GetChars(ordinal: 0, dataOffset: 0, buffer: buffer, bufferOffset: 0, length: buffer.Length);
+
+            resultJsonString.Append(buffer);
+            return buffer.Length;
+        }
+
+        /// <summary>
+        /// This function reads the data from the DbDataReader and returns a JSON string.
+        /// 1. MaxResponseSizeLogicEnabled is used like a feature flag.
+        /// 2. If MaxResponseSize is not specified by the customer or is null,
+        /// getString is used and entire data is read into memory.
+        /// 3. If MaxResponseSize is specified by the customer, getChars is used.
+        /// GetChars tries to read the data in chunks and if the data is more than the specified limit, it throws an exception.
+        /// </summary>
         private async Task<string> GetJsonStringFromDbReader(DbDataReader dbDataReader)
         {
             StringBuilder jsonString = new();
@@ -646,12 +685,41 @@ private async Task<string> GetJsonStringFromDbReader(DbDataReader dbDataReader)
             // 1. https://docs.microsoft.com/en-us/sql/relational-databases/json/format-query-results-as-json-with-for-json-sql-server?view=sql-server-2017#output-of-the-for-json-clause
             // 2. https://stackoverflow.com/questions/54973536/for-json-path-results-in-ssms-truncated-to-2033-characters/54973676
             // 3. https://docs.microsoft.com/en-us/sql/relational-databases/json/use-for-json-output-in-sql-server-and-in-client-apps-sql-server?view=sql-server-2017#use-for-json-output-in-a-c-client-app
-            while (await ReadAsync(dbDataReader))
+
+            if (!ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled())
+            {
+                while (await ReadAsync(dbDataReader))
+                {
+                    jsonString.Append(dbDataReader.GetString(0));
+                }
+            }
+            else
             {
-                jsonString.Append(dbDataReader.GetString(0));
+                long availableSize = _maxResponseSizeBytes;
+                while (await ReadAsync(dbDataReader))
+                {
+                    availableSize -= StreamData(dbDataReader, availableSize, jsonString);
+                }
             }
 
             return jsonString.ToString();
         }
+
+        /// <summary>
+        /// This function validates the size of data being read is within the available size limit.
+        /// </summary>
+        /// <param name="availableSizeBytes">available size in bytes.</param>
+        /// <param name="sizeToBeReadBytes">amount of data trying to be read in bytes</param>
+        /// <exception cref="DataApiBuilderException">exception if size to be read is greater than data to be read.</exception>
+        private void ValidateSize(long availableSizeBytes, long sizeToBeReadBytes)
+        {
+            if (sizeToBeReadBytes > availableSizeBytes)
+            {
+                throw new DataApiBuilderException(
+                    message: $"The JSON result size exceeds max result size of {_maxResponseSizeMB}MB. Please use pagination to reduce size of result.",
+                    statusCode: HttpStatusCode.RequestEntityTooLarge,
+                    subStatusCode: DataApiBuilderException.SubStatusCodes.ErrorProcessingData);
+            }
+        }
     }
 }

src/Service.Tests/Unittests/SqlQueryExecutorUnitTests.cs

@@ -318,6 +318,63 @@ await queryExecutor.Object.ExecuteQueryAsync<object>(
             Assert.AreEqual(3, queryExecutorLogger.Invocations.Count);
         }
 
+        /// <summary>
+        /// Validates streaming logic for QueryExecutor
+        /// In this test the DbDataReader.GetChars method is mocked to return 1024*1024 bytes (1MB) of data.
+        /// Max available size is set to 5 MB.
+        /// Based on number of loops, the data read will be 1MB * readDataLoops.Exception should be thrown in test cases where we go above 5MB.
+        /// This will be in cases where readDataLoops > 5.
+        /// </summary>
+        [DataTestMethod, TestCategory(TestCategory.MSSQL)]
+        [DataRow(4, false,
+            DisplayName = "Max available size is set to 5MB.4 data read loop iterations * 1MB -> should successfully read 4MB because max-db-response-size-mb is 4MB")]
+        [DataRow(5, false,
+            DisplayName = "Max available size is set to 5MB.5 data read loop iterations * 1MB -> should successfully read 5MB because max-db-response-size-mb is 5MB")]
+        [DataRow(6, true,
+            DisplayName = "Max available size is set to 5MB.6 data read loop iterations * 1MB -> Fails to read 6MB because max-db-response-size-mb is 5MB")]
+        public void ValidateStreamingLogicAsync(int readDataLoops, bool exceptionExpected)
+        {
+            TestHelper.SetupDatabaseEnvironment(TestCategory.MSSQL);
+            FileSystem fileSystem = new();
+            FileSystemRuntimeConfigLoader loader = new(fileSystem);
+            RuntimeConfig runtimeConfig = new(
+                Schema: "UnitTestSchema",
+                DataSource: new DataSource(DatabaseType: DatabaseType.MSSQL, "", Options: null),
+                Runtime: new(
+                        Rest: new(),
+                        GraphQL: new(),
+                        Host: new(Cors: null, Authentication: null, MaxResponseSizeMB: 5)
+                    ),
+                Entities: new(new Dictionary<string, Entity>()));
+
+            RuntimeConfigProvider runtimeConfigProvider = TestHelper.GenerateInMemoryRuntimeConfigProvider(runtimeConfig);
+
+            Mock<ILogger<QueryExecutor<SqlConnection>>> queryExecutorLogger = new();
+            Mock<IHttpContextAccessor> httpContextAccessor = new();
+            DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(runtimeConfigProvider);
+
+            // Instantiate the MsSqlQueryExecutor and Setup parameters for the query
+            MsSqlQueryExecutor msSqlQueryExecutor = new(runtimeConfigProvider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object);
+
+            try
+            {
+                Mock<DbDataReader> dbDataReader = new();
+                dbDataReader.Setup(d => d.HasRows).Returns(true);
+                dbDataReader.Setup(x => x.GetChars(It.IsAny<int>(), It.IsAny<long>(), It.IsAny<char[]>(), It.IsAny<int>(), It.IsAny<int>())).Returns(1024 * 1024);
+                int availableSize = (int)runtimeConfig.MaxResponseSizeMB() * 1024 * 1024;
+                for (int i = 0; i < readDataLoops; i++)
+                {
+                    availableSize -= msSqlQueryExecutor.StreamData(dbDataReader: dbDataReader.Object, availableSize: availableSize, resultJsonString: new());
+                }
+            }
+            catch (DataApiBuilderException ex)
+            {
+                Assert.IsTrue(exceptionExpected);
+                Assert.AreEqual(HttpStatusCode.RequestEntityTooLarge, ex.StatusCode);
+                Assert.AreEqual("The JSON result size exceeds max result size of 5MB. Please use pagination to reduce size of result.", ex.Message);
+            }
+        }
+
         [TestCleanup]
         public void CleanupAfterEachTest()
         {