writer_services.go 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575
  1. package services
  2. import (
  3. "bytes"
  4. "fmt"
  5. "io"
  6. "io/ioutil"
  7. "log"
  8. "os"
  9. "path/filepath"
  10. "sort"
  11. "sync"
  12. "time"
  13. "xg_fetl/internal/db_executor"
  14. "xg_fetl/internal/models"
  15. "xg_fetl/module/csv/csvwriter"
  16. )
  17. /*
  18. WriteMain 导出数据库表数据到 CSV 文件,基于分页和文件大小进行分片
  19. 参数:
  20. * executor: 数据库执行器,用于读取数据库表中的数据
  21. * tableNames: 要导出的表名数组
  22. * baseFilePath: 导出文件的基本路径,用于生成不同分片的文件名
  23. * delimiter: CSV 的分隔符
  24. * maxFileSize: 单个 CSV 文件的最大大小
  25. * pageSize: 每次分页读取的行数
  26. */
  27. func WriteMain(
  28. executor db_executor.DBExecutor,
  29. tableNames map[string]models.TableInfo,
  30. baseFilePath string,
  31. delimiter rune,
  32. maxFileSize int64,
  33. pageSize int) error {
  34. {
  35. // --------------
  36. // fmt.Println("--------tableName: 开始导出表数据--------", tableName)
  37. // // 开始分页读取并导出数据
  38. // offset := 0
  39. // for {
  40. // // 从数据库中分页读取数据
  41. // headers, records, err := executor.ReadTableFromDBWithPagination(tableName, offset, pageSize)
  42. // if err != nil {
  43. // log.Fatalf("Failed to read table with pagination: %v", err)
  44. // return err
  45. // }
  46. // fmt.Println("循环")
  47. // // 如果没有更多数据,则结束循环
  48. // if records == nil {
  49. // break
  50. // }
  51. // fmt.Println("baseFilePath: ", baseFilePath)
  52. // csvFilePath := fmt.Sprintf("%s/%s", baseFilePath, tableName)
  53. // // 将数据写入 CSV 文件中
  54. // err = csvwriter.WriteRecordsToCSVWithFileSizeLimit(csvFilePath, headers, records, delimiter, maxFileSize, &fileIndex, &currentFileSize, &currentFilePath)
  55. // if err != nil {
  56. // log.Fatalf("Failed to write records to CSV: %v", err)
  57. // return err
  58. // }
  59. // // 增加偏移量,以获取下一页数据
  60. // offset += pageSize
  61. // }
  62. // fmt.Println("--------tableName 完成--------", tableName)
  63. // }
  64. err := exportTableConcurrently(tableNames, baseFilePath, pageSize, delimiter, &executor, maxFileSize)
  65. if err != nil {
  66. fmt.Println("导出表数据失败:", err)
  67. return err
  68. }
  69. fmt.Println("Table exported successfully with size limit and pagination")
  70. // return nil
  71. }
  72. return nil
  73. }
  74. // 定义用于接收每个表处理结果的结构体
  75. type TableResult struct {
  76. TableName string
  77. Success bool
  78. Error error
  79. GeneratedFiles int
  80. ErrorCount int // 记录该表发生的错误数量
  81. GoroutineCount int // 记录该表使用的协程数量
  82. TotalRows int // 表的总行数
  83. AverageRowSize int64 // 平均行大小
  84. ExportDuration time.Duration // 导出所花费的时间
  85. EstimatedDuration time.Duration // 估算的导出时间
  86. TableSize int64 // 表的容量(字节)
  87. Logs []string // 保存该表的日志信息
  88. }
  89. // 获取当前时间字符串
  90. func currentTime() string {
  91. return time.Now().Format("2006-01-02 15:04:05")
  92. }
  93. // 对表名进行排序,返回排序后的表名列表
  94. func sortedTableNames(tableInfos map[string]models.TableInfo) []string {
  95. tableNames := make([]string, 0, len(tableInfos))
  96. for name := range tableInfos {
  97. tableNames = append(tableNames, name)
  98. }
  99. sort.Strings(tableNames)
  100. return tableNames
  101. }
  102. // 日志消息记录函数
  103. func logMessage(result *TableResult, msg string, logger *log.Logger) {
  104. fmt.Print(msg)
  105. logger.Print(msg)
  106. result.Logs = append(result.Logs, msg)
  107. }
  108. // 估算数据传输速率的函数
  109. func estimateDataTransferRate() (float64, error) {
  110. // 创建临时文件
  111. tmpFile, err := ioutil.TempFile("", "data_transfer_test")
  112. if err != nil {
  113. return 0, err
  114. }
  115. defer os.Remove(tmpFile.Name()) // 确保测试完成后删除临时文件
  116. defer tmpFile.Close()
  117. // 准备测试数据,大小为 50 MB
  118. testDataSize := 50 * 1024 * 1024 // 50 MB
  119. testData := make([]byte, testDataSize)
  120. // 写入测试数据并测量时间
  121. startTime := time.Now()
  122. _, err = tmpFile.Write(testData)
  123. if err != nil {
  124. return 0, err
  125. }
  126. writeDuration := time.Since(startTime)
  127. // 计算写入速率
  128. writeRate := float64(testDataSize) / writeDuration.Seconds()
  129. // 读取测试数据并测量时间
  130. tmpFile.Seek(0, io.SeekStart)
  131. readBuffer := make([]byte, testDataSize)
  132. startTime = time.Now()
  133. _, err = io.ReadFull(tmpFile, readBuffer)
  134. if err != nil && err != io.EOF {
  135. return 0, err
  136. }
  137. readDuration := time.Since(startTime)
  138. // 计算读取速率
  139. readRate := float64(testDataSize) / readDuration.Seconds()
  140. // 取读写速率的平均值作为数据传输速率
  141. dataTransferRate := (writeRate + readRate) / 2
  142. return dataTransferRate, nil
  143. }
  144. // 将总览信息和按表分类的详细日志写入日志文件
  145. func writeLogFile(logFilePath string, resultsMap map[string]TableResult, tableInfos map[string]models.TableInfo) error {
  146. var logBuffer bytes.Buffer
  147. // 构建总览信息
  148. logBuffer.WriteString("========== 导出结果总览 ==========\n")
  149. var successTables []string
  150. var failedTables []TableResult
  151. for _, result := range resultsMap {
  152. if result.Success {
  153. successTables = append(successTables, result.TableName)
  154. } else {
  155. failedTables = append(failedTables, result)
  156. }
  157. }
  158. logBuffer.WriteString(fmt.Sprintf("成功导出的表(%d 个):\n", len(successTables)))
  159. for _, tableName := range successTables {
  160. result := resultsMap[tableName]
  161. avgRowSizeKB := float64(result.AverageRowSize) / 1024.0
  162. tableSizeMB := float64(result.TableSize) / (1024.0 * 1024.0)
  163. exportTime := result.ExportDuration.Seconds()
  164. estimatedTime := result.EstimatedDuration.Seconds()
  165. logBuffer.WriteString(fmt.Sprintf("- %s\n", tableName))
  166. logBuffer.WriteString(fmt.Sprintf(" 总行数:%d\n", result.TotalRows))
  167. logBuffer.WriteString(fmt.Sprintf(" 平均行大小:%.2f KB\n", avgRowSizeKB))
  168. logBuffer.WriteString(fmt.Sprintf(" 表的容量:%.2f MB\n", tableSizeMB))
  169. logBuffer.WriteString(fmt.Sprintf(" 估算导出耗时:%.2f 秒\n", estimatedTime))
  170. logBuffer.WriteString(fmt.Sprintf(" 实际导出耗时:%.2f 秒\n", exportTime))
  171. logBuffer.WriteString(fmt.Sprintf(" 使用协程数:%d\n", result.GoroutineCount))
  172. }
  173. logBuffer.WriteString("\n")
  174. if len(failedTables) > 0 {
  175. logBuffer.WriteString(fmt.Sprintf("导出失败的表(%d 个):\n", len(failedTables)))
  176. for _, failedTable := range failedTables {
  177. logBuffer.WriteString(fmt.Sprintf("- %s\n", failedTable.TableName))
  178. logBuffer.WriteString(fmt.Sprintf(" 错误次数:%d\n", failedTable.ErrorCount))
  179. logBuffer.WriteString(fmt.Sprintf(" 使用协程数:%d\n", failedTable.GoroutineCount))
  180. logBuffer.WriteString(fmt.Sprintf(" 错误信息:%v\n", failedTable.Error))
  181. }
  182. } else {
  183. logBuffer.WriteString("所有表均导出成功。\n")
  184. }
  185. logBuffer.WriteString("\n")
  186. logBuffer.WriteString("========== 详细日志信息 ==========\n")
  187. // 按表名排序,输出每个表的详细日志信息
  188. for _, tableName := range sortedTableNames(tableInfos) {
  189. logBuffer.WriteString(fmt.Sprintf("------------ 表 %s ------------\n", tableName))
  190. result, exists := resultsMap[tableName]
  191. if exists {
  192. for _, msg := range result.Logs {
  193. logBuffer.WriteString(msg)
  194. }
  195. logBuffer.WriteString("\n")
  196. } else {
  197. logBuffer.WriteString(fmt.Sprintf("表 %s 的处理结果不存在。\n", tableName))
  198. }
  199. }
  200. // 将新的日志内容写回日志文件
  201. err := ioutil.WriteFile(logFilePath, logBuffer.Bytes(), 0644)
  202. if err != nil {
  203. return fmt.Errorf("无法写入日志文件: %v", err)
  204. }
  205. return nil
  206. }
  207. /**
  208. * exportTableConcurrently 并发地导出多个表的数据到 CSV 文件。
  209. *
  210. * 参数:
  211. * - tableInfos: 包含表信息的映射,键为表名,值为 TableInfo。
  212. * - baseFilePath: 基础文件路径,用于存储导出的 CSV 文件。
  213. * - pageSize: 分页大小,决定每次从数据库读取的记录数量。
  214. * - delimiter: CSV 文件的分隔符。
  215. * - executor: 数据库执行器,用于执行数据库查询。
  216. * - maxFileSizeMB: 单个 CSV 文件的最大大小(以 MB 为单位)。
  217. *
  218. * 返回值:
  219. * - error: 如果发生错误,返回错误信息。
  220. */
  221. func exportTableConcurrently(
  222. tableInfos map[string]models.TableInfo,
  223. baseFilePath string,
  224. pageSize int,
  225. delimiter rune,
  226. executor *db_executor.DBExecutor,
  227. maxFileSizeMB int64,
  228. ) error {
  229. var (
  230. wg sync.WaitGroup
  231. resultsMap = make(map[string]TableResult)
  232. resultsMutex sync.Mutex // 用于保护 resultsMap 的并发访问
  233. maxConcurrentTableExports = 3
  234. tableSemaphore = make(chan struct{}, maxConcurrentTableExports)
  235. maxConcurrentFileExports = 100
  236. )
  237. // 创建日志文件
  238. logFilePath := filepath.Join(baseFilePath, "export.log")
  239. logFile, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
  240. if err != nil {
  241. return fmt.Errorf("无法打开日志文件: %v", err)
  242. }
  243. defer logFile.Close()
  244. // 创建日志记录器,不添加任何前缀或标志
  245. logger := log.New(logFile, "", 0)
  246. // 估算数据传输速率
  247. dataTransferRate, err := estimateDataTransferRate()
  248. if err != nil {
  249. msg := fmt.Sprintf("%s 无法估算数据传输速率,默认使用 50 MB/s。错误信息:%v\n", currentTime(), err)
  250. fmt.Print(msg)
  251. logger.Print(msg)
  252. dataTransferRate = 50 * 1024 * 1024 // 默认 50 MB/s
  253. } else {
  254. msg := fmt.Sprintf("%s 估算的数据传输速率为:%.2f MB/s\n", currentTime(), dataTransferRate/(1024*1024))
  255. fmt.Print(msg)
  256. logger.Print(msg)
  257. }
  258. // 遍历每个表进行处理,使用协程
  259. for _, tableName := range sortedTableNames(tableInfos) {
  260. tableInfo := tableInfos[tableName]
  261. wg.Add(1)
  262. go func(tableInfo models.TableInfo) {
  263. tableSemaphore <- struct{}{} // 获取表级信号量
  264. defer func() {
  265. <-tableSemaphore // 释放表级信号量
  266. wg.Done()
  267. }()
  268. tableName := tableInfo.Name
  269. totalRecords := tableInfo.Count
  270. // 创建一个 TableResult 来保存该表的结果和日志
  271. result := TableResult{
  272. TableName: tableName,
  273. TotalRows: totalRecords,
  274. GoroutineCount: 1, // 包含当前协程
  275. }
  276. // 创建存储当前表的文件夹
  277. tableDir := filepath.Join(baseFilePath, tableName)
  278. if err := os.MkdirAll(tableDir, 0755); err != nil {
  279. msg := fmt.Sprintf("%s 创建表 %s 文件夹失败: %v\n", currentTime(), tableName, err)
  280. logMessage(&result, msg, logger)
  281. result.Success = false
  282. result.Error = err
  283. result.ErrorCount = 1
  284. resultsMutex.Lock()
  285. resultsMap[tableName] = result
  286. resultsMutex.Unlock()
  287. return
  288. }
  289. // 估算每条记录的平均大小
  290. averageRecordSize, err := estimateAverageRecordSize(*executor, tableName)
  291. if err != nil {
  292. msg := fmt.Sprintf("%s 无法估算表 %s 平均记录大小: %v\n", currentTime(), tableName, err)
  293. logMessage(&result, msg, logger)
  294. result.Success = false
  295. result.Error = err
  296. result.ErrorCount = 1
  297. resultsMutex.Lock()
  298. resultsMap[tableName] = result
  299. resultsMutex.Unlock()
  300. return
  301. }
  302. result.AverageRowSize = averageRecordSize
  303. msg := fmt.Sprintf("%s 表 %s 的平均记录大小: %d 字节\n", currentTime(), tableName, averageRecordSize)
  304. logMessage(&result, msg, logger)
  305. // 计算表的容量
  306. tableSize := averageRecordSize * int64(totalRecords)
  307. result.TableSize = tableSize
  308. tableSizeMB := float64(tableSize) / (1024.0 * 1024.0)
  309. msg = fmt.Sprintf("%s 表 %s 的容量约为:%.2f MB\n", currentTime(), tableName, tableSizeMB)
  310. logMessage(&result, msg, logger)
  311. // 估算导出时间
  312. estimatedDurationSeconds := float64(tableSize) / dataTransferRate
  313. estimatedDuration := time.Duration(estimatedDurationSeconds * float64(time.Second))
  314. result.EstimatedDuration = estimatedDuration
  315. msg = fmt.Sprintf("%s 表 %s 估算导出时间:%.2f 秒\n", currentTime(), tableName, estimatedDurationSeconds)
  316. logMessage(&result, msg, logger)
  317. // 计算每个 CSV 文件应包含的记录数
  318. maxFileSizeBytes := maxFileSizeMB * 1024 * 1024 // 将 MB 转换为字节
  319. recordsPerFile := int(maxFileSizeBytes / averageRecordSize)
  320. if recordsPerFile <= 0 {
  321. recordsPerFile = 1000 // 设置一个最小值,避免 recordsPerFile 为 0
  322. }
  323. msg = fmt.Sprintf("%s 表 %s 的每个文件应包含的记录数: %d\n", currentTime(), tableName, recordsPerFile)
  324. logMessage(&result, msg, logger)
  325. // 计算需要的 CSV 文件数量
  326. totalFiles := (totalRecords + recordsPerFile - 1) / recordsPerFile
  327. msg = fmt.Sprintf("%s 表 %s 需要的 CSV 文件数量: %d\n", currentTime(), tableName, totalFiles)
  328. logMessage(&result, msg, logger)
  329. // 记录导出开始时间
  330. startTime := time.Now()
  331. var fileWg sync.WaitGroup
  332. fileSemaphore := make(chan struct{}, maxConcurrentFileExports)
  333. maxRetries := 3
  334. // 错误计数器
  335. var errorCount int32
  336. // 为每个 CSV 文件启动一个协程
  337. for fileIndex := 0; fileIndex < totalFiles; fileIndex++ {
  338. fileWg.Add(1)
  339. result.GoroutineCount++ // 增加协程计数器
  340. go func(fileIndex int) {
  341. fileSemaphore <- struct{}{} // 获取文件级信号量
  342. defer func() {
  343. <-fileSemaphore // 释放文件级信号量
  344. fileWg.Done()
  345. }()
  346. offset := fileIndex * recordsPerFile
  347. limit := recordsPerFile
  348. if offset+limit > totalRecords {
  349. limit = totalRecords - offset
  350. }
  351. csvFilePath := filepath.Join(tableDir, fmt.Sprintf("%s_%d.csv", tableName, fileIndex+1))
  352. attempt := 0
  353. for {
  354. attempt++
  355. err := exportDataRangeToCSV(*executor, tableName, csvFilePath, delimiter, offset, limit, pageSize)
  356. if err != nil {
  357. if attempt < maxRetries {
  358. msg := fmt.Sprintf("%s 表 %s 导出到文件 %s 时发生错误 (第 %d 次尝试): %v,重试...\n", currentTime(), tableName, csvFilePath, attempt, err)
  359. logMessage(&result, msg, logger)
  360. time.Sleep(2 * time.Second)
  361. continue
  362. } else {
  363. msg := fmt.Sprintf("%s 表 %s 导出到文件 %s 时发生错误 (第 %d 次尝试): %v,已达最大重试次数。\n", currentTime(), tableName, csvFilePath, attempt, err)
  364. logMessage(&result, msg, logger)
  365. // 增加错误计数器
  366. errorCount++
  367. break
  368. }
  369. } else {
  370. msg := fmt.Sprintf("%s 表 %s 的文件 %s 导出完成。\n", currentTime(), tableName, filepath.Base(csvFilePath))
  371. logMessage(&result, msg, logger)
  372. break
  373. }
  374. }
  375. }(fileIndex)
  376. }
  377. // 等待当前表的所有文件导出完成
  378. fileWg.Wait()
  379. // 计算实际导出耗时
  380. exportDuration := time.Since(startTime)
  381. result.ExportDuration = exportDuration
  382. msg = fmt.Sprintf("%s 表 %s 导出完成,共生成 %d 个文件,使用了 %d 个协程。\n", currentTime(), tableName, totalFiles, result.GoroutineCount)
  383. logMessage(&result, msg, logger)
  384. success := errorCount == 0
  385. var finalError error
  386. if !success {
  387. finalError = fmt.Errorf("表 %s 导出过程中发生 %d 个错误", tableName, errorCount)
  388. }
  389. result.Success = success
  390. result.Error = finalError
  391. result.ErrorCount = int(errorCount)
  392. result.GeneratedFiles = totalFiles
  393. resultsMutex.Lock()
  394. resultsMap[tableName] = result
  395. resultsMutex.Unlock()
  396. }(tableInfo)
  397. }
  398. // 等待所有表处理完成
  399. wg.Wait()
  400. // 写入日志文件
  401. err = writeLogFile(logFilePath, resultsMap, tableInfos)
  402. if err != nil {
  403. return fmt.Errorf("无法写入日志文件: %v", err)
  404. }
  405. return nil
  406. }
  407. /**
  408. * estimateAverageRecordSize 估算指定表中每条记录的平均大小(以字节为单位)。
  409. *
  410. * 参数:
  411. * - executor: 数据库执行器,用于执行数据库查询。
  412. * - tableName: 需要估算的表名。
  413. *
  414. * 返回值:
  415. * - int64: 平均每条记录的大小(字节)。
  416. * - error: 如果发生错误,返回错误信息。
  417. */
  418. func estimateAverageRecordSize(executor db_executor.DBExecutor, tableName string) (int64, error) {
  419. // 随机抽取一定数量的记录,计算平均大小
  420. sampleSize := 100 // 抽样的记录数
  421. offset := 0 // 从第 0 条记录开始抽样
  422. // 使用 ReadTableFromDBWithPagination 函数获取样本数据
  423. _, records, err := executor.ReadTableFromDBWithPagination(tableName, offset, sampleSize)
  424. if err != nil {
  425. return 0, fmt.Errorf("failed to query table: %v", err)
  426. }
  427. // 检查是否获取到了数据
  428. if records == nil || len(records) == 0 {
  429. return 0, fmt.Errorf("没有获取到抽样记录,表数量可能为 0")
  430. }
  431. var totalSize int64 = 0 // 总记录大小
  432. var count int64 = 0 // 记录数
  433. // 遍历每一条记录,计算其大小
  434. for _, record := range records {
  435. var recordSize int64 = 0
  436. for _, val := range record {
  437. if val != "" {
  438. recordSize += int64(len(val)) // 累加字段值的长度
  439. }
  440. }
  441. totalSize += recordSize // 累加每条记录的大小
  442. count++
  443. }
  444. if count == 0 {
  445. return 0, fmt.Errorf("没有获取到抽样记录,表数量可能为 0")
  446. }
  447. averageSize := totalSize / count // 计算平均记录大小
  448. return averageSize, nil
  449. }
  450. // exportDataRangeToCSV 导出指定范围的数据到 CSV 文件
  451. // 参数:
  452. // - executor: 数据库执行器
  453. // - tableName: 表名
  454. // - csvFilePath: CSV 文件路径
  455. // - delimiter: CSV 分隔符
  456. // - offset: 数据起始偏移量
  457. // - limit: 数据读取限制数量
  458. // - pageSize: 分页大小
  459. // 返回值:
  460. // - error: 如果发生错误,返回错误信息
  461. func exportDataRangeToCSV(
  462. executor db_executor.DBExecutor,
  463. tableName string,
  464. csvFilePath string,
  465. delimiter rune,
  466. offset int,
  467. limit int,
  468. pageSize int,
  469. ) error {
  470. // 首先确保 CSV 文件存在,创建新文件并写入表头
  471. var headers []string
  472. var headersWritten bool = false
  473. // 计算需要分页的次数
  474. totalPages := (limit + pageSize - 1) / pageSize
  475. for page := 0; page < totalPages; page++ {
  476. pageOffset := offset + page*pageSize
  477. pageLimit := pageSize
  478. if pageOffset+pageLimit > offset+limit {
  479. pageLimit = (offset + limit) - pageOffset
  480. }
  481. // 从数据库读取数据
  482. currentHeaders, records, err := executor.ReadTableFromDBWithPagination(tableName, pageOffset, pageLimit)
  483. if err != nil {
  484. return fmt.Errorf("failed to read data from table %s: %v", tableName, err)
  485. }
  486. // 仅在首次获取到表头时进行赋值
  487. if !headersWritten && currentHeaders != nil {
  488. headers = currentHeaders
  489. headersWritten = true
  490. }
  491. // 判断是否需要写入表头
  492. appendToFile := page > 0 // 第一页时创建新文件,之后的页追加写入
  493. // 写入数据到 CSV 文件
  494. err = csvwriter.WriteRecordsToCSV(csvFilePath, headers, records, delimiter, appendToFile)
  495. if err != nil {
  496. return fmt.Errorf("failed to write data to CSV: %v", err)
  497. }
  498. fmt.Printf("表 %s 文件 %s 已写入 %d/%d 页。\n", tableName, filepath.Base(csvFilePath), page+1, totalPages)
  499. }
  500. return nil
  501. }