MES-ETL/MesETL.App/Services/ETL/MySqlDestination.cs

using System.Text;
using System.Text.RegularExpressions;
using MesETL.App.Const;
using MesETL.App.Helpers;
using MesETL.App.Options;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using MySqlConnector;

namespace MesETL.App.Services.ETL;

/// <summary>
/// Mysql导出
/// </summary>
public partial class MySqlDestination : IDisposable, IAsyncDisposable
{
    private readonly Dictionary<string, IList<DataRecord>> _recordCache;
    private readonly MySqlConnection _conn;
    private readonly ILogger _logger;
    private readonly IOptions<DatabaseOutputOptions> _options;
    private readonly ErrorRecorder.OutputErrorRecorder _outputErrorRecorder;
    private readonly ProcessContext _context;

    public MySqlDestination(
        string connStr, 
        ILogger logger, 
        IOptions<DatabaseOutputOptions> options,
        ErrorRecorder.OutputErrorRecorder outputErrorRecorder, 
        ProcessContext context)
    {
        _conn = new MySqlConnection(connStr);
        _conn.Open();
        _recordCache = new Dictionary<string, IList<DataRecord>>();
        _logger = logger;
        _options = options;
        _outputErrorRecorder = outputErrorRecorder;
        _context = context;
    }
    
    public Task WriteRecordAsync(DataRecord record)
    {
        _recordCache.AddOrUpdate(record.TableName, [record], (_, 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(int maxAllowPacket)
    {
        if (_recordCache.Count == 0)
            return;

        var cmd = _conn.CreateCommand();
        cmd.CommandTimeout = 3 * 60;
       
        try
        {
            var excuseList = GetExcuseList(_recordCache, maxAllowPacket).ToList();
            foreach (var insertSql in excuseList)
            {
                cmd.CommandText = insertSql;
                try
                {
                    await cmd.ExecuteNonQueryAsync();
                }
                catch (Exception e)
                {
                    _logger.LogError(e, "插入数据库时发生错误, sql: {Sql}", cmd.CommandText.Omit(1000));
                    _context.AddException(e);
                    var match = MatchTableName().Match(cmd.CommandText);
                    if (match is { Success: true, Groups.Count: > 1 })
                    {
                        var tableName = match.Groups[1].Value;
                        await _outputErrorRecorder.LogErrorSqlAsync(cmd.CommandText, tableName, e);
                    }
                    else await _outputErrorRecorder.LogErrorSqlAsync(cmd.CommandText, e);
                }
            }
            _recordCache.Clear();
        }
        catch (Exception e)
        {
            _logger.LogError(e, "序列化记录时发生错误");
            throw;
        }
        finally
        {
            await cmd.DisposeAsync();
        }
    }
    
    [GeneratedRegex("INSERT INTO `([^`]+)`")]
    private static partial Regex MatchTableName();

    public IEnumerable<string> GetExcuseList(IDictionary<string, IList<DataRecord>> tableRecords,int maxAllowPacket)
    {
        var sb = new StringBuilder("SET AUTOCOMMIT = 1;\n");
        foreach (var (tableName, records) in tableRecords)
        {
            if (records.Count == 0)
                continue;
            
            var recordIdx = 0;
            StartBuild:
            var noCommas = true;
            
            // INSERT INTO ... VALUES >>>
            sb.Append($"INSERT INTO `{tableName}`(");
            for (var i = 0; i < records[0].Headers.Count; i++)
            {
                var header = records[0].Headers[i];
                sb.Append($"`{header}`");
                if (i != records[0].Headers.Count - 1)
                    sb.Append(',');
            }

            sb.Append(") VALUES ");
            
            // ([FIELDS]), >>>
            for (;recordIdx < records.Count; recordIdx++)
            {
                var record = records[recordIdx];
                var recordSb = new StringBuilder();
                recordSb.Append('(');
                for (var fieldIdx = 0; fieldIdx < record.Fields.Count; fieldIdx++)
                {
                    var field = record.Fields[fieldIdx];

                    // 在这里处理特殊列
                    #region HandleFields

                    if (field.Length == 2 && field == ConstVar.MyDumperNull) // MyDumper导出的NULL为'\N'（'\'不是转义字符）
                    {
                        recordSb.Append(ConstVar.Null);
                        goto Escape;
                    }
                    
                    switch (_options.Value.GetColumnType(record.TableName, record.Headers[fieldIdx]))
                    {
                        case ColumnType.Text:
                            if(string.IsNullOrEmpty(field))
                                recordSb.Append("''");
                            else if (field == ConstVar.Null)
                                recordSb.Append(ConstVar.Null);
                            else recordSb.Append($"_utf8mb4 0x{field}");
                            break;
                        case ColumnType.Blob:
                            if (string.IsNullOrEmpty(field))
                                recordSb.Append("''");
                            else if (field == ConstVar.Null)
                                recordSb.Append(ConstVar.Null);
                            else recordSb.Append($"0x{field}");
                            break;
                        case ColumnType.Json:// 生产库没有JSON列，仅用于测试库进行测试
                            if(string.IsNullOrEmpty(field))
                                recordSb.Append("'[]'"); // JObject or JArray?
                            else if (_options.Value.TreatJsonAsHex)
                                recordSb.Append($"_utf8mb4 0x{field}");
                            else recordSb.AppendLine(field);
                            break;
                        case ColumnType.UnDefine:
                        default:
                            recordSb.Append(field);
                            break;
                    }
                    
                    Escape:

                    #endregion
                    if (fieldIdx != record.Fields.Count - 1)
                        recordSb.Append(',');
                }

                recordSb.Append(')');

                // 若字符数量即将大于限制，则返回SQL，清空StringBuilder，保留当前记录的索引值，然后转到StartBuild标签重新开始一轮INSERT
                if (sb.Length + recordSb.Length + 23 > maxAllowPacket)
                {
                    sb.Append(';').AppendLine();
                    sb.Append("SET AUTOCOMMIT = 1;");
                    yield return sb.ToString();
                    sb.Clear();
                    goto StartBuild;
                }

                if (!noCommas)
                    sb.Append(',').AppendLine();
                noCommas = false;
                sb.Append(recordSb); // StringBuilder.Append(StringBuilder)不会分配多余的内存
            }

            sb.Append(';');
            sb.Append("COMMIT;");
            yield return sb.ToString();
            sb.Clear();
        }
    }


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

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