diff --git a/src/Marten/Storage/MartenTenantAssignmentTable.cs b/src/Marten/Storage/MartenTenantAssignmentTable.cs new file mode 100644 index 0000000000..db6248298a --- /dev/null +++ b/src/Marten/Storage/MartenTenantAssignmentTable.cs @@ -0,0 +1,32 @@ +using Weasel.Postgresql.Tables; + +namespace Marten.Storage; + +/// +/// #4607: Marten-side derivative of Weasel's that +/// adds a disabled boolean column so can support the +/// / +/// lifecycle +/// the same way does (mirrors +/// TenantTable.DisabledColumn). +/// +/// +/// Subclassed on the Marten side rather than added to Weasel so existing Marten pools +/// pick up the new column via Weasel's additive table-delta migration (the NOT NULL +/// DEFAULT false means existing rows default to enabled) without requiring a coordinated +/// Weasel release. Marten owns the only construction site ('s +/// PoolFeatureSchema), so the substitution is fully contained. +/// +/// +internal class MartenTenantAssignmentTable: TenantAssignmentTable +{ + /// The added column's name — used by 's SQL queries. + public const string DisabledColumn = "disabled"; + + public MartenTenantAssignmentTable(string schemaName): base(schemaName) + { + // NOT NULL with a `false` default means the column can be added to an existing + // pool (legacy rows backfill to enabled) without a manual migration. + AddColumn(DisabledColumn).NotNull().DefaultValueByExpression("false"); + } +} diff --git a/src/Marten/Storage/ShardedTenancy.cs b/src/Marten/Storage/ShardedTenancy.cs index 8fc78cdb86..c9bc224766 100644 --- a/src/Marten/Storage/ShardedTenancy.cs +++ b/src/Marten/Storage/ShardedTenancy.cs @@ -192,9 +192,12 @@ public async ValueTask> BuildDatabases() await poolReader.CloseAsync().ConfigureAwait(false); - // Load all tenant assignments + // Load all tenant assignments — only active (non-disabled). Disabled + // tenants are excluded from the in-memory tenant→database cache so + // GetTenantAsync surfaces UnknownTenantIdException for them, mirroring + // MasterTableTenancy's soft-delete semantics (#4607). await using var assignReader = await ((DbCommand)conn - .CreateCommand($"select tenant_id, database_id from {_schemaName}.{TenantAssignmentTable.TableName}")) + .CreateCommand($"select tenant_id, database_id from {_schemaName}.{TenantAssignmentTable.TableName} where {MartenTenantAssignmentTable.DisabledColumn} = false")) .ExecuteReaderAsync().ConfigureAwait(false); while (await assignReader.ReadAsync().ConfigureAwait(false)) @@ -295,9 +298,11 @@ await _dataSource.Value tenantId = _options.TenantIdStyle.MaybeCorrectTenantId(tenantId); await maybeApplyChanges().ConfigureAwait(false); + // #4607: filter out soft-deleted assignments so disabled tenants are not + // resolvable — mirrors MasterTableTenancy's `disabled = false` gate. var result = await _dataSource.Value .CreateCommand( - $"select database_id from {_schemaName}.{TenantAssignmentTable.TableName} where tenant_id = :id") + $"select database_id from {_schemaName}.{TenantAssignmentTable.TableName} where tenant_id = :id and {MartenTenantAssignmentTable.DisabledColumn} = false") .With("id", tenantId) .ExecuteScalarAsync(ct).ConfigureAwait(false); @@ -319,14 +324,18 @@ await conn.CreateCommand($"select pg_advisory_lock({AdvisoryLockKey})") try { + // #4607: explicit assignment also clears the disabled flag so re-assigning + // a soft-deleted tenant via this API reactivates it. Pairs with the + // tolerant DisableTenantAsync / EnableTenantAsync semantics — explicit + // intent overrides prior soft-delete. await conn.CreateCommand( - $"insert into {_schemaName}.{TenantAssignmentTable.TableName} (tenant_id, database_id) values (:tid, :did) on conflict (tenant_id) do update set database_id = :did") + $"insert into {_schemaName}.{TenantAssignmentTable.TableName} (tenant_id, database_id) values (:tid, :did) on conflict (tenant_id) do update set database_id = :did, {MartenTenantAssignmentTable.DisabledColumn} = false") .With("tid", tenantId) .With("did", databaseId) .ExecuteNonQueryAsync(ct).ConfigureAwait(false); await conn.CreateCommand( - $"update {_schemaName}.{DatabasePoolTable.TableName} set tenant_count = (select count(*) from {_schemaName}.{TenantAssignmentTable.TableName} where database_id = :did) where database_id = :did") + $"update {_schemaName}.{DatabasePoolTable.TableName} set tenant_count = (select count(*) from {_schemaName}.{TenantAssignmentTable.TableName} where database_id = :did and {MartenTenantAssignmentTable.DisabledColumn} = false) where database_id = :did") .With("did", databaseId) .ExecuteNonQueryAsync(ct).ConfigureAwait(false); @@ -430,21 +439,87 @@ public async Task AddTenantAsync(string tenantId, CancellationToken toke $"Tenant '{tenantId}' was not assigned to any database after auto-assignment"); } - Task IDynamicTenantSource.DisableTenantAsync(string tenantId) - => throw new NotSupportedException( - "Disable/enable tenant lifecycle is not yet supported on sharded tenancy. " + - "Use RemoveTenantAsync to drop a tenant assignment, or open an issue to request the soft-delete surface."); + /// + /// #4607: soft-delete the tenant — flip disabled = true on its assignment row + /// and evict it from the in-memory tenant→database cache so subsequent tenant + /// resolution surfaces . Mirrors + /// 's lifecycle so the two dynamic sources behave + /// uniformly behind the store-agnostic + /// admin extensions. Idempotent — a no-op for an already-disabled or unknown tenant + /// (no exception; matches MasterTableTenancy's tolerance). + /// + public async Task DisableTenantAsync(string tenantId) + { + tenantId = _options.TenantIdStyle.MaybeCorrectTenantId(tenantId); + await maybeApplyChanges().ConfigureAwait(false); + + await _dataSource.Value + .CreateCommand( + $"update {_schemaName}.{TenantAssignmentTable.TableName} set {MartenTenantAssignmentTable.DisabledColumn} = true where tenant_id = :id") + .With("id", tenantId) + .ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false); + + // Evict from cache (and dispose only if no other tenant is using the same + // shared shard database — sharded tenancy reuses one MartenDatabase per + // assigned shard across tenants, unlike MasterTableTenancy's per-tenant DBs). + _tenantToDatabase = _tenantToDatabase.Remove(tenantId); + } + + /// + /// #4607: re-enable a soft-deleted tenant — flip disabled = false. The next + /// tenant resolution rehydrates the cache via the standard + /// path. Idempotent for already-enabled + /// or unknown tenants. + /// + public async Task EnableTenantAsync(string tenantId) + { + tenantId = _options.TenantIdStyle.MaybeCorrectTenantId(tenantId); + await maybeApplyChanges().ConfigureAwait(false); + + await _dataSource.Value + .CreateCommand( + $"update {_schemaName}.{TenantAssignmentTable.TableName} set {MartenTenantAssignmentTable.DisabledColumn} = false where tenant_id = :id") + .With("id", tenantId) + .ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false); + } Task IDynamicTenantSource.RemoveTenantAsync(string tenantId) => RemoveTenantAsync(tenantId, CancellationToken.None).AsTask(); - Task> IDynamicTenantSource.AllDisabledAsync() - => throw new NotSupportedException( - "Disable/enable tenant lifecycle is not yet supported on sharded tenancy."); + /// + /// #4607: enumerate currently soft-deleted tenants — the rows with + /// disabled = true. Used by the store-agnostic admin extension + /// . + /// + public async Task> AllDisabledAsync() + { + await maybeApplyChanges().ConfigureAwait(false); + + var list = new List(); + await using var conn = _dataSource.Value.CreateConnection(); + await conn.OpenAsync().ConfigureAwait(false); + + try + { + await using var reader = await ((DbCommand)conn + .CreateCommand( + $"select tenant_id from {_schemaName}.{TenantAssignmentTable.TableName} where {MartenTenantAssignmentTable.DisabledColumn} = true")) + .ExecuteReaderAsync().ConfigureAwait(false); - Task IDynamicTenantSource.EnableTenantAsync(string tenantId) - => throw new NotSupportedException( - "Disable/enable tenant lifecycle is not yet supported on sharded tenancy."); + while (await reader.ReadAsync().ConfigureAwait(false)) + { + list.Add(await reader.GetFieldValueAsync(0).ConfigureAwait(false)); + } + + await reader.CloseAsync().ConfigureAwait(false); + } + finally + { + await conn.CloseAsync().ConfigureAwait(false); + } + + return list; + } #endregion @@ -476,12 +551,37 @@ await conn.CreateCommand($"select pg_advisory_lock({AdvisoryLockKey})") try { - // Double-check after acquiring lock - var existingDbId = (string?)await conn + // #4607: under the lock, distinguish three cases: + // (a) tenant has an active assignment → use it + // (b) tenant has a DISABLED assignment → throw UnknownTenantIdException + // (mirrors MasterTableTenancy; auto-assigning here would silently + // resurrect the soft-deleted tenant, possibly onto a different shard) + // (c) no assignment at all → fall through to auto-assign + var existingState = await ((DbCommand)conn .CreateCommand( - $"select database_id from {_schemaName}.{TenantAssignmentTable.TableName} where tenant_id = :id") - .With("id", tenantId) - .ExecuteScalarAsync(CancellationToken.None).ConfigureAwait(false); + $"select database_id, {MartenTenantAssignmentTable.DisabledColumn} from {_schemaName}.{TenantAssignmentTable.TableName} where tenant_id = :id") + .With("id", tenantId)) + .ExecuteReaderAsync(CancellationToken.None).ConfigureAwait(false); + + string? existingDbId = null; + var existingDisabled = false; + try + { + if (await existingState.ReadAsync().ConfigureAwait(false)) + { + existingDbId = await existingState.GetFieldValueAsync(0).ConfigureAwait(false); + existingDisabled = await existingState.GetFieldValueAsync(1).ConfigureAwait(false); + } + } + finally + { + await existingState.CloseAsync().ConfigureAwait(false); + } + + if (existingDisabled) + { + throw new UnknownTenantIdException(tenantId); + } if (existingDbId != null && _databasesById.TryFind(existingDbId, out database)) { @@ -665,7 +765,10 @@ public PoolFeatureSchema(string schemaName, StoreOptions options) protected override IEnumerable schemaObjects() { yield return new DatabasePoolTable(_schemaName); - yield return new TenantAssignmentTable(_schemaName); + // #4607: Marten subclass adds the `disabled` column for soft-delete + // (Disable/Enable lifecycle) without requiring a Weasel release. The + // additive column-add migration upgrades existing pools in place. + yield return new MartenTenantAssignmentTable(_schemaName); } } diff --git a/src/MultiTenancyTests/sharded_tenancy_soft_delete_tests.cs b/src/MultiTenancyTests/sharded_tenancy_soft_delete_tests.cs new file mode 100644 index 0000000000..32890f4611 --- /dev/null +++ b/src/MultiTenancyTests/sharded_tenancy_soft_delete_tests.cs @@ -0,0 +1,286 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using JasperFx; +using JasperFx.MultiTenancy; +using Marten; +using Marten.Storage; +using Marten.Testing.Documents; +using Marten.Testing.Harness; +using Npgsql; +using Shouldly; +using Weasel.Core; +using Weasel.Postgresql; +using Weasel.Postgresql.Tables; +using Xunit; + +namespace MultiTenancyTests; + +/// +/// #4607 — soft-delete (Disable / Enable / AllDisabled) lifecycle on +/// 's implementation. +/// Mirrors 's soft-delete semantics so the two +/// dynamic sources behave uniformly behind the store-agnostic +/// DynamicTenancyAdminExtensions. Backed by a Marten-added +/// disabled boolean not null default false column on the existing +/// mt_tenant_assignments table (). +/// +[Collection("sharded-tenancy")] +public class sharded_tenancy_soft_delete_tests : IAsyncLifetime +{ + private readonly ShardedTenancyFixture _fixture; + private IDocumentStore _store = null!; + + public sharded_tenancy_soft_delete_tests(ShardedTenancyFixture fixture) + { + _fixture = fixture; + } + + public async Task InitializeAsync() + { + await using var conn = new NpgsqlConnection(ConnectionSource.ConnectionString); + await conn.OpenAsync(); + try { await conn.DropSchemaAsync("sharded"); } catch { } + + foreach (var connStr in _fixture.ConnectionStrings.Values) + { + await using var tenantConn = new NpgsqlConnection(connStr); + await tenantConn.OpenAsync(); + try { await tenantConn.DropSchemaAsync("tenants"); } catch { } + await ShardedTenancyFixture.cleanMartenObjectsInPublicSchema(tenantConn); + } + } + + public Task DisposeAsync() + { + _store?.Dispose(); + return Task.CompletedTask; + } + + private IDocumentStore CreateStore() + { + _store = DocumentStore.For(opts => + { + opts.MultiTenantedWithShardedDatabases(x => + { + x.ConnectionString = ConnectionSource.ConnectionString; + x.SchemaName = "sharded"; + x.PartitionSchemaName = "tenants"; + foreach (var (dbName, connStr) in _fixture.ConnectionStrings) + { + x.AddDatabase(dbName, connStr); + } + }); + opts.AutoCreateSchemaObjects = AutoCreate.All; + opts.RegisterDocumentType(); + }); + return _store; + } + + // ---- DisableTenantAsync ---- + + [Fact] + public async Task disabled_tenant_resolution_throws_UnknownTenantIdException() + { + CreateStore(); + var source = (IDynamicTenantSource)_store.Options.Tenancy; + + await source.AddTenantAsync("tenant-a", CancellationToken.None); + // Sanity: tenant is resolvable before the soft-delete. + (await _store.Options.Tenancy.GetTenantAsync("tenant-a")).TenantId.ShouldBe("tenant-a"); + + await source.DisableTenantAsync("tenant-a"); + + await Should.ThrowAsync(async () => + await _store.Options.Tenancy.GetTenantAsync("tenant-a")); + } + + [Fact] + public async Task disabled_tenant_is_filtered_from_FindDatabaseForTenantAsync() + { + CreateStore(); + var sharded = (ShardedTenancy)_store.Options.Tenancy; + var source = (IDynamicTenantSource)sharded; + + var dbId = await source.AddTenantAsync("tenant-b", CancellationToken.None); + (await sharded.FindDatabaseForTenantAsync("tenant-b", CancellationToken.None)).ShouldBe(dbId); + + await source.DisableTenantAsync("tenant-b"); + + (await sharded.FindDatabaseForTenantAsync("tenant-b", CancellationToken.None)) + .ShouldBeNull("FindDatabaseForTenantAsync must hide soft-deleted assignments"); + } + + [Fact] + public async Task auto_assign_on_a_disabled_tenant_throws_UnknownTenantIdException_not_resurrects() + { + // The dangerous case: without the under-lock disabled-check, auto-assign would + // create a fresh assignment for a disabled tenant — silently undoing the + // soft-delete and possibly placing the tenant on a different shard than where + // its data lives. + CreateStore(); + var sharded = (ShardedTenancy)_store.Options.Tenancy; + var source = (IDynamicTenantSource)sharded; + + var originalDbId = await source.AddTenantAsync("tenant-c", CancellationToken.None); + await source.DisableTenantAsync("tenant-c"); + + await Should.ThrowAsync(async () => + await sharded.GetTenantAsync("tenant-c")); + + // No reassignment happened — the disabled row still points at the original shard. + var rawAssignment = await readRawAssignment(_store, "tenant-c"); + rawAssignment.databaseId.ShouldBe(originalDbId, "the disabled row must keep its original database_id"); + rawAssignment.disabled.ShouldBeTrue("the disabled flag must still be set"); + } + + [Fact] + public async Task DisableTenantAsync_is_idempotent_for_already_disabled_or_unknown() + { + CreateStore(); + var source = (IDynamicTenantSource)_store.Options.Tenancy; + + await source.AddTenantAsync("tenant-d", CancellationToken.None); + await source.DisableTenantAsync("tenant-d"); + await source.DisableTenantAsync("tenant-d"); // already disabled — no-op + await source.DisableTenantAsync("unknown-tenant"); // never existed — no-op + } + + // ---- EnableTenantAsync ---- + + [Fact] + public async Task EnableTenantAsync_restores_resolution_for_a_previously_disabled_tenant() + { + CreateStore(); + var source = (IDynamicTenantSource)_store.Options.Tenancy; + + var dbId = await source.AddTenantAsync("tenant-e", CancellationToken.None); + await source.DisableTenantAsync("tenant-e"); + + // Pre-condition: disabled + await Should.ThrowAsync(async () => + await _store.Options.Tenancy.GetTenantAsync("tenant-e")); + + await source.EnableTenantAsync("tenant-e"); + + var tenant = await _store.Options.Tenancy.GetTenantAsync("tenant-e"); + tenant.TenantId.ShouldBe("tenant-e"); + + // Same shard the tenant had before the soft-delete — re-enable doesn't relocate. + (await ((ShardedTenancy)_store.Options.Tenancy).FindDatabaseForTenantAsync("tenant-e", CancellationToken.None)) + .ShouldBe(dbId); + } + + [Fact] + public async Task EnableTenantAsync_is_idempotent_for_already_enabled_or_unknown() + { + CreateStore(); + var source = (IDynamicTenantSource)_store.Options.Tenancy; + + await source.AddTenantAsync("tenant-f", CancellationToken.None); + await source.EnableTenantAsync("tenant-f"); // already enabled + await source.EnableTenantAsync("unknown-tenant"); // never existed + } + + // ---- AllDisabledAsync ---- + + [Fact] + public async Task AllDisabledAsync_returns_only_currently_disabled_tenants() + { + CreateStore(); + var source = (IDynamicTenantSource)_store.Options.Tenancy; + + await source.AddTenantAsync("active-1", CancellationToken.None); + await source.AddTenantAsync("disabled-1", CancellationToken.None); + await source.AddTenantAsync("disabled-2", CancellationToken.None); + await source.AddTenantAsync("active-2", CancellationToken.None); + + await source.DisableTenantAsync("disabled-1"); + await source.DisableTenantAsync("disabled-2"); + + var disabled = await source.AllDisabledAsync(); + + disabled.OrderBy(x => x).ShouldBe(new[] { "disabled-1", "disabled-2" }); + } + + [Fact] + public async Task AllDisabledAsync_is_empty_when_no_tenants_are_disabled() + { + CreateStore(); + var source = (IDynamicTenantSource)_store.Options.Tenancy; + + await source.AddTenantAsync("active-only", CancellationToken.None); + + (await source.AllDisabledAsync()).ShouldBeEmpty(); + } + + // ---- Re-enable via explicit assignment ---- + + [Fact] + public async Task explicit_assign_via_AddTenantAsync_with_databaseId_reactivates_a_disabled_tenant() + { + CreateStore(); + var sharded = (ShardedTenancy)_store.Options.Tenancy; + var source = (IDynamicTenantSource)sharded; + + await source.AddTenantAsync("tenant-g", CancellationToken.None); + await source.DisableTenantAsync("tenant-g"); + + // Caller-supplied overload of AddTenantAsync — semantically "assign tenant + // to this specific database". Stronger intent than DisableTenantAsync, so + // the explicit re-assignment clears the disabled flag. + await source.AddTenantAsync("tenant-g", _fixture.DbNames[0]); + + (await source.AllDisabledAsync()).ShouldNotContain("tenant-g"); + (await sharded.FindDatabaseForTenantAsync("tenant-g", CancellationToken.None)) + .ShouldBe(_fixture.DbNames[0]); + } + + // ---- Cross-source admin extension surface (jasperfx#413) ---- + + [Fact] + public async Task lifecycle_works_through_the_store_agnostic_dynamic_tenant_source_interface() + { + // This is the surface CritterWatch's tenant-management UI uses — it resolves + // IDynamicTenantSource from DI and calls the lifecycle methods without + // sniffing the concrete tenancy type. Sharded must behave like MasterTable. + CreateStore(); + IDynamicTenantSource source = (IDynamicTenantSource)_store.Options.Tenancy; + + await source.AddTenantAsync("uniform-tenant", CancellationToken.None); + await source.DisableTenantAsync("uniform-tenant"); + (await source.AllDisabledAsync()).ShouldContain("uniform-tenant"); + + await source.EnableTenantAsync("uniform-tenant"); + (await source.AllDisabledAsync()).ShouldNotContain("uniform-tenant"); + + await source.RemoveTenantAsync("uniform-tenant"); + (await source.AllDisabledAsync()).ShouldNotContain("uniform-tenant"); + } + + // ---- helpers ---- + + private static async Task<(string databaseId, bool disabled)> readRawAssignment(IDocumentStore store, string tenantId) + { + // Direct read against the assignment table to verify the on-disk state + // independent of the caching paths the production code uses. + await using var conn = new NpgsqlConnection(ConnectionSource.ConnectionString); + await conn.OpenAsync(); + var schemaName = "sharded"; + var tableName = TenantAssignmentTable.TableName; + var disabledColumn = MartenTenantAssignmentTable.DisabledColumn; + await using var reader = await ((System.Data.Common.DbCommand)conn + .CreateCommand( + $"select database_id, {disabledColumn} from {schemaName}.{tableName} where tenant_id = :id") + .With("id", tenantId)) + .ExecuteReaderAsync(); + + if (!await reader.ReadAsync()) throw new InvalidOperationException($"No row for tenant '{tenantId}'"); + var dbId = await reader.GetFieldValueAsync(0); + var disabled = await reader.GetFieldValueAsync(1); + await reader.CloseAsync(); + return (dbId, disabled); + } +}