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 e229ae7239..59fd7a84c4 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 @@ -41,7 +41,6 @@ internal class SqlQueryGenerator : DefaultSqlExpressionVisitor _includeFromCteIds; - private int _curFromCteIndex = -1; private int _tableExpressionCounter = -1; private SqlRootExpression _rootExpression; private readonly SchemaInformation _schemaInfo; @@ -848,7 +847,9 @@ private void HandleTableKindInclude( .Append(")"); // Get FROM ctes - string fromCte = _cteMainSelect; + List fromCte = new List(); + fromCte.Add(_cteMainSelect); + if (includeExpression.Iterate) { // Include Iterate @@ -858,7 +859,7 @@ private void HandleTableKindInclude( // On that case, the fromCte is _cteMainSelect if (TryGetIncludeCtes(includeExpression.SourceResourceType, out _includeFromCteIds)) { - fromCte = _includeFromCteIds[++_curFromCteIndex]; + fromCte = _includeFromCteIds; } } @@ -869,7 +870,7 @@ private void HandleTableKindInclude( { if (TryGetIncludeCtes(includeExpression.TargetResourceType, out _includeFromCteIds)) { - fromCte = _includeFromCteIds[++_curFromCteIndex]; + fromCte = _includeFromCteIds; } } else if (includeExpression.ReferenceSearchParameter?.TargetResourceTypes != null) @@ -884,7 +885,7 @@ private void HandleTableKindInclude( } _includeFromCteIds = _includeFromCteIds.Distinct().ToList(); - fromCte = _includeFromCteIds.Count > 0 ? _includeFromCteIds[++_curFromCteIndex] : fromCte; + fromCte = _includeFromCteIds.Count > 0 ? _includeFromCteIds : fromCte; } } } @@ -895,19 +896,30 @@ private void HandleTableKindInclude( .Append(" = ").Append(Parameters.AddParameter(VLatest.ReferenceSearchParam.ResourceTypeId, Model.GetResourceTypeId(includeExpression.SourceResourceType), true)); } - delimited.BeginDelimitedElement().Append("EXISTS (SELECT * FROM ").Append(fromCte) - .Append(" WHERE ").Append(VLatest.Resource.ResourceTypeId, table).Append(" = T1 AND ") - .Append(VLatest.Resource.ResourceSurrogateId, table).Append(" = Sid1"); - - if (!includeExpression.Iterate) + var scope = delimited.BeginDelimitedElement(); + scope.Append("EXISTS ("); + for (var index = 0; index < fromCte.Count; index++) { - // Limit the join to the main select CTE. - // The main select will have max+1 items in the result set to account for paging, so we only want to join using the max amount. + var cte = fromCte[index]; + scope.Append("SELECT * FROM ").Append(cte) + .Append(" WHERE ").Append(VLatest.Resource.ResourceTypeId, table).Append(" = T1 AND ") + .Append(VLatest.Resource.ResourceSurrogateId, table).Append(" = Sid1"); + + if (!includeExpression.Iterate) + { + // Limit the join to the main select CTE. + // The main select will have max+1 items in the result set to account for paging, so we only want to join using the max amount. - StringBuilder.Append(" AND Row < ").Append(Parameters.AddParameter(context.MaxItemCount + 1, true)); + scope.Append(" AND Row < ").Append(Parameters.AddParameter(context.MaxItemCount + 1, true)); + } + + if (index < fromCte.Count - 1) + { + scope.AppendLine(" UNION ALL "); + } } - StringBuilder.Append(")"); + scope.Append(")"); } if (includeExpression.Reversed) @@ -941,30 +953,9 @@ private void HandleTableKindInclude( } } - // Handle Multiple Results sets to include from - if (count > 1 && _curFromCteIndex >= 0 && _curFromCteIndex < count - 1) + if (includeExpression.WildCard) { - StringBuilder.AppendLine("),"); - - // If it's not the last result set, append a new IncludeLimit cte, since IncludeLimitCte was not created for the current cte - if (_curFromCteIndex < count - 1) - { - var cteToLimit = TableExpressionName(_tableExpressionCounter); - WriteIncludeLimitCte(cteToLimit, context); - } - - // Generate CTE to include from the additional result sets - StringBuilder.Append(TableExpressionName(++_tableExpressionCounter)).AppendLine(" AS").AppendLine("("); - searchParamTableExpression.AcceptVisitor(this, context); - } - else - { - _curFromCteIndex = -1; - - if (includeExpression.WildCard) - { - includeExpression.ReferencedTypes?.ToList().ForEach(t => AddIncludeLimitCte(t, curLimitCte)); - } + includeExpression.ReferencedTypes?.ToList().ForEach(t => AddIncludeLimitCte(t, curLimitCte)); } } @@ -1141,27 +1132,6 @@ private void HandleTableKindSortWithFilter(SearchParamTableExpression searchPara _sortVisited = true; } - private void WriteIncludeLimitCte(string cteToLimit, SearchOptions context) - { - StringBuilder.Append(TableExpressionName(++_tableExpressionCounter)).AppendLine(" AS").AppendLine("("); - - // the related cte is a reverse include, limit the number of returned items and count to - // see if we are over the threshold (to produce a warning to the client) - StringBuilder.Append("SELECT DISTINCT "); - StringBuilder.Append("TOP (").Append(Parameters.AddParameter(context.IncludeCount, true)).Append(") "); - - StringBuilder.Append("T1, Sid1, IsMatch, "); - StringBuilder.Append("CASE WHEN count_big(*) over() > ") - .Append(Parameters.AddParameter(context.IncludeCount, true)) - .AppendLine(" THEN 1 ELSE 0 END AS IsPartial "); - - StringBuilder.Append("FROM ").AppendLine(cteToLimit); - StringBuilder.AppendLine("),"); - - // the 'original' include cte is not in the union, but this new layer is instead - _includeCteIds.Add(TableExpressionName(_tableExpressionCounter)); - } - private SearchParameterQueryGeneratorContext GetContext(string tableAlias = null) { return new SearchParameterQueryGeneratorContext(StringBuilder, Parameters, Model, _schemaInfo, tableAlias); diff --git a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs index edaf6a440c..e8c33eb108 100644 --- a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs +++ b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/SqlServerSearchService.cs @@ -332,6 +332,7 @@ private async Task SearchImpl(SqlSearchOptions sqlSearchOptions, C await CreateStats(expression, cancellationToken); SearchResult searchResult = null; + await _sqlRetryService.ExecuteSql( async (connection, cancellationToken, sqlException) => { @@ -571,6 +572,7 @@ await _sqlRetryService.ExecuteSql( _logger, cancellationToken, true); // this enables reads from replicas + return searchResult; } diff --git a/test/Microsoft.Health.Fhir.Shared.Tests.E2E/Rest/Search/IncludeSearchTests.cs b/test/Microsoft.Health.Fhir.Shared.Tests.E2E/Rest/Search/IncludeSearchTests.cs index af7e1bf92d..895399f91a 100644 --- a/test/Microsoft.Health.Fhir.Shared.Tests.E2E/Rest/Search/IncludeSearchTests.cs +++ b/test/Microsoft.Health.Fhir.Shared.Tests.E2E/Rest/Search/IncludeSearchTests.cs @@ -671,7 +671,7 @@ await SearchAndValidateBundleAsync( public async Task GivenAnIncludeIterateSearchExpressionWithMultitypeArrayReference_WhenSearched_TheIterativeResultsShouldBeAddedToTheBundle() { // Non-recursive iteration - Reference array of multiple target types: CareTeam:participant of type Patient, Practitioner, Organization, etc. - string query = $"_include=CareTeam:participant:Patient&_include:iterate=Patient:general-practitioner&_tag={Fixture.Tag}"; + string query = $"_include=CareTeam:participant&_include:iterate=Patient:general-practitioner&_tag={Fixture.Tag}"; await SearchAndValidateBundleAsync( ResourceType.CareTeam, @@ -682,7 +682,9 @@ await SearchAndValidateBundleAsync( Fixture.TrumanPatient, Fixture.AndersonPractitioner, Fixture.SanchezPractitioner, - Fixture.TaylorPractitioner); + Fixture.TaylorPractitioner, + Fixture.Organization, + Fixture.Practitioner); } [Fact]