diff --git a/src/Config/ObjectModel/HostOptions.cs b/src/Config/ObjectModel/HostOptions.cs index 452c78ef17..7b781577c4 100644 --- a/src/Config/ObjectModel/HostOptions.cs +++ b/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 /// - private const int MAX_RESPONSE_LENGTH_DAB_ENGINE_MB = 158; + public const int MAX_RESPONSE_LENGTH_DAB_ENGINE_MB = 158; /// /// Dab engine default response length. As of now this is same as max response length. /// - 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; } diff --git a/src/Config/ObjectModel/RuntimeConfig.cs b/src/Config/ObjectModel/RuntimeConfig.cs index c09b1c05ec..91b5cbdb35 100644 --- a/src/Config/ObjectModel/RuntimeConfig.cs +++ b/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() diff --git a/src/Core/Resolvers/QueryExecutor.cs b/src/Core/Resolvers/QueryExecutor.cs index 97d41da001..610100ea8d 100644 --- a/src/Core/Resolvers/QueryExecutor.cs +++ b/src/Core/Resolvers/QueryExecutor.cs @@ -36,6 +36,9 @@ public class QueryExecutor : IQueryExecutor private RetryPolicy _retryPolicy; + private int _maxResponseSizeMB; + private long _maxResponseSizeBytes; + /// /// Dictionary that stores dataSourceName to its corresponding connection string builder. /// @@ -51,6 +54,9 @@ public QueryExecutor(DbExceptionParser dbExceptionParser, ConnectionStringBuilders = new Dictionary(); ConfigProvider = configProvider; HttpContextAccessor = httpContextAccessor; + _maxResponseSizeMB = configProvider.GetConfig().MaxResponseSizeMB(); + _maxResponseSizeBytes = _maxResponseSizeMB * 1024 * 1024; + _retryPolicyAsync = Policy .Handle(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); } } - /// public async Task ExtractResultSetFromDbDataReaderAsync(DbDataReader dbDataReader, List? args = null) @@ -636,6 +643,38 @@ public Dictionary GetResultProperties( return resultProperties; } + /// + /// Reads data into jsonString. + /// + /// DbDataReader. + /// Available buffer. + /// jsonString to read into. + /// size of data read in bytes. + 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; + } + + /// + /// 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. + /// private async Task GetJsonStringFromDbReader(DbDataReader dbDataReader) { StringBuilder jsonString = new(); @@ -646,12 +685,41 @@ private async Task 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(); } + + /// + /// This function validates the size of data being read is within the available size limit. + /// + /// available size in bytes. + /// amount of data trying to be read in bytes + /// exception if size to be read is greater than data to be read. + 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); + } + } } } diff --git a/src/Service.Tests/Unittests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/Unittests/SqlQueryExecutorUnitTests.cs index 56a67796d0..d507e849cb 100644 --- a/src/Service.Tests/Unittests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/Unittests/SqlQueryExecutorUnitTests.cs @@ -318,6 +318,63 @@ await queryExecutor.Object.ExecuteQueryAsync( Assert.AreEqual(3, queryExecutorLogger.Invocations.Count); } + /// + /// 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. + /// + [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())); + + RuntimeConfigProvider runtimeConfigProvider = TestHelper.GenerateInMemoryRuntimeConfigProvider(runtimeConfig); + + Mock>> queryExecutorLogger = new(); + Mock 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 = new(); + dbDataReader.Setup(d => d.HasRows).Returns(true); + dbDataReader.Setup(x => x.GetChars(It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny())).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() {