omq/kernel.go
2023-11-20 20:34:09 +05:00

242 lines
5.3 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package main
import (
"bytes"
"database/sql"
"fmt"
"io"
"log/slog"
"os"
"path/filepath"
"strconv"
"strings"
"sync"
"github.com/dimchansky/utfbom"
go_ora "github.com/sijms/go-ora/v2"
)
type Row struct {
isHeader bool
data []any
}
type Job struct {
configFilePath string
scriptFilePath string
exportFileFormat ExportFormat
encoding Encoding
servers []Server
status string
isFinished bool
}
func (j *Job) init() error {
j.status = "Чтение списка серверов..."
servers, err := loadConfig(j.configFilePath)
if err != nil {
return err
}
j.servers = servers
return nil
}
func (j *Job) launch() error {
j.status = "Чтение файла SQL-скрипта..."
sqlBytes, err := readFileIgnoreBOM(j.scriptFilePath)
if err != nil {
return err
}
branchFieldNum := getBranchFieldNumber(string(sqlBytes))
if branchFieldNum <= 0 {
return fmt.Errorf("Некорректное значение номера поля филиала: %v", branchFieldNum)
}
rowsChan := j.iterateServers(string(sqlBytes), branchFieldNum)
err = j.export(rowsChan)
if err != nil {
return err
}
return nil
}
func (j *Job) export(inputRows chan Row) error {
converter, err := j.exportFileFormat.GetExporter(j.encoding)
if err != nil {
return err
}
fileName := filepath.Base(j.scriptFilePath)
fileName = strings.TrimSuffix(fileName, filepath.Ext(fileName))
outputRows := make(chan []any)
go func() {
defer close(outputRows)
gotHeader := false
rowsCache := make([][]any, 0)
rowCount := -1
for row := range inputRows {
rowCount += 1
if gotHeader {
outputRows <- row.data
continue
}
if row.isHeader {
gotHeader = true
outputRows <- row.data
for _, cachedRow := range rowsCache {
outputRows <- cachedRow
}
} else {
rowsCache = append(rowsCache, row.data) // store data row in cache until got a header row
}
}
}()
err = converter.Convert(fileName, outputRows)
j.status = "ЗАВЕРШЕНО."
j.isFinished = true
if err != nil {
return err
}
return nil
}
func (j *Job) iterateServers(sqlStr string, branchFieldNum int) chan Row {
rowsChan := make(chan Row)
go func() {
defer func() {
close(rowsChan)
}()
wg := new(sync.WaitGroup)
wg.Add(len(j.servers))
for i, server := range j.servers {
i := i
server := server
go func() {
defer wg.Done()
j.servers[i].Status = "Подключение к серверу"
db := sql.OpenDB(go_ora.NewConnector(server.Url))
defer db.Close()
err := db.Ping()
if err != nil {
j.servers[i].Error = err
return
}
j.servers[i].Status = "Выполнение SQL-запроса"
rows, err := db.Query(sqlStr)
if err != nil {
j.servers[i].Error = err
return
}
defer rows.Close()
cols, err := rows.Columns()
if err != nil {
j.servers[i].Error = err
return
}
if i == 0 {
rowsChan <- Row{isHeader: true, data: sliceToAnySlice[string](cols)} // Добавление заголовков
}
rowNum := 0
for rows.Next() {
j.servers[i].Status = fmt.Sprintf("Выгружено %d строк", rowNum)
pointers := make([]any, len(cols))
container := make([]any, len(cols))
for i := range pointers {
pointers[i] = &container[i]
}
err = rows.Scan(pointers...)
if err != nil {
j.servers[i].Error = err
break
}
rowsChan <- Row{isHeader: false, data: append(append(container[:branchFieldNum-1], server.Name), container[branchFieldNum:]...)} // Добавление имени сервера
rowNum += 1
}
j.servers[i].Status = fmt.Sprintf("ЗАВЕРШЕНО: Выгружено %d строк.", rowNum)
}()
}
j.status = "Ожидание завершения работы с серверами"
wg.Wait()
}()
return rowsChan
}
func sliceToAnySlice[T string | any](slice []T) []any {
result := make([]any, len(slice))
for i := range result {
result[i] = slice[i]
}
return result
}
func readFileIgnoreBOM(filePath string) ([]byte, error) {
f, err := os.Open(filePath)
if err != nil {
return nil, err
}
defer f.Close()
b, err := io.ReadAll(utfbom.SkipOnly(f))
if err != nil {
return nil, err
}
return bytes.ReplaceAll(b, []byte("\r"), []byte{}), nil
}
// getBranchFieldNumber осуществляет определение номера колонки для подстановки
// наименования филиала. В первой строке SQL-скрипта должен быть комментарий,
// начинающийся на `// ` и содержащий только номер колонки (нумерация с 1).
func getBranchFieldNumber(sqlStr string) int {
lines := strings.Split(sqlStr, "\n")
if len(lines) == 0 {
return 1
}
line := lines[0]
if !strings.HasPrefix(line, "-- ") {
slog.Warn("Не указан номер колонки для вывода филиала, будет перезаписан первый столбец!")
return 1
}
line = strings.TrimPrefix(line, "-- ")
fieldNum, err := strconv.Atoi(line)
if err != nil {
slog.Warn("Неверно указан номер колонки для вывода филиала, будет перезаписан первый столбец!")
return 1
}
return fieldNum
}