Cosmos: Populate __jObject on SaveChanges

src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs

@@ -16,6 +16,8 @@
 using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Infrastructure;
+using Microsoft.EntityFrameworkCore.Metadata;
+using Microsoft.EntityFrameworkCore.Metadata.Conventions;
 using Microsoft.EntityFrameworkCore.Storage;
 using Microsoft.EntityFrameworkCore.Update;
 using Microsoft.EntityFrameworkCore.Utilities;
@@ -465,6 +467,20 @@ private static void ProcessResponse(ResponseMessage response, IUpdateEntry entry
             {
                 entry.SetStoreGeneratedValue(etagProperty, response.Headers.ETag);
             }
+
+            var jObjectProperty = entry.EntityType.FindProperty(StoreKeyConvention.JObjectPropertyName);
+            if (jObjectProperty != null
+                && jObjectProperty.ValueGenerated == ValueGenerated.OnAddOrUpdate
+                && response.Content != null)
+            {
+                using var responseStream = response.Content;
+                using var reader = new StreamReader(responseStream);
+                using var jsonReader = new JsonTextReader(reader);
+
+                var createdDocument = new JsonSerializer().Deserialize<JObject>(jsonReader);
+
+                entry.SetStoreGeneratedValue(jObjectProperty, createdDocument);
+            }
         }
 
         /// <summary>

src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs

@@ -209,7 +209,16 @@ private bool Save(IUpdateEntry entry)
             switch (state)
             {
                 case EntityState.Added:
-                    var newDocument = documentSource.CreateDocument(entry);
+                    var newDocument = documentSource.GetCurrentDocument(entry);
+                    if (newDocument != null)
+                    {
+                        documentSource.UpdateDocument(newDocument, entry);
+                    }
+                    else
+                    {
+                        newDocument = documentSource.CreateDocument(entry);
+                    }
+
                     return _cosmosClient.CreateItem(collectionId, newDocument, entry);
 
                 case EntityState.Modified:
@@ -267,7 +276,16 @@ private Task<bool> SaveAsync(IUpdateEntry entry, CancellationToken cancellationT
             switch (state)
             {
                 case EntityState.Added:
-                    var newDocument = documentSource.CreateDocument(entry);
+                    var newDocument = documentSource.GetCurrentDocument(entry);
+                    if (newDocument != null)
+                    {
+                        documentSource.UpdateDocument(newDocument, entry);
+                    }
+                    else
+                    {
+                        newDocument = documentSource.CreateDocument(entry);
+                    }
+
                     return _cosmosClient.CreateItemAsync(
                         collectionId, newDocument, entry, cancellationToken);
 

test/EFCore.Cosmos.FunctionalTests/EndToEndCosmosTest.cs

@@ -182,6 +182,142 @@ public async Task Can_add_update_delete_detached_entity_end_to_end_async()
             }
         }
 
+        [ConditionalFact]
+        public void Can_add_update_untracked_properties()
+        {
+            var options = Fixture.CreateOptions();
+
+            var customer = new Customer { Id = 42, Name = "Theon" };
+
+            using (var context = new CustomerContext(options))
+            {
+                context.Database.EnsureCreated();
+
+                var entry = context.Add(customer);
+
+                entry.Property<JObject>("__jObject").CurrentValue = new JObject
+                {
+                    ["key1"] = "value1"
+                };
+
+                context.SaveChanges();
+
+                var providersJson = entry.Property<JObject>("__jObject").CurrentValue;
+                Assert.NotNull(providersJson);
+
+                providersJson["key2"] = "value2";
+                entry.State = EntityState.Modified;
+                context.SaveChanges();
+            }
+
+            using (var context = new CustomerContext(options))
+            {
+                var customerFromStore = context.Set<Customer>().Single();
+
+                Assert.Equal(42, customerFromStore.Id);
+                Assert.Equal("Theon", customerFromStore.Name);
+
+                var entry = context.Entry(customerFromStore);
+                var document = entry.Property<JObject>("__jObject").CurrentValue;
+                Assert.Equal("value1", document["key1"]);
+                Assert.Equal("value2", document["key2"]);
+
+                document["key1"] = "value1.1";
+                customerFromStore.Name = "Theon Greyjoy";
+
+                context.SaveChanges();
+            }
+
+            using (var context = new CustomerContext(options))
+            {
+                var customerFromStore = context.Set<Customer>().Single();
+
+                Assert.Equal("Theon Greyjoy", customerFromStore.Name);
+
+                var entry = context.Entry(customerFromStore);
+                var document = entry.Property<JObject>("__jObject").CurrentValue;
+                Assert.Equal("value1.1", document["key1"]);
+                Assert.Equal("value2", document["key2"]);
+
+                context.Remove(customerFromStore);
+
+                context.SaveChanges();
+            }
+
+            using (var context = new CustomerContext(options))
+            {
+                Assert.Empty(context.Set<Customer>().ToList());
+            }
+        }
+
+        [ConditionalFact]
+        public async Task Can_add_update_untracked_properties_async()
+        {
+            var options = Fixture.CreateOptions();
+
+            var customer = new Customer { Id = 42, Name = "Theon" };
+
+            using (var context = new CustomerContext(options))
+            {
+                context.Database.EnsureCreated();
+
+                var entry = context.Add(customer);
+
+                entry.Property<JObject>("__jObject").CurrentValue = new JObject
+                {
+                    ["key1"] = "value1"
+                };
+
+                await context.SaveChangesAsync();
+
+                var providersJson = entry.Property<JObject>("__jObject").CurrentValue;
+                Assert.NotNull(providersJson);
+
+                providersJson["key2"] = "value2";
+                entry.State = EntityState.Modified;
+                await context.SaveChangesAsync();
+            }
+
+            using (var context = new CustomerContext(options))
+            {
+                var customerFromStore = await context.Set<Customer>().SingleAsync();
+
+                Assert.Equal(42, customerFromStore.Id);
+                Assert.Equal("Theon", customerFromStore.Name);
+
+                var entry = context.Entry(customerFromStore);
+                var document = entry.Property<JObject>("__jObject").CurrentValue;
+                Assert.Equal("value1", document["key1"]);
+                Assert.Equal("value2", document["key2"]);
+
+                document["key1"] = "value1.1";
+                customerFromStore.Name = "Theon Greyjoy";
+
+                await context.SaveChangesAsync();
+            }
+
+            using (var context = new CustomerContext(options))
+            {
+                var customerFromStore = await context.Set<Customer>().SingleAsync();
+
+                Assert.Equal("Theon Greyjoy", customerFromStore.Name);
+
+                var entry = context.Entry(customerFromStore);
+                var document = entry.Property<JObject>("__jObject").CurrentValue;
+                Assert.Equal("value1.1", document["key1"]);
+                Assert.Equal("value2", document["key2"]);
+
+                context.Remove(customerFromStore);
+
+                await context.SaveChangesAsync();
+            }
+
+            using (var context = new CustomerContext(options))
+            {
+                Assert.Empty(context.Set<Customer>().ToList());
+            }
+        }
+
         private class Customer
         {
             public int Id { get; set; }

test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs

@@ -284,7 +284,7 @@ public class FakeEntityType : IEntityType
             public IIndex FindIndex(IReadOnlyList<IProperty> properties) => throw new NotImplementedException();
             public IKey FindKey(IReadOnlyList<IProperty> properties) => throw new NotImplementedException();
             public IKey FindPrimaryKey() => throw new NotImplementedException();
-            public IProperty FindProperty(string name) => throw new NotImplementedException();
+            public IProperty FindProperty(string name) => null;
             public IServiceProperty FindServiceProperty(string name) => throw new NotImplementedException();
             public ISkipNavigation FindSkipNavigation(string name) => throw new NotImplementedException();
             public IEnumerable<IAnnotation> GetAnnotations() => throw new NotImplementedException();