Use Linq.AsyncEnumerable library for Mongo

GDATASoftwareAG · Feb 26, 2024 · 625e3a6 · 625e3a6
1 parent 8a2dd47
commit 625e3a6
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 16 deletions.
diff --git a/src/MalwareSampleExchange.Console/Database/MongoMetadataHandler.cs b/src/MalwareSampleExchange.Console/Database/MongoMetadataHandler.cs
@@ -1,12 +1,12 @@
 using System;
 using System.Collections.Generic;
-using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Options;
 using MongoDB.Driver;
 using MalwareSampleExchange.Console.Models;
+using MongoDB.Driver.Linq;
 
 namespace MalwareSampleExchange.Console.Database;
 
@@ -36,25 +36,18 @@ public MongoMetadataHandler(IOptions<MongoMetadataOptions> options)
         _mongoClient = new MongoClient(options.Value.ConnectionString);
     }
 
-    public async IAsyncEnumerable<ExportSample> GetSamplesAsync(DateTime start, DateTime? end, string? sampleSet, [EnumeratorCancellation] CancellationToken token = default)
+    public IAsyncEnumerable<ExportSample> GetSamplesAsync(DateTime start, DateTime? end, string? sampleSet, CancellationToken token = default)
     {
         var mongoDatabase = _mongoClient.GetDatabase(_options.DatabaseName);
         var sampleCollection = mongoDatabase.GetCollection<ExportSample>(_options.CollectionName);
         var list = end == null
-            ? await sampleCollection
-                .FindAsync(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.DoNotUseBefore <= DateTime.Now, cancellationToken: token)
-            : await sampleCollection
-                .FindAsync(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.Imported <= end && sample.DoNotUseBefore <= DateTime.Now, cancellationToken: token);
-
-        while (await list.MoveNextAsync(token))
-        {
-            foreach (var current in list.Current)
-            {
-                yield return current;
-            }
-        }
+            ? sampleCollection
+                .Find(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.DoNotUseBefore <= DateTime.Now).ToAsyncEnumerable(cancellationToken: token)
+            : sampleCollection
+                .Find(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.Imported <= end && sample.DoNotUseBefore <= DateTime.Now).ToAsyncEnumerable(cancellationToken: token);
+        return list;
     }
-    
+
     public async Task InsertSampleAsync(RequestExportSample sample, CancellationToken token = default)
     {
         var mongoDatabase = _mongoClient.GetDatabase(_options.DatabaseName);

diff --git a/src/MalwareSampleExchange.Console/ListRequester/ListRequester.cs b/src/MalwareSampleExchange.Console/ListRequester/ListRequester.cs
@@ -3,7 +3,6 @@
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading;
-using System.Threading.Tasks;
 using JWT.Algorithms;
 using JWT.Builder;
 using Microsoft.Extensions.Logging;

diff --git a/src/MalwareSampleExchange.Console/MalwareSampleExchange.Console.csproj b/src/MalwareSampleExchange.Console/MalwareSampleExchange.Console.csproj
@@ -7,6 +7,7 @@
   <ItemGroup>
     <PackageReference Include="JWT" Version="8.9.0" />
     <PackageReference Include="AWSSDK.S3" Version="3.7.8.5" />
+    <PackageReference Include="MongoDB.Driver.Linq.AsyncEnumerable" Version="2.15.4" />
     <PackageReference Include="Prometheus.Client.AspNetCore" Version="4.5.1" />
     <PackageReference Include="Prometheus.Client.HttpRequestDurations" Version="3.4.1" />
     <PackageReference Include="Microsoft.AspNetCore.Mvc.NewtonsoftJson" Version="6.0.3" />