git.stella-ops.org/src/StellaOps.Excititor.Storage.Mongo/MongoVexConsensusHoldStore.cs

using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using MongoDB.Driver;
using StellaOps.Excititor.Core;

namespace StellaOps.Excititor.Storage.Mongo;

public sealed class MongoVexConsensusHoldStore : IVexConsensusHoldStore
{
    private readonly IMongoCollection<VexConsensusHoldRecord> _collection;

    public MongoVexConsensusHoldStore(IMongoDatabase database)
    {
        ArgumentNullException.ThrowIfNull(database);
        VexMongoMappingRegistry.Register();
        _collection = database.GetCollection<VexConsensusHoldRecord>(VexMongoCollectionNames.ConsensusHolds);
    }

    public async ValueTask<VexConsensusHold?> FindAsync(string vulnerabilityId, string productKey, CancellationToken cancellationToken, IClientSessionHandle? session = null)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(vulnerabilityId);
        ArgumentException.ThrowIfNullOrWhiteSpace(productKey);
        var id = VexConsensusRecord.CreateId(vulnerabilityId, productKey);
        var filter = Builders<VexConsensusHoldRecord>.Filter.Eq(x => x.Id, id);
        var record = session is null
            ? await _collection.Find(filter).FirstOrDefaultAsync(cancellationToken).ConfigureAwait(false)
            : await _collection.Find(session, filter).FirstOrDefaultAsync(cancellationToken).ConfigureAwait(false);
        return record?.ToDomain();
    }

    public async ValueTask SaveAsync(VexConsensusHold hold, CancellationToken cancellationToken, IClientSessionHandle? session = null)
    {
        ArgumentNullException.ThrowIfNull(hold);
        var record = VexConsensusHoldRecord.FromDomain(hold);
        var filter = Builders<VexConsensusHoldRecord>.Filter.Eq(x => x.Id, record.Id);
        if (session is null)
        {
            await _collection.ReplaceOneAsync(filter, record, new ReplaceOptions { IsUpsert = true }, cancellationToken).ConfigureAwait(false);
        }
        else
        {
            await _collection.ReplaceOneAsync(session, filter, record, new ReplaceOptions { IsUpsert = true }, cancellationToken).ConfigureAwait(false);
        }
    }

    public async ValueTask RemoveAsync(string vulnerabilityId, string productKey, CancellationToken cancellationToken, IClientSessionHandle? session = null)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(vulnerabilityId);
        ArgumentException.ThrowIfNullOrWhiteSpace(productKey);
        var id = VexConsensusRecord.CreateId(vulnerabilityId, productKey);
        var filter = Builders<VexConsensusHoldRecord>.Filter.Eq(x => x.Id, id);
        if (session is null)
        {
            await _collection.DeleteOneAsync(filter, cancellationToken).ConfigureAwait(false);
        }
        else
        {
            await _collection.DeleteOneAsync(session, filter, options: null, cancellationToken).ConfigureAwait(false);
        }
    }

    public async IAsyncEnumerable<VexConsensusHold> FindEligibleAsync(DateTimeOffset asOf, int batchSize, [EnumeratorCancellation] CancellationToken cancellationToken, IClientSessionHandle? session = null)
    {
        var cutoff = asOf.UtcDateTime;
        var filter = Builders<VexConsensusHoldRecord>.Filter.Lte(x => x.EligibleAt, cutoff);
        var find = session is null
            ? _collection.Find(filter)
            : _collection.Find(session, filter);

        find = find.SortBy(x => x.EligibleAt);

        if (batchSize > 0)
        {
            find = find.Limit(batchSize);
        }

        using var cursor = await find.ToCursorAsync(cancellationToken).ConfigureAwait(false);
        while (await cursor.MoveNextAsync(cancellationToken).ConfigureAwait(false))
        {
            foreach (var record in cursor.Current)
            {
                yield return record.ToDomain();
            }
        }
    }
}