cyi-box/backend/internal/service/dbm_routes.go

package service

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"github.com/cyi-cc/fun"
	"github.com/xuri/excelize/v2"

	"github.com/cyi-cc/cyi-box/backend/internal/config"
	"github.com/cyi-cc/cyi-box/backend/internal/dbx"
	"github.com/cyi-cc/cyi-box/backend/internal/store"
)

const exportRowCap = 50000

// DbExportHandler POST /dbm/export 表单字段：id, db, sql, format(json|csv|xlsx)
func DbExportHandler(st *store.Store, cfg *config.Config) fun.RouteHandler {
	return func(rc *fun.RouteCtx) error {
		ctx := rc.RequestCtx
		if err := routeAdmin(st, ctx); err != nil {
			ctx.SetBodyString(`{"status":2,"msg":"未授权"}`)
			return nil
		}
		id, _ := strconv.ParseInt(string(ctx.QueryArgs().Peek("id")), 10, 64)
		if id == 0 {
			id, _ = strconv.ParseInt(rc.Data["id"], 10, 64)
		}
		dbName := rc.Data["db"]
		sqlText := strings.TrimSpace(rc.Data["sql"])
		format := rc.Data["format"]
		if format == "" {
			format = "json"
		}
		if id == 0 || sqlText == "" {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"参数不完整"}`)
			return nil
		}
		c, _, err := openConn(st, cfg.VaultKey, id, dbName)
		if err != nil {
			ctx.SetStatusCode(502)
			ctx.SetBodyString("连接失败：" + err.Error())
			return nil
		}
		defer c.Close()

		qctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		res, err := dbx.Query(qctx, c, sqlText, exportRowCap)
		if err != nil {
			ctx.SetStatusCode(500)
			ctx.SetBodyString("执行失败：" + err.Error())
			return nil
		}
		if !res.IsSelect {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"非查询语句没有可导出的结果集"}`)
			return nil
		}

		var body []byte
		var ext, mime string
		switch format {
		case "xlsx":
			body, err = toXLSX(res)
			ext, mime = "xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
		case "csv":
			body, err = toCSV(res)
			ext, mime = "csv", "text/csv; charset=utf-8"
		default:
			ext, mime = "json", "application/json"
			body, err = toJSONRows(res)
		}
		if err != nil {
			ctx.SetStatusCode(500)
			ctx.SetBodyString("导出失败：" + err.Error())
			return nil
		}
		fname := fmt.Sprintf("export_%d.%s", time.Now().Unix(), ext)
		ctx.Response.Header.Set("Content-Disposition", "attachment; filename=\""+fname+"\"")
		ctx.SetContentType(mime)
		ctx.SetBody(body)
		return nil
	}
}

func toJSONRows(res *dbx.Result) ([]byte, error) {
	list := make([]map[string]any, 0, len(res.Rows))
	for _, row := range res.Rows {
		obj := map[string]any{}
		for i, col := range res.Columns {
			obj[col] = row[i]
		}
		list = append(list, obj)
	}
	return json.MarshalIndent(list, "", "  ")
}

func toCSV(res *dbx.Result) ([]byte, error) {
	var buf bytes.Buffer
	buf.Write([]byte{0xEF, 0xBB, 0xBF}) // BOM，Excel 打开不乱码
	w := csv.NewWriter(&buf)
	if err := w.Write(res.Columns); err != nil {
		return nil, err
	}
	for _, row := range res.Rows {
		rec := make([]string, len(row))
		for i, v := range row {
			rec[i] = cellText(v)
		}
		if err := w.Write(rec); err != nil {
			return nil, err
		}
	}
	w.Flush()
	return buf.Bytes(), w.Error()
}

func toXLSX(res *dbx.Result) ([]byte, error) {
	f := excelize.NewFile()
	defer f.Close()
	sheet := "Sheet1"
	for i, c := range res.Columns {
		cell, _ := excelize.CoordinatesToCellName(i+1, 1)
		if err := f.SetCellValue(sheet, cell, c); err != nil {
			return nil, err
		}
	}
	for ri, row := range res.Rows {
		for ci, v := range row {
			cell, _ := excelize.CoordinatesToCellName(ci+1, ri+2)
			if v == nil {
				continue
			}
			if err := f.SetCellValue(sheet, cell, v); err != nil {
				return nil, err
			}
		}
	}
	buf, err := f.WriteToBuffer()
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func cellText(v any) string {
	switch t := v.(type) {
	case nil:
		return ""
	case string:
		return t
	case float64:
		return strconv.FormatFloat(t, 'f', -1, 64)
	case bool:
		if t {
			return "1"
		}
		return "0"
	default:
		return fmt.Sprint(t)
	}
}

// DbImportHandler POST /dbm/import multipart：file + 字段 id/db/table；
// 按扩展名解析 json(对象数组) / csv / xlsx（首行为表头），列名与表列交集后批量插入
func DbImportHandler(st *store.Store, cfg *config.Config) fun.RouteHandler {
	return func(rc *fun.RouteCtx) error {
		ctx := rc.RequestCtx
		if err := routeAdmin(st, ctx); err != nil {
			ctx.SetBodyString(`{"status":2,"msg":"未授权"}`)
			return nil
		}
		id, _ := strconv.ParseInt(string(ctx.QueryArgs().Peek("id")), 10, 64)
		dbName := string(ctx.QueryArgs().Peek("db"))
		table := strings.TrimSpace(string(ctx.QueryArgs().Peek("table")))
		if id == 0 || table == "" {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"参数不完整"}`)
			return nil
		}
		fh, err := ctx.FormFile("file")
		if err != nil {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"缺少文件"}`)
			return nil
		}
		src, err := fh.Open()
		if err != nil {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"读取上传失败"}`)
			return nil
		}
		defer src.Close()
		raw, err := io.ReadAll(io.LimitReader(src, 64<<20))
		if err != nil {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"读取文件失败"}`)
			return nil
		}
		rows, err := parseImport(strings.ToLower(fh.Filename), raw)
		if err != nil {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"` + err.Error() + `"}`)
			return nil
		}
		if len(rows) == 0 {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"文件里没有数据行"}`)
			return nil
		}

		c, e, err := openConn(st, cfg.VaultKey, id, dbName)
		if err != nil {
			ctx.SetStatusCode(502)
			ctx.SetBodyString(`{"status":2,"msg":"连接失败"}`)
			return nil
		}
		defer c.Close()

		// 表名校验：必须在 introspection 列表里，顺便拿到列清单做交集
		tables, cols, err := dbx.TablesAndColumns(c, e.Engine, dbName)
		if err != nil {
			ctx.SetStatusCode(500)
			ctx.SetBodyString(`{"status":2,"msg":"读表结构失败"}`)
			return nil
		}
		validTable := false
		tableCols := map[string]bool{}
		for _, t := range tables {
			if t.Name == table {
				validTable = true
			}
		}
		for _, cl := range cols {
			if cl.Table == table {
				tableCols[cl.Name] = true
			}
		}
		if !validTable {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"表不存在"}`)
			return nil
		}

		// 列名 = 文件表头 ∩ 表列
		fileCols := orderedKeys(rows)
		useCols := []string{}
		for _, k := range fileCols {
			if tableCols[k] {
				useCols = append(useCols, k)
			}
		}
		if len(useCols) == 0 {
			ctx.SetStatusCode(400)
			ctx.SetBodyString(`{"status":2,"msg":"文件列名与表列没有交集"}`)
			return nil
		}

		inserted, skipped, err := bulkInsert(c, e.Engine, table, useCols, rows)
		if err != nil {
			ctx.SetStatusCode(500)
			ctx.SetBodyString(`{"status":2,"msg":"导入失败：` + err.Error() + `"}`)
			return nil
		}
		ctx.SetContentType("application/json")
		fmt.Fprintf(ctx, `{"status":0,"inserted":%d,"skipped":%d}`, inserted, skipped)
		return nil
	}
}

// bulkInsert 逐行插入（不开事务：PG 事务内单行失败会整批回滚，
// 逐行保证「能导多少导多少」，跳过坏行计数返回）。单次最多 1 万行。
func bulkInsert(d *sql.DB, engine, table string, cols []string, rows []map[string]any) (int64, int64, error) {
	if len(rows) > 10000 {
		rows = rows[:10000]
	}
	quote := func(s string) string { return `"` + strings.ReplaceAll(s, `"`, `""`) + `"` }
	if engine == "mysql" {
		quote = func(s string) string { return "`" + strings.ReplaceAll(s, "`", "``") + "`" }
	}
	quoted := make([]string, len(cols))
	placeholders := make([]string, len(cols))
	for i, c := range cols {
		quoted[i] = quote(c)
		if engine == "postgres" {
			placeholders[i] = "$" + strconv.Itoa(i+1)
		} else {
			placeholders[i] = "?"
		}
	}
	stmt := fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s)",
		quote(table), strings.Join(quoted, ","), strings.Join(placeholders, ","))
	ps, err := d.Prepare(stmt)
	if err != nil {
		return 0, 0, err
	}
	defer ps.Close()
	var inserted, skipped int64
	for _, r := range rows {
		vals := make([]any, len(cols))
		for i, c := range cols {
			vals[i] = r[c]
		}
		if _, err := ps.Exec(vals...); err != nil {
			skipped++
			continue
		}
		inserted++
	}
	return inserted, skipped, nil
}

// parseImport 按扩展名分派：.json 对象数组 / .xlsx 首行表头 / 其余按 CSV（兼容 BOM）
func parseImport(filename string, raw []byte) ([]map[string]any, error) {
	switch {
	case strings.HasSuffix(filename, ".json"):
		var rows []map[string]any
		if err := json.Unmarshal(raw, &rows); err != nil {
			return nil, fmt.Errorf("JSON 解析失败：文件应为对象数组 [{...}, ...]")
		}
		return rows, nil
	case strings.HasSuffix(filename, ".xlsx"):
		f, err := excelize.OpenReader(bytes.NewReader(raw))
		if err != nil {
			return nil, fmt.Errorf("Excel 解析失败")
		}
		defer f.Close()
		sheets := f.GetSheetList()
		if len(sheets) == 0 {
			return nil, fmt.Errorf("Excel 文件没有工作表")
		}
		recs, err := f.GetRows(sheets[0])
		if err != nil {
			return nil, fmt.Errorf("读取工作表失败")
		}
		return rowsFromTable(recs), nil
	default:
		r := csv.NewReader(bytes.NewReader(bytes.TrimPrefix(raw, []byte{0xEF, 0xBB, 0xBF})))
		r.FieldsPerRecord = -1 // 容忍长短不一的行
		recs, err := r.ReadAll()
		if err != nil {
			return nil, fmt.Errorf("CSV 解析失败")
		}
		return rowsFromTable(recs), nil
	}
}

// rowsFromTable 二维表（首行表头）→ 对象数组
func rowsFromTable(recs [][]string) []map[string]any {
	if len(recs) < 2 {
		return nil
	}
	head := recs[0]
	out := []map[string]any{}
	for _, rec := range recs[1:] {
		m := map[string]any{}
		for i, h := range head {
			h = strings.TrimSpace(h)
			if h == "" {
				continue
			}
			if i < len(rec) {
				m[h] = rec[i]
			}
		}
		out = append(out, m)
	}
	return out
}

// orderedKeys 取文件列名（按首行出现顺序）
func orderedKeys(rows []map[string]any) []string {
	seen := map[string]bool{}
	out := []string{}
	for _, r := range rows {
		for k := range r {
			if !seen[k] {
				seen[k] = true
				out = append(out, k)
			}
		}
	}
	return out
}