Restore original PKs on failure

Author: fabien.menager

src/PhenX.EntityFrameworkCore.BulkInsert/Graph/GraphBulkInsertOrchestrator.cs

@@ -75,6 +75,9 @@ public async Task<GraphInsertResult<T>> InsertGraph<T>(
 
         var connection = await _context.GetConnection(sync, ctk);
 
+        // Track original primary key values for rollback
+        var originalPkValues = new Dictionary<object, Dictionary<string, object?>>();
+
         try
         {
             // 2. Insert in dependency order (parents first)
@@ -86,6 +89,12 @@ public async Task<GraphInsertResult<T>> InsertGraph<T>(
                     continue;
                 }
 
+                // Save original PK values before any modifications
+                if (options.RestoreOriginalPrimaryKeysOnGraphInsertFailure)
+                {
+                    SaveOriginalPrimaryKeyValues(entitiesToInsert, entityType, graphMetadata, originalPkValues);
+                }
+
                 // Propagate FK values from already-inserted parents
                 PropagateParentForeignKeys(entitiesToInsert, entityType, graphMetadata);
 
@@ -117,6 +126,16 @@ await InsertEntitiesOfType(sync, _context, entityType, entitiesToInsert, options
                 TotalInsertedCount = totalInserted,
             };
         }
+        catch
+        {
+            // Restore original PK values on rollback
+            if (options.RestoreOriginalPrimaryKeysOnGraphInsertFailure)
+            {
+                RestoreOriginalPrimaryKeyValues(originalPkValues, graphMetadata);
+            }
+
+            throw;
+        }
         finally
         {
             await connection.Close(sync, ctk);
@@ -436,14 +455,14 @@ private async Task InsertJoinEntities(
 
         // Use reflection to call the generic BulkInsert method with correctly typed entities
         var method = typeof(GraphBulkInsertOrchestrator)
-            .GetMethod(nameof(InsertJoinEntitiesGenericAsync), BindingFlags.NonPublic | BindingFlags.Instance)!
+            .GetMethod(nameof(InsertJoinEntitiesGeneric), BindingFlags.NonPublic | BindingFlags.Instance)!
             .MakeGenericMethod(joinEntityType);
 
         var task = (Task)method.Invoke(this, [sync, context, tableInfo, joinEntities, options, provider, ctk])!;
         await task;
     }
 
-    private async Task InsertJoinEntitiesGenericAsync<TJoin>(
+    private static async Task InsertJoinEntitiesGeneric<TJoin>(
         bool sync,
         DbContext context,
         TableMetadata tableInfo,
@@ -456,4 +475,66 @@ private async Task InsertJoinEntitiesGenericAsync<TJoin>(
         var typedEntities = joinEntities.Cast<TJoin>().ToList();
         await provider.BulkInsert(sync, context, tableInfo, typedEntities, options, null, ctk);
     }
+
+    private static void SaveOriginalPrimaryKeyValues(
+        List<object> entities,
+        Type entityType,
+        GraphMetadata graphMetadata,
+        Dictionary<object, Dictionary<string, object?>> originalPkValues)
+    {
+        var entityMetadata = graphMetadata.GetEntityMetadata(entityType);
+        if (entityMetadata == null)
+        {
+            return;
+        }
+
+        var efEntityType = graphMetadata.GetEntityType(entityType);
+
+        var pkProperties = efEntityType?.FindPrimaryKey()?.Properties;
+        if (pkProperties == null || !pkProperties.Any())
+        {
+            return;
+        }
+
+        // Only save values for database-generated keys
+        var generatedPkProps = pkProperties
+            .Where(p => p.ValueGenerated != Microsoft.EntityFrameworkCore.Metadata.ValueGenerated.Never)
+            .ToList();
+
+        if (generatedPkProps.Count == 0)
+        {
+            return;
+        }
+
+        foreach (var entity in entities)
+        {
+            var pkValues = new Dictionary<string, object?>();
+            foreach (var pkProp in generatedPkProps)
+            {
+                var value = entityMetadata.GetPropertyValue(entity, pkProp.Name);
+                pkValues[pkProp.Name] = value;
+            }
+            originalPkValues[entity] = pkValues;
+        }
+    }
+
+    private static void RestoreOriginalPrimaryKeyValues(
+        Dictionary<object, Dictionary<string, object?>> originalPkValues,
+        GraphMetadata graphMetadata)
+    {
+        foreach (var (entity, pkValues) in originalPkValues)
+        {
+            var entityType = entity.GetType();
+            var entityMetadata = graphMetadata.GetEntityMetadata(entityType);
+            if (entityMetadata == null)
+            {
+                continue;
+            }
+
+            foreach (var (propertyName, originalValue) in pkValues)
+            {
+                entityMetadata.SetPropertyValue(entity, propertyName, originalValue);
+            }
+        }
+    }
 }

src/PhenX.EntityFrameworkCore.BulkInsert/Options/BulkInsertOptions.cs

@@ -98,6 +98,14 @@ public class BulkInsertOptions
     /// </summary>
     public HashSet<string>? ExcludeNavigations { get; set; }
 
+    /// <summary>
+    /// When enabled, if a graph insert operation fails, the original primary key values of the entities will be restored.
+    /// This ensures that entities in memory remain consistent with the database state after a transaction rollback.
+    /// Can add a little overhead, so it is disabled by default. Enable this option if you need to access the primary
+    /// key values of entities after a failed graph insert operation.
+    /// </summary>
+    public bool RestoreOriginalPrimaryKeysOnGraphInsertFailure { get; set; }
+
     internal int GetCopyTimeoutInSeconds()
     {
         return Math.Max(0, (int)CopyTimeout.TotalSeconds);

tests/PhenX.EntityFrameworkCore.BulkInsert.Tests/Tests/Graph/GraphTestsBase.cs

@@ -751,6 +751,7 @@ public async Task InsertGraph_FailureMidRun_TransactionRolledBack()
         var act = async () => await _context.ExecuteBulkInsertAsync(blogs, options =>
         {
             options.IncludeGraph = true;
+            options.RestoreOriginalPrimaryKeysOnGraphInsertFailure = true; // Ensure original entities are restored on failure
         });
 
         await act.Should().ThrowAsync<Exception>("Insert should fail due to NULL constraint violation");
@@ -768,5 +769,42 @@ public async Task InsertGraph_FailureMidRun_TransactionRolledBack()
         // Verify original entities do NOT have IDs populated (rollback means no database-generated values)
         validBlog.Id.Should().Be(0, "Valid blog should not have ID after rollback");
         invalidBlog.Id.Should().Be(0, "Invalid blog should not have ID after rollback");
+
+        // Act 2 - Fix the invalid data and retry insertion with the same entities
+        // This verifies that entities are properly restored and can be reused after rollback
+        invalidBlog.Posts.First().Title = $"{_run}_FixedPost";
+
+        // Should succeed this time
+        await _context.ExecuteBulkInsertAsync(blogs, options =>
+        {
+            options.IncludeGraph = true;
+        });
+
+        // Assert 2 - Verify that ALL entities are now inserted successfully
+        var insertedBlogsAfterFix = _context.Blogs.Where(b => b.TestRun == _run).ToList();
+        insertedBlogsAfterFix.Should().HaveCount(2, "Both blogs should be inserted after fixing the data");
+
+        var insertedPostsAfterFix = _context.Posts.Where(p => p.TestRun == _run).ToList();
+        insertedPostsAfterFix.Should().HaveCount(2, "Both posts should be inserted after fixing the data");
+
+        var insertedSettingsAfterFix = _context.BlogSettings.Where(s => s.TestRun == _run).ToList();
+        insertedSettingsAfterFix.Should().HaveCount(1, "Blog settings should be inserted after fixing the data");
+
+        // Verify that the original entity references now have IDs populated
+        validBlog.Id.Should().BeGreaterThan(0, "Valid blog should have ID after successful insert");
+        invalidBlog.Id.Should().BeGreaterThan(0, "Fixed blog should have ID after successful insert");
+        validBlog.Posts.First().Id.Should().BeGreaterThan(0, "Valid post should have ID after successful insert");
+        invalidBlog.Posts.First().Id.Should().BeGreaterThan(0, "Fixed post should have ID after successful insert");
+        validBlog.Settings!.Id.Should().BeGreaterThan(0, "Settings should have ID after successful insert");
+
+        // Verify FK relationships are correct
+        validBlog.Posts.First().BlogId.Should().Be(validBlog.Id, "Valid post FK should reference its blog");
+        invalidBlog.Posts.First().BlogId.Should().Be(invalidBlog.Id, "Fixed post FK should reference its blog");
+        validBlog.Settings.BlogId.Should().Be(validBlog.Id, "Settings FK should reference its blog");
+
+        // Verify the corrected title is in the database
+        var fixedPostInDb = _context.Posts.FirstOrDefault(p => p.Id == invalidBlog.Posts.First().Id);
+        fixedPostInDb.Should().NotBeNull();
+        fixedPostInDb!.Title.Should().Be($"{_run}_FixedPost", "Fixed post should have the corrected title");
     }
 }