feat(#335): per-tenant managed partitioning parity with Marten (documents + streams)

Directory.Packages.props

@@ -60,10 +60,16 @@
              databases and retries transient connection failures, with SqlServerMigrator now
              overriding ReleaseConnectionPoolAsync (SqlConnection.ClearPool) and
              IsTransientConnectionFailure over the SQL Server/Azure SQL resource-limit error
-             set. Polecat picks this up through the migrator base class — no code change. -->
-        <PackageVersion Include="Weasel.EntityFrameworkCore" Version="9.17.0" />
-        <PackageVersion Include="Weasel.SqlServer" Version="9.17.0" />
-        <PackageVersion Include="Weasel.Storage" Version="9.17.0" />
+             set. Polecat picks this up through the migrator base class — no code change.
+             Weasel 9.18.0: weasel#362/#374 — ManagedTenantPartitions parity with the
+             PostgreSQL ManagedListPartitions: TenantDropBehavior (RetainData/DeleteData),
+             AllowOrdinalSharing + explicit-ordinal add overloads (tenant bucketing),
+             MigrateAllTablesAsync new-table back-fill, and TenantPartitionAddResult /
+             TablePartitionStatus batch status reporting. Consumed by Polecat #335
+             (per-tenant partitioning parity for documents + streams). -->
+        <PackageVersion Include="Weasel.EntityFrameworkCore" Version="9.18.0" />
+        <PackageVersion Include="Weasel.SqlServer" Version="9.18.0" />
+        <PackageVersion Include="Weasel.Storage" Version="9.18.0" />
 
         <!-- Strongly typed IDs -->
         <PackageVersion Include="StronglyTypedId" Version="1.0.0-beta08" />

docs/documents/partitioning.md

@@ -59,8 +59,71 @@ column/type is reported as a rebuild rather than performed silently.
 
 ## Limitations
 
-- Supported for **single-tenant** document tables only; combining partitioning with conjoined
-  multi-tenancy throws at start-up.
+- Supported for **single-tenant** document tables only; combining member RANGE partitioning with
+  conjoined multi-tenancy throws at start-up. Conjoined document tables can instead use the managed
+  per-tenant partitioning below.
 - Dropping aged partitions for retention (`SWITCH`/`MERGE RANGE`) and externally-managed partition
   rolling are not yet wired into the document API — for now, manage those out of band, or keep pruning
   with a predicate delete.
+
+## Managed per-tenant partitioning (#335)
+
+Conjoined multi-tenanted document tables can be physically partitioned **per tenant** through the
+store's shared managed tenant partitioning — the SQL Server counterpart of Marten's
+`AllDocumentsAreMultiTenantedWithPartitioning` + `PartitionMultiTenantedDocumentsUsingMartenManagement`:
+
+```csharp
+var store = DocumentStore.For(opts =>
+{
+    opts.Connection("...");
+
+    // Make every document conjoined multi-tenanted AND tenant-partitioned:
+    opts.Policies.AllDocumentsAreMultiTenantedWithPartitioning();
+
+    // — or, when the store is already conjoined, enable just the partitioning:
+    // opts.Events.TenancyStyle = TenancyStyle.Conjoined;
+    // opts.Policies.PartitionMultiTenantedDocumentsUsingPolecatManagement();
+
+    // Per-type escape hatch (the [SingleTenanted]/DisablePartitioningIfAny analogue):
+    opts.Policies.ForDocument<AuditRecord>(p => p.DisableTenantPartitioning = true);
+});
+```
+
+Every document table then carries a `tenant_ordinal int` primary-key column and is `RANGE RIGHT`
+partitioned on it, driven by the **one `pc_tenant_partitions` registry per database** — the same
+registry, ordinals, and physical layout as
+[per-tenant event partitioning](/events/multitenancy#per-tenant-event-partitioning), so a store using
+both keeps a single coherent tenant → ordinal map across `pc_events`, `pc_streams`, and every
+document table. The ordinal is resolved server-side from the registry on every write (upsert, insert,
+update, bulk insert, and projection writes), so cross-process ordinal drift cannot mis-route rows.
+
+Tenants are onboarded **lazily on first write** (matching the event-append behavior), or explicitly
+with per-table status reporting:
+
+```csharp
+// Onboard tenants up front — returns Weasel TablePartitionStatus[] per managed table:
+var statuses = await store.Advanced.AddPolecatManagedTenantsAsync(ct, "tenant-a", "tenant-b");
+
+// Tenant bucketing (Weasel 9.18.0): map many small tenants onto one shared partition ordinal.
+// Requires ManagedTenantPartitions.AllowOrdinalSharing:
+await store.Advanced.AddPolecatManagedTenantsAsync(
+    new Dictionary<string, int> { ["small-1"] = 1, ["small-2"] = 1 }, ct);
+
+// Remove a tenant. SQL Server's MERGE RANGE alone would retain the rows —
+// TenantDropBehavior.DeleteData physically purges the tenant's rows from every managed
+// table first (PostgreSQL managed-drop parity):
+await store.Advanced.RemovePolecatManagedTenantsAsync(
+    ["tenant-b"], TenantDropBehavior.DeleteData, ct);
+```
+
+Notes:
+
+- Requires `TenancyStyle.Conjoined` (asserted at store construction) and cannot be combined with
+  member `PartitionByRange` on the same document type — a SQL Server table supports only one
+  partition scheme.
+- The registry and partition function/scheme are database-global objects: one tenant-partitioned
+  store per database.
+- `AddPolecatManagedTenantsAsync` splits the tables of document types **registered** with the store
+  (`opts.Schema.For<T>()` or prior use); a table created later bakes the full boundary set at
+  creation.
+- The daemon's dead-letter document (`pc_doc_deadletterevent`) is always excluded, mirroring Marten.

docs/events/multitenancy.md

@@ -100,9 +100,11 @@ When enabled, Polecat:
   `pc_events_sequence_{ordinal}` object (created on demand) via `NEXT VALUE FOR`, rather than a single
   global `IDENTITY`. `seq_id` is therefore unique only *within* a tenant, so the `pc_events` primary
   key becomes composite `(tenant_ordinal, seq_id)`.
-* **Physically partitions `pc_events` by tenant** — the table is `RANGE RIGHT` partitioned on the
-  tenant `ordinal`, and a new partition is split in as each tenant registers. A tenant's events live
-  in their own physical partition, so per-tenant scans and rebuilds touch only that partition.
+* **Physically partitions `pc_events` and `pc_streams` by tenant** — both tables are `RANGE RIGHT`
+  partitioned on the tenant `ordinal`, and a new partition is split in as each tenant registers. A
+  tenant's events and stream rows live in their own physical partitions, so per-tenant scans and
+  rebuilds touch only those partitions (Marten parity: `mt_streams` rides `mt_events`' tenant
+  partitioning).
 
 ```cs
 // Each tenant's seq_id starts at 1 and advances independently
@@ -145,7 +147,12 @@ path byte-for-byte. The flag **requires** `TenancyStyle.Conjoined` (there is not
 otherwise) and is currently incompatible with `UseArchivedStreamPartitioning` — a SQL Server table
 supports only one partition scheme; both raise an error at store construction.
 
-Physical partitioning applies to `pc_events` (the table that drives the bounded per-tenant scan);
-`pc_streams` is accessed by point lookup and is left unpartitioned. The partition function/scheme are
-database-global objects, so a single database should host one tenant-partitioned event store.
+Physical partitioning applies to `pc_events` and (since #335) `pc_streams`. The partition
+function/scheme are database-global objects, so a single database should host one tenant-partitioned
+event store.
+
+Document tables can join the same managed per-tenant partitioning — including runtime tenant
+onboarding/removal via `store.Advanced.AddPolecatManagedTenantsAsync` /
+`RemovePolecatManagedTenantsAsync` — see
+[Document multi-tenancy partitioning](/documents/partitioning#managed-per-tenant-partitioning-335).
 :::

src/Polecat.Tests/Events/tenant_partitioned_streams_tests.cs

@@ -0,0 +1,139 @@
+using JasperFx;
+using Microsoft.Data.SqlClient;
+using Polecat.Tests.Harness;
+using Polecat.TestUtils;
+
+namespace Polecat.Tests.Events;
+
+/// <summary>
+///     Covers #335 scope 2 — pc_streams is partitioned alongside pc_events under
+///     UseTenantPartitionedEvents (Marten parity: mt_streams rides mt_events' tenant partitioning):
+///     schema shape, physical partition placement of stream rows, and the stream-version update path
+///     against the partitioned table.
+/// </summary>
+[Collection("tenant-partitioning")]
+public class tenant_partitioned_streams_tests : IAsyncLifetime
+{
+    private const string Schema = "pt_streams";
+
+    public async Task InitializeAsync()
+    {
+        await TestSchema.DropSchemaTablesAsync(Schema);
+        await PartitionTestCleanup.DropEventsPartitionObjectsAsync();
+        await TestSchema.DropSequencesAsync(Schema);
+    }
+
+    public Task DisposeAsync() => Task.CompletedTask;
+
+    private static DocumentStore CreateStore()
+    {
+        return DocumentStore.For(opts =>
+        {
+            opts.ConnectionString = ConnectionSource.ConnectionString;
+            opts.DatabaseSchemaName = Schema;
+            opts.AutoCreateSchemaObjects = AutoCreate.All;
+            opts.UseNativeJsonType = ConnectionSource.SupportsNativeJson;
+            opts.Events.TenancyStyle = TenancyStyle.Conjoined;
+            opts.EventGraph.UseTenantPartitionedEvents = true;
+        });
+    }
+
+    [Fact]
+    public async Task streams_table_is_partitioned_alongside_events()
+    {
+        using var store = CreateStore();
+        await store.Database.ApplyAllConfiguredChangesToDatabaseAsync();
+
+        // pc_streams carries the tenant_ordinal partition column and sits on its own
+        // partition scheme, exactly like pc_events.
+        (await TestSchema.ColumnExistsAsync(Schema, "pc_streams", "tenant_ordinal")).ShouldBeTrue();
+        (await TableIsOnPartitionSchemeAsync("pc_streams", "ps_pc_streams_tenant_ordinal")).ShouldBeTrue();
+        (await TableIsOnPartitionSchemeAsync("pc_events", "ps_pc_events_tenant_ordinal")).ShouldBeTrue();
+    }
+
+    [Fact]
+    public async Task stream_rows_land_in_their_tenant_partition()
+    {
+        using var store = CreateStore();
+        await store.Database.ApplyAllConfiguredChangesToDatabaseAsync();
+
+        await using (var red = store.LightweightSession(new SessionOptions { TenantId = "Red" }))
+        {
+            red.Events.StartStream(Guid.NewGuid(), new QuestStarted("Red"), new MonsterSlain("a", 1));
+            await red.SaveChangesAsync();
+        }
+
+        await using (var blue = store.LightweightSession(new SessionOptions { TenantId = "Blue" }))
+        {
+            blue.Events.StartStream(Guid.NewGuid(), new QuestStarted("Blue"));
+            await blue.SaveChangesAsync();
+        }
+
+        // Stream rows carry their tenant's registry ordinal and land in distinct physical
+        // partitions of pc_streams.
+        var rows = await TestSchema.QueryAsync($"""
+            SELECT $PARTITION.pf_pc_streams_tenant_ordinal(st.tenant_ordinal) AS p, st.tenant_id,
+                   st.tenant_ordinal, tp.ordinal
+            FROM [{Schema}].[pc_streams] st
+            JOIN [{Schema}].[pc_tenant_partitions] tp ON tp.tenant_id = st.tenant_id
+            ORDER BY st.tenant_id
+            """);
+
+        rows.Count.ShouldBe(2);
+        rows.Select(r => r[0]).Distinct().Count().ShouldBe(2); // distinct physical partitions
+        foreach (var row in rows)
+        {
+            ((int)row[2]).ShouldBe((int)row[3]); // row ordinal == the tenant's registry ordinal
+        }
+    }
+
+    [Fact]
+    public async Task appending_to_an_existing_stream_updates_the_partitioned_stream_row()
+    {
+        using var store = CreateStore();
+        await store.Database.ApplyAllConfiguredChangesToDatabaseAsync();
+
+        var stream = Guid.NewGuid();
+
+        await using (var session = store.LightweightSession(new SessionOptions { TenantId = "Red" }))
+        {
+            session.Events.StartStream(stream, new QuestStarted("Quest"));
+            await session.SaveChangesAsync();
+        }
+
+        // Second append hits the UPDATE path against the partitioned pc_streams (with the
+        // tenant_ordinal partition-elimination predicate).
+        await using (var session = store.LightweightSession(new SessionOptions { TenantId = "Red" }))
+        {
+            session.Events.Append(stream, new MembersJoined(1, "Town", ["Hero"]), new MonsterSlain("b", 2));
+            await session.SaveChangesAsync();
+        }
+
+        await using var query = store.QuerySession(new SessionOptions { TenantId = "Red" });
+        var state = await query.Events.FetchStreamStateAsync(stream);
+        state.ShouldNotBeNull();
+        state.Version.ShouldBe(3);
+
+        (await query.Events.FetchStreamAsync(stream)).Count.ShouldBe(3);
+    }
+
+    private static async Task<bool> TableIsOnPartitionSchemeAsync(string table, string scheme)
+    {
+        await using var conn = new SqlConnection(ConnectionSource.ConnectionString);
+        await conn.OpenAsync();
+        await using var cmd = conn.CreateCommand();
+        cmd.CommandText = """
+            SELECT COUNT(*)
+            FROM sys.tables t
+            JOIN sys.schemas s ON t.schema_id = s.schema_id
+            JOIN sys.indexes i ON i.object_id = t.object_id AND i.index_id IN (0, 1)
+            JOIN sys.data_spaces ds ON i.data_space_id = ds.data_space_id
+            WHERE s.name = @schema AND t.name = @table AND ds.name = @scheme
+            """;
+        cmd.Parameters.AddWithValue("@schema", Schema);
+        cmd.Parameters.AddWithValue("@table", table);
+        cmd.Parameters.AddWithValue("@scheme", scheme);
+        var count = (int)(await cmd.ExecuteScalarAsync())!;
+        return count == 1;
+    }
+}

src/Polecat.Tests/Harness/TenantPartitioningCollection.cs

@@ -14,8 +14,9 @@ namespace Polecat.Tests.Harness;
 public class TenantPartitioningCollection;
 
 /// <summary>
-///     Drops the database-global partition function/scheme that the managed tenant partitioning leaves
-///     behind for a tenant-partitioned <c>pc_events</c> — they outlive a schema/table drop and must be
+///     Drops the database-global partition functions/schemes that the managed tenant partitioning
+///     leaves behind for tenant-partitioned tables (<c>pc_events</c>, <c>pc_streams</c> since #335,
+///     and any tenant-partitioned <c>pc_doc_*</c>) — they outlive a schema/table drop and must be
 ///     removed explicitly so each test starts clean.
 /// </summary>
 public static class PartitionTestCleanup
@@ -26,21 +27,27 @@ public static async Task DropEventsPartitionObjectsAsync()
         await conn.OpenAsync();
         await using var cmd = conn.CreateCommand();
         cmd.CommandText = """
-            -- The scheme is shared across every test schema's pc_events, so drop ALL tables sitting on
-            -- it (in any schema) before the scheme/function can be removed.
+            -- The schemes are shared across every test schema's tables, so drop ALL tables sitting
+            -- on any pc_* tenant-ordinal scheme (in any schema) before the schemes/functions can be
+            -- removed.
             DECLARE @sql nvarchar(max) = N'';
             SELECT @sql = @sql + 'DROP TABLE [' + s.name + '].[' + t.name + '];'
             FROM sys.tables t
             JOIN sys.schemas s ON t.schema_id = s.schema_id
             JOIN sys.indexes i ON i.object_id = t.object_id AND i.index_id IN (0, 1)
             JOIN sys.data_spaces ds ON i.data_space_id = ds.data_space_id
-            WHERE ds.name = 'ps_pc_events_tenant_ordinal';
+            WHERE ds.name LIKE 'ps[_]pc[_]%[_]tenant[_]ordinal';
             IF @sql <> N'' EXEC sp_executesql @sql;
 
-            IF EXISTS (SELECT 1 FROM sys.partition_schemes WHERE name = 'ps_pc_events_tenant_ordinal')
-                DROP PARTITION SCHEME ps_pc_events_tenant_ordinal;
-            IF EXISTS (SELECT 1 FROM sys.partition_functions WHERE name = 'pf_pc_events_tenant_ordinal')
-                DROP PARTITION FUNCTION pf_pc_events_tenant_ordinal;
+            SET @sql = N'';
+            SELECT @sql = @sql + 'DROP PARTITION SCHEME [' + name + '];'
+            FROM sys.partition_schemes WHERE name LIKE 'ps[_]pc[_]%[_]tenant[_]ordinal';
+            IF @sql <> N'' EXEC sp_executesql @sql;
+
+            SET @sql = N'';
+            SELECT @sql = @sql + 'DROP PARTITION FUNCTION [' + name + '];'
+            FROM sys.partition_functions WHERE name LIKE 'pf[_]pc[_]%[_]tenant[_]ordinal';
+            IF @sql <> N'' EXEC sp_executesql @sql;
             """;
         await cmd.ExecuteNonQueryAsync();
     }

src/Polecat.Tests/Harness/TestSchema.cs

@@ -0,0 +1,84 @@
+using Microsoft.Data.SqlClient;
+using Polecat.TestUtils;
+
+namespace Polecat.Tests.Harness;
+
+/// <summary>
+///     Shared raw-SQL schema helpers for the tenant-partitioning test suites (#335): drop a test
+///     schema's tables (FKs first), drop its sequences, and small introspection/query utilities.
+/// </summary>
+public static class TestSchema
+{
+    public static async Task DropSchemaTablesAsync(string schema)
+    {
+        await using var conn = new SqlConnection(ConnectionSource.ConnectionString);
+        await conn.OpenAsync();
+        await using var cmd = conn.CreateCommand();
+        cmd.CommandText = """
+            DECLARE @sql nvarchar(max) = N'';
+            SELECT @sql = @sql + 'ALTER TABLE [' + s.name + '].[' + t.name + '] DROP CONSTRAINT [' + fk.name + '];'
+            FROM sys.foreign_keys fk
+            JOIN sys.tables t ON fk.parent_object_id = t.object_id
+            JOIN sys.schemas s ON t.schema_id = s.schema_id
+            WHERE s.name = @schema;
+
+            SELECT @sql = @sql + 'DROP TABLE [' + s.name + '].[' + t.name + '];'
+            FROM sys.tables t
+            JOIN sys.schemas s ON t.schema_id = s.schema_id
+            WHERE s.name = @schema;
+            EXEC sp_executesql @sql;
+            """;
+        cmd.Parameters.AddWithValue("@schema", schema);
+        await cmd.ExecuteNonQueryAsync();
+    }
+
+    public static async Task DropSequencesAsync(string schema)
+    {
+        await using var conn = new SqlConnection(ConnectionSource.ConnectionString);
+        await conn.OpenAsync();
+        await using var cmd = conn.CreateCommand();
+        cmd.CommandText = """
+            DECLARE @sql nvarchar(max) = N'';
+            SELECT @sql = @sql + 'DROP SEQUENCE [' + s.name + '].[' + sq.name + '];'
+            FROM sys.sequences sq
+            JOIN sys.schemas s ON sq.schema_id = s.schema_id
+            WHERE s.name = @schema;
+            EXEC sp_executesql @sql;
+            """;
+        cmd.Parameters.AddWithValue("@schema", schema);
+        await cmd.ExecuteNonQueryAsync();
+    }
+
+    public static async Task<bool> ColumnExistsAsync(string schema, string table, string column)
+    {
+        await using var conn = new SqlConnection(ConnectionSource.ConnectionString);
+        await conn.OpenAsync();
+        await using var cmd = conn.CreateCommand();
+        cmd.CommandText = """
+            SELECT COUNT(*) FROM sys.columns c
+            WHERE c.object_id = OBJECT_ID(@table) AND c.name = @column
+            """;
+        cmd.Parameters.AddWithValue("@table", $"[{schema}].[{table}]");
+        cmd.Parameters.AddWithValue("@column", column);
+        var count = (int)(await cmd.ExecuteScalarAsync())!;
+        return count == 1;
+    }
+
+    public static async Task<List<object[]>> QueryAsync(string sql)
+    {
+        await using var conn = new SqlConnection(ConnectionSource.ConnectionString);
+        await conn.OpenAsync();
+        await using var cmd = conn.CreateCommand();
+        cmd.CommandText = sql;
+        var rows = new List<object[]>();
+        await using var reader = await cmd.ExecuteReaderAsync();
+        while (await reader.ReadAsync())
+        {
+            var values = new object[reader.FieldCount];
+            reader.GetValues(values);
+            rows.Add(values);
+        }
+
+        return rows;
+    }
+}