Use sql_select_limit for CommandBehavior.SingleRow. Fixes #679

Author: bgrainger

src/MySqlConnector/Core/IMySqlCommand.cs

@@ -11,6 +11,7 @@ internal interface IMySqlCommand
 	{
 		string? CommandText { get; }
 		CommandType CommandType { get; }
+		CommandBehavior CommandBehavior { get; }
 		MySqlParameterCollection? RawParameters { get; }
 		PreparedStatements? TryGetPreparedStatements();
 		MySqlConnection? Connection { get; }

src/MySqlConnector/Core/SingleCommandPayloadCreator.cs

@@ -191,10 +191,18 @@ private static bool WriteStoredProcedure(IMySqlCommand command, IDictionary<stri
 
 		private static bool WriteCommand(IMySqlCommand command, ByteBufferWriter writer)
 		{
+			var isSingleRow = (command.CommandBehavior & CommandBehavior.SingleRow) != 0;
+			if (isSingleRow)
+				writer.Write(SetSqlSelectLimit);
 			var preparer = new StatementPreparer(command.CommandText!, command.RawParameters, command.CreateStatementPreparerOptions());
-			return preparer.ParseAndBindParameters(writer);
+			var isComplete = preparer.ParseAndBindParameters(writer);
+			if (isComplete && isSingleRow)
+				writer.Write(ClearSqlSelectLimit);
+			return isComplete;
 		}
 
+		static ReadOnlySpan<byte> SetSqlSelectLimit => new byte[] { 83, 69, 84, 32, 115, 113, 108, 95, 115, 101, 108, 101, 99, 116, 95, 108, 105, 109, 105, 116, 61, 49, 59, 10 }; // SET sql_select_limit=1;\n
+		static ReadOnlySpan<byte> ClearSqlSelectLimit => new byte[] { 10, 83, 69, 84, 32, 115, 113, 108, 95, 115, 101, 108, 101, 99, 116, 95, 108, 105, 109, 105, 116, 61, 100, 101, 102, 97, 117, 108, 116, 59 }; // \nSET sql_select_limit=default;
 		static readonly IMySqlConnectorLogger Log = MySqlConnectorLogManager.CreateLogger(nameof(SingleCommandPayloadCreator));
 	}
 }

src/MySqlConnector/MySql.Data.MySqlClient/MySqlCommand.cs

@@ -271,7 +271,8 @@ internal Task<DbDataReader> ExecuteReaderAsync(CommandBehavior behavior, IOBehav
 			if (!IsValid(out var exception))
 				return Utility.TaskFromException<DbDataReader>(exception);
 
-			return CommandExecutor.ExecuteReaderAsync(new IMySqlCommand[] { this }, new SingleCommandPayloadCreator(), behavior, ioBehavior, cancellationToken);
+			m_commandBehavior = behavior;
+			return CommandExecutor.ExecuteReaderAsync(new IMySqlCommand[] { this }, SingleCommandPayloadCreator.Instance, behavior, ioBehavior, cancellationToken);
 		}
 
 		public MySqlCommand Clone() => new MySqlCommand(this);
@@ -354,8 +355,8 @@ private bool IsValid(out Exception exception)
 		PreparedStatements IMySqlCommand.TryGetPreparedStatements() => CommandType == CommandType.Text && !string.IsNullOrWhiteSpace(CommandText) && m_connection is object &&
 			m_connection.State == ConnectionState.Open ? m_connection.Session.TryGetPreparedStatement(CommandText) : null;
 
+		CommandBehavior IMySqlCommand.CommandBehavior => m_commandBehavior;
 		MySqlParameterCollection IMySqlCommand.OutParameters { get; set; }
-
 		MySqlParameter IMySqlCommand.ReturnParameter { get; set; }
 
 		readonly int m_commandId;
@@ -365,6 +366,7 @@ PreparedStatements IMySqlCommand.TryGetPreparedStatements() => CommandType == Co
 		MySqlParameterCollection m_parameterCollection;
 		int? m_commandTimeout;
 		CommandType m_commandType;
+		CommandBehavior m_commandBehavior;
 		Action m_cancelAction;
 	}
 }

tests/SideBySide/BatchTests.cs

@@ -250,6 +250,48 @@ public void ExecuteInvalidSqlBatch()
 				}
 			}
 		}
+
+		[Fact]
+		public void SingleRow()
+		{
+			using (var connection = new MySqlConnection(AppConfig.ConnectionString))
+			{
+				connection.Open();
+				using (var command = new MySqlCommand(@"drop table if exists batch_single_row;
+create table batch_single_row(id integer not null primary key);
+insert into batch_single_row(id) values(1),(2),(3);", connection))
+				{
+					command.ExecuteNonQuery();
+				}
+
+				using (var batch = new MySqlBatch(connection)
+				{
+					BatchCommands =
+					{
+						new MySqlBatchCommand("SELECT id FROM batch_single_row ORDER BY id"),
+						new MySqlBatchCommand("SELECT id FROM batch_single_row ORDER BY id") { CommandBehavior = CommandBehavior.SingleRow },
+					},
+				})
+				using (var reader = batch.ExecuteReader())
+				{
+					Assert.True(reader.Read());
+					Assert.Equal(1, reader.GetInt32(0));
+					Assert.True(reader.Read());
+					Assert.Equal(2, reader.GetInt32(0));
+					Assert.True(reader.Read());
+					Assert.Equal(3, reader.GetInt32(0));
+					Assert.False(reader.Read());
+
+					Assert.True(reader.NextResult());
+					Assert.True(reader.Read());
+					Assert.Equal(1, reader.GetInt32(0));
+					Assert.False(reader.Read());
+
+					Assert.False(reader.NextResult());
+				}
+			}
+		}
+
 		[Fact]
 		public void PrepareNeedsConnection()
 		{

tests/SideBySide/QueryTests.cs

@@ -333,7 +333,7 @@ public async Task MultipleReaders()
 				}
 			}
 		}
-		
+
 		[Fact]
 		public async Task UndisposedReader()
 		{
@@ -573,7 +573,7 @@ public void DapperNullableBoolNullFirst()
 			Assert.True(rows[1].IsBold);
 			Assert.Null(rows[2].IsBold);
 		}
-		
+
 		[SkippableFact(Baseline = "https://bugs.mysql.com/bug.php?id=78760")]
 		public void TabsAndNewLines()
 		{
@@ -1182,6 +1182,26 @@ public void CommandBehaviorCloseConnection()
 			}
 		}
 
+		[Fact]
+		public void CommandBehaviorSingleRow()
+		{
+			using (var connection = new MySqlConnection(AppConfig.ConnectionString))
+			{
+				connection.Open();
+				connection.Execute(@"drop table if exists command_behavior_single_row;
+create table command_behavior_single_row(id integer not null primary key);
+insert into command_behavior_single_row(id) values(1),(2),(3),(4),(5),(6),(7),(8),(9),(10);");
+
+				using (var cmd = new MySqlCommand("SELECT id FROM command_behavior_single_row ORDER BY id;", connection))
+				using (var reader = cmd.ExecuteReader(CommandBehavior.SingleRow))
+				{
+					Assert.True(reader.Read());
+					Assert.Equal(1, reader.GetInt32(0));
+					Assert.False(reader.Read());
+				}
+			}
+		}
+
 		class BoolTest
 		{
 			public int Id { get; set; }