Implement MySqlBatch.Prepare. Fixes #656

mysql-net · Aug 8, 2019 · 0f95533 · 0f95533
1 parent 3befefe
commit 0f95533
Show file tree

Hide file tree

Showing 6 changed files with 164 additions and 74 deletions.
diff --git a/src/MySqlConnector/Core/ConcatenatedCommandPayloadCreator.cs b/src/MySqlConnector/Core/ConcatenatedCommandPayloadCreator.cs
@@ -23,9 +23,6 @@ public bool WriteQueryCommand(ref CommandListPosition commandListPosition, IDict
 			do
 			{
 				var command = commandListPosition.Commands[commandListPosition.CommandIndex];
-				if (command.TryGetPreparedStatements() is object)
-					throw new InvalidOperationException("Can't send prepared statements as part of a concatenated batch.");
-
 				if (Log.IsDebugEnabled())
 					Log.Debug("Session{0} Preparing command payload; CommandText: {1}", command.Connection.Session.Id, command.CommandText);
 

diff --git a/src/MySqlConnector/Core/IMySqlCommand.cs b/src/MySqlConnector/Core/IMySqlCommand.cs
@@ -1,8 +1,6 @@
 using System;
 using System.Data;
-using System.Threading;
 using MySql.Data.MySqlClient;
-using MySqlConnector.Utilities;
 
 namespace MySqlConnector.Core
 {

diff --git a/src/MySqlConnector/Core/ServerSession.cs b/src/MySqlConnector/Core/ServerSession.cs
@@ -133,11 +133,65 @@ public void AbortCancel(ICancellableCommand command)
 
 		public bool IsCancelingQuery => m_state == State.CancelingQuery;
 
-		public void AddPreparedStatement(string commandText, PreparedStatements preparedStatements)
+		public async Task PrepareAsync(IMySqlCommand command, IOBehavior ioBehavior, CancellationToken cancellationToken)
 		{
+			var statementPreparer = new StatementPreparer(command.CommandText, command.RawParameters, command.CreateStatementPreparerOptions());
+			var parsedStatements = statementPreparer.SplitStatements();
+
+			var columnsAndParameters = new ResizableArray<byte>();
+			var columnsAndParametersSize = 0;
+
+			var preparedStatements = new List<PreparedStatement>(parsedStatements.Statements.Count);
+			foreach (var statement in parsedStatements.Statements)
+			{
+				await SendAsync(new PayloadData(statement.StatementBytes), ioBehavior, cancellationToken).ConfigureAwait(false);
+				var payload = await ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
+				var response = StatementPrepareResponsePayload.Create(payload.AsSpan());
+
+				ColumnDefinitionPayload[] parameters = null;
+				if (response.ParameterCount > 0)
+				{
+					parameters = new ColumnDefinitionPayload[response.ParameterCount];
+					for (var i = 0; i < response.ParameterCount; i++)
+					{
+						payload = await ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
+						Utility.Resize(ref columnsAndParameters, columnsAndParametersSize + payload.ArraySegment.Count);
+						Buffer.BlockCopy(payload.ArraySegment.Array, payload.ArraySegment.Offset, columnsAndParameters.Array, columnsAndParametersSize, payload.ArraySegment.Count);
+						parameters[i] = ColumnDefinitionPayload.Create(new ResizableArraySegment<byte>(columnsAndParameters, columnsAndParametersSize, payload.ArraySegment.Count));
+						columnsAndParametersSize += payload.ArraySegment.Count;
+					}
+					if (!SupportsDeprecateEof)
+					{
+						payload = await ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
+						EofPayload.Create(payload.AsSpan());
+					}
+				}
+
+				ColumnDefinitionPayload[] columns = null;
+				if (response.ColumnCount > 0)
+				{
+					columns = new ColumnDefinitionPayload[response.ColumnCount];
+					for (var i = 0; i < response.ColumnCount; i++)
+					{
+						payload = await ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
+						Utility.Resize(ref columnsAndParameters, columnsAndParametersSize + payload.ArraySegment.Count);
+						Buffer.BlockCopy(payload.ArraySegment.Array, payload.ArraySegment.Offset, columnsAndParameters.Array, columnsAndParametersSize, payload.ArraySegment.Count);
+						columns[i] = ColumnDefinitionPayload.Create(new ResizableArraySegment<byte>(columnsAndParameters, columnsAndParametersSize, payload.ArraySegment.Count));
+						columnsAndParametersSize += payload.ArraySegment.Count;
+					}
+					if (!SupportsDeprecateEof)
+					{
+						payload = await ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
+						EofPayload.Create(payload.AsSpan());
+					}
+				}
+
+				preparedStatements.Add(new PreparedStatement(response.StatementId, statement, columns, parameters));
+			}
+
 			if (m_preparedStatements is null)
 				m_preparedStatements = new Dictionary<string, PreparedStatements>();
-			m_preparedStatements.Add(commandText, preparedStatements);
+			m_preparedStatements.Add(command.CommandText, new PreparedStatements(preparedStatements, parsedStatements));
 		}
 
 		public PreparedStatements TryGetPreparedStatement(string commandText)

diff --git a/src/MySqlConnector/MySql.Data.MySqlClient/MySqlBatch.cs b/src/MySqlConnector/MySql.Data.MySqlClient/MySqlBatch.cs
@@ -123,7 +123,7 @@ private Task<DbDataReader> ExecuteReaderAsync(IOBehavior ioBehavior, Cancellatio
 				batchCommand.Batch = this;
 
 			var payloadCreator = Connection.Session.SupportsComMulti ? BatchedCommandPayloadCreator.Instance :
-				// TODO: IsPrepared ? SingleCommandPayloadCreator.Instance :
+				IsPrepared ? SingleCommandPayloadCreator.Instance :
 				ConcatenatedCommandPayloadCreator.Instance;
 			return CommandExecutor.ExecuteReaderAsync(BatchCommands, payloadCreator, CommandBehavior.Default, ioBehavior, cancellationToken);
 		}
@@ -138,9 +138,19 @@ private Task<DbDataReader> ExecuteReaderAsync(IOBehavior ioBehavior, Cancellatio
 
 		public override int Timeout { get; set; }
 
-		public override void Prepare() => throw new NotImplementedException();
+		public override void Prepare()
+		{
+			if (!NeedsPrepare(out var exception))
+			{
+				if (exception is object)
+					throw exception;
+				return;
+			}
+
+			DoPrepareAsync(IOBehavior.Synchronous, default).GetAwaiter().GetResult();
+		}
 
-		public override Task PrepareAsync(CancellationToken cancellationToken = default) => throw new NotImplementedException();
+		public override Task PrepareAsync(CancellationToken cancellationToken = default) => PrepareAsync(AsyncIOBehavior, cancellationToken);
 
 		public override void Cancel() => Connection?.Cancel(this);
 
@@ -235,6 +245,56 @@ private bool IsValid(out Exception exception)
 			return exception is null;
 		}
 
+		private bool NeedsPrepare(out Exception exception)
+		{
+			exception = null;
+			if (Connection is null)
+				exception = new InvalidOperationException("Connection property must be non-null.");
+			else if (Connection.State != ConnectionState.Open)
+				exception = new InvalidOperationException("Connection must be Open; current state is {0}".FormatInvariant(Connection.State));
+			else if (BatchCommands.Count == 0)
+				exception = new InvalidOperationException("BatchCommands must contain a command");
+			else if (Connection?.HasActiveReader ?? false)
+				exception = new InvalidOperationException("Cannot call Prepare when there is an open DataReader for this command; it must be closed first.");
+
+			return exception is null && !Connection.IgnorePrepare;
+		}
+
+		private Task PrepareAsync(IOBehavior ioBehavior, CancellationToken cancellationToken)
+		{
+			if (!NeedsPrepare(out var exception))
+				return exception is null ? Utility.CompletedTask : Utility.TaskFromException(exception);
+
+			return DoPrepareAsync(ioBehavior, cancellationToken);
+		}
+
+		private async Task DoPrepareAsync(IOBehavior ioBehavior, CancellationToken cancellationToken)
+		{
+			foreach (IMySqlCommand batchCommand in BatchCommands)
+			{
+				if (batchCommand.CommandType != CommandType.Text)
+					throw new NotSupportedException("Only CommandType.Text is currently supported by MySqlBatch.Prepare");
+				((MySqlBatchCommand) batchCommand).Batch = this;
+
+				// don't prepare the same SQL twice
+				if (Connection.Session.TryGetPreparedStatement(batchCommand.CommandText) is null)
+					await Connection.Session.PrepareAsync(batchCommand, ioBehavior, cancellationToken).ConfigureAwait(false);
+			}
+		}
+
+		private bool IsPrepared
+		{
+			get
+			{
+				foreach (var command in BatchCommands)
+				{
+					if (Connection.Session.TryGetPreparedStatement(command.CommandText) is null)
+						return false;
+				}
+				return true;
+			}
+		}
+
 		private IOBehavior AsyncIOBehavior => Connection?.AsyncIOBehavior ?? IOBehavior.Asynchronous;
 
 		readonly int m_commandId;

diff --git a/src/MySqlConnector/MySql.Data.MySqlClient/MySqlCommand.cs b/src/MySqlConnector/MySql.Data.MySqlClient/MySqlCommand.cs
@@ -1,12 +1,9 @@
 using System;
-using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
 using System.Threading;
 using System.Threading.Tasks;
 using MySqlConnector.Core;
-using MySqlConnector.Protocol;
-using MySqlConnector.Protocol.Payloads;
 using MySqlConnector.Protocol.Serialization;
 using MySqlConnector.Utilities;
 
@@ -91,7 +88,7 @@ public override void Prepare()
 				return;
 			}
 
-			DoPrepareAsync(IOBehavior.Synchronous, default).GetAwaiter().GetResult();
+			Connection.Session.PrepareAsync(this, IOBehavior.Synchronous, default).GetAwaiter().GetResult();
 		}
 
 #if !NETSTANDARD2_1 && !NETCOREAPP3_0
@@ -115,7 +112,7 @@ private Task PrepareAsync(IOBehavior ioBehavior, CancellationToken cancellationT
 			if (!NeedsPrepare(out var exception))
 				return exception is null ? Utility.CompletedTask : Utility.TaskFromException(exception);
 
-			return DoPrepareAsync(ioBehavior, cancellationToken);
+			return Connection.Session.PrepareAsync(this, ioBehavior, cancellationToken);
 		}
 
 		private bool NeedsPrepare(out Exception exception)
@@ -143,65 +140,6 @@ private bool NeedsPrepare(out Exception exception)
 			return Connection.Session.TryGetPreparedStatement(CommandText) is null;
 		}
 
-		private async Task DoPrepareAsync(IOBehavior ioBehavior, CancellationToken cancellationToken)
-		{
-			var statementPreparer = new StatementPreparer(CommandText, m_parameterCollection, ((IMySqlCommand) this).CreateStatementPreparerOptions());
-			var parsedStatements = statementPreparer.SplitStatements();
-
-			var columnsAndParameters = new ResizableArray<byte>();
-			var columnsAndParametersSize = 0;
-
-			var preparedStatements = new List<PreparedStatement>(parsedStatements.Statements.Count);
-			foreach (var statement in parsedStatements.Statements)
-			{
-				await Connection.Session.SendAsync(new PayloadData(statement.StatementBytes), ioBehavior, cancellationToken).ConfigureAwait(false);
-				var payload = await Connection.Session.ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
-				var response = StatementPrepareResponsePayload.Create(payload.AsSpan());
-
-				ColumnDefinitionPayload[] parameters = null;
-				if (response.ParameterCount > 0)
-				{
-					parameters = new ColumnDefinitionPayload[response.ParameterCount];
-					for (var i = 0; i < response.ParameterCount; i++)
-					{
-						payload = await Connection.Session.ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
-						Utility.Resize(ref columnsAndParameters, columnsAndParametersSize + payload.ArraySegment.Count);
-						Buffer.BlockCopy(payload.ArraySegment.Array, payload.ArraySegment.Offset, columnsAndParameters.Array, columnsAndParametersSize, payload.ArraySegment.Count);
-						parameters[i] = ColumnDefinitionPayload.Create(new ResizableArraySegment<byte>(columnsAndParameters, columnsAndParametersSize, payload.ArraySegment.Count));
-						columnsAndParametersSize += payload.ArraySegment.Count;
-					}
-					if (!Connection.Session.SupportsDeprecateEof)
-					{
-						payload = await Connection.Session.ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
-						EofPayload.Create(payload.AsSpan());
-					}
-				}
-
-				ColumnDefinitionPayload[] columns = null;
-				if (response.ColumnCount > 0)
-				{
-					columns = new ColumnDefinitionPayload[response.ColumnCount];
-					for (var i = 0; i < response.ColumnCount; i++)
-					{
-						payload = await Connection.Session.ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
-						Utility.Resize(ref columnsAndParameters, columnsAndParametersSize + payload.ArraySegment.Count);
-						Buffer.BlockCopy(payload.ArraySegment.Array, payload.ArraySegment.Offset, columnsAndParameters.Array, columnsAndParametersSize, payload.ArraySegment.Count);
-						columns[i] = ColumnDefinitionPayload.Create(new ResizableArraySegment<byte>(columnsAndParameters, columnsAndParametersSize, payload.ArraySegment.Count));
-						columnsAndParametersSize += payload.ArraySegment.Count;
-					}
-					if (!Connection.Session.SupportsDeprecateEof)
-					{
-						payload = await Connection.Session.ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
-						EofPayload.Create(payload.AsSpan());
-					}
-				}
-
-				preparedStatements.Add(new PreparedStatement(response.StatementId, statement, columns, parameters));
-			}
-
-			Connection.Session.AddPreparedStatement(CommandText, new PreparedStatements(preparedStatements, parsedStatements));
-		}
-
 		public override string CommandText
 		{
 			get => m_commandText;

diff --git a/tests/SideBySide/BatchTests.cs b/tests/SideBySide/BatchTests.cs
@@ -250,6 +250,49 @@ public void ExecuteInvalidSqlBatch()
 				}
 			}
 		}
+		[Fact]
+		public void PrepareNeedsConnection()
+		{
+			using (var batch = new MySqlBatch
+			{
+				BatchCommands =
+				{
+					new MySqlBatchCommand("SELECT 1;"),
+				},
+			})
+			{
+				Assert.Throws<InvalidOperationException>(() => batch.Prepare());
+			}
+		}
+
+		[Fact]
+		public void PrepareNeedsOpenConnection()
+		{
+			using (var connection = new MySqlConnection(AppConfig.ConnectionString))
+			using (var batch = new MySqlBatch(connection)
+			{
+				BatchCommands =
+				{
+					new MySqlBatchCommand("SELECT 1;"),
+				},
+			})
+			{
+				Assert.Throws<InvalidOperationException>(() => batch.Prepare());
+			}
+		}
+
+		[Fact]
+		public void PrepareNeedsCommands()
+		{
+			using (var connection = new MySqlConnection(AppConfig.ConnectionString))
+			{
+				connection.Open();
+				using (var batch = new MySqlBatch(connection))
+				{
+					Assert.Throws<InvalidOperationException>(() => batch.Prepare());
+				}
+			}
+		}
 
 		private static string GetIgnoreCommandTransactionConnectionString() =>
 			new MySqlConnectionStringBuilder(AppConfig.ConnectionString)