Add protocol support for query attributes.

This sends an empty set of attributes with each query.

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

Author: bgrainger

src/MySqlConnector/Core/ConcatenatedCommandPayloadCreator.cs

@@ -14,6 +14,18 @@ public bool WriteQueryCommand(ref CommandListPosition commandListPosition, IDict
 			return false;
 
 		writer.Write((byte) CommandKind.Query);
+
+		if (commandListPosition.Commands[commandListPosition.CommandIndex].Connection!.Session.SupportsQueryAttributes)
+		{
+			// attribute count
+			writer.WriteLengthEncodedInteger(0);
+
+			// attribute set count (always 1)
+			writer.Write((byte) 1);
+
+			// TODO: write attributes for all commands
+		}
+
 		bool isComplete;
 		do
 		{

src/MySqlConnector/Core/IMySqlCommand.cs

@@ -10,6 +10,7 @@ internal interface IMySqlCommand
 	bool AllowUserVariables { get; }
 	CommandBehavior CommandBehavior { get; }
 	MySqlParameterCollection? RawParameters { get; }
+	MySqlAttributeCollection? RawAttributes { get; }
 	PreparedStatements? TryGetPreparedStatements();
 	MySqlConnection? Connection { get; }
 	long LastInsertedId { get; }

src/MySqlConnector/Core/ServerSession.cs

@@ -61,6 +61,7 @@ public ServerSession(ConnectionPool? pool, int poolGeneration, int id)
 	public WeakReference<MySqlConnection>? OwningConnection { get; set; }
 	public bool SupportsComMulti => m_supportsComMulti;
 	public bool SupportsDeprecateEof => m_supportsDeprecateEof;
+	public bool SupportsQueryAttributes { get; private set; }
 	public bool SupportsSessionTrack => m_supportsSessionTrack;
 	public bool ProcAccessDenied { get; set; }
 	public ICollection<KeyValuePair<string, object?>> ActivityTags => m_activityTags;
@@ -315,7 +316,7 @@ public void FinishQuerying()
 			// In order to handle this case, we issue a dummy query that will consume the pending cancellation.
 			// See https://bugs.mysql.com/bug.php?id=45679
 			Log.Debug("Session{0} sending 'DO SLEEP(0)' command to clear pending cancellation", m_logArguments);
-			var payload = QueryPayload.Create("DO SLEEP(0);");
+			var payload = QueryPayload.Create(SupportsQueryAttributes, "DO SLEEP(0);");
 #pragma warning disable CA2012 // Safe because method completes synchronously
 			SendAsync(payload, IOBehavior.Synchronous, CancellationToken.None).GetAwaiter().GetResult();
 			payload = ReceiveReplyAsync(IOBehavior.Synchronous, CancellationToken.None).GetAwaiter().GetResult();
@@ -463,10 +464,13 @@ public async Task DisposeAsync(IOBehavior ioBehavior, CancellationToken cancella
 				m_supportsComMulti = (initialHandshake.ProtocolCapabilities & ProtocolCapabilities.MariaDbComMulti) != 0;
 				m_supportsConnectionAttributes = (initialHandshake.ProtocolCapabilities & ProtocolCapabilities.ConnectionAttributes) != 0;
 				m_supportsDeprecateEof = (initialHandshake.ProtocolCapabilities & ProtocolCapabilities.DeprecateEof) != 0;
+				SupportsQueryAttributes = (initialHandshake.ProtocolCapabilities & ProtocolCapabilities.QueryAttributes) != 0;
 				m_supportsSessionTrack = (initialHandshake.ProtocolCapabilities & ProtocolCapabilities.SessionTrack) != 0;
 				var serverSupportsSsl = (initialHandshake.ProtocolCapabilities & ProtocolCapabilities.Ssl) != 0;
 				m_characterSet = ServerVersion.Version >= ServerVersions.SupportsUtf8Mb4 ? CharacterSet.Utf8Mb4GeneralCaseInsensitive : CharacterSet.Utf8GeneralCaseInsensitive;
-				m_setNamesPayload = ServerVersion.Version >= ServerVersions.SupportsUtf8Mb4 ? s_setNamesUtf8mb4Payload : s_setNamesUtf8Payload;
+				m_setNamesPayload = ServerVersion.Version >= ServerVersions.SupportsUtf8Mb4 ?
+					(SupportsQueryAttributes ? s_setNamesUtf8mb4WithAttributesPayload : s_setNamesUtf8mb4NoAttributesPayload) :
+					(SupportsQueryAttributes ? s_setNamesUtf8WithAttributesPayload : s_setNamesUtf8NoAttributesPayload);
 
 				// disable pipelining for RDS MySQL 5.7 (assuming Aurora); otherwise take it from the connection string or default to true
 				if (!cs.Pipelining.HasValue && ServerVersion.Version.Major == 5 && ServerVersion.Version.Minor == 7 && HostName.EndsWith(".rds.amazonaws.com", StringComparison.OrdinalIgnoreCase))
@@ -495,9 +499,9 @@ public async Task DisposeAsync(IOBehavior ioBehavior, CancellationToken cancella
 					}
 				}
 
-				Log.Debug("Session{0} made connection; ServerVersion={1}; ConnectionId={2}; Compression={3}; Attributes={4}; DeprecateEof={5}; Ssl={6}; SessionTrack={7}; Pipelining={8}",
+				Log.Debug("Session{0} made connection; ServerVersion={1}; ConnectionId={2}; Compression={3}; Attributes={4}; DeprecateEof={5}; Ssl={6}; SessionTrack={7}; Pipelining={8}; QueryAttributes={9}",
 					m_logArguments[0], ServerVersion.OriginalString, ConnectionId,
-					m_useCompression, m_supportsConnectionAttributes, m_supportsDeprecateEof, serverSupportsSsl, m_supportsSessionTrack, m_supportsPipelining);
+					m_useCompression, m_supportsConnectionAttributes, m_supportsDeprecateEof, serverSupportsSsl, m_supportsSessionTrack, m_supportsPipelining, SupportsQueryAttributes);
 
 				if (cs.SslMode != MySqlSslMode.None && (cs.SslMode != MySqlSslMode.Preferred || serverSupportsSsl))
 				{
@@ -1633,7 +1637,7 @@ private async Task GetRealServerDetailsAsync(IOBehavior ioBehavior, Cancellation
 		Log.Debug("Session{0} detected proxy; getting CONNECTION_ID(), VERSION() from server", m_logArguments);
 		try
 		{
-			await SendAsync(QueryPayload.Create("SELECT CONNECTION_ID(), VERSION();"), ioBehavior, cancellationToken).ConfigureAwait(false);
+			await SendAsync(QueryPayload.Create(SupportsQueryAttributes, "SELECT CONNECTION_ID(), VERSION();"), ioBehavior, cancellationToken).ConfigureAwait(false);
 
 			// column count: 2
 			await ReceiveReplyAsync(ioBehavior, cancellationToken).ConfigureAwait(false);
@@ -1908,8 +1912,10 @@ protected override void OnStatementBegin(int index)
 
 	static ReadOnlySpan<byte> BeginCertificateBytes => new byte[] { 45, 45, 45, 45, 45, 66, 69, 71, 73, 78, 32, 67, 69, 82, 84, 73, 70, 73, 67, 65, 84, 69, 45, 45, 45, 45, 45 }; // -----BEGIN CERTIFICATE-----
 	static readonly IMySqlConnectorLogger Log = MySqlConnectorLogManager.CreateLogger(nameof(ServerSession));
-	static readonly PayloadData s_setNamesUtf8Payload = QueryPayload.Create("SET NAMES utf8;");
-	static readonly PayloadData s_setNamesUtf8mb4Payload = QueryPayload.Create("SET NAMES utf8mb4;");
+	static readonly PayloadData s_setNamesUtf8NoAttributesPayload = QueryPayload.Create(false, "SET NAMES utf8;");
+	static readonly PayloadData s_setNamesUtf8mb4NoAttributesPayload = QueryPayload.Create(false, "SET NAMES utf8mb4;");
+	static readonly PayloadData s_setNamesUtf8WithAttributesPayload = QueryPayload.Create(true, "SET NAMES utf8;");
+	static readonly PayloadData s_setNamesUtf8mb4WithAttributesPayload = QueryPayload.Create(true, "SET NAMES utf8mb4;");
 	static int s_lastId;
 
 	readonly object m_lock;

src/MySqlConnector/Core/SingleCommandPayloadCreator.cs

@@ -26,6 +26,22 @@ public bool WriteQueryCommand(ref CommandListPosition commandListPosition, IDict
 				Log.Trace("Session{0} Preparing command payload; CommandText: {1}", command.Connection!.Session.Id, command.CommandText);
 
 			writer.Write((byte) CommandKind.Query);
+			var supportsQueryAttributes = command.Connection!.Session.SupportsQueryAttributes;
+			if (supportsQueryAttributes)
+			{
+				// attribute count
+				writer.WriteLengthEncodedInteger((uint) (command.RawAttributes?.Count ?? 0));
+
+				// attribute set count (always 1)
+				writer.Write((byte) 1);
+
+				// TODO: write attributes
+			}
+			else if (command.RawAttributes?.Count > 0)
+			{
+				Log.Warn("Session{0} has query attributes but server doesn't support them; CommandText: {1}", command.Connection!.Session.Id, command.CommandText);
+			}
+
 			WriteQueryPayload(command, cachedProcedures, writer);
 
 			commandListPosition.CommandIndex++;
@@ -62,9 +78,26 @@ private static void WritePreparedStatement(IMySqlCommand command, PreparedStatem
 		if (Log.IsTraceEnabled())
 			Log.Trace("Session{0} Preparing command payload; CommandId: {1}; CommandText: {2}", command.Connection!.Session.Id, preparedStatement.StatementId, command.CommandText);
 
+		var attributes = command.RawAttributes;
+		var supportsQueryAttributes = command.Connection!.Session.SupportsQueryAttributes;
+
 		writer.Write(preparedStatement.StatementId);
-		writer.Write((byte) 0);
+
+		// NOTE: documentation is not updated yet, but due to bugs in MySQL Server 8.0.23-8.0.25, the PARAMETER_COUNT_AVAILABLE (0x08)
+		// flag has to be set in the 'flags' block in order for query attributes to be sent with a prepared statement; we do not version-sniff the
+		// server but assume that it must support this flag
+		writer.Write((byte) (supportsQueryAttributes ? 8 : 0));
 		writer.Write(1);
+
+		if (supportsQueryAttributes)
+		{
+			writer.WriteLengthEncodedInteger((uint) ((attributes?.Count ?? 0) + (preparedStatement.Parameters?.Length ?? 0)));
+		}
+		else if (attributes?.Count > 0)
+		{
+			Log.Warn("Session{0} has attributes for CommandId {1} but the server does not support them", command.Connection!.Session.Id, preparedStatement.StatementId);
+		}
+
 		if (preparedStatement.Parameters?.Length > 0)
 		{
 			// TODO: How to handle incorrect number of parameters?
@@ -114,6 +147,9 @@ private static void WritePreparedStatement(IMySqlCommand command, PreparedStatem
 				}
 
 				writer.Write(TypeMapper.ConvertToColumnTypeAndFlags(mySqlDbType, command.Connection!.GuidFormat));
+
+				if (supportsQueryAttributes)
+					writer.Write((byte) 0); // empty string
 			}
 
 			var options = command.CreateStatementPreparerOptions();

src/MySqlConnector/MySqlAttribute.cs

@@ -0,0 +1,9 @@
+namespace MySqlConnector
+{
+	public sealed class MySqlAttribute
+	{
+		public string AttributeName { get; set; } = "";
+
+		public object? Value { get; set; }
+	}
+}

src/MySqlConnector/MySqlAttributeCollection.cs

@@ -0,0 +1,21 @@
+using System.Collections;
+
+namespace MySqlConnector
+{
+	public sealed class MySqlAttributeCollection : ICollection<MySqlAttribute>
+	{
+		public int Count => m_attributes.Count;
+		public bool IsReadOnly => false;
+		public void Add(MySqlAttribute item) => throw new NotImplementedException();
+		public void Clear() => throw new NotImplementedException();
+		public bool Contains(MySqlAttribute item) => throw new NotImplementedException();
+		public void CopyTo(MySqlAttribute[] array, int arrayIndex) => throw new NotImplementedException();
+		public IEnumerator<MySqlAttribute> GetEnumerator() => m_attributes.GetEnumerator();
+		public bool Remove(MySqlAttribute item) => throw new NotImplementedException();
+		IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+
+		internal MySqlAttributeCollection() => m_attributes = new();
+
+		private readonly List<MySqlAttribute> m_attributes;
+	}
+}

src/MySqlConnector/MySqlBatchCommand.cs

@@ -53,6 +53,8 @@ public MySqlBatchCommand(string? commandText)
 
 	MySqlParameterCollection? IMySqlCommand.RawParameters => m_parameterCollection;
 
+	MySqlAttributeCollection? IMySqlCommand.RawAttributes => null;
+
 	MySqlConnection? IMySqlCommand.Connection => Batch?.Connection;
 
 	long IMySqlCommand.LastInsertedId => m_lastInsertedId;

src/MySqlConnector/MySqlCommand.cs

@@ -75,6 +75,7 @@ private MySqlCommand(MySqlCommand other)
 		DesignTimeVisible = other.DesignTimeVisible;
 		UpdatedRowSource = other.UpdatedRowSource;
 		m_parameterCollection = other.CloneRawParameters();
+		//// TODO: clone attributes
 	}
 
 	/// <summary>
@@ -84,6 +85,13 @@ private MySqlCommand(MySqlCommand other)
 
 	MySqlParameterCollection? IMySqlCommand.RawParameters => m_parameterCollection;
 
+	/// <summary>
+	/// The collection of <see cref="MySqlAttribute"/> objects for this command.
+	/// </summary>
+	public MySqlAttributeCollection Attributes => m_attributeCollection ??= new();
+
+	MySqlAttributeCollection? IMySqlCommand.RawAttributes => m_attributeCollection;
+
 	public new MySqlParameter CreateParameter() => (MySqlParameter) base.CreateParameter();
 
 	/// <inheritdoc/>
@@ -421,6 +429,7 @@ private bool IsValid([NotNullWhen(false)] out Exception? exception)
 	MySqlConnection? m_connection;
 	string m_commandText;
 	MySqlParameterCollection? m_parameterCollection;
+	MySqlAttributeCollection? m_attributeCollection;
 	int? m_commandTimeout;
 	CommandType m_commandType;
 	CommandBehavior m_commandBehavior;

src/MySqlConnector/Protocol/Payloads/HandshakeResponse41Payload.cs

@@ -26,6 +26,7 @@ private static ByteBufferWriter CreateCapabilitiesPayload(ProtocolCapabilities s
 			(serverCapabilities & ProtocolCapabilities.ConnectionAttributes) |
 			(serverCapabilities & ProtocolCapabilities.SessionTrack) |
 			(serverCapabilities & ProtocolCapabilities.DeprecateEof) |
+			(serverCapabilities & ProtocolCapabilities.QueryAttributes) |
 			additionalCapabilities));
 		writer.Write(0x4000_0000);
 		writer.Write((byte) characterSet);

src/MySqlConnector/Protocol/Payloads/QueryPayload.cs

@@ -4,12 +4,14 @@ namespace MySqlConnector.Protocol.Payloads;
 
 internal static class QueryPayload
 {
-	public static PayloadData Create(string query)
+	public static PayloadData Create(bool supportsQueryAttributes, string query)
 	{
 		var length = Encoding.UTF8.GetByteCount(query);
-		var payload = new byte[length + 1];
+		var payload = new byte[length + 1 + (supportsQueryAttributes ? 2 : 0)];
 		payload[0] = (byte) CommandKind.Query;
-		Encoding.UTF8.GetBytes(query, 0, query.Length, payload, 1);
+		if (supportsQueryAttributes)
+			payload[2] = 1;
+		Encoding.UTF8.GetBytes(query, 0, query.Length, payload, supportsQueryAttributes ? 3 : 1);
 		return new PayloadData(payload);
 	}
 }