Fix reusing NpgsqlBatch with autoprepare (#4821)

vonzshik · web-flow · commit 09a8387b4cf3 · 2022-12-21T14:19:54.000+03:00
Fixes #4264
diff --git a/src/Npgsql/NpgsqlBatchCommand.cs b/src/Npgsql/NpgsqlBatchCommand.cs
@@ -20,7 +20,13 @@ public sealed class NpgsqlBatchCommand : DbBatchCommand
     public override string CommandText
     {
         get => _commandText;
-        set => _commandText = value ?? string.Empty;
+        set
+        {
+            _commandText = value ?? string.Empty;
+
+            ResetPreparation();
+            // TODO: Technically should do this also if the parameter list (or type) changes
+        }
     }
 
     /// <inheritdoc/>
@@ -153,6 +159,8 @@ internal PreparedStatement? PreparedStatement
 
     PreparedStatement? _preparedStatement;
 
+    internal NpgsqlConnector? ConnectorPreparedOn { get; set; }
+
     internal bool IsPreparing;
 
     /// <summary>
@@ -248,6 +256,12 @@ internal void ApplyCommandComplete(CommandCompleteMessage msg)
         OID = msg.OID;
     }
 
+    internal void ResetPreparation()
+    {
+        PreparedStatement = null;
+        ConnectorPreparedOn = null;
+    }
+
     /// <summary>
     /// Returns the <see cref="CommandText"/>.
     /// </summary>
diff --git a/src/Npgsql/NpgsqlCommand.cs b/src/Npgsql/NpgsqlCommand.cs
@@ -1346,20 +1346,33 @@ internal virtual async ValueTask<NpgsqlDataReader> ExecuteReader(CommandBehavior
                     {
                     case true:
                         Debug.Assert(_connectorPreparedOn != null);
-                        if (_connectorPreparedOn != connector)
-                        {
-                            // The command was prepared, but since then the connector has changed. Detach all prepared statements.
-                            foreach (var s in InternalBatchCommands)
-                                s.PreparedStatement = null;
-                            ResetPreparation();
-                            goto case false;
-                        }
-
                         if (IsWrappedByBatch)
+                        {
                             foreach (var batchCommand in InternalBatchCommands)
+                            {
+                                if (batchCommand.ConnectorPreparedOn != connector)
+                                {
+                                    foreach (var s in InternalBatchCommands)
+                                        s.ResetPreparation();
+                                    ResetPreparation();
+                                    goto case false;
+                                }
+                                
                                 batchCommand.Parameters.ProcessParameters(dataSource.TypeMapper, validateParameterValues, CommandType);
+                            }
+                        }
                         else
+                        {
+                            if (_connectorPreparedOn != connector)
+                            {
+                                // The command was prepared, but since then the connector has changed. Detach all prepared statements.
+                                foreach (var s in InternalBatchCommands)
+                                    s.PreparedStatement = null;
+                                ResetPreparation();
+                                goto case false;
+                            }
                             Parameters.ProcessParameters(dataSource.TypeMapper, validateParameterValues, CommandType);
+                        }
 
                         NpgsqlEventSource.Log.CommandStartPrepared();
                         break;
@@ -1377,7 +1390,10 @@ internal virtual async ValueTask<NpgsqlDataReader> ExecuteReader(CommandBehavior
                                 ProcessRawQuery(connector.SqlQueryParser, connector.UseConformingStrings, batchCommand);
 
                                 if (connector.Settings.MaxAutoPrepare > 0 && batchCommand.TryAutoPrepare(connector))
+                                {
+                                    batchCommand.ConnectorPreparedOn = connector;
                                     numPrepared++;
+                                }
                             }
                         }
                         else
diff --git a/src/Npgsql/NpgsqlDataReader.cs b/src/Npgsql/NpgsqlDataReader.cs
@@ -563,6 +563,8 @@ async Task<bool> NextResult(bool async, bool isConsuming = false, CancellationTo
                 {
                     preparedStatement.State = PreparedState.Invalidated;
                     Command.ResetPreparation();
+                    foreach (var s in Command.InternalBatchCommands)
+                        s.ResetPreparation();
                 }
             }
 
diff --git a/test/Npgsql.Tests/BatchTests.cs b/test/Npgsql.Tests/BatchTests.cs
@@ -731,6 +731,30 @@ public async Task ExecuteScalar_without_parameters()
         Assert.That(await batch.ExecuteScalarAsync(), Is.EqualTo(1));
     }
 
+    [Test, IssueLink("https://github.com/npgsql/npgsql/issues/4264")]
+    public async Task Batch_with_auto_prepare_reuse()
+    {
+        var csb = new NpgsqlConnectionStringBuilder(ConnectionString)
+        {
+            MaxAutoPrepare = 20
+        };
+
+        await using var conn = await OpenConnectionAsync(csb);
+
+        var tempTableName = await CreateTempTable(conn, "id int");
+
+        await using var batch = new NpgsqlBatch(conn);
+        for (var i = 0; i < 2; ++i)
+        {
+            for (var j = 0; j < 10; ++j)
+            {
+                batch.BatchCommands.Add(new NpgsqlBatchCommand($"DELETE FROM {tempTableName} WHERE 1=0"));
+            }
+            await batch.ExecuteNonQueryAsync();
+            batch.BatchCommands.Clear();
+        }
+    }
+
     #endregion Miscellaneous
 
     #region Logging

Original file line number	Diff line number	Diff line change
`@@ -563,6 +563,8 @@ async Task<bool> NextResult(bool async, bool isConsuming = false, CancellationTo`
`563`	`563`	`{`
`564`	`564`	`preparedStatement.State = PreparedState.Invalidated;`
`565`	`565`	`Command.ResetPreparation();`
	`566`	`+ foreach (var s in Command.InternalBatchCommands)`
	`567`	`+ s.ResetPreparation();`
`566`	`568`	`}`
`567`	`569`	`}`
`568`	`570`