git.stella-ops.org/src/StellaOps.Scanner.Storage/Mongo/MongoBootstrapper.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using MongoDB.Bson;
using MongoDB.Driver;
using StellaOps.Scanner.Storage.Catalog;
using StellaOps.Scanner.Storage.Migrations;

namespace StellaOps.Scanner.Storage.Mongo;

public sealed class MongoBootstrapper
{
    private readonly IMongoDatabase _database;
    private readonly ScannerStorageOptions _options;
    private readonly ILogger<MongoBootstrapper> _logger;
    private readonly MongoMigrationRunner _migrationRunner;

    public MongoBootstrapper(
        IMongoDatabase database,
        IOptions<ScannerStorageOptions> options,
        ILogger<MongoBootstrapper> logger,
        MongoMigrationRunner migrationRunner)
    {
        _database = database ?? throw new ArgumentNullException(nameof(database));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        _migrationRunner = migrationRunner ?? throw new ArgumentNullException(nameof(migrationRunner));
        _options = (options ?? throw new ArgumentNullException(nameof(options))).Value;
    }

    public async Task InitializeAsync(CancellationToken cancellationToken)
    {
        _options.EnsureValid();

        await EnsureCollectionsAsync(cancellationToken).ConfigureAwait(false);
        await EnsureIndexesAsync(cancellationToken).ConfigureAwait(false);
        await _migrationRunner.RunAsync(cancellationToken).ConfigureAwait(false);
    }

    private async Task EnsureCollectionsAsync(CancellationToken cancellationToken)
    {
        var targetCollections = new[]
        {
            ScannerStorageDefaults.Collections.Artifacts,
            ScannerStorageDefaults.Collections.Images,
            ScannerStorageDefaults.Collections.Layers,
            ScannerStorageDefaults.Collections.Links,
            ScannerStorageDefaults.Collections.Jobs,
            ScannerStorageDefaults.Collections.LifecycleRules,
            ScannerStorageDefaults.Collections.RuntimeEvents,
            ScannerStorageDefaults.Collections.Migrations,
        };

        using var cursor = await _database.ListCollectionNamesAsync(cancellationToken: cancellationToken).ConfigureAwait(false);
        var existing = await cursor.ToListAsync(cancellationToken).ConfigureAwait(false);

        foreach (var name in targetCollections)
        {
            if (existing.Contains(name, StringComparer.Ordinal))
            {
                continue;
            }

            _logger.LogInformation("Creating Mongo collection {Collection}", name);
            await _database.CreateCollectionAsync(name, cancellationToken: cancellationToken).ConfigureAwait(false);
        }
    }

    private async Task EnsureIndexesAsync(CancellationToken cancellationToken)
    {
        await EnsureArtifactIndexesAsync(cancellationToken).ConfigureAwait(false);
        await EnsureImageIndexesAsync(cancellationToken).ConfigureAwait(false);
        await EnsureLayerIndexesAsync(cancellationToken).ConfigureAwait(false);
        await EnsureLinkIndexesAsync(cancellationToken).ConfigureAwait(false);
        await EnsureJobIndexesAsync(cancellationToken).ConfigureAwait(false);
        await EnsureLifecycleIndexesAsync(cancellationToken).ConfigureAwait(false);
        await EnsureRuntimeEventIndexesAsync(cancellationToken).ConfigureAwait(false);
    }

    private Task EnsureArtifactIndexesAsync(CancellationToken cancellationToken)
    {
        var collection = _database.GetCollection<ArtifactDocument>(ScannerStorageDefaults.Collections.Artifacts);
        var models = new List<CreateIndexModel<ArtifactDocument>>
        {
            new(
                Builders<ArtifactDocument>.IndexKeys
                    .Ascending(x => x.Type)
                    .Ascending(x => x.BytesSha256),
                new CreateIndexOptions { Name = "artifact_type_bytesSha256", Unique = true }),
            new(
                Builders<ArtifactDocument>.IndexKeys.Ascending(x => x.RefCount),
                new CreateIndexOptions { Name = "artifact_refCount" }),
            new(
                Builders<ArtifactDocument>.IndexKeys.Ascending(x => x.CreatedAtUtc),
                new CreateIndexOptions { Name = "artifact_createdAt" })
        };

        return collection.Indexes.CreateManyAsync(models, cancellationToken);
    }

    private Task EnsureImageIndexesAsync(CancellationToken cancellationToken)
    {
        var collection = _database.GetCollection<ImageDocument>(ScannerStorageDefaults.Collections.Images);
        var models = new List<CreateIndexModel<ImageDocument>>
        {
            new(
                Builders<ImageDocument>.IndexKeys
                    .Ascending(x => x.Repository)
                    .Ascending(x => x.Tag),
                new CreateIndexOptions { Name = "image_repo_tag" }),
            new(
                Builders<ImageDocument>.IndexKeys.Ascending(x => x.LastSeenAtUtc),
                new CreateIndexOptions { Name = "image_lastSeen" })
        };

        return collection.Indexes.CreateManyAsync(models, cancellationToken);
    }

    private Task EnsureLayerIndexesAsync(CancellationToken cancellationToken)
    {
        var collection = _database.GetCollection<LayerDocument>(ScannerStorageDefaults.Collections.Layers);
        var models = new List<CreateIndexModel<LayerDocument>>
        {
            new(
                Builders<LayerDocument>.IndexKeys.Ascending(x => x.LastSeenAtUtc),
                new CreateIndexOptions { Name = "layer_lastSeen" })
        };

        return collection.Indexes.CreateManyAsync(models, cancellationToken);
    }

    private Task EnsureLinkIndexesAsync(CancellationToken cancellationToken)
    {
        var collection = _database.GetCollection<LinkDocument>(ScannerStorageDefaults.Collections.Links);
        var models = new List<CreateIndexModel<LinkDocument>>
        {
            new(
                Builders<LinkDocument>.IndexKeys
                    .Ascending(x => x.FromType)
                    .Ascending(x => x.FromDigest)
                    .Ascending(x => x.ArtifactId),
                new CreateIndexOptions { Name = "link_from_artifact", Unique = true })
        };

        return collection.Indexes.CreateManyAsync(models, cancellationToken);
    }

    private Task EnsureJobIndexesAsync(CancellationToken cancellationToken)
    {
        var collection = _database.GetCollection<JobDocument>(ScannerStorageDefaults.Collections.Jobs);
        var models = new List<CreateIndexModel<JobDocument>>
        {
            new(
                Builders<JobDocument>.IndexKeys
                    .Ascending(x => x.State)
                    .Ascending(x => x.CreatedAtUtc),
                new CreateIndexOptions { Name = "job_state_createdAt" }),
            new(
                Builders<JobDocument>.IndexKeys.Ascending(x => x.HeartbeatAtUtc),
                new CreateIndexOptions { Name = "job_heartbeat" })
        };

        return collection.Indexes.CreateManyAsync(models, cancellationToken);
    }

    private Task EnsureLifecycleIndexesAsync(CancellationToken cancellationToken)
    {
        var collection = _database.GetCollection<LifecycleRuleDocument>(ScannerStorageDefaults.Collections.LifecycleRules);
        var expiresIndex = new CreateIndexModel<LifecycleRuleDocument>(
            Builders<LifecycleRuleDocument>.IndexKeys.Ascending(x => x.ExpiresAtUtc),
            new CreateIndexOptions
            {
                Name = "lifecycle_expiresAt",
                ExpireAfter = TimeSpan.Zero,
            });

        var artifactIndex = new CreateIndexModel<LifecycleRuleDocument>(
            Builders<LifecycleRuleDocument>.IndexKeys
                .Ascending(x => x.ArtifactId)
                .Ascending(x => x.Class),
            new CreateIndexOptions { Name = "lifecycle_artifact_class", Unique = true });

        return collection.Indexes.CreateManyAsync(new[] { expiresIndex, artifactIndex }, cancellationToken);
    }

    private Task EnsureRuntimeEventIndexesAsync(CancellationToken cancellationToken)
    {
        var collection = _database.GetCollection<RuntimeEventDocument>(ScannerStorageDefaults.Collections.RuntimeEvents);
        var models = new List<CreateIndexModel<RuntimeEventDocument>>
        {
            new(
                Builders<RuntimeEventDocument>.IndexKeys.Ascending(x => x.EventId),
                new CreateIndexOptions { Name = "runtime_event_eventId", Unique = true }),
            new(
                Builders<RuntimeEventDocument>.IndexKeys
                    .Ascending(x => x.Tenant)
                    .Ascending(x => x.Node)
                    .Ascending(x => x.When),
                new CreateIndexOptions { Name = "runtime_event_tenant_node_when" }),
            new(
                Builders<RuntimeEventDocument>.IndexKeys
                    .Ascending(x => x.ImageDigest)
                    .Descending(x => x.When),
                new CreateIndexOptions { Name = "runtime_event_imageDigest_when" }),
            new(
                Builders<RuntimeEventDocument>.IndexKeys
                    .Ascending(x => x.BuildId)
                    .Descending(x => x.When),
                new CreateIndexOptions { Name = "runtime_event_buildId_when" }),
            new(
                Builders<RuntimeEventDocument>.IndexKeys.Ascending(x => x.ExpiresAt),
                new CreateIndexOptions
                {
                    Name = "runtime_event_expiresAt",
                    ExpireAfter = TimeSpan.Zero
                })
        };

        return collection.Indexes.CreateManyAsync(models, cancellationToken);
    }
}