opsxdb/server/internal/db/export.go

package db

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

const (
	DefaultExportLimit = 100000
	MaxExportLimit     = 500000
	exportFlushRows    = 512
)

// ExportAllowedQuery 导出专用的只读校验:首个关键词须为 SELECT/WITH/SHOW/DESCRIBE/EXPLAIN。
func ExportAllowedQuery(sqlText string) bool {
	s := stripSQLLead(strings.TrimSpace(sqlText))
	if s == "" {
		return false
	}
	up := strings.ToUpper(s)
	for _, p := range []string{"SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN"} {
		if strings.HasPrefix(up, p+" ") || up == p || strings.HasPrefix(up, p+"\t") || strings.HasPrefix(up, p+"\n") {
			return true
		}
	}
	return false
}

// SelectAllSQL 表模式导出的取数 SQL:按 family 引用(支持 schema.table 多段)。
func SelectAllSQL(driver, table string) string {
	return "SELECT * FROM " + quoteTableNameFor(driver, table)
}

// ExportQuery 把查询结果按 c.Format(csv/json/insert)流式编码写入 w,避免全量载入内存。
// limit <= 0 用默认上限 100000,硬顶 MaxExportLimit。行值经 normalize 保真(同 ExecSQL)。
func ExportQuery(ctx context.Context, c ConnReq, w io.Writer) error {
	ctx = reqCtx(ctx)
	format := strings.ToLower(strings.TrimSpace(c.Format))
	sqlText := strings.TrimSpace(c.SQL)
	if sqlText == "" {
		return fmt.Errorf("缺少导出 SQL")
	}
	limit := c.Limit
	if limit <= 0 {
		limit = DefaultExportLimit
	}
	if limit > MaxExportLimit {
		limit = MaxExportLimit
	}
	return withDB(c, func(db *sql.DB) error {
		rows, err := db.QueryContext(ctx, sqlText)
		if err != nil {
			return err
		}
		defer rows.Close()
		cols, err := rows.Columns()
		if err != nil {
			return err
		}
		bw := bufio.NewWriterSize(w, 64*1024)
		enc := newExportEncoder(bw, format, c.Driver, exportTargetTable(c), cols)
		if err := enc.begin(); err != nil {
			return err
		}
		n := 0
		for rows.Next() {
			if err := ctx.Err(); err != nil {
				return err
			}
			raw := make([]interface{}, len(cols))
			ptrs := make([]interface{}, len(cols))
			for i := range raw {
				ptrs[i] = &raw[i]
			}
			if err := rows.Scan(ptrs...); err != nil {
				return err
			}
			vals := make([]interface{}, len(cols))
			for i := range raw {
				vals[i] = normalize(raw[i])
			}
			if err := enc.row(vals); err != nil {
				return err
			}
			n++
			if n%exportFlushRows == 0 {
				if err := bw.Flush(); err != nil {
					return err
				}
			}
			if n >= limit {
				break
			}
		}
		if err := rows.Err(); err != nil {
			return err
		}
		if err := enc.end(); err != nil {
			return err
		}
		return bw.Flush()
	})
}

// exportTargetTable insert 格式的目标表名:显式 table 优先,否则用通用名。
func exportTargetTable(c ConnReq) string {
	if strings.TrimSpace(c.Table) != "" {
		return c.Table
	}
	return "result"
}

// exportEncoder 按格式把行流式编码进 buffered writer。
type exportEncoder struct {
	bw     *bufio.Writer
	format string
	driver string
	table  string
	cols   []string
	csvW   *csv.Writer
	// jsonPending 上一行 JSON,延迟一行写出,使行间逗号紧跟行尾而非独占一行
	jsonPending string
	n           int
}

func newExportEncoder(bw *bufio.Writer, format, driver, table string, cols []string) *exportEncoder {
	return &exportEncoder{bw: bw, format: format, driver: driver, table: table, cols: cols}
}

func (e *exportEncoder) begin() error {
	switch e.format {
	case "csv":
		e.csvW = csv.NewWriter(e.bw)
		return e.csvW.Write(e.cols)
	case "json":
		_, err := e.bw.WriteString("[\n")
		return err
	case "insert":
		return nil
	default:
		return fmt.Errorf("不支持的导出格式: %s", e.format)
	}
}

func (e *exportEncoder) row(vals []interface{}) error {
	e.n++
	switch e.format {
	case "csv":
		rec := make([]string, len(vals))
		for i, v := range vals {
			rec[i] = csvCell(v)
		}
		return e.csvW.Write(rec)
	case "json":
		m := make(map[string]interface{}, len(e.cols))
		for i, col := range e.cols {
			m[col] = vals[i]
		}
		var rb bytes.Buffer
		if err := json.NewEncoder(&rb).Encode(m); err != nil {
			return err
		}
		line := strings.TrimSuffix(rb.String(), "\n")
		if e.jsonPending != "" {
			if _, err := e.bw.WriteString(e.jsonPending + ",\n"); err != nil {
				return err
			}
		}
		e.jsonPending = line
		return nil
	case "insert":
		var b strings.Builder
		b.WriteString("INSERT INTO ")
		b.WriteString(quoteTableNameFor(e.driver, e.table))
		b.WriteString(" (")
		for i, col := range e.cols {
			if i > 0 {
				b.WriteString(", ")
			}
			b.WriteString(quoteIdentFor(e.driver, col))
		}
		b.WriteString(") VALUES (")
		for i, v := range vals {
			if i > 0 {
				b.WriteString(", ")
			}
			b.WriteString(sqlLiteral(v))
		}
		b.WriteString(");\n")
		_, err := e.bw.WriteString(b.String())
		return err
	}
	return nil
}

func (e *exportEncoder) end() error {
	switch e.format {
	case "csv":
		e.csvW.Flush()
		return e.csvW.Error()
	case "json":
		if e.jsonPending != "" {
			if _, err := e.bw.WriteString(e.jsonPending + "\n"); err != nil {
				return err
			}
			e.jsonPending = ""
		}
		_, err := e.bw.WriteString("]\n")
		return err
	}
	return nil
}

func csvCell(v interface{}) string {
	switch t := v.(type) {
	case nil:
		return ""
	case string:
		return t
	default:
		return fmt.Sprintf("%v", t)
	}
}

// sqlLiteral insert 字面量:NULL/数字原样,字符串单引号 doubling。
func sqlLiteral(v interface{}) string {
	switch t := v.(type) {
	case nil:
		return "NULL"
	case string:
		return "'" + strings.ReplaceAll(t, "'", "''") + "'"
	case bool:
		if t {
			return "1"
		}
		return "0"
	case time.Time:
		return "'" + t.Format(time.RFC3339Nano) + "'"
	case []byte:
		return "'" + strings.ReplaceAll(string(t), "'", "''") + "'"
	default:
		return fmt.Sprintf("%v", t)
	}
}