monorepo/LCU.API.Internal.Document/OwnerNumberMigration/AttributeMigration.cs

using System.Text.RegularExpressions;
using LCU.API.Internal.DocumentLib.Containers;
using LCU.API.Internal.DocumentLib.Utilities;
using MongoDB.Driver;
using OwnerNumberMigration.ValueParsing;

namespace OwnerNumberMigration;

public class AttributeMigration<T> : IAttributeMigration
{
    private readonly IDocumentFilterFactory _filterFactory;
    private readonly IMongoCollectionFactory _collectionFactory;
    private readonly string _attributeName;
    private readonly IValueParsingStrategy<T> _valueParsingStrategy;
    private readonly string _destinationProperty;
    private readonly IObserver<DocumentInformation> _observer;

    public AttributeMigration
    (
        IMongoCollectionFactory collectionFactory,
        string attributeName,
        IValueParsingStrategy<T> valueParsingStrategy,
        string destinationProperty,
        IObserver<DocumentInformation> observer,
        IDocumentFilterFactory filterFactory
    )
    {
        _attributeName = attributeName;
        _valueParsingStrategy = valueParsingStrategy;
        _destinationProperty = destinationProperty;
        _observer = observer;
        _collectionFactory = collectionFactory;
        _filterFactory = filterFactory;
    }

    public async Task Migrate(CancellationToken cancellationToken)
    {
        if (cancellationToken.IsCancellationRequested) return;

        try
        {
            if (typeof(DocumentInformation).GetProperty(_destinationProperty) is null)
            {
                throw new InvalidOperationException();
            }
                           
            var hasAttribute = _filterFactory.AppliesForMigrationFilterFactory(_attributeName, _destinationProperty);
            var collection = _collectionFactory.GetCollection<DocumentInformation>();
            var documents = await RetryOnTooManyRequests(async () => await collection.FindAsync(hasAttribute, cancellationToken: cancellationToken), cancellationToken);

            var done = false;
            while (!done)
            {
                try
                {
                    while (!cancellationToken.IsCancellationRequested && await documents.MoveNextAsync(cancellationToken))
                    {
                        foreach (var document in documents.Current)
                        {
                            if (cancellationToken.IsCancellationRequested) break;

                            _observer.OnNext(document);

                            var value = _valueParsingStrategy.Parse(document.Attributes[_attributeName]);
                            await RetryOnTooManyRequests(async () => await collection.FindOneAndUpdateAsync
                            (
                                d => d.DocumentId == document.DocumentId,
                                Builders<DocumentInformation>.Update.Set(_destinationProperty, value),
                                cancellationToken: cancellationToken
                            ), cancellationToken);

                            _observer.OnCompleted();
                        }
                    }

                    done = true;
                }
                catch (MongoCursorNotFoundException e)
                {
                    _observer.OnError(e);

                    documents = await RetryOnTooManyRequests(async () => await collection.FindAsync(hasAttribute, cancellationToken: cancellationToken), cancellationToken);
                }
            }
        }
        catch (Exception e)
        {
            _observer.OnError(e);
            throw;
        }
    }

    private async Task<TU> RetryOnTooManyRequests<TU>(Func<Task<TU>> action, CancellationToken cancellationToken)
        where TU : class
    {
        TU? result = null;
        while (result is null)
        {
            try
            {
                result = await action();
            }
            catch (MongoCommandException e)
            {
                if (e.Code == 16500)
                {
                    _observer.OnError(e);

                    var wait = int.Parse(Regex.Match(e.Result.GetValue("errmsg").AsString, "\\bRetryAfterMs=(\\d+)\\b").Groups[1].ValueSpan);
                    await Task.Delay(wait, cancellationToken);
                }
                else
                {
                    throw;
                }
            }
        }

        return result;
    }
}