SqlBatchCommand support CommandBehavior

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlBatch.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
@@ -185,13 +185,22 @@ public SqlDataReader ExecuteReader()
             return _batchCommand.ExecuteReaderAsync(cancellationToken);
         }
         /// <include file='../../../../../../doc/snippets/Microsoft.Data.SqlClient/SqlBatch.xml' path='docs/members[@name="SqlBatch"]/ExecuteDbDataReader/*'/>
-        protected override DbDataReader ExecuteDbDataReader(CommandBehavior behavior) => ExecuteReader();
+        protected override DbDataReader ExecuteDbDataReader(CommandBehavior behavior)
+        {
+            ValidateCommandBehavior(nameof(ExecuteDbDataReader), behavior);
+
+            CheckDisposed();
+            SetupBatchCommandExecute();
+            return _batchCommand.ExecuteReader(behavior);
+        }
         /// <include file='../../../../../../doc/snippets/Microsoft.Data.SqlClient/SqlBatch.xml' path='docs/members[@name="SqlBatch"]/ExecuteDbDataReaderAsync/*'/>
         protected override Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBehavior behavior, CancellationToken cancellationToken)
         {
+            ValidateCommandBehavior(nameof(ExecuteDbDataReaderAsync), behavior);
+
             CheckDisposed();
             SetupBatchCommandExecute();
-            return _batchCommand.ExecuteReaderAsync(cancellationToken)
+            return _batchCommand.ExecuteReaderAsync(behavior, cancellationToken)
                 .ContinueWith<DbDataReader>((result) =>
                 {
                     if (result.IsFaulted)
@@ -238,6 +247,15 @@ private void SetupBatchCommandExecute()
             }
             _batchCommand.SetBatchRPCModeReadyToExecute();
         }
+
+        internal static void ValidateCommandBehavior(string method, CommandBehavior behavior)
+        {
+            if (0 != (behavior & ~(CommandBehavior.SequentialAccess | CommandBehavior.CloseConnection)))
+            {
+                ADP.ValidateCommandBehavior(behavior);
+                throw ADP.NotSupportedCommandBehavior(behavior & ~(CommandBehavior.SequentialAccess | CommandBehavior.CloseConnection), method);
+            }
+        }
     }
 }
 

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Batch.cs

@@ -37,7 +37,7 @@ internal void AddBatchCommand(SqlBatchCommand batchCommand)
             {
                 // All batch sql statements must be executed inside sp_executesql, including those
                 // without parameters
-                BuildExecuteSql(CommandBehavior.Default, commandText, batchCommand.Parameters, ref rpc);
+                BuildExecuteSql(batchCommand.CommandBehavior, commandText, batchCommand.Parameters, ref rpc);
             }
 
             _RPCList.Add(rpc);

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Reader.cs

@@ -463,13 +463,13 @@ private void BuildExecuteSql(
             SqlParameter sqlParam;
 
             // @batch_text
-            commandText ??= GetCommandText(behavior);
+            string text = GetCommandText(behavior);
             sqlParam = rpc.systemParams[0];
-            sqlParam.SqlDbType = (commandText.Length << 1) <= TdsEnums.TYPE_SIZE_LIMIT
+            sqlParam.SqlDbType = (text.Length << 1) <= TdsEnums.TYPE_SIZE_LIMIT
                 ? SqlDbType.NVarChar
                 : SqlDbType.NText;
-            sqlParam.Size = commandText.Length;
-            sqlParam.Value = commandText;
+            sqlParam.Size = text.Length;
+            sqlParam.Value = text;
             sqlParam.Direction = ParameterDirection.Input;
 
             // @batch_params
@@ -1383,7 +1383,7 @@ private SqlDataReader RunExecuteReaderTds(
                             $"Command Text '{CommandText}'");
                     }
 
-                    string text = GetCommandText(cmdBehavior) + GetOptionsResetString(cmdBehavior);
+                    string text = GetCommandText(cmdBehavior);
 
                     // If the query requires enclave computations, pass the enclave package in the
                     // SqlBatch TDS stream

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.cs

@@ -2566,7 +2566,7 @@ private string GetCommandText(CommandBehavior behavior)
             // to the sproc.
             Debug.Assert(CommandType is CommandType.Text,
                 "invalid call to GetCommandText for stored proc!");
-            return GetOptionsSetString(behavior) + CommandText;
+            return GetOptionsSetString(behavior) + CommandText + GetOptionsResetString(behavior);
         }
 
         private SqlParameterCollection GetCurrentParameterCollection()

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/Batch/BatchTests.netcore.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
@@ -15,7 +15,6 @@ namespace Microsoft.Data.SqlClient.ManualTesting.Tests
 {
     public static class BatchTests
     {
-
         [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void MissingCommandTextThrows()
         {
@@ -123,14 +122,14 @@ public static void MixedBatchSupported()
 
             using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
             using (var batch = new SqlBatch
-                   {
-                       Connection = connection,
-                       BatchCommands =
+            {
+                Connection = connection,
+                BatchCommands =
                        {
                            new SqlBatchCommand("select @@SPID", CommandType.Text),
                            new SqlBatchCommand("sp_help", CommandType.StoredProcedure, new List<SqlParameter> { new("@objname", "sys.indexes") })
                        }
-                   })
+            })
             {
                 connection.RetryLogicProvider = prov;
                 connection.Open();
@@ -631,6 +630,64 @@ public static async Task ExecuteReaderAsyncMultiple()
             Assert.Equal(10, resultRowCount);
         }
 
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        public static void ExecuteReaderCommandCommandBehaviorSchemaOnlyKeyInfo()
+        {
+            System.Collections.ObjectModel.ReadOnlyCollection<DbColumn> schema;
+
+            using (SqlConnection conn = new SqlConnection(DataTestUtility.TCPConnectionString))
+            using (SqlBatch batch = new SqlBatch(conn))
+            {
+                conn.Open();
+
+                var cmd = new SqlBatchCommand("SELECT * FROM Categories");
+                cmd.CommandBehavior = CommandBehavior.SchemaOnly | CommandBehavior.KeyInfo;
+                batch.BatchCommands.Add(cmd);
+
+                using var reader = batch.ExecuteReader();
+
+                Assert.False(reader.Read());
+
+                schema = reader.GetColumnSchema();
+            }
+
+            Assert.Equal(4, schema.Count);
+            Assert.True(schema[0].IsKey);
+        }
+
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
+        public static void ExecuteReaderCommandBehaviorCloseConnection()
+        {
+            int resultSetCount = 0;
+            int resultRowCount = 0;
+
+            using (SqlConnection conn = new SqlConnection(DataTestUtility.TCPConnectionString))
+            using (SqlBatch batch = new SqlBatch(conn))
+            {
+                conn.Open();
+
+                batch.BatchCommands.Add(new SqlBatchCommand("SELECT 1"));
+                batch.BatchCommands.Add(new SqlBatchCommand("SELECT 2"));
+
+                using (var reader = batch.ExecuteReader(CommandBehavior.CloseConnection))
+                {
+                    do
+                    {
+                        resultSetCount += 1;
+                        while (reader.Read())
+                        {
+                            resultRowCount += 1;
+                        }
+                    } while (reader.NextResult());
+                }
+
+                Assert.Equal(ConnectionState.Closed, conn.State);
+            }
+
+            Assert.Equal(2, resultSetCount);
+            Assert.Equal(2, resultRowCount);
+        }
+
         private static SqlParameter CreateParameter<T>(string name, SqlDbType type, T value, ParameterDirection direction = ParameterDirection.Input)
         {
             var parameter = new SqlParameter(name, type);

src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlBatchTest.cs

@@ -0,0 +1,39 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#if NET
+
+using System;
+using System.Data;
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.UnitTests;
+
+/// <summary>
+/// Provides unit tests for verifying the behavior of the SqlBatch class.
+/// </summary>
+public class SqlBatchTest
+{
+    /// <summary>
+    /// Verifies that SqlBatch.ValidateCommandBehavior throws an ArgumentOutOfRangeException when an invalid CommandBehavior is specified.
+    /// </summary>
+    [Fact]
+    public void InvalidCommandBehaviorValidateCommandBehavior_Throws()
+    {
+        ArgumentOutOfRangeException ex = Assert.Throws<ArgumentOutOfRangeException>(() => SqlBatch.ValidateCommandBehavior("ExecuteNonQuery", (CommandBehavior)64));
+        Assert.Contains("CommandBehavior", ex.Message, StringComparison.OrdinalIgnoreCase);
+    }
+
+    /// <summary>
+    /// Verifies that SqlBatch.ValidateCommandBehavior throws an ArgumentOutOfRangeException when a valid but unsupported CommandBehavior is specified.
+    /// </summary>
+    [Fact]
+    public void NotSupportedCommandBehaviorValidateCommandBehavior_Throws()
+    {
+        ArgumentOutOfRangeException ex = Assert.Throws<ArgumentOutOfRangeException>(() => SqlBatch.ValidateCommandBehavior("ExecuteNonQuery", CommandBehavior.KeyInfo));
+        Assert.Contains("not supported", ex.Message, StringComparison.OrdinalIgnoreCase);
+    }
+}
+
+#endif