git.stella-ops.org/src/StellaOps.Feedser.Storage.Mongo/Migrations/EnsureDocumentExpiryIndexesMigration.cs

using System;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Options;
using MongoDB.Bson;
using MongoDB.Driver;

namespace StellaOps.Feedser.Storage.Mongo.Migrations;

internal sealed class EnsureDocumentExpiryIndexesMigration : IMongoMigration
{
    private readonly MongoStorageOptions _options;

    public EnsureDocumentExpiryIndexesMigration(IOptions<MongoStorageOptions> options)
    {
        ArgumentNullException.ThrowIfNull(options);
        _options = options.Value;
    }

    public string Id => "20241005_document_expiry_indexes";

    public string Description => "Ensure document.expiresAt index matches configured retention";

    public async Task ApplyAsync(IMongoDatabase database, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(database);

        var needsTtl = _options.RawDocumentRetention > TimeSpan.Zero;
        var collection = database.GetCollection<BsonDocument>(MongoStorageDefaults.Collections.Document);

        using var cursor = await collection.Indexes.ListAsync(cancellationToken).ConfigureAwait(false);
        var indexes = await cursor.ToListAsync(cancellationToken).ConfigureAwait(false);

        var ttlIndex = indexes.FirstOrDefault(x => TryGetName(x, out var name) && string.Equals(name, "document_expiresAt_ttl", StringComparison.Ordinal));
        var nonTtlIndex = indexes.FirstOrDefault(x => TryGetName(x, out var name) && string.Equals(name, "document_expiresAt", StringComparison.Ordinal));

        if (needsTtl)
        {
            var shouldRebuild = ttlIndex is null || !IndexMatchesTtlExpectations(ttlIndex);
            if (shouldRebuild)
            {
                if (ttlIndex is not null)
                {
                    await collection.Indexes.DropOneAsync("document_expiresAt_ttl", cancellationToken).ConfigureAwait(false);
                }

                if (nonTtlIndex is not null)
                {
                    await collection.Indexes.DropOneAsync("document_expiresAt", cancellationToken).ConfigureAwait(false);
                }

                var options = new CreateIndexOptions<BsonDocument>
                {
                    Name = "document_expiresAt_ttl",
                    ExpireAfter = TimeSpan.Zero,
                    PartialFilterExpression = Builders<BsonDocument>.Filter.Exists("expiresAt", true),
                };

                var keys = Builders<BsonDocument>.IndexKeys.Ascending("expiresAt");
                await collection.Indexes.CreateOneAsync(new CreateIndexModel<BsonDocument>(keys, options), cancellationToken: cancellationToken).ConfigureAwait(false);
            }
            else if (nonTtlIndex is not null)
            {
                await collection.Indexes.DropOneAsync("document_expiresAt", cancellationToken).ConfigureAwait(false);
            }
        }
        else
        {
            if (ttlIndex is not null)
            {
                await collection.Indexes.DropOneAsync("document_expiresAt_ttl", cancellationToken).ConfigureAwait(false);
            }

            var shouldRebuild = nonTtlIndex is null || !IndexMatchesNonTtlExpectations(nonTtlIndex);
            if (shouldRebuild)
            {
                if (nonTtlIndex is not null)
                {
                    await collection.Indexes.DropOneAsync("document_expiresAt", cancellationToken).ConfigureAwait(false);
                }

                var options = new CreateIndexOptions<BsonDocument>
                {
                    Name = "document_expiresAt",
                    PartialFilterExpression = Builders<BsonDocument>.Filter.Exists("expiresAt", true),
                };

                var keys = Builders<BsonDocument>.IndexKeys.Ascending("expiresAt");
                await collection.Indexes.CreateOneAsync(new CreateIndexModel<BsonDocument>(keys, options), cancellationToken: cancellationToken).ConfigureAwait(false);
            }
        }
    }

    private static bool IndexMatchesTtlExpectations(BsonDocument index)
    {
        if (!index.TryGetValue("expireAfterSeconds", out var expireAfter) || expireAfter.ToDouble() != 0)
        {
            return false;
        }

        if (!index.TryGetValue("partialFilterExpression", out var partialFilter) || partialFilter is not BsonDocument partialDoc)
        {
            return false;
        }

        if (!partialDoc.TryGetValue("expiresAt", out var expiresAtRule) || expiresAtRule is not BsonDocument expiresAtDoc)
        {
            return false;
        }

        return expiresAtDoc.Contains("$exists") && expiresAtDoc["$exists"].ToBoolean();
    }

    private static bool IndexMatchesNonTtlExpectations(BsonDocument index)
    {
        if (index.Contains("expireAfterSeconds"))
        {
            return false;
        }

        if (!index.TryGetValue("partialFilterExpression", out var partialFilter) || partialFilter is not BsonDocument partialDoc)
        {
            return false;
        }

        if (!partialDoc.TryGetValue("expiresAt", out var expiresAtRule) || expiresAtRule is not BsonDocument expiresAtDoc)
        {
            return false;
        }

        return expiresAtDoc.Contains("$exists") && expiresAtDoc["$exists"].ToBoolean();
    }

    private static bool TryGetName(BsonDocument index, out string name)
    {
        if (index.TryGetValue("name", out var value) && value.IsString)
        {
            name = value.AsString;
            return true;
        }

        name = string.Empty;
        return false;
    }
}