From 30c0e50e4d27eda4d871605b2015c8bc6b7277b6 Mon Sep 17 00:00:00 2001
From: SergeyGaluzo <95932081+SergeyGaluzo@users.noreply.github.com>
Date: Mon, 23 Dec 2024 12:48:32 -0800
Subject: [PATCH] Sort value logic without reader field count (#4759)

* Simple sort value logic

* plus
---
 .../Visitors/QueryGenerators/SqlQueryGenerator.cs   |  2 +-
 .../Features/Search/SqlServerSearchService.cs       | 13 +++++--------
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/QueryGenerators/SqlQueryGenerator.cs b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/QueryGenerators/SqlQueryGenerator.cs
index 61695e6567..dcb1a7e073 100644
--- a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/QueryGenerators/SqlQueryGenerator.cs
+++ b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/QueryGenerators/SqlQueryGenerator.cs
@@ -1364,7 +1364,7 @@ private static bool IsPrimaryKeySort(SearchOptions searchOptions)
             return searchOptions.Sort.All(s => s.searchParameterInfo.Name is SearchParameterNames.ResourceType or SearchParameterNames.LastUpdated);
         }
 
-        private bool IsSortValueNeeded(SearchOptions context)
+        internal bool IsSortValueNeeded(SearchOptions context)
         {
             if (context.Sort.Count == 0)
             {
diff --git a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs
index 5b61d12661..c3deb5e0d3 100644
--- a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs
+++ b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs
@@ -68,7 +68,6 @@ internal class SqlServerSearchService : SearchService
         private readonly SchemaInformation _schemaInformation;
         private readonly ICompressedRawResourceConverter _compressedRawResourceConverter;
         private readonly RequestContextAccessor<IFhirRequestContext> _requestContextAccessor;
-        private const int _defaultNumberOfColumnsReadFromResult = 11;
         private readonly SearchParameterInfo _fakeLastUpdate = new SearchParameterInfo(SearchParameterNames.LastUpdated, SearchParameterNames.LastUpdated);
         private readonly IParameterStore _parameterStore;
         private static ResourceSearchParamStats _resourceSearchParamStats;
@@ -336,6 +335,7 @@ await _sqlRetryService.ExecuteSql(
                     using (SqlCommand sqlCommand = connection.CreateCommand()) // WARNING, this code will not set sqlCommand.Transaction. Sql transactions via C#/.NET are not supported in this method.
                     {
                         sqlCommand.CommandTimeout = (int)_sqlServerDataStoreConfiguration.CommandTimeout.TotalSeconds;
+                        var isSortValueNeeded = false;
 
                         var exportTimeTravel = clonedSearchOptions.QueryHints != null && ContainsGlobalEndSurrogateId(clonedSearchOptions);
                         if (exportTimeTravel)
@@ -358,6 +358,7 @@ await _sqlRetryService.ExecuteSql(
                                 sqlException);
 
                             expression.AcceptVisitor(queryGenerator, clonedSearchOptions);
+                            isSortValueNeeded = queryGenerator.IsSortValueNeeded(clonedSearchOptions);
 
                             SqlCommandSimplifier.RemoveRedundantParameters(stringBuilder, sqlCommand.Parameters, _logger);
 
@@ -415,7 +416,6 @@ await _sqlRetryService.ExecuteSql(
 
                             string sortValue = null;
                             var isResultPartial = false;
-                            int numberOfColumnsRead = 0;
 
                             while (await reader.ReadAsync(cancellationToken))
                             {
@@ -439,8 +439,6 @@ await _sqlRetryService.ExecuteSql(
                                     continue;
                                 }
 
-                                numberOfColumnsRead = reader.FieldCount;
-
                                 // If we get to this point, we know there are more results so we need a continuation token
                                 // Additionally, this resource shouldn't be included in the results
                                 if (matchCount >= clonedSearchOptions.MaxItemCount && isMatch)
@@ -477,10 +475,9 @@ await _sqlRetryService.ExecuteSql(
                                     newContinuationType = resourceTypeId;
                                     newContinuationId = resourceSurrogateId;
 
-                                    // For normal queries, we select _defaultNumberOfColumnsReadFromResult number of columns.
-                                    // If we have more, that means we have an extra column tracking sort value.
+                                    // If sort value needed, that means we have an extra column tracking sort value.
                                     // Keep track of sort value if this is the last row.
-                                    if (matchCount == clonedSearchOptions.MaxItemCount - 1 && reader.FieldCount > _defaultNumberOfColumnsReadFromResult)
+                                    if (matchCount == clonedSearchOptions.MaxItemCount - 1 && isSortValueNeeded)
                                     {
                                         var tempSortValue = reader.GetValue(SortValueColumnName);
                                         if ((tempSortValue as DateTime?) != null)
@@ -547,7 +544,7 @@ await _sqlRetryService.ExecuteSql(
                                 clonedSearchOptions.Sort[0].searchParameterInfo.Code != KnownQueryParameterNames.LastUpdated)
                             {
                                 // If there is an extra column for sort value, we know we have searched for sort values. If no results were returned, we don't know if we have searched for sort values so we need to assume we did so we run the second phase.
-                                sqlSearchOptions.DidWeSearchForSortValue = numberOfColumnsRead > _defaultNumberOfColumnsReadFromResult;
+                                sqlSearchOptions.DidWeSearchForSortValue = isSortValueNeeded;
                             }
 
                             // This value is set inside the SortRewriter. If it is set, we need to pass