From c3a7ce2b36589c789389a50976d109276ef4826f Mon Sep 17 00:00:00 2001
From: rajithaalurims <110048715+rajithaalurims@users.noreply.github.com>
Date: Tue, 6 Feb 2024 14:51:11 -0600
Subject: [PATCH] Removed resource type expansion for global search (#3701)

* Removed resource type expansion for global search

* Removing resource type expansions only for non Smart global search
---
 .../PartitionEliminationRewriterTests.cs      |  2 +-
 .../Visitors/PartitionEliminationRewriter.cs  | 29 +++++++++++++++----
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/Expressions/PartitionEliminationRewriterTests.cs b/src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/Expressions/PartitionEliminationRewriterTests.cs
index f669907c7..3fef49be6 100644
--- a/src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/Expressions/PartitionEliminationRewriterTests.cs
+++ b/src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/Expressions/PartitionEliminationRewriterTests.cs
@@ -70,7 +70,7 @@ namespace Microsoft.Health.Fhir.SqlServer.UnitTests.Features.Search.Expressions
                 .AcceptVisitor(_rewriter);
 
             Assert.Equal(
-                "(SqlRoot (SearchParamTables:) (ResourceTable: (Param _id (StringEquals TokenCode 'foo')) (Param _type (TokenCode IN (AllergyIntolerance, Claim, Condition, Device, DiagnosticReport)))))",
+                "(SqlRoot (SearchParamTables:) (ResourceTable: (Param _id (StringEquals TokenCode 'foo'))))",
                 rewritten.ToString());
         }
 
diff --git a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/PartitionEliminationRewriter.cs b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/PartitionEliminationRewriter.cs
index 4c1c25205..cc6fc17ab 100644
--- a/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/PartitionEliminationRewriter.cs
+++ b/src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/PartitionEliminationRewriter.cs
@@ -74,6 +74,7 @@ namespace Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions.Visitors
 
             int primaryKeyValueIndex = -1;
             bool hasTypeRestriction = false;
+            bool needTypeRestriction = false;
             for (var i = 0; i < expression.ResourceTableExpressions.Count; i++)
             {
                 SearchParameterInfo parameter = expression.ResourceTableExpressions[i].Parameter;
@@ -88,6 +89,20 @@ namespace Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions.Visitors
                 }
             }
 
+            // We still need this resource expansion for Smart requests that does system wide search and returns only resources that are part of same
+            // compartment along with universal resources.
+            // Refer to this test case GivenFhirUserClaimPractitioner_WhenAllResourcesRequested_ResourcesInTheSameComparementAndUniversalResourcesAlsoReturned
+            if (!hasTypeRestriction)
+            {
+                for (var i = 0; i < expression.SearchParamTableExpressions.Count; i++)
+                {
+                    if (expression.SearchParamTableExpressions[i].ToString().Contains(_resourceTypeSearchParameter.Code, StringComparison.Ordinal))
+                    {
+                        needTypeRestriction = true;
+                    }
+                }
+            }
+
             if (primaryKeyValueIndex < 0)
             {
                 // no continuation token
@@ -97,14 +112,18 @@ namespace Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions.Visitors
                     // This is already constrained to be one or more resource types.
                     return expression;
                 }
+                else if (needTypeRestriction)
+                {
+                    // Explicitly allow all resource types. SQL tends to create far better query plans than when there is no filter on ResourceTypeId.
 
-                // Explicitly allow all resource types. SQL tends to create far better query plans than when there is no filter on ResourceTypeId.
+                    var updatedResourceTableExpressions = new List<SearchParameterExpressionBase>(expression.ResourceTableExpressions.Count + 1);
+                    updatedResourceTableExpressions.AddRange(expression.ResourceTableExpressions);
+                    updatedResourceTableExpressions.Add(GetAllTypesExpression());
 
-                var updatedResourceTableExpressions = new List<SearchParameterExpressionBase>(expression.ResourceTableExpressions.Count + 1);
-                updatedResourceTableExpressions.AddRange(expression.ResourceTableExpressions);
-                updatedResourceTableExpressions.Add(GetAllTypesExpression());
+                    return new SqlRootExpression(expression.SearchParamTableExpressions, updatedResourceTableExpressions);
+                }
 
-                return new SqlRootExpression(expression.SearchParamTableExpressions, updatedResourceTableExpressions);
+                return expression;
             }
 
             // There is a primary key continuation token.