git.stella-ops.org/src/StellaOps.Scheduler.Storage.Mongo/Repositories/ImpactSnapshotRepository.cs

using System;
using MongoDB.Bson;
using MongoDB.Driver;
using StellaOps.Scheduler.Models;
using StellaOps.Scheduler.Storage.Mongo.Internal;
using StellaOps.Scheduler.Storage.Mongo.Serialization;

namespace StellaOps.Scheduler.Storage.Mongo.Repositories;

internal sealed class ImpactSnapshotRepository : IImpactSnapshotRepository
{
    private static readonly FilterDefinitionBuilder<BsonDocument> Filter = Builders<BsonDocument>.Filter;
    private static readonly SortDefinitionBuilder<BsonDocument> Sort = Builders<BsonDocument>.Sort;

    private readonly IMongoCollection<BsonDocument> _collection;

    public ImpactSnapshotRepository(SchedulerMongoContext context)
    {
        if (context is null)
        {
            throw new ArgumentNullException(nameof(context));
        }

        _collection = context.Database.GetCollection<BsonDocument>(context.Options.ImpactSnapshotsCollection);
    }

    public async Task UpsertAsync(
        ImpactSet snapshot,
        IClientSessionHandle? session = null,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(snapshot);

        var document = ImpactSetDocumentMapper.ToBsonDocument(snapshot);
        var filter = Filter.Eq("_id", document["_id"]);
        var options = new ReplaceOptions { IsUpsert = true };

        if (session is null)
        {
            await _collection.ReplaceOneAsync(filter, document, options, cancellationToken).ConfigureAwait(false);
        }
        else
        {
            await _collection.ReplaceOneAsync(session, filter, document, options, cancellationToken).ConfigureAwait(false);
        }
    }

    public async Task<ImpactSet?> GetBySnapshotIdAsync(
        string snapshotId,
        IClientSessionHandle? session = null,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(snapshotId))
        {
            throw new ArgumentException("Snapshot id must be provided.", nameof(snapshotId));
        }

        var filter = Filter.Eq("_id", snapshotId);
        var query = session is null
            ? _collection.Find(filter)
            : _collection.Find(session, filter);

        var document = await query.FirstOrDefaultAsync(cancellationToken).ConfigureAwait(false);
        return document is null ? null : ImpactSetDocumentMapper.FromBsonDocument(document);
    }

    public async Task<ImpactSet?> GetLatestBySelectorAsync(
        Selector selector,
        IClientSessionHandle? session = null,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(selector);
        if (string.IsNullOrWhiteSpace(selector.TenantId))
        {
            throw new ArgumentException("Selector tenantId is required to resolve impact snapshots.", nameof(selector));
        }

        var digest = ImpactSetDocumentMapper.ComputeSelectorDigest(selector);
        var filters = Filter.And(
            Filter.Eq("selectorDigest", digest),
            Filter.Eq("selector.tenantId", selector.TenantId));

        var find = session is null
            ? _collection.Find(filters)
            : _collection.Find(session, filters);

        var document = await find
            .Sort(Sort.Descending("generatedAt"))
            .FirstOrDefaultAsync(cancellationToken)
            .ConfigureAwait(false);

        return document is null ? null : ImpactSetDocumentMapper.FromBsonDocument(document);
    }
}