Hardening the multi-tenancy, marten managed partitioning to really, r…

…eally respect single tenanted documents mixed in. Closes GH-3442
JasperFx · Sep 25, 2024 · d385992 · d385992
1 parent 546583c
commit d385992
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 7 deletions.
diff --git a/src/Marten/DocumentStore.cs b/src/Marten/DocumentStore.cs
@@ -48,7 +48,6 @@ public DocumentStore(StoreOptions options)
         _logger = options.Logger();
         Serializer = options.Serializer();
 
-
         // Workaround to make database creation lazy so all StoreOptions
         // customizations can be done first
         if (Tenancy is DefaultTenancy d)

diff --git a/src/Marten/Schema/MartenManagedTenantListPartitions.cs b/src/Marten/Schema/MartenManagedTenantListPartitions.cs
@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using JasperFx.Core.Reflection;
 using Marten.Storage;
 using Marten.Storage.Metadata;
 using Npgsql;
@@ -34,7 +35,7 @@ public MartenManagedTenantListPartitions(StoreOptions options, string? schemaNam
 
     public void Apply(DocumentMapping mapping)
     {
-        if (mapping.TenancyStyle == TenancyStyle.Conjoined)
+        if (mapping.TenancyStyle == TenancyStyle.Conjoined || mapping.DocumentType.HasAttribute<SingleTenantedAttribute>())
         {
             mapping.Partitioning =
                 new ListPartitioning { Columns = [TenantIdColumn.Name] }.UsePartitionManager(Partitions);

diff --git a/src/Marten/Schema/SingleTenantedAttribute.cs b/src/Marten/Schema/SingleTenantedAttribute.cs
@@ -1,6 +1,7 @@
 #nullable enable
 using System;
 using Marten.Storage;
+using Weasel.Postgresql.Tables.Partitioning;
 
 namespace Marten.Schema;
 
@@ -15,5 +16,10 @@ public class SingleTenantedAttribute: MartenAttribute
     public override void Modify(DocumentMapping mapping)
     {
         mapping.TenancyStyle = TenancyStyle.Conjoined;
+        if (mapping.Partitioning != null && mapping.Partitioning is ListPartitioning lp &&
+            lp.PartitionManager is ManagedListPartitions)
+        {
+            mapping.Partitioning = null;
+        }
     }
 }
diff --git a/src/Marten/Storage/DocumentTable.cs b/src/Marten/Storage/DocumentTable.cs
@@ -81,7 +81,17 @@ public DocumentTable(DocumentMapping mapping): base(mapping.TableName)
         Indexes.AddRange(mapping.Indexes);
         ForeignKeys.AddRange(mapping.ForeignKeys);
 
-        Partitioning = mapping.Partitioning;
+        if (mapping.Partitioning != null)
+        {
+            if (mapping.Partitioning.Columns.All(HasColumn))
+            {
+                Partitioning = mapping.Partitioning;
+            }
+            else
+            {
+                Console.WriteLine($"Warning: Table {Identifier} is missing columns specified in the Partitioning scheme. This is probably an error in configuration");
+            }
+        }
 
         // Any column referred to in the partitioning has to be
         // part of the primary key
@@ -90,7 +100,7 @@ public DocumentTable(DocumentMapping mapping): base(mapping.TableName)
             IgnorePartitionsInMigration = mapping.IgnorePartitions;
             foreach (var columnName in Partitioning.Columns)
             {
-                var column = this.ModifyColumn(columnName);
+                var column = ModifyColumn(columnName);
                 column.AsPrimaryKey();
             }
         }

diff --git a/src/Marten/StoreOptions.cs b/src/Marten/StoreOptions.cs
@@ -63,6 +63,7 @@ public enum TenantIdStyle
 public partial class StoreOptions: IReadOnlyStoreOptions, IMigrationLogger, IDocumentSchemaResolver
 {
     public const int DefaultTimeout = 5;
+    internal const string? NoConnectionMessage = "No tenancy is configured! Ensure that you provided connection string in `AddMarten` method or called `UseNpgsqlDataSource`";
 
     internal static readonly Func<string, NpgsqlDataSourceBuilder> DefaultNpgsqlDataSourceBuilderFactory =
         connectionString => new NpgsqlDataSourceBuilder(connectionString);
@@ -396,7 +397,7 @@ void IReadOnlyStoreOptions.AssertDocumentTypeIsSoftDeleted(Type documentType)
     public ITenancy Tenancy
     {
         get => (_tenancy ?? throw new InvalidOperationException(
-                "No tenancy is configured! Ensure that you provided connection string in `AddMarten` method or called `UseNpgsqlDataSource`"))
+                NoConnectionMessage))
             .Value;
         set => _tenancy = new Lazy<ITenancy>(() => value);
     }
@@ -689,7 +690,7 @@ internal void Validate()
         if (_tenancy == null)
         {
             throw new InvalidOperationException(
-                "Tenancy not specified - provide either connection string or connection factory through Connection(..)");
+                NoConnectionMessage);
         }
     }
 
@@ -833,7 +834,11 @@ public PoliciesExpression ForAllDocuments(Action<DocumentMapping> configure)
         /// <returns></returns>
         public PoliciesExpression AllDocumentsAreMultiTenanted()
         {
-            return ForAllDocuments(_ => _.TenancyStyle = TenancyStyle.Conjoined);
+            return ForAllDocuments(mapping =>
+            {
+                if (mapping.DocumentType.HasAttribute<SingleTenantedAttribute>()) return;
+                mapping.TenancyStyle = TenancyStyle.Conjoined;
+            });
         }
 
         /// <summary>
@@ -848,6 +853,7 @@ public PoliciesExpression AllDocumentsAreMultiTenantedWithPartitioning(Action<Pa
             return ForAllDocuments(mapping =>
             {
                 if (mapping.DocumentType == typeof(DeadLetterEvent)) return;
+                if (mapping.DocumentType.HasAttribute<SingleTenantedAttribute>()) return;
 
                 mapping.PrimaryKeyTenancyOrdering = PrimaryKeyTenancyOrdering.Id_Then_TenantId;