Handle column-pruning for cached data

Author: samsharma2700

src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlBulkCopy.cs

@@ -238,6 +238,10 @@ private int RowNumber
         // Metadata caching fields for CacheMetadata option
         private BulkCopySimpleResultSet _cachedMetadata;
         private string _cachedDestinationTableName;
+        // Per-operation clone of the destination table metadata, used when CacheMetadata is
+        // enabled so that column-pruning in AnalyzeTargetAndCreateUpdateBulkCommand does not
+        // mutate the cached BulkCopySimpleResultSet.
+        private _SqlMetaDataSet _operationMetaData;
 
 #if DEBUG
         internal static bool s_setAlwaysTaskOnWrite; //when set and in DEBUG mode, TdsParser::WriteBulkCopyValue will always return a task
@@ -612,7 +616,17 @@ private string AnalyzeTargetAndCreateUpdateBulkCommand(BulkCopySimpleResultSet i
             bool appendComma = false;
 
             // Loop over the metadata for each result column.
+            // When using cached metadata, clone the metadata set so that null-pruning of
+            // unmatched/rejected columns does not mutate the shared cache. Without this,
+            // changing ColumnMappings between WriteToServer calls (e.g. mapping fewer columns
+            // on the first call, then more on the second) would permanently lose metadata
+            // entries from the cache.
             _SqlMetaDataSet metaDataSet = internalResults[MetaDataResultId].MetaData;
+            if (IsCopyOption(SqlBulkCopyOptions.CacheMetadata) && _cachedMetadata != null)
+            {
+                metaDataSet = metaDataSet.Clone();
+            }
+            _operationMetaData = metaDataSet;
             _sortedColumnMappings = new List<_ColumnMapping>(metaDataSet.Length);
             for (int i = 0; i < metaDataSet.Length; i++)
             {
@@ -909,7 +923,7 @@ private void WriteMetaData(BulkCopySimpleResultSet internalResults)
         {
             _stateObj.SetTimeoutSeconds(BulkCopyTimeout);
 
-            _SqlMetaDataSet metadataCollection = internalResults[MetaDataResultId].MetaData;
+            _SqlMetaDataSet metadataCollection = _operationMetaData ?? internalResults[MetaDataResultId].MetaData;
             _stateObj._outputMessageType = TdsEnums.MT_BULK;
             _parser.WriteBulkCopyMetaData(metadataCollection, _sortedColumnMappings.Count, _stateObj);
         }
@@ -944,6 +958,7 @@ private void Dispose(bool disposing)
                 _parser = null;
                 _cachedMetadata = null;
                 _cachedDestinationTableName = null;
+                _operationMetaData = null;
                 try
                 {
                     // Just in case there is a lingering transaction (which there shouldn't be)
@@ -2711,7 +2726,7 @@ private Task CopyBatchesAsyncContinued(BulkCopySimpleResultSet internalResults,
 
                 // Load encryption keys now (if needed)
                 _parser.LoadColumnEncryptionKeys(
-                    internalResults[MetaDataResultId].MetaData,
+                    _operationMetaData ?? internalResults[MetaDataResultId].MetaData,
                     _connection);
 
                 Task task = CopyRowsAsync(0, _savedBatchSize, cts); // This is copying 1 batch of rows and setting _hasMoreRowToCopy = true/false.
@@ -3238,6 +3253,7 @@ private void ResetWriteToServerGlobalVariables()
             _dataTableSource = null;
             _dbDataReaderRowSource = null;
             _isAsyncBulkCopy = false;
+            _operationMetaData = null;
             _rowEnumerator = null;
             _rowSource = null;
             _rowSourceType = ValueSourceType.Unspecified;

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/SqlBulkCopyTest/CacheMetadata.cs

@@ -339,6 +339,66 @@ public static void Test(string dstConstr, string dstTable)
         }
     }
 
+    public class CacheMetadataColumnSubsetChange
+    {
+        private static readonly string initialQueryTemplate = "create table {0} (col1 int, col2 nvarchar(50), col3 nvarchar(50))";
+
+        // Test that mapping a subset of columns on the first call, then all columns on the
+        // second call, works correctly with CacheMetadata. This verifies that null-pruning of
+        // unmatched columns in AnalyzeTargetAndCreateUpdateBulkCommand does not mutate the
+        // cached metadata, which would cause a NullReferenceException on the second call.
+        public static void Test(string dstConstr, string dstTable)
+        {
+            string initialQuery = string.Format(initialQueryTemplate, dstTable);
+
+            using DataTable sourceData = new DataTable();
+            sourceData.Columns.Add("id", typeof(int));
+            sourceData.Columns.Add("firstName", typeof(string));
+            sourceData.Columns.Add("lastName", typeof(string));
+            sourceData.Rows.Add(1, "Alice", "Smith");
+            sourceData.Rows.Add(2, "Bob", "Jones");
+
+            using SqlConnection dstConn = new(dstConstr);
+            using SqlCommand dstCmd = dstConn.CreateCommand();
+            dstConn.Open();
+
+            try
+            {
+                Helpers.TryExecute(dstCmd, initialQuery);
+
+                using SqlBulkCopy bulkcopy = new(dstConn, SqlBulkCopyOptions.CacheMetadata, null);
+                bulkcopy.DestinationTableName = dstTable;
+
+                // First write: map only col1 and col2 (col3 is unmatched and will be pruned).
+                bulkcopy.ColumnMappings.Add("id", "col1");
+                bulkcopy.ColumnMappings.Add("firstName", "col2");
+                bulkcopy.WriteToServer(sourceData);
+                Helpers.VerifyResults(dstConn, dstTable, 3, 2);
+
+                // Second write: map all three columns including col3.
+                // Without the clone fix, this would fail because col3 metadata was
+                // permanently nulled in the cache during the first call.
+                bulkcopy.ColumnMappings.Clear();
+                bulkcopy.ColumnMappings.Add("id", "col1");
+                bulkcopy.ColumnMappings.Add("firstName", "col2");
+                bulkcopy.ColumnMappings.Add("lastName", "col3");
+                bulkcopy.WriteToServer(sourceData);
+                Helpers.VerifyResults(dstConn, dstTable, 3, 4);
+
+                // Verify col3 has the expected data from the second write.
+                using (SqlCommand verifyCmd = new("select col3 from " + dstTable + " where col1 = 1 and col3 is not null", dstConn))
+                {
+                    object result = verifyCmd.ExecuteScalar();
+                    Assert.Equal("Smith", result);
+                }
+            }
+            finally
+            {
+                Helpers.TryExecute(dstCmd, "drop table " + dstTable);
+            }
+        }
+    }
+
     public class CacheMetadataCombinedWithKeepNulls
     {
         private static readonly string initialQueryTemplate = "create table {0} (col1 int, col2 nvarchar(50) default 'DefaultVal', col3 nvarchar(50))";

src/Microsoft.Data.SqlClient/tests/ManualTests/SQL/SqlBulkCopyTest/SqlBulkCopyTest.cs

@@ -345,6 +345,12 @@ public void CacheMetadataColumnMappingsChangeTest()
             CacheMetadataColumnMappingsChange.Test(_connStr, AddGuid("SqlBulkCopyTest_CacheMetadataColMap"));
         }
 
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureServer))]
+        public void CacheMetadataColumnSubsetChangeTest()
+        {
+            CacheMetadataColumnSubsetChange.Test(_connStr, AddGuid("SqlBulkCopyTest_CacheMetadataSubset"));
+        }
+
         [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), nameof(DataTestUtility.IsNotAzureServer))]
         public void CacheMetadataCombinedWithKeepNullsTest()
         {

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

@@ -0,0 +1,126 @@
+// 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.
+
+using Xunit;
+
+namespace Microsoft.Data.SqlClient.UnitTests
+{
+    /// <summary>
+    /// Tests that verify _SqlMetaDataSet.Clone() produces independent copies,
+    /// ensuring that null-pruning of unmatched columns in AnalyzeTargetAndCreateUpdateBulkCommand
+    /// does not corrupt the cached metadata when CacheMetadata is enabled.
+    /// </summary>
+    public class SqlBulkCopyCacheMetadataTest
+    {
+        [Fact]
+        public void SqlMetaDataSet_Clone_ProducesIndependentCopy()
+        {
+            // Arrange: create a metadata set with 3 columns simulating a destination table
+            _SqlMetaDataSet original = new _SqlMetaDataSet(3);
+            original[0].column = "col1";
+            original[1].column = "col2";
+            original[2].column = "col3";
+
+            // Act: clone and then null out an entry in the clone (simulating column pruning)
+            _SqlMetaDataSet clone = original.Clone();
+            clone[2] = null;
+
+            // Assert: the original is not affected by the mutation of the clone
+            Assert.NotNull(original[0]);
+            Assert.NotNull(original[1]);
+            Assert.NotNull(original[2]);
+            Assert.Equal("col1", original[0].column);
+            Assert.Equal("col2", original[1].column);
+            Assert.Equal("col3", original[2].column);
+        }
+
+        [Fact]
+        public void SqlMetaDataSet_Clone_NullingMultipleEntries_OriginalRetainsAll()
+        {
+            // Arrange: simulate a table with 4 columns
+            _SqlMetaDataSet original = new _SqlMetaDataSet(4);
+            original[0].column = "id";
+            original[1].column = "name";
+            original[2].column = "email";
+            original[3].column = "phone";
+
+            // Act: clone and null out entries 1 and 3 (simulating mapping only id and email)
+            _SqlMetaDataSet clone = original.Clone();
+            clone[1] = null;
+            clone[3] = null;
+
+            // Assert: clone has nulls where expected
+            Assert.NotNull(clone[0]);
+            Assert.Null(clone[1]);
+            Assert.NotNull(clone[2]);
+            Assert.Null(clone[3]);
+
+            // Assert: original retains all entries
+            for (int i = 0; i < 4; i++)
+            {
+                Assert.NotNull(original[i]);
+            }
+            Assert.Equal("name", original[1].column);
+            Assert.Equal("phone", original[3].column);
+        }
+
+        [Fact]
+        public void SqlMetaDataSet_Clone_RepeatedCloneAndPrune_OriginalSurvives()
+        {
+            // Arrange: simulate the scenario where multiple WriteToServer calls each
+            // clone and prune different subsets of columns
+            _SqlMetaDataSet original = new _SqlMetaDataSet(3);
+            original[0].column = "col1";
+            original[1].column = "col2";
+            original[2].column = "col3";
+
+            // First operation: map only col1 and col2 (prune col3)
+            _SqlMetaDataSet clone1 = original.Clone();
+            clone1[2] = null;
+
+            // Second operation: map only col1 and col3 (prune col2)
+            _SqlMetaDataSet clone2 = original.Clone();
+            clone2[1] = null;
+
+            // Third operation: map all columns (no pruning needed)
+            _SqlMetaDataSet clone3 = original.Clone();
+
+            // Assert: original is fully intact after all operations
+            Assert.NotNull(original[0]);
+            Assert.NotNull(original[1]);
+            Assert.NotNull(original[2]);
+            Assert.Equal("col1", original[0].column);
+            Assert.Equal("col2", original[1].column);
+            Assert.Equal("col3", original[2].column);
+
+            // Assert: each clone reflects its own pruning
+            Assert.Null(clone1[2]);
+            Assert.NotNull(clone1[1]);
+
+            Assert.Null(clone2[1]);
+            Assert.NotNull(clone2[2]);
+
+            Assert.NotNull(clone3[0]);
+            Assert.NotNull(clone3[1]);
+            Assert.NotNull(clone3[2]);
+        }
+
+        [Fact]
+        public void SqlMetaDataSet_Clone_PreservesOrdinals()
+        {
+            // Verify that cloned entries maintain correct ordinal values,
+            // which are used for column matching in AnalyzeTargetAndCreateUpdateBulkCommand
+            _SqlMetaDataSet original = new _SqlMetaDataSet(3);
+            original[0].column = "col1";
+            original[1].column = "col2";
+            original[2].column = "col3";
+
+            _SqlMetaDataSet clone = original.Clone();
+
+            Assert.Equal(original[0].ordinal, clone[0].ordinal);
+            Assert.Equal(original[1].ordinal, clone[1].ordinal);
+            Assert.Equal(original[2].ordinal, clone[2].ordinal);
+        }
+    }
+}