Implement MySqlBulkCopy.RowsCopied event. Fixes #769

bgrainger · bgrainger · commit b7bf205a715c · 2020-02-08T11:31:13.000-08:00
diff --git a/src/MySqlConnector/MySql.Data.MySqlClient/MySqlBulkCopy.cs b/src/MySqlConnector/MySql.Data.MySqlClient/MySqlBulkCopy.cs
@@ -26,6 +26,20 @@ public MySqlBulkCopy(MySqlConnection connection, MySqlTransaction? transaction =
 
 		public string? DestinationTableName { get; set; }
 
+		/// <summary>
+		/// Defines the number of rows to be processed before generating a notification event.
+		/// </summary>
+		public int NotifyAfter { get; set; }
+
+		/// <summary>
+		/// Occurs every time that the number of rows specified by the <see cref="NotifyAfter"/> property have been processed,
+		/// and once after all rows have been copied (if <see cref="NotifyAfter"/> is non-zero).
+		/// </summary>
+		/// <remarks>
+		/// Receipt of a RowsCopied event does not imply that any rows have been sent to the server or committed.
+		/// </remarks>
+		public event MySqlRowsCopiedEventHandler? RowsCopied;
+
 #if !NETSTANDARD1_3
 		public void WriteToServer(DataTable dataTable)
 		{
@@ -167,6 +181,12 @@ internal async Task SendDataReaderAsync(IOBehavior ioBehavior, CancellationToken
 			var buffer = ArrayPool<byte>.Shared.Rent(maxLength + 1);
 			var outputIndex = 0;
 
+			// allocate a reusable MySqlRowsCopiedEventArgs if event notification is necessary
+			var rowsCopied = 0;
+			MySqlRowsCopiedEventArgs? eventArgs = null;
+			if (NotifyAfter > 0 && RowsCopied is object)
+				eventArgs = new MySqlRowsCopiedEventArgs();
+
 			try
 			{
 				var values = new object?[m_valuesEnumerator!.FieldCount];
@@ -210,14 +230,30 @@ await m_valuesEnumerator.MoveNextAsync().ConfigureAwait(false) :
 					else
 					{
 						buffer[outputIndex++] = (byte) '\n';
+
+						rowsCopied++;
+						if (eventArgs is object && rowsCopied % NotifyAfter == 0)
+						{
+							eventArgs.RowsCopied = rowsCopied;
+							RowsCopied!(this, eventArgs);
+							if (eventArgs.Abort)
+								break;
+						}
 					}
 				}
 
-				if (outputIndex != 0)
+				if (outputIndex != 0 && !(eventArgs?.Abort ?? false))
 				{
 					var payload2 = new PayloadData(new ArraySegment<byte>(buffer, 0, outputIndex));
 					await m_connection.Session.SendReplyAsync(payload2, ioBehavior, cancellationToken).ConfigureAwait(false);
 				}
+
+				// send final RowsCopied event (if it wasn't already sent)
+				if (eventArgs is object && rowsCopied % NotifyAfter != 0)
+				{
+					eventArgs.RowsCopied = rowsCopied;
+					RowsCopied!(this, eventArgs);
+				}
 			}
 			finally
 			{
diff --git a/src/MySqlConnector/MySql.Data.MySqlClient/MySqlRowsCopiedEventArgs.cs b/src/MySqlConnector/MySql.Data.MySqlClient/MySqlRowsCopiedEventArgs.cs
@@ -0,0 +1,26 @@
+using System;
+
+namespace MySql.Data.MySqlClient
+{
+	public sealed class MySqlRowsCopiedEventArgs : EventArgs
+	{
+		/// <summary>
+		/// Gets or sets a value that indicates whether the bulk copy operation should be aborted.
+		/// </summary>
+		public bool Abort { get; set; }
+
+		/// <summary>
+		/// Gets a value that returns the number of rows copied during the current bulk copy operation.
+		/// </summary>
+		public long RowsCopied { get; internal set; }
+
+		internal MySqlRowsCopiedEventArgs()
+		{
+		}
+	}
+
+	/// <summary>
+	/// Represents the method that handles the <see cref="MySqlBulkCopy.RowsCopied"/> event of a <see cref="MySqlBulkCopy"/>.
+	/// </summary>
+	public delegate void MySqlRowsCopiedEventHandler(object sender, MySqlRowsCopiedEventArgs e);
+}
diff --git a/tests/SideBySide/BulkLoaderAsync.cs b/tests/SideBySide/BulkLoaderAsync.cs
@@ -2,6 +2,7 @@
 using System.Data;
 using System.Data.Common;
 using System.IO;
+using System.Linq;
 using System.Threading.Tasks;
 using MySql.Data.MySqlClient;
 using Xunit;
@@ -501,6 +502,92 @@ public async Task BulkLoadDataTableWithTooLongData()
 			};
 			await Assert.ThrowsAsync<MySqlException>(async () => await bulkCopy.WriteToServerAsync(dataTable));
 		}
+
+		[Theory]
+		[InlineData(0, 15, 0, 0)]
+		[InlineData(5, 15, 3, 15)]
+		[InlineData(5, 16, 4, 16)]
+		[InlineData(int.MaxValue, 15, 1, 15)]
+		public async Task BulkCopyNotifyAfter(int notifyAfter, int rowCount, int expectedEventCount, int expectedRowsCopied)
+		{
+			using var connection = new MySqlConnection(GetLocalConnectionString());
+			await connection.OpenAsync();
+			using (var cmd = new MySqlCommand(@"drop table if exists bulk_copy_notify_after;
+				create table bulk_copy_notify_after(value int);", connection))
+			{
+				await cmd.ExecuteNonQueryAsync();
+			}
+
+			var bulkCopy = new MySqlBulkCopy(connection)
+			{
+				NotifyAfter = notifyAfter,
+				DestinationTableName = "bulk_copy_notify_after",
+			};
+			int eventCount = 0;
+			long rowsCopied = 0;
+			bulkCopy.RowsCopied += (s, e) =>
+			{
+				eventCount++;
+				rowsCopied = e.RowsCopied;
+			};
+
+			var dataTable = new DataTable()
+			{
+				Columns = { new DataColumn("value", typeof(int)) },
+			};
+			foreach (var x in Enumerable.Range(1, rowCount))
+				dataTable.Rows.Add(new object[] { x });
+
+			await bulkCopy.WriteToServerAsync(dataTable);
+			Assert.Equal(expectedEventCount, eventCount);
+			Assert.Equal(expectedRowsCopied, rowsCopied);
+		}
+
+		[Theory]
+		[InlineData(0, 40, 0, 0, 0, 40)]
+		[InlineData(5, 40, 15, 3, 15, 0)]
+		[InlineData(5, 40, 20, 4, 20, 16)]
+		[InlineData(int.MaxValue, 20, 0, 1, 20, 20)]
+		public async Task BulkCopyAbort(int notifyAfter, int rowCount, int abortAfter, int expectedEventCount, int expectedRowsCopied, long expectedCount)
+		{
+			using var connection = new MySqlConnection(GetLocalConnectionString());
+			await connection.OpenAsync();
+			using (var cmd = new MySqlCommand(@"drop table if exists bulk_copy_abort;
+				create table bulk_copy_abort(value longtext);", connection))
+			{
+				await cmd.ExecuteNonQueryAsync();
+			}
+
+			var bulkCopy = new MySqlBulkCopy(connection)
+			{
+				NotifyAfter = notifyAfter,
+				DestinationTableName = "bulk_copy_abort",
+			};
+			int eventCount = 0;
+			long rowsCopied = 0;
+			bulkCopy.RowsCopied += (s, e) =>
+			{
+				eventCount++;
+				rowsCopied = e.RowsCopied;
+				if (e.RowsCopied >= abortAfter)
+					e.Abort = true;
+			};
+
+			var dataTable = new DataTable()
+			{
+				Columns = { new DataColumn("value", typeof(string)) },
+			};
+			var str = new string('a', 1_000_000);
+			foreach (var x in Enumerable.Range(1, rowCount))
+				dataTable.Rows.Add(new object[] { str });
+
+			await bulkCopy.WriteToServerAsync(dataTable);
+			Assert.Equal(expectedEventCount, eventCount);
+			Assert.Equal(expectedRowsCopied, rowsCopied);
+
+			using (var cmd = new MySqlCommand("select count(value) from bulk_copy_abort;", connection))
+				Assert.Equal(expectedCount, await cmd.ExecuteScalarAsync());
+		}
 #endif
 
 		[Fact]
diff --git a/tests/SideBySide/BulkLoaderSync.cs b/tests/SideBySide/BulkLoaderSync.cs
@@ -2,6 +2,7 @@
 using System.Data;
 using System.Data.Common;
 using System.IO;
+using System.Linq;
 using MySql.Data.MySqlClient;
 using Xunit;
 using Dapper;
@@ -697,6 +698,92 @@ public void BulkLoadDataTableWithTooLongString()
 			};
 			Assert.Throws<MySqlException>(() => bulkCopy.WriteToServer(dataTable));
 		}
+
+		[Theory]
+		[InlineData(0, 15, 0, 0)]
+		[InlineData(5, 15, 3, 15)]
+		[InlineData(5, 16, 4, 16)]
+		[InlineData(int.MaxValue, 15, 1, 15)]
+		public void BulkCopyNotifyAfter(int notifyAfter, int rowCount, int expectedEventCount, int expectedRowsCopied)
+		{
+			using var connection = new MySqlConnection(GetLocalConnectionString());
+			connection.Open();
+			using (var cmd = new MySqlCommand(@"drop table if exists bulk_copy_notify_after;
+				create table bulk_copy_notify_after(value int);", connection))
+			{
+				cmd.ExecuteNonQuery();
+			}
+
+			var bulkCopy = new MySqlBulkCopy(connection)
+			{
+				NotifyAfter = notifyAfter,
+				DestinationTableName = "bulk_copy_notify_after",
+			};
+			int eventCount = 0;
+			long rowsCopied = 0;
+			bulkCopy.RowsCopied += (s, e) =>
+			{
+				eventCount++;
+				rowsCopied = e.RowsCopied;
+			};
+
+			var dataTable = new DataTable()
+			{
+				Columns = { new DataColumn("value", typeof(int)) },
+			};
+			foreach (var x in Enumerable.Range(1, rowCount))
+				dataTable.Rows.Add(new object[] { x });
+
+			bulkCopy.WriteToServer(dataTable);
+			Assert.Equal(expectedEventCount, eventCount);
+			Assert.Equal(expectedRowsCopied, rowsCopied);
+		}
+
+		[Theory]
+		[InlineData(0, 40, 0, 0, 0, 40)]
+		[InlineData(5, 40, 15, 3, 15, 0)]
+		[InlineData(5, 40, 20, 4, 20, 16)]
+		[InlineData(int.MaxValue, 20, 0, 1, 20, 20)]
+		public void BulkCopyAbort(int notifyAfter, int rowCount, int abortAfter, int expectedEventCount, int expectedRowsCopied, long expectedCount)
+		{
+			using var connection = new MySqlConnection(GetLocalConnectionString());
+			connection.Open();
+			using (var cmd = new MySqlCommand(@"drop table if exists bulk_copy_abort;
+				create table bulk_copy_abort(value longtext);", connection))
+			{
+				cmd.ExecuteNonQuery();
+			}
+
+			var bulkCopy = new MySqlBulkCopy(connection)
+			{
+				NotifyAfter = notifyAfter,
+				DestinationTableName = "bulk_copy_abort",
+			};
+			int eventCount = 0;
+			long rowsCopied = 0;
+			bulkCopy.RowsCopied += (s, e) =>
+			{
+				eventCount++;
+				rowsCopied = e.RowsCopied;
+				if (e.RowsCopied >= abortAfter)
+					e.Abort = true;
+			};
+
+			var dataTable = new DataTable()
+			{
+				Columns = { new DataColumn("value", typeof(string)) },
+			};
+			var str = new string('a', 1_000_000);
+			foreach (var x in Enumerable.Range(1, rowCount))
+				dataTable.Rows.Add(new object[] { str });
+
+			bulkCopy.WriteToServer(dataTable);
+			Assert.Equal(expectedEventCount, eventCount);
+			Assert.Equal(expectedRowsCopied, rowsCopied);
+
+			using (var cmd = new MySqlCommand("select count(value) from bulk_copy_abort;", connection))
+				Assert.Equal(expectedCount, cmd.ExecuteScalar());
+		}
 #endif
 
 		[Fact]
