[Client encryption] Adds streaming JSON processing support to feed iterators

Microsoft.Azure.Cosmos.Encryption.Custom/src/Common/PooledMemoryStream.cs

@@ -303,6 +303,9 @@ public override void Write(byte[] buffer, int offset, int count)
 
 #if NET8_0_OR_GREATER
         public override void Write(ReadOnlySpan<byte> buffer)
+#else
+        public void Write(ReadOnlySpan<byte> buffer)
+#endif
         {
             this.EnsureNotDisposed();
 
@@ -332,7 +335,6 @@ public override void Write(ReadOnlySpan<byte> buffer)
                 this.length = this.position;
             }
         }
-#endif
 
         public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {

Microsoft.Azure.Cosmos.Encryption.Custom/src/DecryptableItem.cs

@@ -4,6 +4,7 @@
 
 namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
+    using System;
     using System.Threading.Tasks;
 
     /// <summary>
@@ -71,13 +72,23 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
     /// ]]>
     /// </code>
     /// </example>
-    public abstract class DecryptableItem
+    public abstract class DecryptableItem : IAsyncDisposable
     {
         /// <summary>
         /// Decrypts and deserializes the content.
         /// </summary>
         /// <typeparam name="T">The type of item to be returned.</typeparam>
         /// <returns>The requested item and the decryption related context.</returns>
         public abstract Task<(T, DecryptionContext)> GetItemAsync<T>();
+
+        /// <summary>
+        /// Disposes any resources held by the decryptable item.
+        /// Default implementation does nothing. Override in derived classes that hold disposable resources.
+        /// </summary>
+        /// <returns>A ValueTask representing the asynchronous dispose operation.</returns>
+        public virtual ValueTask DisposeAsync()
+        {
+            return default;
+        }
     }
 }

Microsoft.Azure.Cosmos.Encryption.Custom/src/EncryptionContainer.cs

@@ -16,6 +16,8 @@ internal sealed class EncryptionContainer : Container
     {
         private readonly Container container;
 
+        internal JsonProcessor DefaultJsonProcessor { get; private set; } = JsonProcessor.Newtonsoft;
+
         public CosmosSerializer CosmosSerializer { get; }
 
         public Encryptor Encryptor { get; }
@@ -631,7 +633,10 @@ public override FeedIterator<T> GetItemQueryIterator<T>(
                     queryDefinition,
                     continuationToken,
                     requestOptions),
-                this.ResponseFactory);
+                this.ResponseFactory,
+                this.Encryptor,
+                this.CosmosSerializer,
+                requestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override FeedIterator<T> GetItemQueryIterator<T>(
@@ -644,7 +649,10 @@ public override FeedIterator<T> GetItemQueryIterator<T>(
                     queryText,
                     continuationToken,
                     requestOptions),
-                this.ResponseFactory);
+                this.ResponseFactory,
+                this.Encryptor,
+                this.CosmosSerializer,
+                requestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override Task<ContainerResponse> ReadContainerAsync(
@@ -724,7 +732,7 @@ public override FeedIterator GetItemQueryStreamIterator(
                     continuationToken,
                     requestOptions),
                 this.Encryptor,
-                this.CosmosSerializer);
+                requestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override FeedIterator GetItemQueryStreamIterator(
@@ -738,7 +746,7 @@ public override FeedIterator GetItemQueryStreamIterator(
                     continuationToken,
                     requestOptions),
                 this.Encryptor,
-                this.CosmosSerializer);
+                requestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override Task<ThroughputResponse> ReplaceThroughputAsync(
@@ -771,7 +779,7 @@ public override FeedIterator GetItemQueryStreamIterator(
                     continuationToken,
                     requestOptions),
                 this.Encryptor,
-                this.CosmosSerializer);
+                requestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override FeedIterator<T> GetItemQueryIterator<T>(
@@ -786,7 +794,10 @@ public override FeedIterator<T> GetItemQueryIterator<T>(
                     queryDefinition,
                     continuationToken,
                     requestOptions),
-                this.ResponseFactory);
+                this.ResponseFactory,
+                this.Encryptor,
+                this.CosmosSerializer,
+                requestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override ChangeFeedEstimator GetChangeFeedEstimator(
@@ -807,7 +818,7 @@ public override FeedIterator GetChangeFeedStreamIterator(
                     changeFeedMode,
                     changeFeedRequestOptions),
                 this.Encryptor,
-                this.CosmosSerializer);
+                changeFeedRequestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override FeedIterator<T> GetChangeFeedIterator<T>(
@@ -820,7 +831,10 @@ public override FeedIterator<T> GetChangeFeedIterator<T>(
                     changeFeedStartFrom,
                     changeFeedMode,
                     changeFeedRequestOptions),
-                this.ResponseFactory);
+                this.ResponseFactory,
+                this.Encryptor,
+                this.CosmosSerializer,
+                changeFeedRequestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor);
         }
 
         public override Task<ItemResponse<T>> PatchItemAsync<T>(
@@ -917,6 +931,7 @@ public override ChangeFeedProcessorBuilder GetChangeFeedProcessorBuilder(
                     Stream decryptedChanges = await EncryptionProcessor.DeserializeAndDecryptResponseAsync(
                         changes,
                         this.Encryptor,
+                        this.DefaultJsonProcessor,
                         cancellationToken);
 
                     // Call the original passed in delegate
@@ -939,6 +954,7 @@ public override ChangeFeedProcessorBuilder GetChangeFeedProcessorBuilderWithManu
                     Stream decryptedChanges = await EncryptionProcessor.DeserializeAndDecryptResponseAsync(
                         changes,
                         this.Encryptor,
+                        this.DefaultJsonProcessor,
                         cancellationToken);
 
                     // Call the original passed in delegate
@@ -1008,6 +1024,13 @@ public override Task<bool> IsFeedRangePartOfAsync(
         }
 #endif
 
+#if NET8_0_OR_GREATER
+        public void UseStreamingJsonProcessingByDefault()
+        {
+            this.DefaultJsonProcessor = JsonProcessor.Stream;
+        }
+#endif
+
 #if PREVIEW && SDKPROJECTREF
         public override Task<SemanticRerankResult> SemanticRerankAsync(
             string rerankContext,
@@ -1036,6 +1059,7 @@ private async Task<ResponseMessage> ReadManyItemsHelperAsync(
             Stream decryptedContent = await EncryptionProcessor.DeserializeAndDecryptResponseAsync(
                 responseMessage.Content,
                 this.Encryptor,
+                readManyRequestOptions?.GetJsonProcessor(this.DefaultJsonProcessor) ?? this.DefaultJsonProcessor,
                 cancellationToken);
 
             return new DecryptedResponseMessage(responseMessage, decryptedContent);

Microsoft.Azure.Cosmos.Encryption.Custom/src/EncryptionContainerExtensions.cs

@@ -28,6 +28,26 @@ public static Container WithEncryptor(
                 encryptor);
         }
 
+#if NET8_0_OR_GREATER
+        /// <summary>
+        /// Configures the specified <see cref="Container"/> to use streaming JSON processing by default.
+        /// </summary>
+        /// <param name="container">The <see cref="Container"/> instance to configure. Must be an <see cref="EncryptionContainer"/>.</param>
+        /// <returns>The configured <see cref="EncryptionContainer"/> instance.</returns>
+        /// <exception cref="NotSupportedException">Thrown if <paramref name="container"/> is not an <see cref="EncryptionContainer"/>.</exception>
+        public static Container UseStreamingJsonProcessingByDefault(this Container container)
+        {
+            if (container is not EncryptionContainer encryptionContainer)
+            {
+                throw new NotSupportedException($"{nameof(UseStreamingJsonProcessingByDefault)} is only supported with {nameof(EncryptionContainer)}.");
+            }
+
+            encryptionContainer.UseStreamingJsonProcessingByDefault();
+
+            return encryptionContainer;
+        }
+#endif
+
         /// <summary>
         /// This method gets the FeedIterator from LINQ IQueryable to execute query asynchronously.
         /// This will create the fresh new FeedIterator when called which will support decryption.
@@ -57,7 +77,10 @@ public static FeedIterator<T> ToEncryptionFeedIterator<T>(
 
             return new EncryptionFeedIterator<T>(
                 (EncryptionFeedIterator)encryptionContainer.ToEncryptionStreamIterator(query),
-                encryptionContainer.ResponseFactory);
+                encryptionContainer.ResponseFactory,
+                encryptionContainer.Encryptor,
+                encryptionContainer.CosmosSerializer,
+                encryptionContainer.DefaultJsonProcessor);
         }
 
         /// <summary>
@@ -90,7 +113,7 @@ public static FeedIterator ToEncryptionStreamIterator<T>(
             return new EncryptionFeedIterator(
                 query.ToStreamIterator(),
                 encryptionContainer.Encryptor,
-                encryptionContainer.CosmosSerializer);
+                encryptionContainer.DefaultJsonProcessor);
         }
     }
 }

Microsoft.Azure.Cosmos.Encryption.Custom/src/EncryptionFeedIterator.cs

@@ -4,27 +4,32 @@
 
 namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
-    using System;
-    using System.Collections.Generic;
     using System.IO;
     using System.Threading;
     using System.Threading.Tasks;
-    using Newtonsoft.Json.Linq;
 
     internal sealed class EncryptionFeedIterator : FeedIterator
     {
         private readonly FeedIterator feedIterator;
         private readonly Encryptor encryptor;
-        private readonly CosmosSerializer cosmosSerializer;
+        private readonly JsonProcessor jsonProcessor;
 
         public EncryptionFeedIterator(
             FeedIterator feedIterator,
             Encryptor encryptor,
-            CosmosSerializer cosmosSerializer)
+            JsonProcessor jsonProcessor)
+        {
+            this.feedIterator = feedIterator ?? throw new System.ArgumentNullException(nameof(feedIterator));
+            this.encryptor = encryptor ?? throw new System.ArgumentNullException(nameof(encryptor));
+            this.jsonProcessor = jsonProcessor;
+        }
+
+        public EncryptionFeedIterator(
+            FeedIterator feedIterator,
+            Encryptor encryptor,
+            RequestOptions requestOptions)
+            : this(feedIterator, encryptor, requestOptions?.GetJsonProcessor() ?? JsonProcessor.Newtonsoft)
         {
-            this.feedIterator = feedIterator;
-            this.encryptor = encryptor;
-            this.cosmosSerializer = cosmosSerializer;
         }
 
         public override bool HasMoreResults => this.feedIterator.HasMoreResults;
@@ -41,6 +46,7 @@ public override async Task<ResponseMessage> ReadNextAsync(CancellationToken canc
                     Stream decryptedContent = await EncryptionProcessor.DeserializeAndDecryptResponseAsync(
                         responseMessage.Content,
                         this.encryptor,
+                        this.jsonProcessor,
                         cancellationToken);
 
                     return new DecryptedResponseMessage(responseMessage, decryptedContent);
@@ -50,49 +56,9 @@ public override async Task<ResponseMessage> ReadNextAsync(CancellationToken canc
             }
         }
 
-        public async Task<(ResponseMessage, List<T>)> ReadNextWithoutDecryptionAsync<T>(CancellationToken cancellationToken = default)
-        {
-            CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(options: null);
-            using (diagnosticsContext.CreateScope("FeedIterator.ReadNextWithoutDecryption"))
-            {
-                ResponseMessage responseMessage = await this.feedIterator.ReadNextAsync(cancellationToken);
-                List<T> decryptableContent = null;
-
-                if (responseMessage.IsSuccessStatusCode && responseMessage.Content != null)
-                {
-                    decryptableContent = this.ConvertResponseToDecryptableItems<T>(
-                        responseMessage.Content);
-
-                    return (responseMessage, decryptableContent);
-                }
-
-                return (responseMessage, decryptableContent);
-            }
-        }
-
-        private List<T> ConvertResponseToDecryptableItems<T>(
-            Stream content)
+        internal Task<ResponseMessage> ReadNextRawResponseAsync(CancellationToken cancellationToken = default)
         {
-            JObject contentJObj = EncryptionProcessor.BaseSerializer.FromStream<JObject>(content);
-
-            if (contentJObj.SelectToken(Constants.DocumentsResourcePropertyName) is not JArray documents)
-            {
-                throw new InvalidOperationException("Feed Response body contract was violated. Feed Response did not have an array of Documents.");
-            }
-
-            List<T> decryptableItems = new (documents.Count);
-
-            foreach (JToken value in documents)
-            {
-                DecryptableItemCore item = new (
-                    value,
-                    this.encryptor,
-                    this.cosmosSerializer);
-
-                decryptableItems.Add((T)(object)item);
-            }
-
-            return decryptableItems;
+            return this.feedIterator.ReadNextAsync(cancellationToken);
         }
     }
 }