Add support for custom delete statements

jbogard · Apr 29, 2024 · a3ac94e · a3ac94e
1 parent 2ad376d
commit a3ac94e
Show file tree

Hide file tree

Showing 9 changed files with 77 additions and 24 deletions.
diff --git a/Respawn.DatabaseTests/SqlServerTests.cs b/Respawn.DatabaseTests/SqlServerTests.cs
@@ -101,6 +101,47 @@ public async Task ShouldDeleteData()
             _database.ExecuteScalar<int>("SELECT COUNT(1) FROM Foo").ShouldBe(0);
         }
 
+        [Fact]
+        public async Task ShouldDeleteDataUsingCustomDeleteStatements()
+        {
+            await _database.ExecuteAsync("create table Foo (Value [int])");
+            await _database.ExecuteAsync("create table Bar (Value [int])");
+
+            await _database.InsertBulkAsync(Enumerable.Range(0, 100).Select(i => new Foo { Value = i }));
+            await _database.InsertBulkAsync(Enumerable.Range(0, 100).Select(i => new Bar { Value = i }));
+
+            _database.ExecuteScalar<int>("SELECT COUNT(1) FROM Foo").ShouldBe(100);
+            _database.ExecuteScalar<int>("SELECT COUNT(1) FROM Bar").ShouldBe(100);
+
+            var checkpoint = await Respawner.CreateAsync(_connection, new RespawnerOptions()
+            {
+                FormatDeleteStatement = table =>
+                {
+                    if (table.Name == "Foo")
+                    {
+                        return $"DELETE FROM {table.GetFullName('"')} WHERE Value > 20;";
+                    }
+                    else
+                    {
+                        return $"DELETE FROM {table.GetFullName('"')} WHERE Value > 30;";
+                    }
+                }
+            });
+
+            try
+            {
+                await checkpoint.ResetAsync(_connection);
+            }
+            catch
+            {
+                _output.WriteLine(checkpoint.DeleteSql);
+                throw;
+            }
+
+            _database.ExecuteScalar<int>("SELECT COUNT(1) FROM Foo").ShouldBe(21);
+            _database.ExecuteScalar<int>("SELECT COUNT(1) FROM Bar").ShouldBe(31);
+        }
+
         [Fact]
         public async Task ShouldHandleRelationships()
         {
@@ -262,8 +303,8 @@ public async Task ShouldIgnoreTables()
             {
                 _output.WriteLine(checkpoint.DeleteSql);
                 throw;
-            } 
-            
+            }
+
             _output.WriteLine(checkpoint.DeleteSql);
             _database.ExecuteScalar<int>("SELECT COUNT(1) FROM Foo").ShouldBe(100);
             _database.ExecuteScalar<int>("SELECT COUNT(1) FROM Bar").ShouldBe(0);
@@ -293,7 +334,7 @@ public async Task ShouldIgnoreTablesWithSchema()
             {
                 TablesToIgnore = new[]
                 {
-                    new Table("A", "Foo"), 
+                    new Table("A", "Foo"),
                     new Table("A", "FooWithBrackets")
                 }
             });

diff --git a/Respawn/IDbAdapter.cs b/Respawn/IDbAdapter.cs
@@ -10,7 +10,7 @@ public interface IDbAdapter
         string BuildTableCommandText(RespawnerOptions options);
         string BuildTemporalTableCommandText(RespawnerOptions options);
         string BuildRelationshipCommandText(RespawnerOptions options);
-        string BuildDeleteCommandText(GraphBuilder builder);
+        string BuildDeleteCommandText(GraphBuilder builder, RespawnerOptions options);
         string BuildReseedSql(IEnumerable<Table> tablesToDelete);
         string BuildTurnOffSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOffSystemVersioning);
         string BuildTurnOnSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOnSystemVersioning);

diff --git a/Respawn/InformixDbAdapter.cs b/Respawn/InformixDbAdapter.cs
@@ -177,7 +177,7 @@ INNER JOIN systables T2
             return commandText;
         }
 
-        public string BuildDeleteCommandText(GraphBuilder graph)
+        public string BuildDeleteCommandText(GraphBuilder graph, RespawnerOptions options)
         {
             var builder = new StringBuilder();
 
@@ -187,7 +187,7 @@ public string BuildDeleteCommandText(GraphBuilder graph)
             }
             foreach (var table in graph.ToDelete)
             {
-                builder.AppendLine($"DELETE FROM {table.GetFullName(QuoteCharacter)};");
+                builder.AppendLine(options.FormatDeleteStatement?.Invoke(table) ?? $"DELETE FROM {table.GetFullName(QuoteCharacter)};");
             }
             foreach (var table in graph.CyclicalTableRelationships)
             {
@@ -204,7 +204,7 @@ public string BuildDeleteCommandText(GraphBuilder graph)
         public string BuildTurnOffSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOffSystemVersioning) => throw new System.NotImplementedException();
 
         public string BuildTurnOnSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOnSystemVersioning) => throw new System.NotImplementedException();
-        
+
         public Task<bool> CheckSupportsTemporalTables(DbConnection connection)
         {
             return Task.FromResult(false);

diff --git a/Respawn/MySqlAdapter.cs b/Respawn/MySqlAdapter.cs
@@ -178,14 +178,14 @@ public string BuildRelationshipCommandText(RespawnerOptions options)
             return commandText;
         }
 
-        public string BuildDeleteCommandText(GraphBuilder graph)
+        public string BuildDeleteCommandText(GraphBuilder graph, RespawnerOptions options)
         {
             var builder = new StringBuilder();
 
             builder.AppendLine("SET FOREIGN_KEY_CHECKS=0;");
             foreach (var table in graph.ToDelete)
             {
-                builder.AppendLine($"DELETE FROM {table.GetFullName(QuoteCharacter)};");
+                builder.AppendLine(options.FormatDeleteStatement?.Invoke(table) ?? $"DELETE FROM {table.GetFullName(QuoteCharacter)};");
             }
             builder.AppendLine("SET FOREIGN_KEY_CHECKS=1;");
 
@@ -208,7 +208,7 @@ public string BuildReseedSql(IEnumerable<Table> tablesToDelete)
         public string BuildTurnOffSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOffSystemVersioning) => throw new System.NotImplementedException();
 
         public string BuildTurnOnSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOnSystemVersioning) => throw new System.NotImplementedException();
-        
+
         public Task<bool> CheckSupportsTemporalTables(DbConnection connection)
         {
             return Task.FromResult(false);

diff --git a/Respawn/OracleDbAdapter.cs b/Respawn/OracleDbAdapter.cs
@@ -170,21 +170,22 @@ from all_CONSTRAINTS     a
             return commandText;
         }
 
-        public string BuildDeleteCommandText(GraphBuilder graph)
+        public string BuildDeleteCommandText(GraphBuilder graph, RespawnerOptions options)
         {
-            var deleteSql = string.Join("\n", BuildCommands(graph));
+            var deleteSql = string.Join("\n", BuildCommands(graph, options));
             return $"BEGIN\n{deleteSql}\nEND;";
         }
 
-        private static IEnumerable<string> BuildCommands(GraphBuilder graph)
+        private static IEnumerable<string> BuildCommands(GraphBuilder graph, RespawnerOptions options)
         {
             foreach (var rel in graph.CyclicalTableRelationships)
             {
                 yield return $"EXECUTE IMMEDIATE 'ALTER TABLE {rel.ParentTable.GetFullName(QuoteCharacter)} DISABLE CONSTRAINT {QuoteCharacter}{rel.Name}{QuoteCharacter}';";
             }
             foreach (var table in graph.ToDelete)
             {
-                yield return $"EXECUTE IMMEDIATE 'delete from {table.GetFullName(QuoteCharacter)}';";
+                var deleteCommand = options.FormatDeleteStatement?.Invoke(table) ?? $"delete from {table.GetFullName(QuoteCharacter)}";
+                yield return $"EXECUTE IMMEDIATE '{deleteCommand}';";
             }
             foreach (var rel in graph.CyclicalTableRelationships)
             {
@@ -198,7 +199,7 @@ private static IEnumerable<string> BuildCommands(GraphBuilder graph)
         public string BuildTurnOffSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOffSystemVersioning) => throw new System.NotImplementedException();
 
         public string BuildTurnOnSystemVersioningCommandText(IEnumerable<TemporalTable> tablesToTurnOnSystemVersioning) => throw new System.NotImplementedException();
-        
+
         public Task<bool> CheckSupportsTemporalTables(DbConnection connection)
         {
             return Task.FromResult(false);

diff --git a/Respawn/PostgresDbAdapter.cs b/Respawn/PostgresDbAdapter.cs
@@ -183,18 +183,28 @@ from INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS rc
             return commandText;
         }
 
-        public string BuildDeleteCommandText(GraphBuilder graph)
+        public string BuildDeleteCommandText(GraphBuilder graph, RespawnerOptions options)
         {
             var builder = new StringBuilder();
 
             foreach (var table in graph.CyclicalTableRelationships.Select(rel => rel.ParentTable))
             {
                 builder.AppendLine($"ALTER TABLE {table.GetFullName(QuoteCharacter)} DISABLE TRIGGER ALL;");
             }
-            if (graph.ToDelete.Any())
+            if (graph.ToDelete.Count > 0)
             {
-                var allTables = graph.ToDelete.Select(table => table.GetFullName(QuoteCharacter));
-                builder.AppendLine($"truncate table {string.Join(",", allTables)} cascade;");
+                if (options.FormatDeleteStatement != null)
+                {
+                    foreach (var table in graph.ToDelete)
+                    {
+                        builder.AppendLine(options.FormatDeleteStatement(table) ?? $"truncate table {table.GetFullName(QuoteCharacter)} cascade;");
+                    }
+                }
+                else
+                {
+                    var allTables = graph.ToDelete.Select(table => table.GetFullName(QuoteCharacter));
+                    builder.AppendLine($"truncate table {string.Join(",", allTables)} cascade;");
+                }
             }
             foreach (var table in graph.CyclicalTableRelationships.Select(rel => rel.ParentTable))
             {

diff --git a/Respawn/Respawner.cs b/Respawn/Respawner.cs
@@ -48,7 +48,7 @@ public static async Task<Respawner> CreateAsync(string nameOrConnectionString, R
         }
 
         /// <summary>
-        /// Creates a <see cref="Respawner"/> based on the supplied connection and options. 
+        /// Creates a <see cref="Respawner"/> based on the supplied connection and options.
         /// </summary>
         /// <param name="connection">Connection object for your target database</param>
         /// <param name="options">Options</param>
@@ -149,7 +149,7 @@ private async Task BuildDeleteTables(DbConnection connection)
 
             var graphBuilder = new GraphBuilder(allTables, allRelationships);
 
-            DeleteSql = Options.DbAdapter.BuildDeleteCommandText(graphBuilder);
+            DeleteSql = Options.DbAdapter.BuildDeleteCommandText(graphBuilder, Options);
             ReseedSql = Options.WithReseed ? Options.DbAdapter.BuildReseedSql(graphBuilder.ToDelete) : null;
         }
 

diff --git a/Respawn/RespawnerOptions.cs b/Respawn/RespawnerOptions.cs
@@ -12,6 +12,7 @@ public class RespawnerOptions
     public bool CheckTemporalTables { get; init; }
     public bool WithReseed { get; init; }
     public int? CommandTimeout { get; init; }
-    public IDbAdapter DbAdapter { get; init; } = Respawn.DbAdapter.SqlServer;
+    public Func<Table, string>? FormatDeleteStatement { get; init; }
 
+    public IDbAdapter DbAdapter { get; init; } = Respawn.DbAdapter.SqlServer;
 }
diff --git a/Respawn/SqlServerDbAdapter.cs b/Respawn/SqlServerDbAdapter.cs
@@ -181,7 +181,7 @@ sys.foreign_keys sfk
             return commandText;
         }
 
-        public string BuildDeleteCommandText(GraphBuilder graph)
+        public string BuildDeleteCommandText(GraphBuilder graph, RespawnerOptions options)
         {
             var builder = new StringBuilder();
 
@@ -191,7 +191,7 @@ public string BuildDeleteCommandText(GraphBuilder graph)
             }
             foreach (var table in graph.ToDelete)
             {
-                builder.AppendLine($"DELETE {table.GetFullName(QuoteCharacter)};");
+                builder.AppendLine(options.FormatDeleteStatement?.Invoke(table) ?? $"DELETE {table.GetFullName(QuoteCharacter)};");
             }
             foreach (var table in graph.CyclicalTableRelationships.Select(rel => rel.ParentTable))
             {