From 9b513c70d4d224ae43beeac3888c9e86974c93b4 Mon Sep 17 00:00:00 2001 From: Erik O'Leary Date: Mon, 18 Mar 2024 12:51:39 -0500 Subject: [PATCH 1/6] Implement IAsyncEnumerable on CosmosLinqQuery --- .../src/Linq/CosmosLinqExtensions.cs | 26 +++++++++++++++++++ .../src/Linq/CosmosLinqQuery.cs | 20 +++++++++++++- .../CosmosItemLinqTests.cs | 25 ++++++++++++++++++ 3 files changed, 70 insertions(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs index 7f7a22b466..f3bd146190 100644 --- a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs +++ b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs @@ -284,6 +284,32 @@ public static QueryDefinition ToQueryDefinition(this IQueryable query) throw new ArgumentException("ToQueryDefinition is only supported on Cosmos LINQ query operations", nameof(query)); } + /// + /// This extension method returns the query as an asynchronous enumerable. + /// + /// the type of object to query. + /// the IQueryable{T} to be converted. + /// An asynchronous enumerable to go through the items. + /// + /// This example shows how to get the query as an asynchronous enumerable. + /// + /// + /// linqQueryable = this.Container.GetItemLinqQueryable(); + /// IAsyncEnumerable asyncEnumerable = linqQueryable.Where(item => (item.taskNum < 100)).AsAsyncEnumerable(); + /// ]]> + /// + /// + public static IAsyncEnumerable AsAsyncEnumerable(this IQueryable query) + { + if (query is IAsyncEnumerable asyncEnumerable) + { + return asyncEnumerable; + } + + throw new ArgumentException("AsAsyncEnumerable is only supported on Cosmos LINQ query operations", nameof(query)); + } + /// /// This extension method gets the FeedIterator from LINQ IQueryable to execute query asynchronously. /// This will create the fresh new FeedIterator when called. diff --git a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs index 6676d096c9..8d465afd6c 100644 --- a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs +++ b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs @@ -22,7 +22,7 @@ namespace Microsoft.Azure.Cosmos.Linq /// This is the entry point for LINQ query creation/execution, it generate query provider, implements IOrderedQueryable. /// /// - internal sealed class CosmosLinqQuery : IDocumentQuery, IOrderedQueryable + internal sealed class CosmosLinqQuery : IDocumentQuery, IOrderedQueryable, IAsyncEnumerable { private readonly CosmosLinqQueryProvider queryProvider; private readonly Guid correlatedActivityId; @@ -283,5 +283,23 @@ private FeedIteratorInlineCore CreateFeedIterator(bool isContinuationExpected this.responseFactory.CreateQueryFeedUserTypeResponse), this.container.ClientContext); } + + public async IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + using FeedIteratorInlineCore localFeedIterator = this.CreateFeedIterator(isContinuationExpected: false, out ScalarOperationKind scalarOperationKind); + Debug.Assert( + scalarOperationKind == ScalarOperationKind.None, + "CosmosLinqQuery Assert!", + $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'"); + + while (localFeedIterator.HasMoreResults) + { + FeedResponse response = await localFeedIterator.ReadNextAsync(cancellationToken); + foreach (T item in response) + { + yield return item; + } + } + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs index efffeae1cb..206895a727 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs @@ -141,6 +141,31 @@ public void LinqQueryToIteratorBlockTest(bool isStreamIterator) } } + [TestMethod] + public async Task LinqQueryToAsyncEnumerable() + { + ToDoActivity toDoActivity = ToDoActivity.CreateRandomToDoActivity(); + toDoActivity.taskNum = 20; + toDoActivity.id = "minTaskNum"; + await this.Container.CreateItemAsync(toDoActivity, new PartitionKey(toDoActivity.pk)); + toDoActivity.taskNum = 100; + toDoActivity.id = "maxTaskNum"; + await this.Container.CreateItemAsync(toDoActivity, new PartitionKey(toDoActivity.pk)); + + IAsyncEnumerable query = this.Container.GetItemLinqQueryable() + .OrderBy(p => p.cost) + .AsAsyncEnumerable(); + + int found = 0; + await foreach (ToDoActivity item in query) + { + Assert.IsNotNull(item); + ++found; + } + + Assert.IsTrue(found > 0); + } + [TestMethod] [DataRow(false)] [DataRow(true)] From 2256fa28a8983f6e1ef0fb65e0c9a57421f7fc8d Mon Sep 17 00:00:00 2001 From: Erik O'Leary Date: Thu, 21 Mar 2024 12:02:36 -0500 Subject: [PATCH 2/6] Moved `GetAsyncEnumerator` closer to `GetEnumerator`, added doccomments. Dispose feed iterator in GetEnumerator --- .../src/Linq/CosmosLinqQuery.cs | 47 +++++++++++-------- .../CosmosItemLinqTests.cs | 2 +- 2 files changed, 29 insertions(+), 20 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs index 8d465afd6c..d445eea3c2 100644 --- a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs +++ b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs @@ -109,7 +109,7 @@ public IEnumerator GetEnumerator() " use GetItemQueryIterator to execute asynchronously"); } - FeedIterator localFeedIterator = this.CreateFeedIterator(false, out ScalarOperationKind scalarOperationKind); + using FeedIterator localFeedIterator = this.CreateFeedIterator(false, out ScalarOperationKind scalarOperationKind); Debug.Assert( scalarOperationKind == ScalarOperationKind.None, "CosmosLinqQuery Assert!", @@ -128,6 +128,33 @@ public IEnumerator GetEnumerator() } } + /// + /// Retrieves an object that can iterate through the individual results of the query asynchronously. + /// + /// + /// This triggers an asynchronous multi-page load. + /// + /// Cancellation token + /// IEnumerator + public async IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + using FeedIteratorInlineCore localFeedIterator = this.CreateFeedIterator(isContinuationExpected: false, out ScalarOperationKind scalarOperationKind); + Debug.Assert( + scalarOperationKind == ScalarOperationKind.None, + "CosmosLinqQuery Assert!", + $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'"); + + while (localFeedIterator.HasMoreResults) + { + FeedResponse response = await localFeedIterator.ReadNextAsync(cancellationToken); + + foreach (T item in response) + { + yield return item; + } + } + } + /// /// Synchronous Multi-Page load /// @@ -283,23 +310,5 @@ private FeedIteratorInlineCore CreateFeedIterator(bool isContinuationExpected this.responseFactory.CreateQueryFeedUserTypeResponse), this.container.ClientContext); } - - public async IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) - { - using FeedIteratorInlineCore localFeedIterator = this.CreateFeedIterator(isContinuationExpected: false, out ScalarOperationKind scalarOperationKind); - Debug.Assert( - scalarOperationKind == ScalarOperationKind.None, - "CosmosLinqQuery Assert!", - $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'"); - - while (localFeedIterator.HasMoreResults) - { - FeedResponse response = await localFeedIterator.ReadNextAsync(cancellationToken); - foreach (T item in response) - { - yield return item; - } - } - } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs index 206895a727..356b296165 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs @@ -163,7 +163,7 @@ public async Task LinqQueryToAsyncEnumerable() ++found; } - Assert.IsTrue(found > 0); + Assert.AreEqual(2, found); } [TestMethod] From 2670b811e1dec8555572f269ac2eb71bfd964601 Mon Sep 17 00:00:00 2001 From: Erik O'Leary Date: Fri, 19 Apr 2024 08:17:16 -0500 Subject: [PATCH 3/6] Refactor `AsAsyncEnumerable` method in `CosmosLinqExtensions.cs` to use `CosmosLinqQuery` instead of `IAsyncEnumerable` --- Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs index f3bd146190..771059e5be 100644 --- a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs +++ b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs @@ -302,7 +302,7 @@ public static QueryDefinition ToQueryDefinition(this IQueryable query) /// public static IAsyncEnumerable AsAsyncEnumerable(this IQueryable query) { - if (query is IAsyncEnumerable asyncEnumerable) + if (query is CosmosLinqQuery asyncEnumerable) { return asyncEnumerable; } From 5ed5f5e0c842d15e52c661ec16445a816bc81d27 Mon Sep 17 00:00:00 2001 From: Sidney Andrews Date: Fri, 25 Oct 2024 20:17:32 +0000 Subject: [PATCH 4/6] Initial implementation --- .../src/Linq/CosmosLinqExtensions.cs | 58 ++++++++++--- .../src/Linq/CosmosLinqQuery.cs | 27 +++---- .../FeedEnumerators/FeedAsyncEnumerator.cs | 52 ++++++++++++ .../FeedAsyncEnumeratorExtensions.cs | 74 +++++++++++++++++ .../FeedIterators/FeedIteratorExtensions.cs | 81 +++++++++++++++++++ 5 files changed, 266 insertions(+), 26 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumerator.cs create mode 100644 Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumeratorExtensions.cs create mode 100644 Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIteratorExtensions.cs diff --git a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs index 3c1c237f10..e1d3582996 100644 --- a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs +++ b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs @@ -19,7 +19,7 @@ namespace Microsoft.Azure.Cosmos.Linq /// This class provides extension methods for cosmos LINQ code. /// public static class CosmosLinqExtensions - { + { /// /// Returns the integer identifier corresponding to a specific item within a physical partition. /// This method is to be used in LINQ expressions only and will be evaluated on server. @@ -37,7 +37,7 @@ public static class CosmosLinqExtensions public static int DocumentId(this object obj) { throw new NotImplementedException(ClientResources.TypeCheckExtensionFunctionsNotImplemented); - } + } /// /// Returns a Boolean value indicating if the type of the specified expression is an array. @@ -304,22 +304,58 @@ public static QueryDefinition ToQueryDefinition(this IQueryable query) } /// - /// This extension method returns the query as an asynchronous enumerable. - /// - /// the type of object to query. - /// the IQueryable{T} to be converted. - /// An asynchronous enumerable to go through the items. + /// This extension method returns the current LINQ-derived NoSQL query for items within an Azure Cosmos DB for NoSQL container as an . + /// + /// The enables iteration over pages of results with sets of items. + /// + /// + /// The generic type to deserialize items to. + /// + /// + /// The to be converted to an asynchornous enumerable. + /// + /// + /// + /// + /// This method returns an asynchronous enumerable that enables iteration over pages of results and sets of items for each page. + /// /// - /// This example shows how to get the query as an asynchronous enumerable. + /// This example shows how to: + /// + /// 1. Create a from an existing . + /// 1. Use this extension method to convert the iterator into an . /// /// /// linqQueryable = this.Container.GetItemLinqQueryable(); - /// IAsyncEnumerable asyncEnumerable = linqQueryable.Where(item => (item.taskNum < 100)).AsAsyncEnumerable(); + /// IQueryable query = container.GetItemLinqQueryable() + /// .Where(item => item.partitionKey == "example-partition-key") + /// .OrderBy(item => item.value); + /// + /// IAsyncEnumerable> pages = query.AsAsyncEnumerable(); + /// + /// List items = new(); + /// double requestCharge = 0.0; + /// await foreach(var page in pages) + /// { + /// requestCharge += page.RequestCharge; + /// foreach (var item in page) + /// { + /// items.Add(item); + /// } + /// } + /// + /// + /// + /// /// /// - public static IAsyncEnumerable AsAsyncEnumerable(this IQueryable query) + public static IAsyncEnumerable> AsAsyncEnumerable(this IQueryable query) { if (query is CosmosLinqQuery asyncEnumerable) { diff --git a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs index d445eea3c2..4949140281 100644 --- a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs +++ b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs @@ -22,7 +22,7 @@ namespace Microsoft.Azure.Cosmos.Linq /// This is the entry point for LINQ query creation/execution, it generate query provider, implements IOrderedQueryable. /// /// - internal sealed class CosmosLinqQuery : IDocumentQuery, IOrderedQueryable, IAsyncEnumerable + internal sealed class CosmosLinqQuery : IDocumentQuery, IOrderedQueryable, IAsyncEnumerable> { private readonly CosmosLinqQueryProvider queryProvider; private readonly Guid correlatedActivityId; @@ -129,14 +129,19 @@ public IEnumerator GetEnumerator() } /// - /// Retrieves an object that can iterate through the individual results of the query asynchronously. + /// This method retrieves an that enables iteration, asynchronously, over pages of results with sets of items. /// + /// + /// (optional) The cancellation token that could be used to cancel the operation. + /// + /// + /// This method returns an asynchronous enumerator that enables iteration over pages of results and sets of items for each page. + /// /// - /// This triggers an asynchronous multi-page load. + /// > [!IMPORTANT] + /// > This method triggers an asynchronous multi-page load. /// - /// Cancellation token - /// IEnumerator - public async IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) + public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) { using FeedIteratorInlineCore localFeedIterator = this.CreateFeedIterator(isContinuationExpected: false, out ScalarOperationKind scalarOperationKind); Debug.Assert( @@ -144,15 +149,7 @@ public async IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellati "CosmosLinqQuery Assert!", $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'"); - while (localFeedIterator.HasMoreResults) - { - FeedResponse response = await localFeedIterator.ReadNextAsync(cancellationToken); - - foreach (T item in response) - { - yield return item; - } - } + return localFeedIterator.BuildFeedAsyncEnumerator(cancellationToken); } /// diff --git a/Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumerator.cs new file mode 100644 index 0000000000..84c9c4fa5a --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumerator.cs @@ -0,0 +1,52 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Collections.Generic; + using System.Threading; + + /// + /// This class encapsulates the asynchronous enumerator for an instance of . + /// + /// The generic type to deserialize items to. + public sealed class FeedAsyncEnumerator : IAsyncEnumerable>, IDisposable + { + private readonly FeedIterator feedIterator; + + /// + /// This constructor creates an instance of . + /// + /// + /// The target feed iterator instance. + /// + public FeedAsyncEnumerator(FeedIterator feedIterator) + { + this.feedIterator = feedIterator; + } + + /// + /// This method is used to get the asynchronous enumerator for an instance of . + /// + /// + /// (optional) The cancellation token that could be used to cancel the operation. + /// + /// + /// The asynchronous enumerator for an instance of . + /// + public IAsyncEnumerator> GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + return this.feedIterator.BuildFeedAsyncEnumerator(cancellationToken); + } + + /// + /// Releases the unmanaged resources used by the and optionally releases the managed resources. + /// + public void Dispose() + { + this.feedIterator.Dispose(); + } + } +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumeratorExtensions.cs b/Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumeratorExtensions.cs new file mode 100644 index 0000000000..5e3db9013c --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Resource/FeedEnumerators/FeedAsyncEnumeratorExtensions.cs @@ -0,0 +1,74 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System.Collections.Generic; + using System.Runtime.CompilerServices; + using System.Threading; + + /// + /// This class provides extension methods for Azure Cosmos DB for NoSQL's asynchronous feed iterator. + /// + public static class FeedAsyncEnumeratorExtensions + { + /// + /// This extension method returns the current feed iterator for items within an Azure Cosmos DB for NoSQL container as an . + /// + /// The enables iteration over pages of results with sets of items. + /// + /// + /// The generic type to deserialize items to. + /// + /// + /// The feed iterator to be converted to an asynchronous enumerable. + /// + /// + /// + /// + /// This method returns an asynchronous enumerable that enables iteration over pages of results and sets of items for each page. + /// + /// + /// This example shows how to: + /// + /// 1. Create a from an existing . + /// 1. Use this extension method to convert the iterator into an . + /// + /// + /// feedIterator = container.GetItemQueryIterator( + /// queryText: "SELECT * FROM items" + /// ); + /// + /// IAsyncEnumerable> pages = query.AsAsyncEnumerable(); + /// + /// List items = new(); + /// double requestCharge = 0.0; + /// await foreach(var page in pages) + /// { + /// requestCharge += page.RequestCharge; + /// foreach (var item in page) + /// { + /// items.Add(item); + /// } + /// } + /// + /// + /// + /// + /// + /// + public static IAsyncEnumerable> AsAsyncEnumerable( + this FeedIterator feedIterator) + { + return new FeedAsyncEnumerator(feedIterator); + } + } +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIteratorExtensions.cs b/Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIteratorExtensions.cs new file mode 100644 index 0000000000..3bec781ad3 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIteratorExtensions.cs @@ -0,0 +1,81 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System.Collections.Generic; + using System.Threading; + + /// + /// This class provides extension methods for Azure Cosmos DB for NoSQL's feed iterator. + /// + public static class FeedIteratorExtensions + { + /// + /// This extension method takes an existing and builds an that enables iteration over pages of results with sets of items. + /// + /// The generic type to deserialize items to. + /// + /// The target feed iterator instance. + /// + /// + /// (optional) The cancellation token that could be used to cancel the operation. + /// + /// + /// This extension method returns an asynchronous enumerator that enables iteration over pages of results and sets of items for each page. + /// + /// + /// + /// + /// + /// This example shows how to: + /// + /// 1. Create a from an existing . + /// 1. Use this extension method to convert the iterator into an . + /// + /// Finally, use the `await foreach` statement to iterate over the pages and items. + /// + /// feedIterator = container.GetItemQueryIterator( + /// queryText: "SELECT * FROM items" + /// ); + /// + /// IAsyncEnumerator> pages = feedIterator.BuildFeedAsyncEnumerator(); + /// + /// List items = new(); + /// double requestCharge = 0.0; + /// await foreach(var page in pages) + /// { + /// requestCharge += page.RequestCharge; + /// foreach (var item in page) + /// { + /// items.Add(item); + /// } + /// } + /// ]]> + /// + /// + /// + /// + /// + /// + public static async IAsyncEnumerator> BuildFeedAsyncEnumerator( + this FeedIterator feedIterator, + CancellationToken cancellationToken = default) + { + while (feedIterator.HasMoreResults) + { + FeedResponse response = await feedIterator.ReadNextAsync(cancellationToken); + + yield return response; + } + } + } +} \ No newline at end of file From 67b223c5ce817279ed04abf3328537279131e705 Mon Sep 17 00:00:00 2001 From: Sidney Andrews Date: Fri, 25 Oct 2024 20:18:57 +0000 Subject: [PATCH 5/6] Initial emulator tests --- .../CosmosBasicQueryTests.cs | 52 +++++++++++++++++++ .../CosmosItemLinqTests.cs | 45 +++++++++++----- 2 files changed, 84 insertions(+), 13 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs index 61f917bbe2..3ad9e9d721 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs @@ -839,6 +839,58 @@ public async Task TesOdeTokenCompatibilityWithNonOdePipeline() Assert.IsTrue(responseMessage.CosmosException.ToString().Contains(expectedErrorMessage)); } + [TestMethod] + public async Task ItemQueryIteratorToAsyncEnumerable() + { + CosmosClient client = DirectCosmosClient; + Container container = client.GetContainer(DatabaseId, ContainerId); + + int expected = 3; + string partitionKey = "example-partition-key"; + + List insertionTasks = new(); + foreach(int index in Enumerable.Range(1, expected)) + { + insertionTasks.Add( + container.CreateItemAsync( + item: ToDoActivity.CreateRandomToDoActivity( + id: $"{index:0000}", + pk: partitionKey + ) + ) + ); + } + await Task.WhenAll(insertionTasks); + + QueryDefinition query = new QueryDefinition( + query: "SELECT * FROM items i WHERE i.pk = @pk ORDER BY i.cost" + ).WithParameter("@pk", partitionKey); + + FeedIterator feedIterator = container.GetItemQueryIterator( + queryDefinition: query + ); + + IAsyncEnumerable> pages = feedIterator.AsAsyncEnumerable(); + + double requestUnits = 0.0; + int found = 0; + await foreach (FeedResponse page in pages) + { + Assert.IsNotNull(page); + Assert.AreEqual(page.StatusCode, System.Net.HttpStatusCode.OK); + requestUnits += page.RequestCharge; + found += page.Count; + foreach (ToDoActivity item in page) + { + Assert.IsNotNull(item); + Assert.AreEqual(item.pk, partitionKey); + } + } + + Assert.IsTrue(requestUnits > 0); + Assert.AreEqual(expected, found); + } + private class CustomHandler : RequestHandler { string correlatedActivityId; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs index 356b296165..de7801573c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemLinqTests.cs @@ -144,26 +144,45 @@ public void LinqQueryToIteratorBlockTest(bool isStreamIterator) [TestMethod] public async Task LinqQueryToAsyncEnumerable() { - ToDoActivity toDoActivity = ToDoActivity.CreateRandomToDoActivity(); - toDoActivity.taskNum = 20; - toDoActivity.id = "minTaskNum"; - await this.Container.CreateItemAsync(toDoActivity, new PartitionKey(toDoActivity.pk)); - toDoActivity.taskNum = 100; - toDoActivity.id = "maxTaskNum"; - await this.Container.CreateItemAsync(toDoActivity, new PartitionKey(toDoActivity.pk)); + int expected = 3; + string partitionKey = "example-partition-key"; + + List insertionTasks = new(); + foreach(int index in Enumerable.Range(1, expected)) + { + insertionTasks.Add( + this.Container.CreateItemAsync( + item: ToDoActivity.CreateRandomToDoActivity( + id: $"{index:0000}", + pk: partitionKey + ) + ) + ); + } + await Task.WhenAll(insertionTasks); - IAsyncEnumerable query = this.Container.GetItemLinqQueryable() - .OrderBy(p => p.cost) + IAsyncEnumerable> pages = this.Container.GetItemLinqQueryable() + .Where(i => i.pk == partitionKey) + .OrderBy(i => i.cost) .AsAsyncEnumerable(); + double requestUnits = 0.0; int found = 0; - await foreach (ToDoActivity item in query) + await foreach (FeedResponse page in pages) { - Assert.IsNotNull(item); - ++found; + Assert.IsNotNull(page); + Assert.AreEqual(page.StatusCode, System.Net.HttpStatusCode.OK); + requestUnits += page.RequestCharge; + found += page.Count; + foreach (ToDoActivity item in page) + { + Assert.IsNotNull(item); + Assert.AreEqual(item.pk, partitionKey); + } } - Assert.AreEqual(2, found); + Assert.IsTrue(requestUnits > 0); + Assert.AreEqual(expected, found); } [TestMethod] From afb699c9763170e5e32635a7011bc7c6ed4c7a82 Mon Sep 17 00:00:00 2001 From: Sidney Andrews Date: Fri, 25 Oct 2024 18:24:39 -0400 Subject: [PATCH 6/6] Initial unit tests --- .../src/Linq/CosmosLinqExtensions.cs | 2 +- .../Async/CosmosAsyncEnumerableTests.cs | 84 +++++++++++++++++++ 2 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Async/CosmosAsyncEnumerableTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs index e1d3582996..d74c16e04d 100644 --- a/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs +++ b/Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs @@ -362,7 +362,7 @@ public static IAsyncEnumerable> AsAsyncEnumerable(this IQuery return asyncEnumerable; } - throw new ArgumentException("AsAsyncEnumerable is only supported on Cosmos LINQ query operations", nameof(query)); + throw new NotSupportedException("AsAsyncEnumerable is only supported for IQueryable instances created by the Azure Cosmos DB LINQ provider."); } /// diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Async/CosmosAsyncEnumerableTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Async/CosmosAsyncEnumerableTests.cs new file mode 100644 index 0000000000..991f5988a8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Async/CosmosAsyncEnumerableTests.cs @@ -0,0 +1,84 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Linq +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Linq.Expressions; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Scenarios; + using Microsoft.Azure.Cosmos.Serializer; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; + + [TestClass] + public class CosmosAsyncEnumerableTests + { + [TestMethod] + public async Task FeedIteratorAsAsyncEnumerableTest() + { + // Create list of expected results + List expected = new List + { + new Item("1", "A"), + new Item("2", "A"), + new Item("3", "A"), + new Item("4", "B"), + new Item("5", "B"), + new Item("6", "C") + }; + + // Generate dictionary by grouping by partition key + Dictionary> itemsByPartition = expected.GroupBy(item => item.PartitionKey).ToDictionary(group => group.Key, group => group.ToList()); + IEnumerator enumerator = itemsByPartition.Keys.GetEnumerator(); + + // Create mock FeedIterator that returns the expected results + Mock> mockFeedIterator = new Mock>(); + mockFeedIterator.Setup(i => i.HasMoreResults).Returns(() => enumerator.MoveNext()); + mockFeedIterator.Setup(i => i.ReadNextAsync(It.IsAny())).Returns((CancellationToken _) => + { + Mock> mockFeedResponse = new Mock>(); + mockFeedResponse.Setup(i => i.GetEnumerator()).Returns(() => itemsByPartition[enumerator.Current].GetEnumerator()); + return Task.FromResult(mockFeedResponse.Object); + }); + FeedIterator feedIterator = mockFeedIterator.Object; + + // Method under test: Convert FeedIterator to IAsyncEnumerable> + IAsyncEnumerable> asyncEnumerable = feedIterator.AsAsyncEnumerable(); + + // Parse results using asynchronous foreach iteration + List actual = new List {}; + await foreach (FeedResponse page in asyncEnumerable) + { + foreach(Item item in page) + { + actual.Add(item); + } + } + + // Verify results + CollectionAssert.AreEqual(expected, actual); + } + + [TestMethod] + public void IQueryableAsAsyncEnumerableTest() + { + // Create mock Container that returns the expected results + Mock> mockQueryable = new Mock>(); + IQueryable queryable = mockQueryable.Object; + + // Method under test: Convert IQueryable to IAsyncEnumerable> with expected exception + Assert.ThrowsException(() => queryable.AsAsyncEnumerable()); + } + + public record Item( + string Id, + string PartitionKey + ); + } +} \ No newline at end of file