CBDDC/samples/ZB.MOM.WW.CBDDC.Sample.Console/SampleDocumentStore.cs

using System.Text.Json;
using Microsoft.Extensions.Logging;
using ZB.MOM.WW.CBDDC.Core.Network;
using ZB.MOM.WW.CBDDC.Core.Storage;
using ZB.MOM.WW.CBDDC.Core.Sync;
using ZB.MOM.WW.CBDDC.Persistence.Surreal;

namespace ZB.MOM.WW.CBDDC.Sample.Console;

/// <summary>
///     Surreal-backed document store for the sample app.
/// </summary>
public class SampleDocumentStore : SurrealDocumentStore<SampleDbContext>
{
    private const string UsersCollection = "Users";
    private const string TodoListsCollection = "TodoLists";

    public SampleDocumentStore(
        SampleDbContext context,
        IPeerNodeConfigurationProvider configProvider,
        IVectorClockService vectorClockService,
        ILogger<SampleDocumentStore>? logger = null)
        : base(
            context,
            context.SurrealEmbeddedClient,
            context.SchemaInitializer,
            configProvider,
            vectorClockService,
            new LastWriteWinsConflictResolver(),
            null,
            null,
            logger)
    {
        WatchCollection(UsersCollection, context.Users, u => u.Id);
        WatchCollection(TodoListsCollection, context.TodoLists, t => t.Id);
    }

    protected override async Task ApplyContentToEntityAsync(
        string collection,
        string key,
        JsonElement content,
        CancellationToken cancellationToken)
    {
        await UpsertEntityAsync(collection, key, content, cancellationToken);
    }

    protected override async Task ApplyContentToEntitiesBatchAsync(
        IEnumerable<(string Collection, string Key, JsonElement Content)> documents,
        CancellationToken cancellationToken)
    {
        foreach ((string collection, string key, var content) in documents)
            await UpsertEntityAsync(collection, key, content, cancellationToken);
    }

    protected override async Task<JsonElement?> GetEntityAsJsonAsync(
        string collection,
        string key,
        CancellationToken cancellationToken)
    {
        return collection switch
        {
            UsersCollection => SerializeEntity(await _context.Users.FindByIdAsync(key, cancellationToken)),
            TodoListsCollection => SerializeEntity(await _context.TodoLists.FindByIdAsync(key, cancellationToken)),
            _ => null
        };
    }

    protected override async Task RemoveEntityAsync(
        string collection,
        string key,
        CancellationToken cancellationToken)
    {
        await DeleteEntityAsync(collection, key, cancellationToken);
    }

    protected override async Task RemoveEntitiesBatchAsync(
        IEnumerable<(string Collection, string Key)> documents,
        CancellationToken cancellationToken)
    {
        foreach ((string collection, string key) in documents)
            await DeleteEntityAsync(collection, key, cancellationToken);
    }

    protected override async Task<IEnumerable<(string Key, JsonElement Content)>> GetAllEntitiesAsJsonAsync(
        string collection,
        CancellationToken cancellationToken)
    {
        return collection switch
        {
            UsersCollection => (await _context.Users.FindAllAsync(cancellationToken))
                .Select(u => (u.Id, SerializeEntity(u)!.Value))
                .ToList(),
            TodoListsCollection => (await _context.TodoLists.FindAllAsync(cancellationToken))
                .Select(t => (t.Id, SerializeEntity(t)!.Value))
                .ToList(),
            _ => []
        };
    }

    private async Task UpsertEntityAsync(
        string collection,
        string key,
        JsonElement content,
        CancellationToken cancellationToken)
    {
        switch (collection)
        {
            case UsersCollection:
                var user = content.Deserialize<User>() ?? throw new InvalidOperationException("Failed to deserialize user.");
                user.Id = key;
                if (await _context.Users.FindByIdAsync(key, cancellationToken) == null)
                    await _context.Users.InsertAsync(user, cancellationToken);
                else
                    await _context.Users.UpdateAsync(user, cancellationToken);
                break;

            case TodoListsCollection:
                var todo = content.Deserialize<TodoList>() ??
                           throw new InvalidOperationException("Failed to deserialize todo list.");
                todo.Id = key;
                if (await _context.TodoLists.FindByIdAsync(key, cancellationToken) == null)
                    await _context.TodoLists.InsertAsync(todo, cancellationToken);
                else
                    await _context.TodoLists.UpdateAsync(todo, cancellationToken);
                break;

            default:
                throw new NotSupportedException($"Collection '{collection}' is not supported for sync.");
        }
    }

    private async Task DeleteEntityAsync(string collection, string key, CancellationToken cancellationToken)
    {
        switch (collection)
        {
            case UsersCollection:
                await _context.Users.DeleteAsync(key, cancellationToken);
                break;
            case TodoListsCollection:
                await _context.TodoLists.DeleteAsync(key, cancellationToken);
                break;
            default:
                _logger.LogWarning("Attempted to remove entity from unsupported collection: {Collection}", collection);
                break;
        }
    }

    private static JsonElement? SerializeEntity<T>(T? entity) where T : class
    {
        return entity == null ? null : JsonSerializer.SerializeToElement(entity);
    }
}