|
1 | 1 | using System; |
2 | 2 | using System.Collections.Generic; |
3 | | -using System.Runtime.CompilerServices; |
4 | 3 | using System.Threading; |
5 | 4 | using System.Threading.Tasks; |
6 | 5 | using Microsoft.Extensions.Hosting; |
7 | 6 | using Microsoft.Extensions.Options; |
8 | 7 | using MongoDB.Driver; |
9 | 8 | using MalwareSampleExchange.Console.Models; |
| 9 | +using MongoDB.Driver.Linq; |
10 | 10 |
|
11 | 11 | namespace MalwareSampleExchange.Console.Database; |
12 | 12 |
|
@@ -36,25 +36,18 @@ public MongoMetadataHandler(IOptions<MongoMetadataOptions> options) |
36 | 36 | _mongoClient = new MongoClient(options.Value.ConnectionString); |
37 | 37 | } |
38 | 38 |
|
39 | | - public async IAsyncEnumerable<ExportSample> GetSamplesAsync(DateTime start, DateTime? end, string? sampleSet, [EnumeratorCancellation] CancellationToken token = default) |
| 39 | + public IAsyncEnumerable<ExportSample> GetSamplesAsync(DateTime start, DateTime? end, string? sampleSet, CancellationToken token = default) |
40 | 40 | { |
41 | 41 | var mongoDatabase = _mongoClient.GetDatabase(_options.DatabaseName); |
42 | 42 | var sampleCollection = mongoDatabase.GetCollection<ExportSample>(_options.CollectionName); |
43 | 43 | var list = end == null |
44 | | - ? await sampleCollection |
45 | | - .FindAsync(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.DoNotUseBefore <= DateTime.Now, cancellationToken: token) |
46 | | - : await sampleCollection |
47 | | - .FindAsync(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.Imported <= end && sample.DoNotUseBefore <= DateTime.Now, cancellationToken: token); |
48 | | - |
49 | | - while (await list.MoveNextAsync(token)) |
50 | | - { |
51 | | - foreach (var current in list.Current) |
52 | | - { |
53 | | - yield return current; |
54 | | - } |
55 | | - } |
| 44 | + ? sampleCollection |
| 45 | + .Find(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.DoNotUseBefore <= DateTime.Now).ToAsyncEnumerable(cancellationToken: token) |
| 46 | + : sampleCollection |
| 47 | + .Find(sample => sample.SampleSet == sampleSet && sample.Imported >= start && sample.Imported <= end && sample.DoNotUseBefore <= DateTime.Now).ToAsyncEnumerable(cancellationToken: token); |
| 48 | + return list; |
56 | 49 | } |
57 | | - |
| 50 | + |
58 | 51 | public async Task InsertSampleAsync(RequestExportSample sample, CancellationToken token = default) |
59 | 52 | { |
60 | 53 | var mongoDatabase = _mongoClient.GetDatabase(_options.DatabaseName); |
|
0 commit comments