fix: repopulate natural-key lookup table on projection rebuild (#259)

src/Polecat.Tests/Events/Bug_259_natural_key_rebuild.cs

@@ -0,0 +1,198 @@
+using JasperFx;
+using JasperFx.Events.Daemon;
+using JasperFx.Events.Projections;
+using Microsoft.Data.SqlClient;
+using Polecat.Projections;
+using Polecat.TestUtils;
+using Shouldly;
+
+namespace Polecat.Tests.Events;
+
+/// <summary>
+/// #259 (parity with marten#4788/#4793): the pc_natural_key_X lookup table is maintained by the
+/// auto-registered NaturalKeyProjection on the inline-append path, which drives off newly-appended
+/// StreamActions. The async-daemon rebuild replays persisted events without appending streams, so
+/// before the fix the table was never repopulated on rebuild and FetchForWriting by natural key
+/// missed. The fix wipes the table on teardown and re-emits the upserts per rebuild page (routed
+/// through the projection batch's per-tenant session). Reuses the OrderAggregate/OrderNumber/
+/// NkOrderCreated natural-key types from natural_key_tests.cs.
+/// </summary>
+public class Bug_259_natural_key_rebuild : IAsyncLifetime
+{
+    private const string Schema = "natural_key_rebuild";
+    private const string Table = "pc_natural_key_orderaggregate";
+
+    public async Task InitializeAsync() => await DropSchemaTablesAsync(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.Projections.Add<SingleStreamProjection<OrderAggregate, Guid>>(ProjectionLifecycle.Inline);
+        });
+    }
+
+    [Fact]
+    public async Task natural_key_table_is_repopulated_on_rebuild()
+    {
+        using var store = CreateStore();
+
+        var streamId = Guid.NewGuid();
+        var orderNumber = new OrderNumber("ORD-1");
+        await using (var session = store.LightweightSession())
+        {
+            session.Events.StartStream(streamId, new NkOrderCreated(orderNumber, "Alice"));
+            await session.SaveChangesAsync();
+        }
+
+        // Inline append populated the lookup table + the natural-key fetch works.
+        (await CountRowsAsync()).ShouldBe(1);
+        await using (var query = store.LightweightSession())
+        {
+            (await query.Events.FetchForWriting<OrderAggregate, OrderNumber>(orderNumber)).Aggregate.ShouldNotBeNull();
+        }
+
+        // Simulate corruption / pre-rebuild state.
+        await ExecuteAsync($"DELETE FROM [{Schema}].[{Table}];");
+        (await CountRowsAsync()).ShouldBe(0);
+
+        // Rebuild the snapshot projection through the async daemon.
+        using var daemon = (IProjectionDaemon)await store.BuildProjectionDaemonAsync();
+        var projectionName = store.Options.Projections.All.Single().Name;
+        await daemon.RebuildProjectionAsync(projectionName, CancellationToken.None);
+
+        // The fix: the lookup table is repopulated, so the natural-key fetch succeeds again.
+        (await CountRowsAsync()).ShouldBe(1);
+        await using (var query = store.LightweightSession())
+        {
+            var stream = await query.Events.FetchForWriting<OrderAggregate, OrderNumber>(orderNumber);
+            stream.Aggregate.ShouldNotBeNull();
+            stream.Aggregate!.OrderNum.ShouldBe(orderNumber);
+        }
+    }
+
+    [Fact]
+    public async Task natural_key_table_is_repopulated_for_multiple_streams_on_rebuild()
+    {
+        using var store = CreateStore();
+
+        var orders = Enumerable.Range(0, 5)
+            .Select(i => (Id: Guid.NewGuid(), Number: new OrderNumber($"ORD-{i}")))
+            .ToList();
+
+        await using (var session = store.LightweightSession())
+        {
+            foreach (var (oid, num) in orders)
+            {
+                session.Events.StartStream(oid, new NkOrderCreated(num, "Customer"));
+            }
+
+            await session.SaveChangesAsync();
+        }
+
+        (await CountRowsAsync()).ShouldBe(5);
+
+        await ExecuteAsync($"DELETE FROM [{Schema}].[{Table}];");
+
+        using var daemon = (IProjectionDaemon)await store.BuildProjectionDaemonAsync();
+        var projectionName = store.Options.Projections.All.Single().Name;
+        await daemon.RebuildProjectionAsync(projectionName, CancellationToken.None);
+
+        (await CountRowsAsync()).ShouldBe(5);
+
+        await using var query = store.LightweightSession();
+        foreach (var (_, num) in orders)
+        {
+            (await query.Events.FetchForWriting<OrderAggregate, OrderNumber>(num)).Aggregate.ShouldNotBeNull();
+        }
+    }
+
+    [Fact]
+    public async Task rebuild_teardown_wipes_table_and_excludes_archived_streams()
+    {
+        // Teardown wipe + the archived-stream delta in one test: an active stream's natural-key row is
+        // re-emitted on rebuild, while an archived stream's row is NOT (the rebuild event loader filters
+        // is_archived = 0, so its events are never replayed). Proving the archived row disappears also
+        // proves teardown wiped the table first — rebuild only ever INSERTs, so a surviving row could
+        // only mean teardown left it.
+        using var store = CreateStore();
+
+        var activeId = Guid.NewGuid();
+        var archivedId = Guid.NewGuid();
+        var activeNumber = new OrderNumber("ORD-active");
+
+        await using (var session = store.LightweightSession())
+        {
+            session.Events.StartStream(activeId, new NkOrderCreated(activeNumber, "Alice"));
+            session.Events.StartStream(archivedId, new NkOrderCreated(new OrderNumber("ORD-archived"), "Bob"));
+            await session.SaveChangesAsync();
+        }
+
+        (await CountRowsAsync()).ShouldBe(2); // both rows present after inline append
+
+        await using (var session = store.LightweightSession())
+        {
+            session.Events.ArchiveStream(archivedId);
+            await session.SaveChangesAsync();
+        }
+
+        using var daemon = (IProjectionDaemon)await store.BuildProjectionDaemonAsync();
+        var projectionName = store.Options.Projections.All.Single().Name;
+        await daemon.RebuildProjectionAsync(projectionName, CancellationToken.None);
+
+        // Only the active stream's row remains: teardown wiped both, rebuild re-emitted only the
+        // non-archived stream's events.
+        (await CountRowsAsync()).ShouldBe(1);
+        await using var query = store.LightweightSession();
+        (await query.Events.FetchForWriting<OrderAggregate, OrderNumber>(activeNumber)).Aggregate.ShouldNotBeNull();
+    }
+
+    private Task<int> CountRowsAsync() => ScalarAsync($"SELECT COUNT(*) FROM [{Schema}].[{Table}];");
+
+    private static async Task<int> ScalarAsync(string sql)
+    {
+        await using var conn = new SqlConnection(ConnectionSource.ConnectionString);
+        await conn.OpenAsync();
+        await using var cmd = conn.CreateCommand();
+        cmd.CommandText = sql;
+        return Convert.ToInt32(await cmd.ExecuteScalarAsync());
+    }
+
+    private static async Task ExecuteAsync(string sql)
+    {
+        await using var conn = new SqlConnection(ConnectionSource.ConnectionString);
+        await conn.OpenAsync();
+        await using var cmd = conn.CreateCommand();
+        cmd.CommandText = sql;
+        await cmd.ExecuteNonQueryAsync();
+    }
+
+    private 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();
+    }
+}

src/Polecat/DocumentStore.EventStore.cs

@@ -128,6 +128,32 @@ async ValueTask<IProjectionBatch<IDocumentSession, IQuerySession>>
         var connStr = database is PolecatDatabase pdb ? pdb.ConnectionString : Options.ConnectionString;
         var batch = new PolecatProjectionBatch(this, Events, connStr);
         await batch.RecordProgress(range);
+
+        // #259: when rebuilding a snapshot whose aggregate has a [NaturalKey], re-emit the natural-key
+        // upserts for this page's events. NaturalKeyProjection.ApplyAsync drives off newly-appended
+        // StreamActions on the inline path and never fires during rebuild — without this hook the
+        // pc_natural_key_X table stays empty after teardown. Routing each tenant's upserts through
+        // batch.SessionForTenant lands them on a session whose work-tracker flushes inside this batch's
+        // transaction, alongside the rebuilt snapshots.
+        if (mode == ShardExecutionMode.Rebuild && range.Events.Any())
+        {
+            var naturalKeySource = Options.Projections.All
+                .FirstOrDefault(s => string.Equals(s.Name, range.ShardName.Name, StringComparison.OrdinalIgnoreCase))
+                as JasperFx.Events.Aggregation.IAggregateProjection;
+
+            if (naturalKeySource?.NaturalKeyDefinition != null)
+            {
+                var naturalKeyProjection = new Events.Projections.NaturalKeyProjection(
+                    naturalKeySource.NaturalKeyDefinition!, Events);
+
+                foreach (var byTenant in range.Events.GroupBy(e => e.TenantId ?? JasperFx.StorageConstants.DefaultTenantId))
+                {
+                    var ops = batch.SessionForTenant(byTenant.Key);
+                    naturalKeyProjection.QueueUpsertsForEvents(ops, byTenant);
+                }
+            }
+        }
+
         return batch;
     }
 
@@ -530,9 +556,22 @@ private async Task TeardownProjectionStateAsync(IEventDatabase database, string
         var publishedTableNames = Array.Empty<string>();
         if (Options.Projections.TryFindProjection(subscriptionName, out var source))
         {
-            publishedTableNames = source.PublishedTypes()
+            var tables = source.PublishedTypes()
                 .Select(t => GetProvider(t).QualifiedTableName)
-                .ToArray();
+                .ToList();
+
+            // #259: a [NaturalKey] aggregate maintains its pc_natural_key_X lookup table via the
+            // auto-registered NaturalKeyProjection on the inline-append path. Teardown of the parent
+            // projection must also wipe that table so the rebuild repopulates it from scratch
+            // (the rebuild re-emits the upserts via StartProjectionBatchAsync). DELETE FROM works the
+            // same as for the doc tables, so it rides the same loop below.
+            if (source is JasperFx.Events.Aggregation.IAggregateProjection { NaturalKeyDefinition: not null } natural)
+            {
+                var aggregateName = natural.NaturalKeyDefinition.AggregateType.Name.ToLowerInvariant();
+                tables.Add($"[{Events.DatabaseSchemaName}].[pc_natural_key_{aggregateName}]");
+            }
+
+            publishedTableNames = tables.ToArray();
         }
 
         await Options.ResiliencePipeline.ExecuteAsync(static async (state, ct) =>

src/Polecat/Events/Projections/NaturalKeyProjection.cs

@@ -1,3 +1,4 @@
+using JasperFx;
 using JasperFx.Events;
 using JasperFx.Events.Projections;
 using Polecat.Internal;
@@ -39,34 +40,68 @@ public Task ApplyAsync(IDocumentSession operations, IEnumerable<StreamAction> st
 
             foreach (var e in stream.Events)
             {
-                // Check for Archived event to mark natural key as archived
-                if (e.EventType == typeof(Archived))
-                {
-                    sessionBase.WorkTracker.Add(
-                        new NaturalKeyArchiveOperation(_qualifiedTableName, streamId, _isGuidStream,
-                            _isConjoined, tenantId));
-                    continue;
-                }
-
-                // Check if this event type has a natural key mapping
-                var mapping = _definition.EventMappings
-                    .FirstOrDefault(m => m.EventType.IsAssignableFrom(e.Data.GetType()));
+                QueueOperationForEvent(sessionBase, streamId, tenantId, e);
+            }
+        }
 
-                if (mapping == null) continue;
+        return Task.CompletedTask;
+    }
 
-                var naturalKeyValue = mapping.Extractor(e.Data);
-                if (naturalKeyValue == null) continue;
+    /// <summary>
+    ///     #259: rebuild-time counterpart to <see cref="ApplyAsync" />. The async-daemon rebuild path
+    ///     replays already-persisted events without appending streams, so <see cref="ApplyAsync" />'s
+    ///     StreamAction-driven dispatch never fires and the <c>pc_natural_key_X</c> table stays empty
+    ///     after teardown. This entry-point feeds raw <see cref="IEvent" />s through the same operation
+    ///     builder, pulling stream id/key + tenant id off the event itself (events read from pc_events
+    ///     always carry these). Called from <c>StartProjectionBatchAsync</c> per rebuild page with
+    ///     <paramref name="operations" /> routed through <c>PolecatProjectionBatch.SessionForTenant</c>,
+    ///     so the operations land on a session whose work-tracker flushes inside the batch transaction
+    ///     alongside the rebuilt snapshots. Archived events are replayed too, so a stream archived
+    ///     before the rebuild keeps its natural-key row marked <c>is_archived = 1</c> afterward.
+    /// </summary>
+    internal void QueueUpsertsForEvents(IDocumentSession operations, IEnumerable<IEvent> events)
+    {
+        if (operations is not DocumentSessionBase sessionBase) return;
 
-                // Unwrap strong-typed id to primitive value
-                var unwrapped = _definition.Unwrap(naturalKeyValue);
-                if (unwrapped == null) continue;
+        foreach (var e in events)
+        {
+            var streamId = _isGuidStream ? (object)e.StreamId : e.StreamKey!;
+            var tenantId = e.TenantId ?? StorageConstants.DefaultTenantId;
+            QueueOperationForEvent(sessionBase, streamId, tenantId, e);
+        }
+    }
 
-                sessionBase.WorkTracker.Add(
-                    new NaturalKeyUpsertOperation(_qualifiedTableName, unwrapped, streamId, _isGuidStream,
-                        _isConjoined, tenantId));
-            }
+    /// <summary>
+    ///     Shared operation builder used by both the inline append path (<see cref="ApplyAsync" />) and
+    ///     the rebuild path (<see cref="QueueUpsertsForEvents" />): an Archived event marks the
+    ///     natural-key row archived; an event carrying a mapped natural key upserts the mapping.
+    /// </summary>
+    private void QueueOperationForEvent(DocumentSessionBase sessionBase, object streamId, string tenantId, IEvent e)
+    {
+        // Check for Archived event to mark natural key as archived
+        if (e.EventType == typeof(Archived))
+        {
+            sessionBase.WorkTracker.Add(
+                new NaturalKeyArchiveOperation(_qualifiedTableName, streamId, _isGuidStream,
+                    _isConjoined, tenantId));
+            return;
         }
 
-        return Task.CompletedTask;
+        // Check if this event type has a natural key mapping
+        var mapping = _definition.EventMappings
+            .FirstOrDefault(m => m.EventType.IsAssignableFrom(e.Data.GetType()));
+
+        if (mapping == null) return;
+
+        var naturalKeyValue = mapping.Extractor(e.Data);
+        if (naturalKeyValue == null) return;
+
+        // Unwrap strong-typed id to primitive value
+        var unwrapped = _definition.Unwrap(naturalKeyValue);
+        if (unwrapped == null) return;
+
+        sessionBase.WorkTracker.Add(
+            new NaturalKeyUpsertOperation(_qualifiedTableName, unwrapped, streamId, _isGuidStream,
+                _isConjoined, tenantId));
     }
 }