MES-ETL/ConsoleApp2/Services/MySqlDestination.cs

using System.Text;
using ConsoleApp2.Helpers;
using Microsoft.Extensions.Logging;
using MySqlConnector;

namespace ConsoleApp2.Services;

/// <summary>
/// Mysql导出
/// </summary>
public class MySqlDestination : IDisposable, IAsyncDisposable
{
    private readonly Dictionary<string, IList<DataRecord>> _recordCache;
    private readonly MySqlConnection _conn;
    private readonly ILogger _logger;
    private readonly bool _prettyOutput;

    public MySqlDestination(string connStr, ILogger logger, bool prettyOutput = false)
    {
        _conn = new MySqlConnection(connStr);
        _conn.Open();
        _recordCache = new Dictionary<string, IList<DataRecord>>();
        _logger = logger;
        _prettyOutput = prettyOutput;
    }

    public Task WriteRecordAsync(DataRecord record)
    {
        _recordCache.AddOrUpdate(record.TableName, [record], (key, value) =>
        {
            value.Add(record);
            return value;
        });
        return Task.CompletedTask;
    }

    public async Task WriteRecordsAsync(IEnumerable<DataRecord> records)
    {
        foreach (var record in records)
        {
            await WriteRecordAsync(record);
        }
    }

    public async Task FlushAsync()
    {
        if (_recordCache.Count == 0)
            return;
        
        var cmd = _conn.CreateCommand();
        cmd.CommandText = SerializeRecords(_recordCache, _prettyOutput);
        try
        {
            await cmd.ExecuteNonQueryAsync();
            _recordCache.Clear();
        }
        catch (Exception e)
        {
            _logger.LogCritical(e, "Error when flushing records, sql: {Sql}", cmd.CommandText.Omit(1000));
            throw;
        }
        finally
        {
            await cmd.DisposeAsync();
        }
    }

    public static string SerializeRecords(IDictionary<string, IList<DataRecord>> tableRecords,
        bool prettyOutput = false)
    {
        var sb = new StringBuilder();

        foreach (var (tableName, records) in tableRecords)
        {
            if (records.Count == 0)
                continue;
            sb.Append($"INSERT INTO `{tableName}`(");
            for (var i = 0; i < records[0].Headers.Length; i++)
            {
                var header = records[0].Headers[i];
                sb.Append($"`{header}`");
                if (i != records[0].Headers.Length - 1)
                    sb.Append(',');
            }

            sb.Append(") VALUES ");
            if (prettyOutput)
                sb.AppendLine();

            for (var i = 0; i < records.Count; i++)
            {
                var record = records[i];
                sb.Append('(');
                for (var j = 0; j < record.Fields.Length; j++)
                {
                    var field = record.Fields[j];
                    sb.Append(field);
                    if (j != record.Fields.Length - 1)
                        sb.Append(',');
                }

                sb.Append(')');

                if (i != records.Count - 1) // not last field
                    sb.Append(',');
                if (prettyOutput) sb.AppendLine();
            }

            sb.AppendLine(";");
        }

        return sb.ToString();
    }


    public void Dispose()
    {
        _conn.Close();
        _conn.Dispose();
    }

    public async ValueTask DisposeAsync()
    {
        await _conn.CloseAsync();
        await _conn.DisposeAsync();
    }
}