Hotfix: Skip Multiary Chained Expression For Rewriter

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

@@ -81,13 +81,13 @@ private static SearchParameterInfo BuildReferenceParam()
                 targetResourceTypes: new[] { "Patient" });
         }
 
-        private static ChainedExpression BuildChainedExpression(Expression inner)
+        private static ChainedExpression BuildChainedExpression(Expression inner, bool reversed = false)
         {
             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.Reversed), reversed);
             SetBackingField(expression, nameof(ChainedExpression.Expression), inner);
             return expression;
         }
@@ -138,6 +138,32 @@ public void GivenAllowListedBirthdateExactDayInChainedExpression_WhenRewritten_T
             Assert.Same(inner, result.Expression);
         }
 
+        [Fact]
+        public void GivenRootExpressionContainsReverseChainAndAllowListedBirthdateExactDay_WhenRewritten_ThenPassThroughBirthdate()
+        {
+            var birthdate = new SearchParameterExpression(BuildBirthdateParam(), EqualityPattern(StartOfDay, EndOfDay));
+            var reverseChain = BuildChainedExpression(Expression.GreaterThanOrEqual(FieldName.DateTimeStart, null, StartOfDay), reversed: true);
+            var root = Expression.And(reverseChain, birthdate);
+
+            var result = Assert.IsType<MultiaryExpression>(root.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance));
+            var rewrittenBirthdate = Assert.IsType<SearchParameterExpression>(result.Expressions[1]);
+
+            Assert.Same(birthdate, rewrittenBirthdate);
+        }
+
+        [Fact]
+        public void GivenRootExpressionContainsForwardChainAndAllowListedBirthdateExactDay_WhenRewritten_ThenPassThroughBirthdate()
+        {
+            var birthdate = new SearchParameterExpression(BuildBirthdateParam(), EqualityPattern(StartOfDay, EndOfDay));
+            var forwardChain = BuildChainedExpression(Expression.GreaterThanOrEqual(FieldName.DateTimeStart, null, StartOfDay), reversed: false);
+            var root = Expression.And(forwardChain, birthdate);
+
+            var result = Assert.IsType<MultiaryExpression>(root.AcceptVisitor(ScalarTemporalEqualityRewriter.Instance));
+            var rewrittenBirthdate = Assert.IsType<SearchParameterExpression>(result.Expressions[1]);
+
+            Assert.Same(birthdate, rewrittenBirthdate);
+        }
+
         [Theory]
         [MemberData(nameof(NonRewritableExpressions))]
         public void GivenAllowListedBirthdateWithNonExactDayExpression_WhenRewritten_ThenPassThrough(Expression inner)

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

@@ -5,6 +5,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using Microsoft.Health.Fhir.Core.Features.Search.Expressions;
 using Microsoft.Health.Fhir.Core.Models;
 using Microsoft.Health.Fhir.SqlServer.Features.Search.Expressions;
@@ -50,8 +51,23 @@ private enum Precision
             ExactDay,
         }
 
+        public override Expression VisitMultiary(MultiaryExpression expression, bool context)
+        {
+            // Don't run rewriter for chained or reverse chained expression. Once we remove the union in this rewriter below and
+            // adhere to the FHIR spec for dates, we can remove this method. Our current SQL generator has difficulty with unions
+            // emitted by this rewriter due to FindRestrictingPredecessorTableExpressionIndex() has a catch all that returns
+            // currentIndex - 1 for unions.
+            if (!context && ContainsChain(expression))
+            {
+                return expression;
+            }
+
+            return base.VisitMultiary(expression, context);
+        }
+
         public override Expression VisitChained(ChainedExpression expression, bool context)
         {
+            // Logic here is the same as VisitMultiary.
             Expression visitedExpression = expression.Expression.AcceptVisitor(this, context: true);
             if (ReferenceEquals(visitedExpression, expression.Expression))
             {
@@ -208,5 +224,25 @@ private static bool IsEndLe(BinaryExpression be) =>
         private static bool IsUtc(DateTimeOffset value) => value.Offset == TimeSpan.Zero;
 
         private static bool IsUtcMidnight(DateTimeOffset value) => IsUtc(value) && value.TimeOfDay == TimeSpan.Zero;
+
+        private static bool ContainsChain(Expression expression)
+        {
+            if (expression is ChainedExpression)
+            {
+                return true;
+            }
+
+            if (expression is MultiaryExpression multiary)
+            {
+                return multiary.Expressions.Any(ContainsChain);
+            }
+
+            if (expression is UnionExpression union)
+            {
+                return union.Expressions.Any(ContainsChain);
+            }
+
+            return false;
+        }
     }
 }

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

@@ -96,6 +96,28 @@ public async Task GivenAChainedSearchExpressionOverBirthdate_WhenSearched_ThenCo
             ValidateBundle(bundle, Fixture.SmithSnomedDiagnosticReport, Fixture.SmithLoincDiagnosticReport);
         }
 
+        [HttpIntegrationFixtureArgumentSets(DataStore.SqlServer, Format.Json)]
+        [Fact]
+        public async Task GivenAChainedSearchExpressionOverBirthdateMonthPrecision_WhenSearched_ThenCorrectBundleShouldBeReturned()
+        {
+            string query = $"_tag={Fixture.Tag}&subject:Patient.birthdate=1990-05";
+
+            Bundle bundle = await Client.SearchAsync(ResourceType.DiagnosticReport, query);
+
+            ValidateBundle(bundle, Fixture.SmithSnomedDiagnosticReport, Fixture.SmithLoincDiagnosticReport, Fixture.TrumanSnomedDiagnosticReport, Fixture.TrumanLoincDiagnosticReport);
+        }
+
+        [HttpIntegrationFixtureArgumentSets(DataStore.SqlServer, Format.Json)]
+        [Fact]
+        public async Task GivenAChainedSearchExpressionOverBirthdateDayWithAdditionalCriteriaOnTheSameTarget_WhenSearched_ThenCorrectBundleShouldBeReturned()
+        {
+            string query = $"_tag={Fixture.Tag}&subject:Patient.birthdate={Fixture.SmithPatientBirthDate}&subject:Patient.family=Smith";
+
+            Bundle bundle = await Client.SearchAsync(ResourceType.DiagnosticReport, query);
+
+            ValidateBundle(bundle, Fixture.SmithSnomedDiagnosticReport, Fixture.SmithLoincDiagnosticReport);
+        }
+
         [Fact]
         public async Task GivenAChainedSearchExpressionOverASimpleParameter_WhenSearchedWithPaging_ThenCorrectBundleShouldBeReturned()
         {
@@ -144,6 +166,86 @@ public async Task GivenAReverseChainSearchExpressionOverASimpleParameter_WhenSea
             ValidateBundle(bundle, Fixture.TrumanPatient);
         }
 
+        [HttpIntegrationFixtureArgumentSets(DataStore.SqlServer, Format.Json)]
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public async Task GivenAReverseChainSearchExpressionCombinedWithAnExactDayBirthdate_WhenSearched_ThenCorrectBundleShouldBeReturned(bool includeAccurateTotal)
+        {
+            string query = $"_tag={Fixture.Tag}&birthdate={Fixture.SmithPatientBirthDate}&_has:Observation:patient:code={Fixture.SnomedCode}";
+            if (includeAccurateTotal)
+            {
+                query += "&_total=accurate";
+            }
+
+            Bundle bundle = await Client.SearchAsync(ResourceType.Patient, query);
+
+            if (includeAccurateTotal)
+            {
+                Assert.Equal(1, bundle.Total.GetValueOrDefault());
+            }
+
+            ValidateBundle(bundle, Fixture.SmithPatient);
+        }
+
+#if !Stu3
+        [HttpIntegrationFixtureArgumentSets(DataStore.SqlServer, Format.Json)]
+        [Fact]
+        public async Task GivenAReverseChainSearchExpressionOverImagingStudyStartedCombinedWithAnExactDayBirthdateAndTag_WhenSearchedWithPost_ThenCorrectBundleShouldBeReturned()
+        {
+            string tenantTagCode = Guid.NewGuid().ToString("N");
+            var tenantMeta = new Meta
+            {
+                Tag = new List<Coding>
+                {
+                    new Coding(Fixture.TenantTagSystem, tenantTagCode),
+                },
+            };
+
+            Patient imagingStudyExactDayPatient = (await Client.CreateAsync(
+                new Patient
+                {
+                    Meta = tenantMeta,
+                    BirthDate = Fixture.ImagingStudyPatientBirthDate,
+                })).Resource;
+
+            Patient imagingStudyMonthPrecisionPatient = (await Client.CreateAsync(
+                new Patient
+                {
+                    Meta = tenantMeta,
+                    BirthDate = Fixture.ImagingStudyMonthPrecisionPatientBirthDate,
+                })).Resource;
+
+            await Client.CreateAsync(
+                new ImagingStudy
+                {
+                    Meta = tenantMeta,
+                    Status = ImagingStudy.ImagingStudyStatus.Available,
+                    Subject = new ResourceReference($"Patient/{imagingStudyExactDayPatient.Id}"),
+                    Started = Fixture.ImagingStudyStarted,
+                });
+
+            await Client.CreateAsync(
+                new ImagingStudy
+                {
+                    Meta = tenantMeta,
+                    Status = ImagingStudy.ImagingStudyStatus.Available,
+                    Subject = new ResourceReference($"Patient/{imagingStudyMonthPrecisionPatient.Id}"),
+                    Started = Fixture.ImagingStudyStarted,
+                });
+
+            Bundle bundle = await Client.SearchPostAsync(
+                ResourceType.Patient.ToString(),
+                null,
+                default,
+                ("_has:ImagingStudy:patient:started", Fixture.ImagingStudyStarted),
+                ("birthdate", Fixture.ImagingStudyPatientBirthDate),
+                ("_tag", $"{Fixture.TenantTagSystem}|{tenantTagCode}"));
+
+            ValidateBundle(bundle, imagingStudyExactDayPatient, imagingStudyMonthPrecisionPatient);
+        }
+#endif
+
         [Fact]
         public async Task GivenAReverseChainSearchExpressionWithMultipleTargetTypes_WhenSearched_ThenCorrectBundleShouldBeReturned()
         {
@@ -400,6 +502,16 @@ public ClassFixture(DataStore dataStore, Format format, TestFhirServerFactory te
 
             public string OrganizationIdentifier { get; } = Guid.NewGuid().ToString();
 
+#if !Stu3
+            public string TenantTagSystem { get; } = "urn:tenantId";
+
+            public string ImagingStudyPatientBirthDate { get; } = "2018-06-06";
+
+            public string ImagingStudyMonthPrecisionPatientBirthDate { get; } = "2018-06";
+
+            public string ImagingStudyStarted { get; } = "2018-02-02T05:00:00.000";
+#endif
+
             public Patient SmithPatient { get; private set; }
 
             public DiagnosticReport SmithSnomedDiagnosticReport { get; private set; }