From c81a3dd6dc13922d8a5c17bd2f5c74d9870ee696 Mon Sep 17 00:00:00 2001
From: "Jeremy D. Miller" <jeremydmiller@yahoo.com>
Date: Sun, 8 Dec 2024 14:14:29 -0600
Subject: [PATCH] Duplicates the tenant id in a sub query when using ctid in
 the case of using conjoined tenancy and partition by tenant

---
 ...artitioning_with_tenant_id_and_subquery.cs | 58 +++++++++++++++++++
 .../DocumentQueryableMemberCollection.cs      |  4 ++
 .../SqlGeneration/Filters/SubQueryFilter.cs   | 11 ++++
 src/Marten/Schema/DocumentMapping.cs          |  2 +
 4 files changed, 75 insertions(+)
 create mode 100644 src/CoreTests/Partitioning/querying_against_conjoined_partitioning_with_tenant_id_and_subquery.cs

diff --git a/src/CoreTests/Partitioning/querying_against_conjoined_partitioning_with_tenant_id_and_subquery.cs b/src/CoreTests/Partitioning/querying_against_conjoined_partitioning_with_tenant_id_and_subquery.cs
new file mode 100644
index 0000000000..404367f268
--- /dev/null
+++ b/src/CoreTests/Partitioning/querying_against_conjoined_partitioning_with_tenant_id_and_subquery.cs
@@ -0,0 +1,58 @@
+using System.Linq;
+using System.Threading.Tasks;
+using Marten;
+using Marten.Testing.Documents;
+using Marten.Testing.Harness;
+using Shouldly;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace CoreTests.Partitioning;
+
+public class querying_against_conjoined_partitioning_with_tenant_id_and_subquery : OneOffConfigurationsContext
+{
+    private readonly ITestOutputHelper _output;
+
+    public querying_against_conjoined_partitioning_with_tenant_id_and_subquery(ITestOutputHelper output)
+    {
+        _output = output;
+    }
+
+    [Fact]
+    public async Task do_not_bleed_tenant_data_because_of_select_queries()
+    {
+        var reds = Target.GenerateRandomData(100).ToArray();
+        var blues = Target.GenerateRandomData(1000).ToArray();
+
+        StoreOptions(opts =>
+        {
+            opts.Policies.AllDocumentsAreMultiTenantedWithPartitioning(x =>
+            {
+                x.ByList()
+                    .AddPartition("red", "red")
+                    .AddPartition("blue", "blue");
+            });
+        });
+
+        await theStore.BulkInsertAsync("red", reds);
+        await theStore.BulkInsertAsync("blue", blues);
+
+        using var session = theStore.LightweightSession("red");
+        session.Logger = new TestOutputMartenLogger(_output);
+
+        var matching = await session.Query<Target>()
+            .Where(x => x.Children.Any(c => c.Number > 8))
+            .OrderBy(x => x.Id)
+            .Select(x => x.Id)
+            .ToListAsync();
+
+        var expected = reds.Where(x => x.Children.Any(c => c.Number > 8))
+            .OrderBy(x => x.Id)
+            .Select(x => x.Id)
+            .ToList();
+
+        matching.Count.ShouldBe(expected.Count);
+
+        matching.ShouldBe(expected);
+    }
+}
diff --git a/src/Marten/Linq/Members/DocumentQueryableMemberCollection.cs b/src/Marten/Linq/Members/DocumentQueryableMemberCollection.cs
index 6b523f1dbb..cf193f8a3c 100644
--- a/src/Marten/Linq/Members/DocumentQueryableMemberCollection.cs
+++ b/src/Marten/Linq/Members/DocumentQueryableMemberCollection.cs
@@ -6,8 +6,10 @@
 using System.Linq.Expressions;
 using System.Reflection;
 using JasperFx.Core;
+using JasperFx.Core.Reflection;
 using Marten.Linq.Parsing.Operators;
 using Marten.Schema;
+using Marten.Storage;
 using Weasel.Postgresql;
 using Weasel.Postgresql.SqlGeneration;
 
@@ -24,6 +26,8 @@ public DocumentQueryableMemberCollection(IDocumentMapping mapping, StoreOptions
         ElementType = mapping.DocumentType;
     }
 
+    public TenancyStyle TenancyStyle { get; set; } = TenancyStyle.Single;
+
     public string MemberName => string.Empty;
 
     public string JsonPathSegment => "";
diff --git a/src/Marten/Linq/SqlGeneration/Filters/SubQueryFilter.cs b/src/Marten/Linq/SqlGeneration/Filters/SubQueryFilter.cs
index 504b476580..c35acdf6a3 100644
--- a/src/Marten/Linq/SqlGeneration/Filters/SubQueryFilter.cs
+++ b/src/Marten/Linq/SqlGeneration/Filters/SubQueryFilter.cs
@@ -1,6 +1,7 @@
 #nullable enable
 using Marten.Internal;
 using Marten.Linq.Members;
+using Marten.Storage;
 using Weasel.Postgresql;
 using Weasel.Postgresql.SqlGeneration;
 
@@ -43,6 +44,16 @@ void ISqlFragment.Apply(ICommandBuilder builder)
             builder.Append("NOT(");
         }
 
+        if (builder.TenantId != Tenancy.DefaultTenantId && Member is ChildCollectionMember child)
+        {
+            if (child.Ancestors[0] is DocumentQueryableMemberCollection c && c.TenancyStyle == TenancyStyle.Conjoined)
+            {
+                builder.Append("d.tenant_id = ");
+                builder.AppendParameter(builder.TenantId);
+                builder.Append(" and ");
+            }
+        }
+
         builder.Append("d.ctid in (select ctid from ");
         builder.Append(_exportName!);
         builder.Append(")");
diff --git a/src/Marten/Schema/DocumentMapping.cs b/src/Marten/Schema/DocumentMapping.cs
index d1c5be3a02..395db9a4d0 100644
--- a/src/Marten/Schema/DocumentMapping.cs
+++ b/src/Marten/Schema/DocumentMapping.cs
@@ -128,6 +128,8 @@ public DocumentMapping(Type documentType, StoreOptions storeOptions)
 
         StoreOptions.applyPostPolicies(this);
 
+        QueryMembers.TenancyStyle = TenancyStyle;
+
         _schema = new Lazy<DocumentSchema>(() => new DocumentSchema(this));
 
         if (DisablePartitioningIfAny)