package services import ( "fmt" "io" "io/ioutil" "log" "os" "path/filepath" "sort" "sync" "time" "xg_fetl/internal/db_executor" "xg_fetl/internal/models" "xg_fetl/module/csv/csvwriter" ) /* WriteMain 导出数据库表数据到 CSV 文件,基于分页和文件大小进行分片 参数: * executor: 数据库执行器,用于读取数据库表中的数据 * tableNames: 要导出的表名数组 * baseFilePath: 导出文件的基本路径,用于生成不同分片的文件名 * delimiter: CSV 的分隔符 * maxFileSize: 单个 CSV 文件的最大大小 * pageSize: 每次分页读取的行数 */ func WriteMain( executor db_executor.DBExecutor, tableNames map[string]models.TableInfo, baseFilePath string, delimiter rune, maxFileSize int64, pageSize int) error { { // -------------- // fmt.Println("--------tableName: 开始导出表数据--------", tableName) // // 开始分页读取并导出数据 // offset := 0 // for { // // 从数据库中分页读取数据 // headers, records, err := executor.ReadTableFromDBWithPagination(tableName, offset, pageSize) // if err != nil { // log.Fatalf("Failed to read table with pagination: %v", err) // return err // } // fmt.Println("循环") // // 如果没有更多数据,则结束循环 // if records == nil { // break // } // fmt.Println("baseFilePath: ", baseFilePath) // csvFilePath := fmt.Sprintf("%s/%s", baseFilePath, tableName) // // 将数据写入 CSV 文件中 // err = csvwriter.WriteRecordsToCSVWithFileSizeLimit(csvFilePath, headers, records, delimiter, maxFileSize, &fileIndex, ¤tFileSize, ¤tFilePath) // if err != nil { // log.Fatalf("Failed to write records to CSV: %v", err) // return err // } // // 增加偏移量,以获取下一页数据 // offset += pageSize // } // fmt.Println("--------tableName 完成--------", tableName) // } err := exportTableConcurrently(tableNames, baseFilePath, pageSize, delimiter, &executor, maxFileSize) if err != nil { fmt.Println("导出表数据失败:", err) return err } fmt.Println("Table exported successfully with size limit and pagination") // return nil } return nil } // 定义用于接收每个表处理结果的结构体 type TableResult struct { TableName string Success bool Error error GeneratedFiles int ErrorCount int // 记录该表发生的错误数量 GoroutineCount int // 记录该表使用的协程数量 TotalRows int // 表的总行数 AverageRowSize int64 // 平均行大小 ExportDuration time.Duration // 导出所花费的时间 EstimatedDuration time.Duration // 估算的导出时间 TableSize int64 // 表的容量(字节) Logs []string // 保存该表的日志信息 } // 获取当前时间字符串 func currentTime() string { return time.Now().Format("2006-01-02 15:04:05") } // 对表名进行排序,返回排序后的表名列表 func sortedTableNames(tableInfos map[string]models.TableInfo) []string { tableNames := make([]string, 0, len(tableInfos)) for name := range tableInfos { tableNames = append(tableNames, name) } sort.Strings(tableNames) return tableNames } /** * exportTableConcurrently 并发地导出多个表的数据到 CSV 文件。 * * 参数: * - tableInfos: 包含表信息的映射,键为表名,值为 TableInfo。 * - baseFilePath: 基础文件路径,用于存储导出的 CSV 文件。 * - pageSize: 分页大小,决定每次从数据库读取的记录数量。 * - delimiter: CSV 文件的分隔符。 * - executor: 数据库执行器,用于执行数据库查询。 * - maxFileSizeMB: 单个 CSV 文件的最大大小(以 MB 为单位)。 * * 返回值: * - error: 如果发生错误,返回错误信息。 */ func exportTableConcurrently( tableInfos map[string]models.TableInfo, baseFilePath string, pageSize int, delimiter rune, executor *db_executor.DBExecutor, maxFileSizeMB int64, ) error { var ( wg sync.WaitGroup resultsMap = make(map[string]TableResult) resultsMutex sync.Mutex // 用于保护 resultsMap 的并发访问 maxConcurrentTableExports = 3 tableSemaphore = make(chan struct{}, maxConcurrentTableExports) maxConcurrentFileExports = 100 ) // ! 创建日志文件 logFilePath := filepath.Join(baseFilePath, "export.log") logFile, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) if err != nil { return fmt.Errorf("无法打开日志文件: %v", err) } defer logFile.Close() // 创建日志记录器,不添加任何前缀或标志 logger := log.New(logFile, "", 0) // ! 估算数据传输速率 dataTransferRate, err := estimateDataTransferRate() if err != nil { msg := fmt.Sprintf("%s 无法估算数据传输速率,默认使用 50 MB/s。错误信息:%v\n", currentTime(), err) fmt.Print(msg) logger.Print(msg) dataTransferRate = 50 * 1024 * 1024 // 默认 50 MB/s } else { msg := fmt.Sprintf("%s 估算的数据传输速率为:%.2f MB/s\n", currentTime(), dataTransferRate/(1024*1024)) fmt.Print(msg) logger.Print(msg) } // ! 遍历每个表进行处理,使用协程 for _, tableName := range sortedTableNames(tableInfos) { tableInfo := tableInfos[tableName] wg.Add(1) go func(tableInfo models.TableInfo) { tableSemaphore <- struct{}{} // 获取表级信号量 defer func() { <-tableSemaphore // 释放表级信号量 wg.Done() }() tableName := tableInfo.Name totalRecords := tableInfo.Count // 创建一个 TableResult 来保存该表的结果和日志 result := TableResult{ TableName: tableName, TotalRows: totalRecords, GoroutineCount: 1, // 包含当前协程 } // !创建存储当前表的文件夹 tableDir := filepath.Join(baseFilePath, tableName) if err := os.MkdirAll(tableDir, 0755); err != nil { msg := fmt.Sprintf("%s 创建表 %s 文件夹失败: %v\n", currentTime(), tableName, err) logMessage(&result, msg, logger) result.Success = false result.Error = err result.ErrorCount = 1 resultsMutex.Lock() resultsMap[tableName] = result resultsMutex.Unlock() return } // !估算每条记录的平均大小 averageRecordSize, err := estimateAverageRecordSize(*executor, tableName) if err != nil { msg := fmt.Sprintf("%s 无法估算表 %s 平均记录大小: %v\n", currentTime(), tableName, err) logMessage(&result, msg, logger) result.Success = false result.Error = err result.ErrorCount = 1 resultsMutex.Lock() resultsMap[tableName] = result resultsMutex.Unlock() return } result.AverageRowSize = averageRecordSize msg := fmt.Sprintf("%s 表 %s 的平均记录大小: %d 字节\n", currentTime(), tableName, averageRecordSize) logMessage(&result, msg, logger) // !计算表的容量 tableSize := averageRecordSize * int64(totalRecords) result.TableSize = tableSize tableSizeMB := float64(tableSize) / (1024.0 * 1024.0) msg = fmt.Sprintf("%s 表 %s 的容量约为:%.2f MB\n", currentTime(), tableName, tableSizeMB) logMessage(&result, msg, logger) // !估算导出时间 estimatedDurationSeconds := float64(tableSize) / dataTransferRate estimatedDuration := time.Duration(estimatedDurationSeconds * float64(time.Second)) result.EstimatedDuration = estimatedDuration msg = fmt.Sprintf("%s 表 %s 估算导出时间:%.2f 秒\n", currentTime(), tableName, estimatedDurationSeconds) logMessage(&result, msg, logger) // !计算每个 CSV 文件应包含的记录数 maxFileSizeBytes := maxFileSizeMB * 1024 * 1024 // 将 MB 转换为字节 recordsPerFile := int(maxFileSizeBytes / averageRecordSize) if recordsPerFile <= 0 { recordsPerFile = 1000 // 设置一个最小值,避免 recordsPerFile 为 0 } msg = fmt.Sprintf("%s 表 %s 的每个文件应包含的记录数: %d\n", currentTime(), tableName, recordsPerFile) logMessage(&result, msg, logger) // !计算需要的 CSV 文件数量 totalFiles := (totalRecords + recordsPerFile - 1) / recordsPerFile msg = fmt.Sprintf("%s 表 %s 需要的 CSV 文件数量: %d\n", currentTime(), tableName, totalFiles) logMessage(&result, msg, logger) //! 记录导出开始时间 startTime := time.Now() var fileWg sync.WaitGroup fileSemaphore := make(chan struct{}, maxConcurrentFileExports) maxRetries := 3 // 错误计数器 var errorCount int32 // 为每个 CSV 文件启动一个协程 for fileIndex := 0; fileIndex < totalFiles; fileIndex++ { fileWg.Add(1) result.GoroutineCount++ // 增加协程计数器 go func(fileIndex int) { fileSemaphore <- struct{}{} // 获取文件级信号量 defer func() { <-fileSemaphore // 释放文件级信号量 fileWg.Done() }() offset := fileIndex * recordsPerFile limit := recordsPerFile if offset+limit > totalRecords { limit = totalRecords - offset } csvFilePath := filepath.Join(tableDir, fmt.Sprintf("%s_%d.csv", tableName, fileIndex+1)) attempt := 0 for { attempt++ err := exportDataRangeToCSV(*executor, tableName, csvFilePath, delimiter, offset, limit, pageSize) if err != nil { if attempt < maxRetries { msg := fmt.Sprintf("%s 表 %s 导出到文件 %s 时发生错误 (第 %d 次尝试): %v,重试...\n", currentTime(), tableName, csvFilePath, attempt, err) logMessage(&result, msg, logger) time.Sleep(2 * time.Second) continue } else { msg := fmt.Sprintf("%s 表 %s 导出到文件 %s 时发生错误 (第 %d 次尝试): %v,已达最大重试次数。\n", currentTime(), tableName, csvFilePath, attempt, err) logMessage(&result, msg, logger) // 增加错误计数器 errorCount++ break } } else { msg := fmt.Sprintf("%s 表 %s 的文件 %s 导出完成。\n", currentTime(), tableName, filepath.Base(csvFilePath)) logMessage(&result, msg, logger) break } } }(fileIndex) } // 等待当前表的所有文件导出完成 fileWg.Wait() // 计算实际导出耗时 exportDuration := time.Since(startTime) result.ExportDuration = exportDuration msg = fmt.Sprintf("%s 表 %s 导出完成,共生成 %d 个文件,使用了 %d 个协程。\n", currentTime(), tableName, totalFiles, result.GoroutineCount) logMessage(&result, msg, logger) success := errorCount == 0 var finalError error if !success { finalError = fmt.Errorf("表 %s 导出过程中发生 %d 个错误", tableName, errorCount) } result.Success = success result.Error = finalError result.ErrorCount = int(errorCount) result.GeneratedFiles = totalFiles resultsMutex.Lock() resultsMap[tableName] = result resultsMutex.Unlock() }(tableInfo) } // 等待所有表处理完成 wg.Wait() // 写入日志文件 err = WriteExportReportToLogFile(logFilePath, resultsMap, tableInfos) if err != nil { return fmt.Errorf("无法写入日志文件: %v", err) } return nil } /** * estimateAverageRecordSize 估算指定表中每条记录的平均大小(以字节为单位)。 * * 参数: * - executor: 数据库执行器,用于执行数据库查询。 * - tableName: 需要估算的表名。 * * 返回值: * - int64: 平均每条记录的大小(字节)。 * - error: 如果发生错误,返回错误信息。 */ func estimateAverageRecordSize(executor db_executor.DBExecutor, tableName string) (int64, error) { // 随机抽取一定数量的记录,计算平均大小 sampleSize := 100 // 抽样的记录数 offset := 0 // 从第 0 条记录开始抽样 // 使用 ReadTableFromDBWithPagination 函数获取样本数据 _, records, err := executor.ReadTableFromDBWithPagination(tableName, offset, sampleSize) if err != nil { return 0, fmt.Errorf("failed to query table: %v", err) } // 检查是否获取到了数据 if records == nil || len(records) == 0 { return 0, fmt.Errorf("没有获取到抽样记录,表数量可能为 0") } var totalSize int64 = 0 // 总记录大小 var count int64 = 0 // 记录数 // 遍历每一条记录,计算其大小 for _, record := range records { var recordSize int64 = 0 for _, val := range record { if val != "" { recordSize += int64(len(val)) // 累加字段值的长度 } } totalSize += recordSize // 累加每条记录的大小 count++ } if count == 0 { return 0, fmt.Errorf("没有获取到抽样记录,表数量可能为 0") } averageSize := totalSize / count // 计算平均记录大小 return averageSize, nil } // exportDataRangeToCSV 导出指定范围的数据到 CSV 文件 // 参数: // - executor: 数据库执行器 // - tableName: 表名 // - csvFilePath: CSV 文件路径 // - delimiter: CSV 分隔符 // - offset: 数据起始偏移量 // - limit: 数据读取限制数量 // - pageSize: 分页大小 // 返回值: // - error: 如果发生错误,返回错误信息 func exportDataRangeToCSV( executor db_executor.DBExecutor, tableName string, csvFilePath string, delimiter rune, offset int, limit int, pageSize int, ) error { // !首先确保 CSV 文件存在,创建新文件并写入表头 var headers []string var headersWritten bool = false // !计算需要分页的次数 totalPages := (limit + pageSize - 1) / pageSize for page := 0; page < totalPages; page++ { pageOffset := offset + page*pageSize pageLimit := pageSize if pageOffset+pageLimit > offset+limit { pageLimit = (offset + limit) - pageOffset } // !从数据库读取数据 currentHeaders, records, err := executor.ReadTableFromDBWithPagination(tableName, pageOffset, pageLimit) if err != nil { return fmt.Errorf("failed to read data from table %s: %v", tableName, err) } // !仅在首次获取到表头时进行赋值 if !headersWritten && currentHeaders != nil { headers = currentHeaders headersWritten = true } // !判断是否需要写入表头 appendToFile := page > 0 // 第一页时创建新文件,之后的页追加写入 // !写入数据到 CSV 文件 err = csvwriter.WriteRecordsToCSV(csvFilePath, headers, records, delimiter, appendToFile) if err != nil { return fmt.Errorf("failed to write data to CSV: %v", err) } fmt.Printf("表 %s 文件 %s 已写入 %d/%d 页。\n", tableName, filepath.Base(csvFilePath), page+1, totalPages) } return nil } // 估算数据传输速率的函数 func estimateDataTransferRate() (float64, error) { // 创建临时文件 tmpFile, err := ioutil.TempFile("", "data_transfer_test") if err != nil { return 0, err } defer os.Remove(tmpFile.Name()) // 确保测试完成后删除临时文件 defer tmpFile.Close() // 准备测试数据,大小为 50 MB testDataSize := 50 * 1024 * 1024 // 50 MB testData := make([]byte, testDataSize) // 写入测试数据并测量时间 startTime := time.Now() _, err = tmpFile.Write(testData) if err != nil { return 0, err } writeDuration := time.Since(startTime) // 计算写入速率 writeRate := float64(testDataSize) / writeDuration.Seconds() // 读取测试数据并测量时间 tmpFile.Seek(0, io.SeekStart) readBuffer := make([]byte, testDataSize) startTime = time.Now() _, err = io.ReadFull(tmpFile, readBuffer) if err != nil && err != io.EOF { return 0, err } readDuration := time.Since(startTime) // 计算读取速率 readRate := float64(testDataSize) / readDuration.Seconds() // 取读写速率的平均值作为数据传输速率 dataTransferRate := (writeRate + readRate) / 2 return dataTransferRate, nil }