opsxsql/backend/internal/handler/redis_export.go

package handler

import (
	"bufio"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"sql-platform/config"
	"sql-platform/internal/handler/resp"
	"sql-platform/internal/model"
)

// Redis 大 Key 导出的两段式设计：
//
//	① POST /redis/:id/keys/:key/export-ticket  （走 Bearer 鉴权）→ 一次性短时效票据
//	② GET  /redis/:id/keys/:key/export?ticket=  （不走 Bearer，凭证在票据里）
//
// 为什么不给导出接口直接挂 Bearer：浏览器原生下载（<a download>/window.open）带不了
// Authorization 头；若前端用 axios 取 blob，则整份导出物会进浏览器内存，与「大 key
// 不爆内存」的目标相悖。改为短时效、绑定具体 key 的票据后，浏览器可以边收边写磁盘。

const redisExportTicketTTL = 2 * time.Minute

const (
	redisExportMaxRows    = 2_000_000
	redisExportBatchRows  = 1000
	redisExportChunkBytes = 256 << 10 // 超大 string 按 256KB 分块 GETRANGE
)

func redisExportSignature(instanceID, db, key, exp string) string {
	mac := hmac.New(sha256.New, []byte(config.C.JWTSecret))
	mac.Write([]byte(strings.Join([]string{"redis-export", instanceID, db, key, exp}, "\x00")))
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

func buildRedisExportTicket(instanceID, db, key string) (string, int64) {
	exp := strconv.FormatInt(time.Now().Add(redisExportTicketTTL).Unix(), 10)
	return exp + "." + redisExportSignature(instanceID, db, key, exp), int64(redisExportTicketTTL.Seconds())
}

func verifyRedisExportTicket(ticket, instanceID, db, key string) error {
	parts := strings.SplitN(ticket, ".", 2)
	if len(parts) != 2 {
		return fmt.Errorf("导出票据格式不正确，请重新点击导出")
	}
	expUnix, err := strconv.ParseInt(parts[0], 10, 64)
	if err != nil {
		return fmt.Errorf("导出票据格式不正确，请重新点击导出")
	}
	if time.Now().Unix() > expUnix {
		return fmt.Errorf("导出链接已过期（有效期 %d 秒），请重新点击导出", int(redisExportTicketTTL.Seconds()))
	}
	want := redisExportSignature(instanceID, db, key, parts[0])
	if !hmac.Equal([]byte(want), []byte(parts[1])) {
		return fmt.Errorf("导出票据无效")
	}
	return nil
}

// POST /api/v1/redis/:id/keys/:key/export-ticket
func RedisCreateExportTicket(c *gin.Context) {
	id := c.Param("id")
	key := c.Param("key")
	db := c.DefaultQuery("db", "0")
	if key == "" {
		resp.BadRequest(c, "缺少 key")
		return
	}
	ticket, ttl := buildRedisExportTicket(id, db, key)
	resp.OK(c, gin.H{"ticket": ticket, "url": c.Request.URL.Path, "expiresIn": ttl})
}

// GET /api/v1/redis/:id/keys/:key/export?db=0&format=csv|jsonl&ticket=xxx
//
// 全程流式：按批从 Redis 取数 → 立即写响应 → flush。
// 服务端内存只占用「一批 + 64KB 缓冲」，与 key 总大小无关。
func RedisExportKeyStream(c *gin.Context) {
	id := c.Param("id")
	key := c.Param("key")
	dbParam := c.DefaultQuery("db", "0")
	format := strings.ToLower(c.DefaultQuery("format", "csv"))
	if format == "json" {
		format = "jsonl"
	}
	if format != "csv" && format != "jsonl" {
		resp.BadRequest(c, "format 只支持 csv / jsonl")
		return
	}
	if err := verifyRedisExportTicket(c.Query("ticket"), id, dbParam, key); err != nil {
		resp.Forbidden(c, err.Error())
		return
	}

	db, _ := strconv.Atoi(dbParam)
	var instance model.Instance
	if err := model.DB.First(&instance, id).Error; err != nil {
		resp.NotFound(c, "Instance not found")
		return
	}
	client, err := getRedisClient(&instance, db)
	if err != nil {
		resp.InternalErr(c, err)
		return
	}
	defer client.Close()

	ctx := c.Request.Context()
	keyType, _ := client.Type(ctx, key).Result()
	if keyType == "none" {
		resp.NotFound(c, "Key 不存在")
		return
	}

	filename := sanitizeExportFilename(key) + exportExt(format)
	c.Header("Content-Type", exportContentType(format))
	c.Header("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filename))
	c.Header("X-Content-Type-Options", "nosniff")
	c.Header("Cache-Control", "no-store")
	c.Status(http.StatusOK)

	w := bufio.NewWriterSize(c.Writer, 64<<10)
	flusher, _ := c.Writer.(http.Flusher)
	flush := func() {
		_ = w.Flush()
		if flusher != nil {
			flusher.Flush()
		}
	}

	ew := &redisExportWriter{w: w, format: format}
	ew.begin()
	if err := streamRedisKey(ctx, client, key, keyType, ew, flush); err != nil {
		// 响应头已发出，无法改状态码；把错误写进文件尾部，方便使用者与排障。
		ew.errorRow("导出中断: " + err.Error())
	}
	ew.end()
	flush()
}

// redisExportWriter 统一的 CSV / JSONL 输出器
type redisExportWriter struct {
	w        *bufio.Writer
	format   string
	cw       *csv.Writer
	rowCount int
}

func (e *redisExportWriter) begin() {
	if e.format == "csv" {
		e.cw = csv.NewWriter(e.w)
		// Excel 打开 UTF-8 CSV 的正确方式：写 BOM
		e.w.Write([]byte{0xEF, 0xBB, 0xBF})
	}
}

func (e *redisExportWriter) header(cols ...string) {
	if e.format == "csv" {
		_ = e.cw.Write(cols)
		e.cw.Flush()
		return
	}
	e.writeJSON(map[string]any{"_type": "header", "columns": cols})
}

func (e *redisExportWriter) row(cols ...string) {
	e.rowCount++
	if e.format == "csv" {
		_ = e.cw.Write(cols)
		if e.rowCount%redisExportBatchRows == 0 {
			e.cw.Flush()
		}
		return
	}
	obj := make(map[string]any, 2)
	obj["value"] = cols[len(cols)-1]
	e.writeJSON(obj)
}

func (e *redisExportWriter) rowKV(k, v string, extra ...string) {
	e.rowCount++
	if e.format == "csv" {
		_ = e.cw.Write(append([]string{k, v}, extra...))
		if e.rowCount%redisExportBatchRows == 0 {
			e.cw.Flush()
		}
		return
	}
	obj := map[string]any{"key": k, "value": v}
	if len(extra) > 0 {
		obj["score"] = extra[0]
	}
	e.writeJSON(obj)
}

func (e *redisExportWriter) writeJSON(v any) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	_, _ = e.w.Write(b)
	_ = e.w.WriteByte('\n')
}

func (e *redisExportWriter) errorRow(msg string) {
	if e.format == "csv" {
		_ = e.cw.Write([]string{"__error__", msg})
		e.cw.Flush()
		return
	}
	e.writeJSON(map[string]any{"_error": msg})
}

func (e *redisExportWriter) end() {
	if e.format == "csv" && e.cw != nil {
		e.cw.Flush()
	}
}

// streamRedisKey 分批读取并写出，任何时刻只有一批数据在内存里。
func streamRedisKey(
	ctx context.Context,
	client *redisClientWrapper,
	key, keyType string,
	ew *redisExportWriter,
	flush func(),
) error {
	rows := 0
	budget := func() bool { return rows < redisExportMaxRows }

	switch keyType {
	case "string":
		ew.header("offset", "content")
		total, _ := client.StrLen(ctx, key).Result()
		var offset int64
		for offset < total && budget() {
			chunk, err := client.GetRange(ctx, key, offset, offset+redisExportChunkBytes-1).Result()
			if err != nil {
				return err
			}
			if chunk == "" {
				break
			}
			ew.rowKV(strconv.FormatInt(offset, 10), chunk)
			rows++
			offset += int64(len(chunk))
			flush()
		}

	case "list":
		ew.header("index", "value")
		total, _ := client.LLen(ctx, key).Result()
		offset := int64(0)
		for offset < total && budget() {
			items, err := client.LRange(ctx, key, offset, offset+redisExportBatchRows-1).Result()
			if err != nil {
				return err
			}
			if len(items) == 0 {
				break
			}
			for i, v := range items {
				ew.rowKV(strconv.FormatInt(offset+int64(i), 10), v)
				rows++
			}
			offset += int64(len(items))
			flush()
		}

	case "set":
		ew.header("member")
		var cursor uint64
		for budget() {
			members, next, err := client.SScan(ctx, key, cursor, "", redisExportBatchRows).Result()
			if err != nil {
				return err
			}
			for _, m := range members {
				ew.row(m)
				rows++
			}
			flush()
			if next == 0 {
				break
			}
			cursor = next
		}

	case "zset":
		ew.header("rank", "member", "score")
		total, _ := client.ZCard(ctx, key).Result()
		offset := int64(0)
		for offset < total && budget() {
			items, err := client.ZRangeWithScores(ctx, key, offset, offset+redisExportBatchRows-1).Result()
			if err != nil {
				return err
			}
			if len(items) == 0 {
				break
			}
			for i, z := range items {
				ew.rowKV(strconv.FormatInt(offset+int64(i), 10),
					fmt.Sprint(z.Member), strconv.FormatFloat(z.Score, 'f', -1, 64))
				rows++
			}
			offset += int64(len(items))
			flush()
		}

	case "hash":
		ew.header("field", "value")
		var cursor uint64
		for budget() {
			flat, next, err := client.HScan(ctx, key, cursor, "", redisExportBatchRows).Result()
			if err != nil {
				return err
			}
			for i := 0; i+1 < len(flat); i += 2 {
				ew.rowKV(flat[i], flat[i+1])
				rows++
			}
			flush()
			if next == 0 {
				break
			}
			cursor = next
		}

	default:
		return fmt.Errorf("不支持导出的类型: %s", keyType)
	}

	if rows >= redisExportMaxRows {
		ew.errorRow(fmt.Sprintf("已达单次导出上限 %d 行，后续数据未包含", redisExportMaxRows))
	}
	return nil
}

func exportExt(format string) string {
	if format == "jsonl" {
		return ".jsonl"
	}
	return ".csv"
}

func exportContentType(format string) string {
	if format == "jsonl" {
		return "application/x-ndjson; charset=utf-8"
	}
	return "text/csv; charset=utf-8"
}

// sanitizeExportFilename 把 key 转成安全文件名（去掉路径分隔符等）
func sanitizeExportFilename(key string) string {
	var b strings.Builder
	for _, r := range key {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == ':' || r == '-' || r == '_' || r == '.':
			b.WriteRune(r)
		default:
			b.WriteRune('_')
		}
	}
	s := b.String()
	if s == "" {
		s = "redis-key"
	}
	if len(s) > 120 {
		s = s[:120]
	}
	return s
}