git.stella-ops.org/src/__Libraries/StellaOps.Provcache.Postgres/PostgresEvidenceChunkRepository.cs

using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
using StellaOps.Determinism;
using StellaOps.Provcache.Entities;

namespace StellaOps.Provcache.Postgres;

/// <summary>
/// PostgreSQL implementation of <see cref="IEvidenceChunkRepository"/>.
/// </summary>
public sealed class PostgresEvidenceChunkRepository : IEvidenceChunkRepository
{
    private readonly ProvcacheDbContext _context;
    private readonly ILogger<PostgresEvidenceChunkRepository> _logger;
    private readonly TimeProvider _timeProvider;
    private readonly IGuidProvider _guidProvider;

    public PostgresEvidenceChunkRepository(
        ProvcacheDbContext context,
        ILogger<PostgresEvidenceChunkRepository> logger,
        TimeProvider? timeProvider = null,
        IGuidProvider? guidProvider = null)
    {
        _context = context ?? throw new ArgumentNullException(nameof(context));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        _timeProvider = timeProvider ?? TimeProvider.System;
        _guidProvider = guidProvider ?? SystemGuidProvider.Instance;
    }

    /// <inheritdoc />
    public async Task<IReadOnlyList<EvidenceChunk>> GetChunksAsync(
        string proofRoot,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);

        var entities = await _context.EvidenceChunks
            .Where(e => e.ProofRoot == proofRoot)
            .OrderBy(e => e.ChunkIndex)
            .AsNoTracking()
            .ToListAsync(cancellationToken)
            .ConfigureAwait(false);

        _logger.LogDebug("Retrieved {Count} chunks for proof root {ProofRoot}", entities.Count, proofRoot);
        return entities.Select(MapToModel).ToList();
    }

    /// <inheritdoc />
    public async Task<EvidenceChunk?> GetChunkAsync(
        string proofRoot,
        int chunkIndex,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);

        var entity = await _context.EvidenceChunks
            .Where(e => e.ProofRoot == proofRoot && e.ChunkIndex == chunkIndex)
            .AsNoTracking()
            .FirstOrDefaultAsync(cancellationToken)
            .ConfigureAwait(false);

        return entity is null ? null : MapToModel(entity);
    }

    /// <inheritdoc />
    public async Task<IReadOnlyList<EvidenceChunk>> GetChunkRangeAsync(
        string proofRoot,
        int startIndex,
        int count,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);

        if (startIndex < 0)
        {
            throw new ArgumentOutOfRangeException(nameof(startIndex), "Start index must be non-negative.");
        }

        if (count <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(count), "Count must be positive.");
        }

        var entities = await _context.EvidenceChunks
            .Where(e => e.ProofRoot == proofRoot && e.ChunkIndex >= startIndex)
            .OrderBy(e => e.ChunkIndex)
            .Take(count)
            .AsNoTracking()
            .ToListAsync(cancellationToken)
            .ConfigureAwait(false);

        return entities.Select(MapToModel).ToList();
    }

    /// <inheritdoc />
    public async Task<ChunkManifest?> GetManifestAsync(
        string proofRoot,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);

        // Get metadata without loading blobs
        var chunks = await _context.EvidenceChunks
            .Where(e => e.ProofRoot == proofRoot)
            .OrderBy(e => e.ChunkIndex)
            .Select(e => new
            {
                e.ChunkId,
                e.ChunkIndex,
                e.ChunkHash,
                e.BlobSize,
                e.ContentType,
                e.CreatedAt
            })
            .AsNoTracking()
            .ToListAsync(cancellationToken)
            .ConfigureAwait(false);

        if (chunks.Count == 0)
        {
            return null;
        }

        var metadata = chunks
            .Select(c => new ChunkMetadata
            {
                ChunkId = c.ChunkId,
                Index = c.ChunkIndex,
                Hash = c.ChunkHash,
                Size = c.BlobSize,
                ContentType = c.ContentType
            })
            .ToList();

        return new ChunkManifest
        {
            ProofRoot = proofRoot,
            TotalChunks = chunks.Count,
            TotalSize = chunks.Sum(c => (long)c.BlobSize),
            Chunks = metadata,
            GeneratedAt = _timeProvider.GetUtcNow()
        };
    }

    /// <inheritdoc />
    public async Task StoreChunksAsync(
        string proofRoot,
        IEnumerable<EvidenceChunk> chunks,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);
        ArgumentNullException.ThrowIfNull(chunks);

        var chunkList = chunks.ToList();

        if (chunkList.Count == 0)
        {
            _logger.LogDebug("No chunks to store for proof root {ProofRoot}", proofRoot);
            return;
        }

        // Update proof root in chunks if not set
        var entities = chunkList.Select(c => MapToEntity(c, proofRoot)).ToList();

        _context.EvidenceChunks.AddRange(entities);
        await _context.SaveChangesAsync(cancellationToken).ConfigureAwait(false);

        _logger.LogDebug("Stored {Count} chunks for proof root {ProofRoot}", chunkList.Count, proofRoot);
    }

    /// <inheritdoc />
    public async Task<int> DeleteChunksAsync(
        string proofRoot,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);

        var deleted = await _context.EvidenceChunks
            .Where(e => e.ProofRoot == proofRoot)
            .ExecuteDeleteAsync(cancellationToken)
            .ConfigureAwait(false);

        _logger.LogDebug("Deleted {Count} chunks for proof root {ProofRoot}", deleted, proofRoot);
        return deleted;
    }

    /// <inheritdoc />
    public async Task<int> GetChunkCountAsync(
        string proofRoot,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);

        return await _context.EvidenceChunks
            .CountAsync(e => e.ProofRoot == proofRoot, cancellationToken)
            .ConfigureAwait(false);
    }

    /// <inheritdoc />
    public async Task<long> GetTotalSizeAsync(
        string proofRoot,
        CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proofRoot);

        return await _context.EvidenceChunks
            .Where(e => e.ProofRoot == proofRoot)
            .SumAsync(e => (long)e.BlobSize, cancellationToken)
            .ConfigureAwait(false);
    }

    /// <summary>
    /// Gets total storage across all proof roots.
    /// </summary>
    public async Task<long> GetTotalStorageAsync(CancellationToken cancellationToken = default)
    {
        return await _context.EvidenceChunks
            .SumAsync(e => (long)e.BlobSize, cancellationToken)
            .ConfigureAwait(false);
    }

    /// <summary>
    /// Prunes chunks older than the specified date.
    /// </summary>
    public async Task<int> PruneOldChunksAsync(
        DateTimeOffset olderThan,
        CancellationToken cancellationToken = default)
    {
        return await _context.EvidenceChunks
            .Where(e => e.CreatedAt < olderThan)
            .ExecuteDeleteAsync(cancellationToken)
            .ConfigureAwait(false);
    }

    private static EvidenceChunk MapToModel(ProvcacheEvidenceChunkEntity entity)
    {
        return new EvidenceChunk
        {
            ChunkId = entity.ChunkId,
            ProofRoot = entity.ProofRoot,
            ChunkIndex = entity.ChunkIndex,
            ChunkHash = entity.ChunkHash,
            Blob = entity.Blob,
            BlobSize = entity.BlobSize,
            ContentType = entity.ContentType,
            CreatedAt = entity.CreatedAt
        };
    }

    private ProvcacheEvidenceChunkEntity MapToEntity(EvidenceChunk chunk, string proofRoot)
    {
        return new ProvcacheEvidenceChunkEntity
        {
            ChunkId = chunk.ChunkId == Guid.Empty ? _guidProvider.NewGuid() : chunk.ChunkId,
            ProofRoot = proofRoot,
            ChunkIndex = chunk.ChunkIndex,
            ChunkHash = chunk.ChunkHash,
            Blob = chunk.Blob,
            BlobSize = chunk.BlobSize,
            ContentType = chunk.ContentType,
            CreatedAt = chunk.CreatedAt
        };
    }
}