Fix auto-preparation with SchemaOnly (#4420)

roji · roji · commit 12f167706dee · 2022-04-17T22:48:39.000+02:00
Fixes #4404 (cherry picked from commit e914ed7)
diff --git a/src/Npgsql/NpgsqlCommand.cs b/src/Npgsql/NpgsqlCommand.cs
@@ -1034,9 +1034,8 @@ async Task WriteExecuteSchemaOnly(NpgsqlConnector connector, bool async, bool fl
 
                 if (batchCommand.PreparedStatement?.State == PreparedState.Prepared)
                     continue;   // Prepared, we already have the RowDescription
-                Debug.Assert(batchCommand.PreparedStatement == null);
 
-                await connector.WriteParse(batchCommand.FinalCommandText!, string.Empty, batchCommand.PositionalParameters, async, cancellationToken);
+                await connector.WriteParse(batchCommand.FinalCommandText!, batchCommand.StatementName, batchCommand.PositionalParameters, async, cancellationToken);
                 await connector.WriteDescribe(StatementOrPortal.Statement, batchCommand.StatementName, async, cancellationToken);
                 wroteSomething = true;
             }
diff --git a/src/Npgsql/NpgsqlDataReader.cs b/src/Npgsql/NpgsqlDataReader.cs
@@ -623,7 +623,27 @@ async Task<bool> NextResultSchemaOnly(bool async, bool isConsuming = false, Canc
                 }
                 else
                 {
+                    var pStatement = statement.PreparedStatement;
+                    if (pStatement != null)
+                    {
+                        Debug.Assert(!pStatement.IsPrepared);
+                        if (pStatement.StatementBeingReplaced != null)
+                        {
+                            Expect<CloseCompletedMessage>(await Connector.ReadMessage(async), Connector);
+                            pStatement.StatementBeingReplaced.CompleteUnprepare();
+                            pStatement.StatementBeingReplaced = null;
+                        }
+                    }
+
                     Expect<ParseCompleteMessage>(await Connector.ReadMessage(async), Connector);
+
+                    if (statement.IsPreparing)
+                    {
+                        pStatement!.State = PreparedState.Prepared;
+                        Connector.PreparedStatementManager.NumPrepared++;
+                        statement.IsPreparing = false;
+                    }
+
                     Expect<ParameterDescriptionMessage>(await Connector.ReadMessage(async), Connector);
                     var msg = await Connector.ReadMessage(async);
                     switch (msg.Code)
@@ -659,12 +679,31 @@ async Task<bool> NextResultSchemaOnly(bool async, bool isConsuming = false, Canc
         {
             State = ReaderState.Consumed;
 
-            // Reference the triggering statement from the exception (for batching)
-            if (e is PostgresException postgresException &&
-                Command.IsWrappedByBatch &&
-                StatementIndex >= 0 && StatementIndex < _statements.Count)
+            // Reference the triggering statement from the exception
+            if (e is PostgresException postgresException && StatementIndex >= 0 && StatementIndex < _statements.Count)
             {
                 postgresException.BatchCommand = _statements[StatementIndex];
+
+                // Prevent the command or batch from by recycled (by the connection) when it's disposed. This is important since
+                // the exception is very likely to escape the using statement of the command, and by that time some other user may
+                // already be using the recycled instance.
+                if (!Command.IsWrappedByBatch)
+                {
+                    Command.IsCached = false;
+                }
+            }
+
+            // An error means all subsequent statements were skipped by PostgreSQL.
+            // If any of them were being prepared, we need to update our bookkeeping to put
+            // them back in unprepared state.
+            for (; StatementIndex < _statements.Count; StatementIndex++)
+            {
+                var statement = _statements[StatementIndex];
+                if (statement.IsPreparing)
+                {
+                    statement.IsPreparing = false;
+                    statement.PreparedStatement!.AbortPrepare();
+                }
             }
 
             throw;
diff --git a/test/Npgsql.Tests/AutoPrepareTests.cs b/test/Npgsql.Tests/AutoPrepareTests.cs
@@ -1,4 +1,5 @@
 ﻿using System;
+using System.Data;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -556,6 +557,25 @@ CREATE OR REPLACE FUNCTION pg_temp.emit_exception() RETURNS VOID AS
         Assert.That(conn.ExecuteScalar("SELECT 3"), Is.EqualTo(3));
     }
 
+    [Test, IssueLink("https://github.com/npgsql/npgsql/issues/4404")]
+    public async Task SchemaOnly()
+    {
+        var csb = new NpgsqlConnectionStringBuilder(ConnectionString)
+        {
+            AutoPrepareMinUsages = 2,
+            MaxAutoPrepare = 10,
+        };
+
+        using var _ = CreateTempPool(csb, out var connString);
+        await using var conn = await OpenConnectionAsync(connString);
+        await using var cmd = new NpgsqlCommand("SELECT 1", conn);
+
+        for (var i = 0; i < 5; i++)
+        {
+            await using var reader = await cmd.ExecuteReaderAsync(CommandBehavior.SchemaOnly);
+        }
+    }
+
     void DumpPreparedStatements(NpgsqlConnection conn)
     {
         using var cmd = new NpgsqlCommand("SELECT name,statement FROM pg_prepared_statements", conn);

Original file line number	Diff line number	Diff line change
`@@ -1034,9 +1034,8 @@ async Task WriteExecuteSchemaOnly(NpgsqlConnector connector, bool async, bool fl`
`1034`	`1034`
`1035`	`1035`	`if (batchCommand.PreparedStatement?.State == PreparedState.Prepared)`
`1036`	`1036`	`continue; // Prepared, we already have the RowDescription`
`1037`		`- Debug.Assert(batchCommand.PreparedStatement == null);`
`1038`	`1037`
`1039`		`- await connector.WriteParse(batchCommand.FinalCommandText!, string.Empty, batchCommand.PositionalParameters, async, cancellationToken);`
	`1038`	`+ await connector.WriteParse(batchCommand.FinalCommandText!, batchCommand.StatementName, batchCommand.PositionalParameters, async, cancellationToken);`
`1040`	`1039`	`await connector.WriteDescribe(StatementOrPortal.Statement, batchCommand.StatementName, async, cancellationToken);`
`1041`	`1040`	`wroteSomething = true;`
`1042`	`1041`	`}`