Support multiple prepared statements in one command.

Signed-off-by: Bradley Grainger <[email protected]>

Author: bgrainger

src/MySqlConnector/Core/CommandExecutor.cs

@@ -39,6 +39,8 @@ public static async ValueTask<MySqlDataReader> ExecuteReaderAsync(CommandListPos
 				}
 			}
 
+			await payloadCreator.WritePrologueAsync(connection, commandListPosition, ioBehavior, cancellationToken).ConfigureAwait(false);
+
 			var writer = new ByteBufferWriter();
 			//// cachedProcedures will be non-null if there is a stored procedure, which is also the only time it will be read
 			if (!payloadCreator.WriteQueryCommand(ref commandListPosition, cachedProcedures!, writer, false))

src/MySqlConnector/Core/ConcatenatedCommandPayloadCreator.cs

@@ -8,6 +8,9 @@ internal sealed class ConcatenatedCommandPayloadCreator : ICommandPayloadCreator
 {
 	public static ICommandPayloadCreator Instance { get; } = new ConcatenatedCommandPayloadCreator();
 
+	public ValueTask WritePrologueAsync(MySqlConnection connection, CommandListPosition commandListPosition, IOBehavior ioBehavior, CancellationToken cancellationToken) =>
+		throw new NotSupportedException();
+
 	public bool WriteQueryCommand(ref CommandListPosition commandListPosition, IDictionary<string, CachedProcedure?> cachedProcedures, ByteBufferWriter writer, bool appendSemicolon)
 	{
 		if (commandListPosition.CommandIndex == commandListPosition.CommandCount)

src/MySqlConnector/Core/ICommandPayloadCreator.cs

@@ -7,6 +7,16 @@ namespace MySqlConnector.Core;
 /// </summary>
 internal interface ICommandPayloadCreator
 {
+	/// <summary>
+	/// Writes any prologue data that needs to be sent before the current command in the command list.
+	/// </summary>
+	/// <param name="connection">The <see cref="MySqlConnection"/>.</param>
+	/// <param name="commandListPosition">The <see cref="CommandListPosition"/> giving the current command and current prepared statement.</param>
+	/// <param name="ioBehavior">The IO behavior.</param>
+	/// <param name="cancellationToken">A cancellation token to cancel the asynchronous operation.</param>
+	/// <returns>A <see cref="ValueTask"/> representing the potentially-asynchronous operation.</returns>
+	ValueTask WritePrologueAsync(MySqlConnection connection, CommandListPosition commandListPosition, IOBehavior ioBehavior, CancellationToken cancellationToken);
+
 	/// <summary>
 	/// Writes the payload for an "execute query" command to <paramref name="writer"/>.
 	/// </summary>

src/MySqlConnector/Core/SingleCommandPayloadCreator.cs

@@ -1,3 +1,4 @@
+using System.Runtime.InteropServices;
 using MySqlConnector.Logging;
 using MySqlConnector.Protocol;
 using MySqlConnector.Protocol.Serialization;
@@ -12,6 +13,74 @@ internal sealed class SingleCommandPayloadCreator : ICommandPayloadCreator
 	// with this as the first column name, the result set will be treated as 'out' parameters for the previous command.
 	public static string OutParameterSentinelColumnName => "\uE001\b\x0B";
 
+	public async ValueTask WritePrologueAsync(MySqlConnection connection, CommandListPosition commandListPosition, IOBehavior ioBehavior, CancellationToken cancellationToken)
+	{
+		// get the current command and check for prepared statements
+		var command = commandListPosition.CommandAt(commandListPosition.CommandIndex);
+		var preparedStatements = commandListPosition.PreparedStatements ?? command.TryGetPreparedStatements();
+		if (preparedStatements is not null)
+		{
+			// get the current prepared statement; WriteQueryCommand will advance this
+			var preparedStatement = preparedStatements.Statements[commandListPosition.PreparedStatementIndex];
+			if (preparedStatement.Parameters is { } parameters)
+			{
+				// check each parameter
+				for (var i = 0; i < parameters.Length; i++)
+				{
+					// look up this parameter in the command's parameter collection and check if it is a Stream
+					var parameterName = preparedStatement.Statement.NormalizedParameterNames![i];
+					var parameterIndex = parameterName is not null ? (command.RawParameters?.UnsafeIndexOf(parameterName) ?? -1) : preparedStatement.Statement.ParameterIndexes[i];
+					if (parameterIndex != -1 && command.RawParameters![parameterIndex] is { Value: Stream stream and not MemoryStream })
+					{
+						// send almost-full packets, but don't send exactly ProtocolUtility.MaxPacketSize bytes in one payload (as that's ambiguous to whether another packet follows).
+						const int maxDataSize = 16_000_000;
+						int totalBytesRead;
+						while (true)
+						{
+							// write seven-byte COM_STMT_SEND_LONG_DATA header: https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_stmt_send_long_data.html
+							var writer = new ByteBufferWriter(maxDataSize);
+							writer.Write((byte) CommandKind.StatementSendLongData);
+							writer.Write(preparedStatement.StatementId);
+							writer.Write((ushort) i);
+
+							// keep reading from the stream until we've filled the buffer to send
+#if NET7_0_OR_GREATER
+							if (ioBehavior == IOBehavior.Synchronous)
+								totalBytesRead = stream.ReadAtLeast(writer.GetSpan(maxDataSize).Slice(0, maxDataSize), maxDataSize, throwOnEndOfStream: false);
+							else
+								totalBytesRead = await stream.ReadAtLeastAsync(writer.GetMemory(maxDataSize).Slice(0, maxDataSize), maxDataSize, throwOnEndOfStream: false, cancellationToken).ConfigureAwait(false);
+							writer.Advance(totalBytesRead);
+#else
+							totalBytesRead = 0;
+							int bytesRead;
+							do
+							{
+								var sizeToRead = maxDataSize - totalBytesRead;
+								ReadOnlyMemory<byte> bufferMemory = writer.GetMemory(sizeToRead);
+								if (!MemoryMarshal.TryGetArray(bufferMemory, out var arraySegment))
+									throw new InvalidOperationException("Failed to get array segment from buffer memory.");
+								if (ioBehavior == IOBehavior.Synchronous)
+									bytesRead = stream.Read(arraySegment.Array!, arraySegment.Offset, sizeToRead);
+								else
+									bytesRead = await stream.ReadAsync(arraySegment.Array!, arraySegment.Offset, sizeToRead, cancellationToken).ConfigureAwait(false);
+								totalBytesRead += bytesRead;
+								writer.Advance(bytesRead);
+							} while (bytesRead > 0);
+#endif
+
+							if (totalBytesRead == 0)
+								break;
+
+							// send StatementSendLongData; MySQL Server will keep appending the sent data to the parameter value
+							using var payload = writer.ToPayloadData();
+							await connection.Session.SendAsync(payload, ioBehavior, cancellationToken).ConfigureAwait(false);
+						}
+					}
+				}
+			}
+		}
+	}
+
 	public bool WriteQueryCommand(ref CommandListPosition commandListPosition, IDictionary<string, CachedProcedure?> cachedProcedures, ByteBufferWriter writer, bool appendSemicolon)
 	{
 		if (commandListPosition.CommandIndex == commandListPosition.CommandCount)
@@ -58,6 +127,7 @@ public bool WriteQueryCommand(ref CommandListPosition commandListPosition, IDict
 			{
 				commandListPosition.CommandIndex++;
 				commandListPosition.PreparedStatementIndex = 0;
+				commandListPosition.PreparedStatements = null;
 			}
 		}
 		return true;

src/MySqlConnector/MySqlCommand.cs

@@ -359,68 +359,15 @@ internal async ValueTask<MySqlDataReader> ExecuteReaderAsync(CommandBehavior beh
 		return await ExecuteReaderNoResetTimeoutAsync(behavior, ioBehavior, cancellationToken).ConfigureAwait(false);
 	}
 
-	internal async ValueTask<MySqlDataReader> ExecuteReaderNoResetTimeoutAsync(CommandBehavior behavior, IOBehavior ioBehavior, CancellationToken cancellationToken)
+	internal ValueTask<MySqlDataReader> ExecuteReaderNoResetTimeoutAsync(CommandBehavior behavior, IOBehavior ioBehavior, CancellationToken cancellationToken)
 	{
 		if (!IsValid(out var exception))
-			throw exception;
-
-		if (((IMySqlCommand) this).TryGetPreparedStatements() is { Statements.Count: 1 } statements)
-		{
-			for (var i = 0; i < Parameters.Count; i++)
-			{
-				if (Parameters[i].Value is Stream stream and not MemoryStream)
-				{
-					// send almost-full packets, but don't send exactly ProtocolUtility.MaxPacketSize bytes in one payload (as that's ambiguous to whether another packet follows).
-					const int maxDataSize = 16_000_000;
-					int totalBytesRead;
-					while (true)
-					{
-						// write seven-byte COM_STMT_SEND_LONG_DATA header: https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_stmt_send_long_data.html
-						var writer = new ByteBufferWriter(maxDataSize);
-						writer.Write((byte) CommandKind.StatementSendLongData);
-						writer.Write(statements.Statements[0].StatementId);
-						writer.Write((ushort) i);
-
-						// keep reading from the stream until we've filled the buffer to send
-#if NET7_0_OR_GREATER
-						if (ioBehavior == IOBehavior.Synchronous)
-							totalBytesRead = stream.ReadAtLeast(writer.GetSpan(maxDataSize).Slice(0, maxDataSize), maxDataSize, throwOnEndOfStream: false);
-						else
-							totalBytesRead = await stream.ReadAtLeastAsync(writer.GetMemory(maxDataSize).Slice(0, maxDataSize), maxDataSize, throwOnEndOfStream: false, cancellationToken).ConfigureAwait(false);
-						writer.Advance(totalBytesRead);
-#else
-						totalBytesRead = 0;
-						int bytesRead;
-						do
-						{
-							var sizeToRead = maxDataSize - totalBytesRead;
-							ReadOnlyMemory<byte> bufferMemory = writer.GetMemory(sizeToRead);
-							if (!MemoryMarshal.TryGetArray(bufferMemory, out var arraySegment))
-								throw new InvalidOperationException("Failed to get array segment from buffer memory.");
-							if (ioBehavior == IOBehavior.Synchronous)
-								bytesRead = stream.Read(arraySegment.Array!, arraySegment.Offset, sizeToRead);
-							else
-								bytesRead = await stream.ReadAsync(arraySegment.Array!, arraySegment.Offset, sizeToRead, cancellationToken).ConfigureAwait(false);
-							totalBytesRead += bytesRead;
-							writer.Advance(bytesRead);
-						} while (bytesRead > 0);
-#endif
-
-						if (totalBytesRead == 0)
-							break;
-
-						// send StatementSendLongData; MySQL Server will keep appending the sent data to the parameter value
-						using var payload = writer.ToPayloadData();
-						await Connection!.Session.SendAsync(payload, ioBehavior, cancellationToken).ConfigureAwait(false);
-					}
-				}
-			}
-		}
+			return ValueTaskExtensions.FromException<MySqlDataReader>(exception);
 
 		var activity = NoActivity ? null : Connection!.Session.StartActivity(ActivitySourceHelper.ExecuteActivityName,
 			ActivitySourceHelper.DatabaseStatementTagName, CommandText);
 		m_commandBehavior = behavior;
-		return await CommandExecutor.ExecuteReaderAsync(new(this), SingleCommandPayloadCreator.Instance, behavior, activity, ioBehavior, cancellationToken).ConfigureAwait(false);
+		return CommandExecutor.ExecuteReaderAsync(new(this), SingleCommandPayloadCreator.Instance, behavior, activity, ioBehavior, cancellationToken);
 	}
 
 	public MySqlCommand Clone() => new(this);

tests/IntegrationTests/InsertTests.cs

@@ -520,22 +520,21 @@ public async Task SendLongData(int dataLength, int chunkLength, bool isAsync)
 
 		using var chunkStream = new ChunkStream(data, chunkLength);
 
-		using var writeCommand = new MySqlCommand("insert into insert_mysql_long_data(value) values(@value);", connection);
+		using var writeCommand = new MySqlCommand("""
+			insert into insert_mysql_long_data(value) values(@value);
+			select length(value) from insert_mysql_long_data order by rowid;
+			""", connection);
 		writeCommand.Parameters.AddWithValue("@value", chunkStream);
 		writeCommand.Prepare();
-		if (isAsync)
-			await writeCommand.ExecuteNonQueryAsync().ConfigureAwait(true);
-		else
-			writeCommand.ExecuteNonQuery();
-
-		using var readCommand = new MySqlCommand("select length(value) from insert_mysql_long_data order by rowid;", connection);
-		using (var reader = readCommand.ExecuteReader())
+		using (var reader = isAsync ? await writeCommand.ExecuteReaderAsync().ConfigureAwait(true) : writeCommand.ExecuteReader())
 		{
 			Assert.True(reader.Read());
-			Assert.Equal(chunkStream.Length, reader.GetInt32(0));
+			Assert.Equal(1, reader.FieldCount);
+			Assert.Equal(dataLength, reader.GetInt32(0));
+			Assert.False(reader.Read());
 		}
 
-		readCommand.CommandText = "select value from insert_mysql_long_data order by rowid;";
+		using var readCommand = new MySqlCommand("select value from insert_mysql_long_data order by rowid;", connection);
 		using (var reader = readCommand.ExecuteReader())
 		{
 			Assert.True(reader.Read());