Merge pull request #225 from AzureCosmosDB/copilot/fix-nested-elements-json-sink

Fix Dictionary serialization in JSON sink for MongoDB nested documents

Author: philnach

Extensions/Json/Cosmos.DataTransfer.JsonExtension.UnitTests/JsonFileSinkTests.cs

@@ -2,6 +2,7 @@
 using Cosmos.DataTransfer.Common.UnitTests;
 using Microsoft.Extensions.Logging.Abstractions;
 using Newtonsoft.Json;
+using Newtonsoft.Json.Linq;
 
 namespace Cosmos.DataTransfer.JsonExtension.UnitTests
 {
@@ -45,5 +46,68 @@ public async Task WriteAsync_WithFlatObjects_WritesToValidFile()
             Assert.IsTrue(outputData.Any(o => o.Id == 2 && o.Name == "Two"));
             Assert.IsTrue(outputData.Any(o => o.Id == 3 && o.Name == "Three"));
         }
+
+        [TestMethod]
+        public async Task WriteAsync_WithNestedDictionaries_SerializesCorrectly()
+        {
+            // Test case for the MongoDB nested elements issue
+            var sink = new JsonFileSink();
+
+            var data = new List<DictionaryDataItem>
+            {
+                new(new Dictionary<string, object?>
+                {
+                    { "_id", new Dictionary<string, object?> { { "$oid", "some_id" } } },
+                    { "thread_id", "thread_id" },
+                    { "content", new List<Dictionary<string, object?>>
+                        {
+                            new Dictionary<string, object?>
+                            {
+                                { "text", "a message text" },
+                                { "type", "text" }
+                            }
+                        }
+                    },
+                    { "role", "user" }
+                })
+            };
+            
+            string outputFile = $"{DateTime.Now:yy-MM-dd}_FS_Nested_Output.json";
+            var config = TestHelpers.CreateConfig(new Dictionary<string, string>
+            {
+                { "FilePath", outputFile }
+            });
+
+            await sink.WriteAsync(data.ToAsyncEnumerable(), config, new JsonFileSource(), NullLogger.Instance);
+
+            var jsonContent = await File.ReadAllTextAsync(outputFile);
+            var outputArray = JArray.Parse(jsonContent);
+            
+            Assert.AreEqual(1, outputArray.Count);
+            
+            var doc = outputArray[0] as JObject;
+            Assert.IsNotNull(doc);
+            
+            // Verify _id is an object with $oid field
+            var idObj = doc["_id"] as JObject;
+            Assert.IsNotNull(idObj, "_id should be an object");
+            Assert.AreEqual("some_id", idObj["$oid"]?.ToString());
+            
+            // Verify thread_id is a string
+            Assert.AreEqual("thread_id", doc["thread_id"]?.ToString());
+            
+            // Verify content is an array of objects
+            var contentArray = doc["content"] as JArray;
+            Assert.IsNotNull(contentArray, "content should be an array");
+            Assert.AreEqual(1, contentArray.Count);
+            
+            var contentItem = contentArray[0] as JObject;
+            Assert.IsNotNull(contentItem, "content item should be an object");
+            Assert.AreEqual("a message text", contentItem["text"]?.ToString());
+            Assert.AreEqual("text", contentItem["type"]?.ToString());
+            
+            // Verify role is a string
+            Assert.AreEqual("user", doc["role"]?.ToString());
+        }
     }
 }

Interfaces/Cosmos.DataTransfer.Common.UnitTests/DataItemJsonConverterTests.cs

@@ -194,5 +194,81 @@ public void Test_AsJsonString(bool includeNullFields) {
         var json = DataItemJsonConverter.AsJsonString(obj, false, includeNullFields);
         Assert.AreEqual(expected, json);
     }
+
+    [TestMethod]
+    [DataRow(false)]
+    [DataRow(true)]
+    public void Test_WriteFieldValue_DictionaryAsNestedObject(bool includeNullFields)
+    {
+        // Test that Dictionary<string, object?> is properly serialized as nested object
+        var nestedDict = new Dictionary<string, object?>
+        {
+            { "text", "a message text" },
+            { "type", "text" },
+            { "NULL", null }
+        };
+        
+        var expected = "\"x\":{\"text\":\"a message text\",\"type\":\"text\",\"NULL\":null}";
+        if (!includeNullFields)
+        {
+            expected = expected.Replace(",\"NULL\":null", "");
+        }
+        
+        var (writer, readFunc) = CreateUtf8JsonWriter();
+        DataItemJsonConverter.WriteFieldValue(writer, "x", nestedDict, includeNullFields: includeNullFields);
+        Assert.AreEqual(expected, readFunc(), $"includeNullFields: {includeNullFields}");
+    }
+
+    [TestMethod]
+    [DataRow(false)]
+    [DataRow(true)]
+    public void Test_WriteFieldValue_ArrayOfDictionaries(bool includeNullFields)
+    {
+        // Test array of dictionaries (simulating MongoDB nested array scenario)
+        var arrayOfDicts = new List<Dictionary<string, object?>>
+        {
+            new Dictionary<string, object?>
+            {
+                { "text", "a message text" },
+                { "type", "text" }
+            },
+            new Dictionary<string, object?>
+            {
+                { "text", "another message" },
+                { "type", "text" }
+            }
+        };
+        
+        var expected = "\"x\":[{\"text\":\"a message text\",\"type\":\"text\"},{\"text\":\"another message\",\"type\":\"text\"}]";
+        
+        var (writer, readFunc) = CreateUtf8JsonWriter();
+        DataItemJsonConverter.WriteFieldValue(writer, "x", arrayOfDicts, includeNullFields: includeNullFields);
+        Assert.AreEqual(expected, readFunc(), $"includeNullFields: {includeNullFields}");
+    }
+
+    [TestMethod]
+    public void Test_AsJsonString_CompleteMongoScenario()
+    {
+        // Test complete scenario from the issue: nested _id object and array of content dictionaries
+        var mongoStyleDoc = new DictionaryDataItem(new Dictionary<string, object?>
+        {
+            { "_id", new Dictionary<string, object?> { { "$oid", "some_id" } } },
+            { "thread_id", "thread_id" },
+            { "content", new List<Dictionary<string, object?>>
+                {
+                    new Dictionary<string, object?>
+                    {
+                        { "text", "a message text" },
+                        { "type", "text" }
+                    }
+                }
+            },
+            { "role", "user" }
+        });
+
+        var expected = "{\"_id\":{\"$oid\":\"some_id\"},\"thread_id\":\"thread_id\",\"content\":[{\"text\":\"a message text\",\"type\":\"text\"}],\"role\":\"user\"}";
+        var json = DataItemJsonConverter.AsJsonString(mongoStyleDoc, false, false);
+        Assert.AreEqual(expected, json);
+    }
 }
 

Interfaces/Cosmos.DataTransfer.Common/DataItemJsonConverter.cs

@@ -104,6 +104,12 @@ internal static void WriteFieldValue(Utf8JsonWriter writer, string fieldName, ob
             {
                 WriteDataItem(writer, child, includeNullFields, propertyName);
             }
+            else if (fieldValue is IDictionary<string, object?> dict)
+            {
+                // Handle dictionaries (e.g., from MongoDB BsonDocument conversion) as nested objects
+                var dictItem = new DictionaryDataItem(dict);
+                WriteDataItem(writer, dictItem, includeNullFields, propertyName);
+            }
             else if (fieldValue is not string && fieldValue is IEnumerable children)
             {
                 writer.WriteStartArray(propertyName);
@@ -113,6 +119,12 @@ internal static void WriteFieldValue(Utf8JsonWriter writer, string fieldName, ob
                     {
                         WriteDataItem(writer, arrayChild, includeNullFields);
                     }
+                    else if (arrayItem is IDictionary<string, object?> arrayDict)
+                    {
+                        // Handle dictionaries (e.g., from MongoDB BsonDocument conversion) as nested objects
+                        var arrayDictItem = new DictionaryDataItem(arrayDict);
+                        WriteDataItem(writer, arrayDictItem, includeNullFields);
+                    }
                     else if (TryGetLong(arrayItem, out var longValue))
                     {
                         writer.WriteNumberValue(longValue);