Fix explicit batch preparation reset (#5459)

Fixes #5458

(cherry picked from commit 9e3a8df)

src/Npgsql/NpgsqlCommand.cs

@@ -663,6 +663,7 @@ Task Prepare(bool async, CancellationToken cancellationToken = default)
                 ProcessRawQuery(connector.SqlQueryParser, connector.UseConformingStrings, batchCommand);
 
                 needToPrepare = batchCommand.ExplicitPrepare(connector) || needToPrepare;
+                batchCommand.ConnectorPreparedOn = connector;
             }
 
             if (logger.IsEnabled(LogLevel.Debug) && needToPrepare)

test/Npgsql.Tests/PrepareTests.cs

@@ -5,6 +5,9 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Npgsql.BackendMessages;
+using Npgsql.Internal.Postgres;
+using Npgsql.Tests.Support;
 using NpgsqlTypes;
 using NUnit.Framework;
 using static Npgsql.Tests.TestUtil;
@@ -13,6 +16,8 @@ namespace Npgsql.Tests;
 
 public class PrepareTests: TestBase
 {
+    static uint Int4Oid => PostgresMinimalDatabaseInfo.DefaultTypeCatalog.GetOid(DataTypeNames.Int4).Value;
+
     [Test]
     public void Basic()
     {
@@ -771,6 +776,127 @@ public async Task Explicit_prepare_unprepare_many_queries()
         await cmd.UnprepareAsync();
     }
 
+    [Test]
+    public async Task Explicitly_prepared_batch_sends_prepared_queries()
+    {
+        await using var postmaster = PgPostmasterMock.Start(ConnectionString);
+        await using var dataSource = CreateDataSource(postmaster.ConnectionString);
+
+        await using var conn = await dataSource.OpenConnectionAsync();
+        var server = await postmaster.WaitForServerConnection();
+
+        await using var batch = new NpgsqlBatch(conn)
+        {
+            BatchCommands = { new("SELECT 1"), new("SELECT 2") }
+        };
+
+        var prepareTask = batch.PrepareAsync();
+
+        await server.ExpectMessages(
+            FrontendMessageCode.Parse, FrontendMessageCode.Describe,
+            FrontendMessageCode.Parse, FrontendMessageCode.Describe,
+            FrontendMessageCode.Sync);
+
+        await server
+            .WriteParseComplete()
+            .WriteParameterDescription(new FieldDescription(Int4Oid))
+            .WriteRowDescription(new FieldDescription(Int4Oid))
+            .WriteParseComplete()
+            .WriteParameterDescription(new FieldDescription(Int4Oid))
+            .WriteRowDescription(new FieldDescription(Int4Oid))
+            .WriteReadyForQuery()
+            .FlushAsync();
+
+        await prepareTask;
+
+        for (var i = 0; i < 2; i++)
+            await ExecutePreparedBatch(batch, server);
+
+        async Task ExecutePreparedBatch(NpgsqlBatch batch, PgServerMock server)
+        {
+            var executeBatchTask = batch.ExecuteNonQueryAsync();
+
+            await server.ExpectMessages(
+                FrontendMessageCode.Bind, FrontendMessageCode.Execute,
+                FrontendMessageCode.Bind, FrontendMessageCode.Execute,
+                FrontendMessageCode.Sync);
+
+            await server
+                .WriteBindComplete()
+                .WriteCommandComplete()
+                .WriteBindComplete()
+                .WriteCommandComplete()
+                .WriteReadyForQuery()
+                .FlushAsync();
+
+            await executeBatchTask;
+        }
+    }
+
+    [Test]
+    public async Task Auto_prepared_batch_sends_prepared_queries()
+    {
+        var csb = new NpgsqlConnectionStringBuilder(ConnectionString)
+        {
+            AutoPrepareMinUsages = 1,
+            MaxAutoPrepare = 10
+        };
+        await using var postmaster = PgPostmasterMock.Start(csb.ConnectionString);
+        await using var dataSource = CreateDataSource(postmaster.ConnectionString);
+
+        await using var conn = await dataSource.OpenConnectionAsync();
+        var server = await postmaster.WaitForServerConnection();
+
+        await using var batch = new NpgsqlBatch(conn)
+        {
+            BatchCommands = { new("SELECT 1"), new("SELECT 2") }
+        };
+
+        var firstBatchExecuteTask = batch.ExecuteNonQueryAsync();
+
+        await server.ExpectMessages(
+            FrontendMessageCode.Parse, FrontendMessageCode.Bind, FrontendMessageCode.Describe, FrontendMessageCode.Execute,
+            FrontendMessageCode.Parse, FrontendMessageCode.Bind, FrontendMessageCode.Describe, FrontendMessageCode.Execute,
+            FrontendMessageCode.Sync);
+
+        await server
+            .WriteParseComplete()
+            .WriteBindComplete()
+            .WriteRowDescription(new FieldDescription(Int4Oid))
+            .WriteCommandComplete()
+            .WriteParseComplete()
+            .WriteBindComplete()
+            .WriteRowDescription(new FieldDescription(Int4Oid))
+            .WriteCommandComplete()
+            .WriteReadyForQuery()
+            .FlushAsync();
+
+        await firstBatchExecuteTask;
+
+        for (var i = 0; i < 2; i++)
+            await ExecutePreparedBatch(batch, server);
+
+        async Task ExecutePreparedBatch(NpgsqlBatch batch, PgServerMock server)
+        {
+            var executeBatchTask = batch.ExecuteNonQueryAsync();
+
+            await server.ExpectMessages(
+                FrontendMessageCode.Bind, FrontendMessageCode.Execute,
+                FrontendMessageCode.Bind, FrontendMessageCode.Execute,
+                FrontendMessageCode.Sync);
+
+            await server
+                .WriteBindComplete()
+                .WriteCommandComplete()
+                .WriteBindComplete()
+                .WriteCommandComplete()
+                .WriteReadyForQuery()
+                .FlushAsync();
+
+            await executeBatchTask;
+        }
+    }
+
     NpgsqlConnection OpenConnectionAndUnprepare()
     {
         var conn = OpenConnection();

test/Npgsql.Tests/Support/PgServerMock.cs

@@ -228,6 +228,20 @@ internal PgServerMock WriteRowDescription(params FieldDescription[] fields)
         return this;
     }
 
+    internal PgServerMock WriteParameterDescription(params FieldDescription[] fields)
+    {
+        CheckDisposed();
+
+        _writeBuffer.WriteByte((byte)BackendMessageCode.ParameterDescription);
+        _writeBuffer.WriteInt32(1 + 4 + 2 + fields.Length * 4);
+        _writeBuffer.WriteUInt16((ushort)fields.Length);
+
+        foreach (var field in fields)
+            _writeBuffer.WriteUInt32(field.TypeOID);
+
+        return this;
+    }
+
     internal PgServerMock WriteNoData()
     {
         CheckDisposed();