Disable scalar temporal rewrite for chained searches

src/Microsoft.Health.Fhir.SqlServer.UnitTests/Features/Search/Expressions/ScalarTemporalEqualityRewriterTests.cs

@@ -4,6 +4,8 @@
 // -------------------------------------------------------------------------------------------------
 
 using System;
+using System.Reflection;
+using System.Runtime.CompilerServices;
 using Microsoft.Health.Fhir.Core.Features.Search.Expressions;
 using Microsoft.Health.Fhir.Core.Models;
 using Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions;
@@ -67,6 +69,35 @@ private static SearchParameterInfo BuildBirthdateParam(Uri url = null, SearchPar
                 baseResourceTypes: new[] { "Patient" });
         }
 
+        private static SearchParameterInfo BuildReferenceParam()
+        {
+            return new SearchParameterInfo(
+                "Observation-patient",
+                "patient",
+                SearchParamType.Reference,
+                new Uri("http://hl7.org/fhir/SearchParameter/Observation-patient"),
+                expression: "Observation.subject",
+                baseResourceTypes: new[] { "Observation" },
+                targetResourceTypes: new[] { "Patient" });
+        }
+
+        private static ChainedExpression BuildChainedExpression(Expression inner)
+        {
+            var expression = (ChainedExpression)RuntimeHelpers.GetUninitializedObject(typeof(ChainedExpression));
+            SetBackingField(expression, nameof(ChainedExpression.ResourceTypes), new[] { "Observation" });
+            SetBackingField(expression, nameof(ChainedExpression.ReferenceSearchParameter), BuildReferenceParam());
+            SetBackingField(expression, nameof(ChainedExpression.TargetResourceTypes), new[] { "Patient" });
+            SetBackingField(expression, nameof(ChainedExpression.Reversed), false);
+            SetBackingField(expression, nameof(ChainedExpression.Expression), inner);
+            return expression;
+        }
+
+        private static void SetBackingField<T>(ChainedExpression expression, string propertyName, T value)
+        {
+            FieldInfo field = typeof(ChainedExpression).GetField($"<{propertyName}>k__BackingField", BindingFlags.Instance | BindingFlags.NonPublic);
+            field.SetValue(expression, value);
+        }
+
         private static MultiaryExpression EqualityPattern(DateTimeOffset start, DateTimeOffset end) =>
             Expression.And(
                 Expression.GreaterThanOrEqual(FieldName.DateTimeStart, null, start),
@@ -78,7 +109,7 @@ public void GivenAllowListedBirthdateExactDay_WhenRewritten_ThenEmitsDaySplitUni
         {
             var expr = new SearchParameterExpression(BuildBirthdateParam(), EqualityPattern(start, end));
 
-            var result = expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance, null);
+            var result = expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance);
 
             AssertDaySplitUnion(result, start, end);
         }
@@ -91,18 +122,29 @@ public void GivenAllowListedBirthdateExactDayReversedOperandOrder_WhenRewritten_
                 Expression.GreaterThanOrEqual(FieldName.DateTimeStart, null, StartOfDay));
             var expr = new SearchParameterExpression(BuildBirthdateParam(), reversedPattern);
 
-            var result = expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance, null);
+            var result = expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance);
 
             AssertDaySplitUnion(result, StartOfDay, EndOfDay);
         }
 
+        [Fact]
+        public void GivenAllowListedBirthdateExactDayInChainedExpression_WhenRewritten_ThenPassThrough()
+        {
+            var inner = new SearchParameterExpression(BuildBirthdateParam(), EqualityPattern(StartOfDay, EndOfDay));
+            var expr = BuildChainedExpression(inner);
+
+            var result = Assert.IsType<ChainedExpression>(expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance));
+
+            Assert.Same(inner, result.Expression);
+        }
+
         [Theory]
         [MemberData(nameof(NonRewritableExpressions))]
         public void GivenAllowListedBirthdateWithNonExactDayExpression_WhenRewritten_ThenPassThrough(Expression inner)
         {
             var expr = new SearchParameterExpression(BuildBirthdateParam(), inner);
 
-            var result = Assert.IsType<SearchParameterExpression>(expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance, null));
+            var result = Assert.IsType<SearchParameterExpression>(expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance));
 
             Assert.Same(expr, result);
         }
@@ -113,7 +155,7 @@ public void GivenNonAllowListedParameter_WhenEqualityPatternMatched_ThenPassThro
         {
             var expr = new SearchParameterExpression(param, EqualityPattern(StartOfDay, EndOfDay));
 
-            var result = Assert.IsType<SearchParameterExpression>(expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance, null));
+            var result = Assert.IsType<SearchParameterExpression>(expr.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance));
 
             Assert.Same(expr, result);
         }

src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/ScalarTemporalEqualityRewriter.cs

@@ -33,7 +33,7 @@ namespace Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions.Visitors
     /// This rewriter must run BEFORE <see cref="DateTimeEqualityRewriter"/> so the input pattern still has only two
     /// predicates. Composite parameters and range operators are out of scope and pass through unchanged.
     /// </summary>
-    internal class ScalarTemporalEqualityRewriter : SqlExpressionRewriterWithInitialContext<object>
+    internal class ScalarTemporalEqualityRewriter : SqlExpressionRewriterWithInitialContext<bool>
     {
         internal static readonly ScalarTemporalEqualityRewriter Instance = new ScalarTemporalEqualityRewriter();
 
@@ -50,8 +50,24 @@ private enum Precision
             ExactDay,
         }
 
-        public override Expression VisitSearchParameter(SearchParameterExpression expression, object context)
+        public override Expression VisitChained(ChainedExpression expression, bool context)
         {
+            Expression visitedExpression = expression.Expression.AcceptVisitor(this, context: true);
+            if (ReferenceEquals(visitedExpression, expression.Expression))
+            {
+                return expression;
+            }
+
+            return new ChainedExpression(expression.ResourceTypes, expression.ReferenceSearchParameter, expression.TargetResourceTypes, expression.Reversed, visitedExpression);
+        }
+
+        public override Expression VisitSearchParameter(SearchParameterExpression expression, bool context)
+        {
+            if (context)
+            {
+                return expression;
+            }
+
             // 1. Only allow-listed scalar date parameters are eligible.
             if (!IsActivatedScalarTemporalParameter(expression))
             {

test/Microsoft.Health.Fhir.Shared.Tests.E2E/Rest/Search/ChainingSearchTests.cs

@@ -85,6 +85,17 @@ public async Task GivenAChainedSearchExpressionOverASimpleParameter_WhenSearched
             ValidateBundle(bundle, Fixture.SmithSnomedDiagnosticReport, Fixture.SmithLoincDiagnosticReport, Fixture.TrumanSnomedDiagnosticReport, Fixture.TrumanLoincDiagnosticReport);
         }
 
+        [HttpIntegrationFixtureArgumentSets(DataStore.SqlServer, Format.Json)]
+        [Fact]
+        public async Task GivenAChainedSearchExpressionOverBirthdate_WhenSearched_ThenCorrectBundleShouldBeReturned()
+        {
+            string query = $"_tag={Fixture.Tag}&subject:Patient.birthdate={Fixture.SmithPatientBirthDate}";
+
+            Bundle bundle = await Client.SearchAsync(ResourceType.DiagnosticReport, query);
+
+            ValidateBundle(bundle, Fixture.SmithSnomedDiagnosticReport, Fixture.SmithLoincDiagnosticReport);
+        }
+
         [Fact]
         public async Task GivenAChainedSearchExpressionOverASimpleParameter_WhenSearchedWithPaging_ThenCorrectBundleShouldBeReturned()
         {
@@ -379,6 +390,8 @@ public ClassFixture(DataStore dataStore, Format format, TestFhirServerFactory te
 
             public string SmithPatientGivenName { get; } = Guid.NewGuid().ToString();
 
+            public string SmithPatientBirthDate { get; } = "1990-05-15";
+
             public string TrumanPatientGivenName { get; } = Guid.NewGuid().ToString();
 
             public string SnomedCode { get; } = Guid.NewGuid().ToString();
@@ -426,8 +439,8 @@ protected override async Task OnInitializedAsync()
 #endif
 
                 AdamsPatient = (await TestFhirClient.CreateAsync(new Patient { Meta = meta, Gender = AdministrativeGender.Female, Name = new List<HumanName> { new HumanName { Family = "Adams" } } })).Resource;
-                SmithPatient = (await TestFhirClient.CreateAsync(new Patient { Meta = meta, Gender = AdministrativeGender.Male, Name = new List<HumanName> { new HumanName { Given = new[] { SmithPatientGivenName }, Family = "Smith" } }, ManagingOrganization = new ResourceReference($"Organization/{organization.Id}") })).Resource;
-                TrumanPatient = (await TestFhirClient.CreateAsync(new Patient { Meta = meta, Gender = AdministrativeGender.Male, Name = new List<HumanName> { new HumanName { Given = new[] { TrumanPatientGivenName }, Family = "Truman" } } })).Resource;
+                SmithPatient = (await TestFhirClient.CreateAsync(new Patient { Meta = meta, Gender = AdministrativeGender.Male, BirthDate = SmithPatientBirthDate, Name = new List<HumanName> { new HumanName { Given = new[] { SmithPatientGivenName }, Family = "Smith" } }, ManagingOrganization = new ResourceReference($"Organization/{organization.Id}") })).Resource;
+                TrumanPatient = (await TestFhirClient.CreateAsync(new Patient { Meta = meta, Gender = AdministrativeGender.Male, BirthDate = "1990-05-16", Name = new List<HumanName> { new HumanName { Given = new[] { TrumanPatientGivenName }, Family = "Truman" } } })).Resource;
 
                 DeviceLoincSubject = (await TestFhirClient.CreateAsync(new Device { Meta = meta })).Resource;
                 DeviceSnomedSubject = (await TestFhirClient.CreateAsync(new Device { Meta = meta })).Resource;