New API to drop Marten managed tenant id partitions at runtime. Updat…

…ed Weasel. Closes GH-3583
JasperFx · Jan 14, 2025 · 130bd28 · 130bd28
1 parent 82d51f6
commit 130bd28
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 5 deletions.
diff --git a/src/Marten/AdvancedOperations.cs b/src/Marten/AdvancedOperations.cs
@@ -235,7 +235,7 @@ public async Task RebuildSingleStreamAsync<T>(Guid id, CancellationToken token =
     /// </summary>
     /// <param name="token"></param>
     /// <param name="tenantIds"></param>
-    public Task AddMartenManagedTenantsAsync(CancellationToken token, params string[] tenantIds)
+    public Task<TablePartitionStatus[]> AddMartenManagedTenantsAsync(CancellationToken token, params string[] tenantIds)
     {
         var dict = new Dictionary<string, string>();
         foreach (var tenantId in tenantIds)
@@ -275,6 +275,33 @@ public async Task<TablePartitionStatus[]> AddMartenManagedTenantsAsync(Cancellat
             token).ConfigureAwait(false);
     }
 
+    /// <summary>
+    /// Drop a tenant partition from all tables that use the Marten managed tenant partitioning. NOTE: you have to supply
+    /// the partition suffix for the tenant, not necessarily the tenant id. In most cases we think this will probably
+    /// be the same value, but you may have to "sanitize" the suffix name
+    /// </summary>
+    /// <param name="suffixes"></param>
+    /// <param name="token"></param>
+    /// <exception cref="InvalidOperationException"></exception>
+    public async Task RemoveMartenManagedTenantsAsync(string[] suffixes, CancellationToken token)
+    {
+        if (_store.Options.TenantPartitions == null)
+        {
+            throw new InvalidOperationException(
+                $"Marten-managed per-tenant partitioning is not active in this store. Did you miss a call to {nameof(StoreOptions)}.{nameof(StoreOptions.Policies)}.{nameof(StoreOptions.PoliciesExpression.PartitionMultiTenantedDocumentsUsingMartenManagement)}()?");
+        }
+
+        if (_store.Tenancy is not DefaultTenancy)
+            throw new InvalidOperationException(
+                "This option is not (yet) supported in combination with database per tenant multi-tenancy");
+        var database = (PostgresqlDatabase)_store.Tenancy.Default.Database;
+
+
+        var logger = _store.Options.LogFactory?.CreateLogger<DocumentStore>() ?? NullLogger<DocumentStore>.Instance;
+        await _store.Options.TenantPartitions.Partitions.DropPartitionFromAllTables(database, logger, suffixes,
+            token).ConfigureAwait(false);
+    }
+
     /// <summary>
     /// Configure and execute a batch masking of protected data for a subset of the events
     /// in the event store

diff --git a/src/Marten/Marten.csproj b/src/Marten/Marten.csproj
@@ -61,7 +61,7 @@
         <PackageReference Include="Polly.Core" Version="8.5.0" />
         <PackageReference Include="System.Threading.Tasks.Dataflow" Version="8.0.0" />
         <PackageReference Include="System.Text.Json" Version="8.0.5" />
-        <PackageReference Include="Weasel.Postgresql" Version="7.13.2" />
+        <PackageReference Include="Weasel.Postgresql" Version="7.13.3" />
     </ItemGroup>
 
     <!--SourceLink specific settings-->

diff --git a/src/MultiTenancyTests/marten_managed_tenant_id_partitioning.cs b/src/MultiTenancyTests/marten_managed_tenant_id_partitioning.cs
@@ -27,8 +27,10 @@ public async Task InitializeAsync()
         await conn.OpenAsync();
         try
         {
-            await conn.CreateCommand($"delete from tenants.{MartenManagedTenantListPartitions.TableName}")
-                .ExecuteNonQueryAsync();
+            await conn.DropSchemaAsync("tenants");
+
+            // await conn.CreateCommand($"delete from tenants.{MartenManagedTenantListPartitions.TableName}")
+            //     .ExecuteNonQueryAsync();
         }
         catch (Exception)
         {
@@ -75,6 +77,41 @@ await theStore
 
     }
 
+    [Fact]
+    public async Task add_then_remove_tenants_at_runtime()
+    {
+        StoreOptions(opts =>
+        {
+            opts.Policies.AllDocumentsAreMultiTenanted();
+            opts.Policies.PartitionMultiTenantedDocumentsUsingMartenManagement("tenants");
+
+            opts.Schema.For<Target>();
+            opts.Schema.For<User>();
+        }, true);
+
+        var statuses = await theStore
+            .Advanced
+            // This is ensuring that there are tenant id partitions for all multi-tenanted documents
+            // with the named tenant ids
+            .AddMartenManagedTenantsAsync(CancellationToken.None, "a1", "a2", "a3");
+
+        foreach (var status in statuses)
+        {
+            status.Status.ShouldBe(PartitionMigrationStatus.Complete);
+        }
+
+        await theStore.Storage.ApplyAllConfiguredChangesToDatabaseAsync();
+        await theStore.Storage.Database.AssertDatabaseMatchesConfigurationAsync();
+
+        await theStore.Advanced.RemoveMartenManagedTenantsAsync(["a2"], CancellationToken.None);
+
+        var targetTable = await theStore.Storage.Database.ExistingTableFor(typeof(Target));
+        assertTableHasTenantPartitions(targetTable, "a1", "a3");
+
+        var userTable = await theStore.Storage.Database.ExistingTableFor(typeof(User));
+        assertTableHasTenantPartitions(userTable, "a1", "a3");
+    }
+
 
 
     [Fact]
@@ -166,7 +203,11 @@ public async Task can_build_then_add_additive_partitions_later()
         await theStore.Storage.ApplyAllConfiguredChangesToDatabaseAsync();
 
         // Little overlap to prove it's idempotent
-        await theStore.Advanced.AddMartenManagedTenantsAsync(CancellationToken.None, "a1", "b1", "b2");
+        var statuses = await theStore.Advanced.AddMartenManagedTenantsAsync(CancellationToken.None, "a1", "b1", "b2");
+        foreach (var status in statuses)
+        {
+            status.Status.ShouldBe(PartitionMigrationStatus.Complete);
+        }
 
         var targetTable = await theStore.Storage.Database.ExistingTableFor(typeof(Target));
         assertTableHasTenantPartitions(targetTable, "a1", "a2", "a3", "b1", "b2");
@@ -175,6 +216,8 @@ public async Task can_build_then_add_additive_partitions_later()
         assertTableHasTenantPartitions(userTable, "a1", "a2", "a3", "b1", "b2");
     }
 
+
+
     private void assertTableHasTenantPartitions(Table table, params string[] tenantIds)
     {
         var partitioning = table.Partitioning.ShouldBeOfType<ListPartitioning>();